<commit_before>use std::io::Write;
use time::get_time;

use {Sender};

impl<'a> Sender<'a> {
    pub fn add_u64(&mut self, name: &str, value: u64) {
        let ts = get_time().sec;
        writeln!(self.0.buffer, "{} {} {}", name, value, ts).unwrap();
    }
    pub fn add_int<I: Into<i64>=i64>(&mut self, name: &str, value: I) {
        let ts = get_time().sec;
        let val = value.into();
        writeln!(self.0.buffer, "{} {} {}", name, val, ts).unwrap();
    }
}

impl<'a> Drop for Sender<'a> {
    fn drop(&mut self) {
        if self.1 {
            self.0.notify.wakeup().unwrap()
        }
    }
}
<commit_msg>Add timestamped notifiers<commit_after>use std::io::Write;
use time::get_time;

use {Sender};

impl<'a> Sender<'a> {
    /// Send a generic convertible to i64
    pub fn add_int<I: Into<i64>=i64>(&mut self, name: &str, value: I) {
        let ts = get_time().sec;
        let val = value.into();
        writeln!(self.0.buffer, "{} {} {}", name, val, ts).unwrap();
    }
    /// The special case for u64, as it's the only type doesn't fit in i64
    pub fn add_u64(&mut self, name: &str, value: u64) {
        let ts = get_time().sec;
        writeln!(self.0.buffer, "{} {} {}", name, value, ts).unwrap();
    }
    /// Send a generic convertible to i64 with a timestamp
    pub fn add_int_at<I=i64>(&mut self, name: &str, value: I, ts: u64)
        where I: Into<i64>
    {
        let val = value.into();
        writeln!(self.0.buffer, "{} {} {}", name, val, ts).unwrap();
    }
    /// Send an u64 with a timestamp
    pub fn add_u64_at(&mut self, name: &str, value: u64, ts: u64) {
        writeln!(self.0.buffer, "{} {} {}", name, value, ts).unwrap();
    }
}

impl<'a> Drop for Sender<'a> {
    fn drop(&mut self) {
        if self.1 {
            self.0.notify.wakeup().unwrap()
        }
    }
}
<|endoftext|>


<commit_before>//! Serde integration support.

use std::fmt;
use std::marker::PhantomData;

use serde::*;
use serde::de::{Visitor, Error, Unexpected};

fn i64_to_u64<'d, V: Visitor<'d>, E: Error>(v: V, n: i64) -> Result<V::Value, E> {
	if n >= 0 {
		v.visit_u64(n as u64)
	} else {
		Err(E::invalid_value(Unexpected::Signed(n), &v))
	}
}

/// Ignore deserialization errors and revert to default.
pub fn ignore_errors<'d, T: Deserialize<'d> + Default, D: Deserializer<'d>>(d: D) -> Result<T, D::Error> {
	Ok(T::deserialize(d).ok().unwrap_or_default())
}

/// Deserialize a maybe-string ID into a u64.
pub fn deserialize_id<'d, D: Deserializer<'d>>(d: D) -> Result<u64, D::Error> {
	struct IdVisitor;
	impl<'d> Visitor<'d> for IdVisitor {
		type Value = u64;

		fn expecting(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
			write!(fmt, "a u64 or parseable string")
		}

		fn visit_i64<E: Error>(self, v: i64) -> Result<u64, E> {
			i64_to_u64(self, v)
		}

		fn visit_u64<E: Error>(self, v: u64) -> Result<u64, E> {
			Ok(v)
		}

		fn visit_str<E: Error>(self, v: &str) -> Result<u64, E> {
			v.parse::<u64>().map_err(|_| E::invalid_value(Unexpected::Str(v), &self))
		}
	}

	d.deserialize_u64(IdVisitor)
}

/// Deserialize a maybe-string discriminator into a u16.
/// Also enforces 0 <= N <= 9999.
#[allow(unused_comparisons)]
pub fn deserialize_discrim<'d, D: Deserializer<'d>>(d: D) -> Result<u16, D::Error> {
	macro_rules! check {
		($self:ident, $v:ident, $wrong:expr) => {
			if $v >= 0 && $v <= 9999 {
				Ok($v as u16)
			} else {
				Err(E::invalid_value($wrong, &$self))
			}
		}
	}

	struct DiscrimVisitor;
	impl<'d> Visitor<'d> for DiscrimVisitor {
		type Value = u16;

		fn expecting(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
			write!(fmt, "a u16 in [0, 9999] or parseable string")
		}

		fn visit_i64<E: Error>(self, v: i64) -> Result<u16, E> {
			check!(self, v, Unexpected::Signed(v))
		}

		fn visit_u64<E: Error>(self, v: u64) -> Result<u16, E> {
			check!(self, v, Unexpected::Unsigned(v))
		}

		fn visit_str<E: Error>(self, v: &str) -> Result<u16, E> {
			v.parse::<u16>()
				.map_err(|_| E::invalid_value(Unexpected::Str(v), &self))
				.and_then(|v| self.visit_u16(v))
		}
	}

	d.deserialize_u16(DiscrimVisitor)
}

/// Deserialize a single-field struct like a newtype struct.
macro_rules! serial_single_field {
	($typ:ident as $field:ident: $inner:path) => {
		impl ::serde::Serialize for $typ {
			fn serialize<S: ::serde::ser::Serializer>(&self, s: S) -> ::std::result::Result<S::Ok, S::Error> {
				self.$field.serialize(s)
			}
		}

		impl<'d> ::serde::Deserialize<'d> for $typ {
			fn deserialize<D: ::serde::de::Deserializer<'d>>(d: D) -> ::std::result::Result<$typ, D::Error> {
				<$inner as ::serde::de::Deserialize>::deserialize(d).map(|v| $typ { $field: v })
			}
		}
	}
}

/// Special support for the oddly complex `ReactionEmoji`.
pub mod reaction_emoji {
	use super::*;
	use model::{ReactionEmoji, EmojiId};

	#[derive(Serialize)]
	struct EmojiSer<'s> {
		name: &'s str,
		id: Option<EmojiId>
	}

	#[derive(Deserialize)]
	struct EmojiDe {
		name: String,
		id: Option<EmojiId>,
	}

	pub fn serialize<S: Serializer>(v: &ReactionEmoji, s: S) -> Result<S::Ok, S::Error> {
		(match *v {
			ReactionEmoji::Unicode(ref name) => EmojiSer { name: name, id: None },
			ReactionEmoji::Custom { ref name, id } => EmojiSer { name: name, id: Some(id) },
		}).serialize(s)
	}

	pub fn deserialize<'d, D: Deserializer<'d>>(d: D) -> Result<ReactionEmoji, D::Error> {
		Ok(match try!(EmojiDe::deserialize(d)) {
			EmojiDe { name, id: None } => ReactionEmoji::Unicode(name),
			EmojiDe { name, id: Some(id) } => ReactionEmoji::Custom { name: name, id: id },
		})
	}
}

/// Support for named enums.
pub mod named {
	use super::*;

	pub trait NamedEnum: Sized {
		fn name(&self) -> &'static str;
		fn from_name(name: &str) -> Option<Self>;
		fn typename() -> &'static str;
	}

	pub fn serialize<T: NamedEnum, S: Serializer>(v: &T, s: S) -> Result<S::Ok, S::Error> {
		v.name().serialize(s)
	}

	pub fn deserialize<'d, T: NamedEnum, D: Deserializer<'d>>(d: D) -> Result<T, D::Error> {
		struct NameVisitor<T>(PhantomData<T>);
		impl<'d, T: NamedEnum> Visitor<'d> for NameVisitor<T> {
			type Value = T;

			fn expecting(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
				write!(fmt, "a valid {} name", T::typename())
			}

			fn visit_str<E: Error>(self, v: &str) -> Result<T, E> {
				T::from_name(v).ok_or_else(|| E::invalid_value(Unexpected::Str(v), &self))
			}
		}

		d.deserialize_string(NameVisitor(PhantomData))
	}
}
macro_rules! serial_names {
	($typ:ident; $($entry:ident, $value:expr;)*) => {
		impl $typ {
			pub fn name(&self) -> &'static str {
				match *self {
					$($typ::$entry => $value,)*
				}
			}

			pub fn from_name(name: &str) -> Option<Self> {
				match name {
					$($value => Some($typ::$entry),)*
					_ => None,
				}
			}
		}

		impl ::serial::named::NamedEnum for $typ {
			fn name(&self) -> &'static str {
				self.name()
			}

			fn from_name(name: &str) -> Option<Self> {
				Self::from_name(name)
			}

			fn typename() -> &'static str {
				stringify!($typ)
			}
		}
	}
}

/// Support for numeric enums.
pub mod numeric {
	use super::*;

	pub trait NumericEnum: Sized {
		fn num(&self) -> u64;
		fn from_num(num: u64) -> Option<Self>;
		fn typename() -> &'static str;
	}

	pub fn serialize<T: NumericEnum, S: Serializer>(v: &T, s: S) -> Result<S::Ok, S::Error> {
		v.num().serialize(s)
	}

	pub fn deserialize<'d, T: NumericEnum, D: Deserializer<'d>>(d: D) -> Result<T, D::Error> {
		struct NumVisitor<T>(PhantomData<T>);
		impl<'d, T: NumericEnum> Visitor<'d> for NumVisitor<T> {
			type Value = T;

			fn expecting(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
				write!(fmt, "a valid {} number", T::typename())
			}

			fn visit_i64<E: Error>(self, v: i64) -> Result<T, E> {
				i64_to_u64(self, v)
			}

			fn visit_u64<E: Error>(self, v: u64) -> Result<T, E> {
				T::from_num(v).ok_or_else(|| E::invalid_value(Unexpected::Unsigned(v), &self))
			}
		}

		d.deserialize_string(NumVisitor(PhantomData))
	}
}
macro_rules! serial_numbers {
	($typ:ident; $($entry:ident, $value:expr;)*) => {
		impl $typ {
			pub fn num(&self) -> u64 {
				match *self {
					$($typ::$entry => $value,)*
				}
			}

			pub fn from_num(num: u64) -> Option<Self> {
				match num {
					$($value => Some($typ::$entry),)*
					_ => None,
				}
			}
		}
		impl ::serial::numeric::NumericEnum for $typ {
			fn num(&self) -> u64 {
				self.num()
			}

			fn from_num(num: u64) -> Option<Self> {
				Self::from_num(num)
			}

			fn typename() -> &'static str {
				stringify!($typ)
			}
		}
	}
}

/// Support for using "named" or "numeric" as the default ser/de impl.
macro_rules! serial_use_mapping {
	($typ:ident, $which:ident) => {
		impl ::serde::Serialize for $typ {
			#[inline]
			fn serialize<S: ::serde::ser::Serializer>(&self, s: S) -> ::std::result::Result<S::Ok, S::Error> {
				::serial::$which::serialize(self, s)
			}
		}

		impl<'d> ::serde::Deserialize<'d> for $typ {
			#[inline]
			fn deserialize<D: ::serde::de::Deserializer<'d>>(d: D) -> ::std::result::Result<$typ, D::Error> {
				::serial::$which::deserialize(d)
			}
		}
	}
}
<commit_msg>Fixed serialization<commit_after>//! Serde integration support.

use std::fmt;
use std::marker::PhantomData;

use serde::*;
use serde::de::{Visitor, Error, Unexpected};

fn i64_to_u64<'d, V: Visitor<'d>, E: Error>(v: V, n: i64) -> Result<V::Value, E> {
	if n >= 0 {
		v.visit_u64(n as u64)
	} else {
		Err(E::invalid_value(Unexpected::Signed(n), &v))
	}
}

/// Ignore deserialization errors and revert to default.
pub fn ignore_errors<'d, T: Deserialize<'d> + Default, D: Deserializer<'d>>(d: D) -> Result<T, D::Error> {
	Ok(T::deserialize(d).ok().unwrap_or_default())
}

/// Deserialize a maybe-string ID into a u64.
pub fn deserialize_id<'d, D: Deserializer<'d>>(d: D) -> Result<u64, D::Error> {
	struct IdVisitor;
	impl<'d> Visitor<'d> for IdVisitor {
		type Value = u64;

		fn expecting(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
			write!(fmt, "a u64 or parseable string")
		}

		fn visit_i64<E: Error>(self, v: i64) -> Result<u64, E> {
			i64_to_u64(self, v)
		}

		fn visit_u64<E: Error>(self, v: u64) -> Result<u64, E> {
			Ok(v)
		}

		fn visit_str<E: Error>(self, v: &str) -> Result<u64, E> {
			v.parse::<u64>().map_err(|_| E::invalid_value(Unexpected::Str(v), &self))
		}
	}

	d.deserialize_any(IdVisitor)
}

/// Deserialize a maybe-string discriminator into a u16.
/// Also enforces 0 <= N <= 9999.
#[allow(unused_comparisons)]
pub fn deserialize_discrim<'d, D: Deserializer<'d>>(d: D) -> Result<u16, D::Error> {
	macro_rules! check {
		($self:ident, $v:ident, $wrong:expr) => {
			if $v >= 0 && $v <= 9999 {
				Ok($v as u16)
			} else {
				Err(E::invalid_value($wrong, &$self))
			}
		}
	}

	struct DiscrimVisitor;
	impl<'d> Visitor<'d> for DiscrimVisitor {
		type Value = u16;

		fn expecting(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
			write!(fmt, "a u16 in [0, 9999] or parseable string")
		}

		fn visit_i64<E: Error>(self, v: i64) -> Result<u16, E> {
			check!(self, v, Unexpected::Signed(v))
		}

		fn visit_u64<E: Error>(self, v: u64) -> Result<u16, E> {
			check!(self, v, Unexpected::Unsigned(v))
		}

		fn visit_str<E: Error>(self, v: &str) -> Result<u16, E> {
			v.parse::<u16>()
				.map_err(|_| E::invalid_value(Unexpected::Str(v), &self))
				.and_then(|v| self.visit_u16(v))
		}
	}

	d.deserialize_any(DiscrimVisitor)
}

/// Deserialize a single-field struct like a newtype struct.
macro_rules! serial_single_field {
	($typ:ident as $field:ident: $inner:path) => {
		impl ::serde::Serialize for $typ {
			fn serialize<S: ::serde::ser::Serializer>(&self, s: S) -> ::std::result::Result<S::Ok, S::Error> {
				self.$field.serialize(s)
			}
		}

		impl<'d> ::serde::Deserialize<'d> for $typ {
			fn deserialize<D: ::serde::de::Deserializer<'d>>(d: D) -> ::std::result::Result<$typ, D::Error> {
				<$inner as ::serde::de::Deserialize>::deserialize(d).map(|v| $typ { $field: v })
			}
		}
	}
}

/// Special support for the oddly complex `ReactionEmoji`.
pub mod reaction_emoji {
	use super::*;
	use model::{ReactionEmoji, EmojiId};

	#[derive(Serialize)]
	struct EmojiSer<'s> {
		name: &'s str,
		id: Option<EmojiId>
	}

	#[derive(Deserialize)]
	struct EmojiDe {
		name: String,
		id: Option<EmojiId>,
	}

	pub fn serialize<S: Serializer>(v: &ReactionEmoji, s: S) -> Result<S::Ok, S::Error> {
		(match *v {
			ReactionEmoji::Unicode(ref name) => EmojiSer { name: name, id: None },
			ReactionEmoji::Custom { ref name, id } => EmojiSer { name: name, id: Some(id) },
		}).serialize(s)
	}

	pub fn deserialize<'d, D: Deserializer<'d>>(d: D) -> Result<ReactionEmoji, D::Error> {
		Ok(match try!(EmojiDe::deserialize(d)) {
			EmojiDe { name, id: None } => ReactionEmoji::Unicode(name),
			EmojiDe { name, id: Some(id) } => ReactionEmoji::Custom { name: name, id: id },
		})
	}
}

/// Support for named enums.
pub mod named {
	use super::*;

	pub trait NamedEnum: Sized {
		fn name(&self) -> &'static str;
		fn from_name(name: &str) -> Option<Self>;
		fn typename() -> &'static str;
	}

	pub fn serialize<T: NamedEnum, S: Serializer>(v: &T, s: S) -> Result<S::Ok, S::Error> {
		v.name().serialize(s)
	}

	pub fn deserialize<'d, T: NamedEnum, D: Deserializer<'d>>(d: D) -> Result<T, D::Error> {
		struct NameVisitor<T>(PhantomData<T>);
		impl<'d, T: NamedEnum> Visitor<'d> for NameVisitor<T> {
			type Value = T;

			fn expecting(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
				write!(fmt, "a valid {} name", T::typename())
			}

			fn visit_str<E: Error>(self, v: &str) -> Result<T, E> {
				T::from_name(v).ok_or_else(|| E::invalid_value(Unexpected::Str(v), &self))
			}
		}

		d.deserialize_any(NameVisitor(PhantomData))
	}
}
macro_rules! serial_names {
	($typ:ident; $($entry:ident, $value:expr;)*) => {
		impl $typ {
			pub fn name(&self) -> &'static str {
				match *self {
					$($typ::$entry => $value,)*
				}
			}

			pub fn from_name(name: &str) -> Option<Self> {
				match name {
					$($value => Some($typ::$entry),)*
					_ => None,
				}
			}
		}

		impl ::serial::named::NamedEnum for $typ {
			fn name(&self) -> &'static str {
				self.name()
			}

			fn from_name(name: &str) -> Option<Self> {
				Self::from_name(name)
			}

			fn typename() -> &'static str {
				stringify!($typ)
			}
		}
	}
}

/// Support for numeric enums.
pub mod numeric {
	use super::*;

	pub trait NumericEnum: Sized {
		fn num(&self) -> u64;
		fn from_num(num: u64) -> Option<Self>;
		fn typename() -> &'static str;
	}

	pub fn serialize<T: NumericEnum, S: Serializer>(v: &T, s: S) -> Result<S::Ok, S::Error> {
		v.num().serialize(s)
	}

	pub fn deserialize<'d, T: NumericEnum, D: Deserializer<'d>>(d: D) -> Result<T, D::Error> {
		struct NumVisitor<T>(PhantomData<T>);
		impl<'d, T: NumericEnum> Visitor<'d> for NumVisitor<T> {
			type Value = T;

			fn expecting(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
				write!(fmt, "a valid {} number", T::typename())
			}

			fn visit_i64<E: Error>(self, v: i64) -> Result<T, E> {
				i64_to_u64(self, v)
			}

			fn visit_u64<E: Error>(self, v: u64) -> Result<T, E> {
				T::from_num(v).ok_or_else(|| E::invalid_value(Unexpected::Unsigned(v), &self))
			}
		}

		d.deserialize_any(NumVisitor(PhantomData))
	}
}
macro_rules! serial_numbers {
	($typ:ident; $($entry:ident, $value:expr;)*) => {
		impl $typ {
			pub fn num(&self) -> u64 {
				match *self {
					$($typ::$entry => $value,)*
				}
			}

			pub fn from_num(num: u64) -> Option<Self> {
				match num {
					$($value => Some($typ::$entry),)*
					_ => None,
				}
			}
		}
		impl ::serial::numeric::NumericEnum for $typ {
			fn num(&self) -> u64 {
				self.num()
			}

			fn from_num(num: u64) -> Option<Self> {
				Self::from_num(num)
			}

			fn typename() -> &'static str {
				stringify!($typ)
			}
		}
	}
}

/// Support for using "named" or "numeric" as the default ser/de impl.
macro_rules! serial_use_mapping {
	($typ:ident, $which:ident) => {
		impl ::serde::Serialize for $typ {
			#[inline]
			fn serialize<S: ::serde::ser::Serializer>(&self, s: S) -> ::std::result::Result<S::Ok, S::Error> {
				::serial::$which::serialize(self, s)
			}
		}

		impl<'d> ::serde::Deserialize<'d> for $typ {
			#[inline]
			fn deserialize<D: ::serde::de::Deserializer<'d>>(d: D) -> ::std::result::Result<$typ, D::Error> {
				::serial::$which::deserialize(d)
			}
		}
	}
}
<|endoftext|>


<commit_before>use std::collections::HashMap;
use std::net::{Shutdown, SocketAddr};
use std::io;

use byteorder::{ByteOrder, BigEndian};
use mio::deprecated::{EventLoop, Handler, TryWrite};
use mio::*;
use mio::tcp::*;
use bytes::buf::{Buf, ByteBuf};

use constants::*;
use packet::{Packet, PacketMagic, PTYPES, EofError};
use queues::QueueHolder;
use worker::Worker;

struct GearmanRemote {
    socket: TcpStream,
    packet: Packet,
    queues: QueueHolder,
    worker: Worker,
    sendqueue: Vec<ByteBuf>,
    interest: Ready,
}

pub struct GearmanServer {
    pub socket: TcpListener,
    remotes: HashMap<Token, GearmanRemote>,
    token_counter: usize,
    queues: QueueHolder,
}

const SERVER_TOKEN: Token = Token(0);

impl GearmanRemote {
    pub fn new(socket: TcpStream, queues: QueueHolder) -> GearmanRemote {
        GearmanRemote {
            socket: socket,
            packet: Packet::new(),
            queues: queues,
            worker: Worker::new(),
            sendqueue: Vec::with_capacity(2), // Don't need many bufs
            interest: Ready::readable(),
        }
    }

    pub fn read(&mut self) -> Result<(), EofError> {
        match self.packet.from_socket(&mut self.socket, &mut self.worker, self.queues.clone())? {
            None => debug!("Done reading"),
            Some(p) => {
                info!("Queueing {:?}", &p);
                self.sendqueue.push(ByteBuf::from_slice(&p.to_byteslice()));
                self.interest.remove(Ready::readable());
                self.interest.insert(Ready::writable());
            }
        }
        Ok(())
    }

    pub fn write(&mut self) -> Result<(), io::Error> {
        while !self.sendqueue.is_empty() {
            if !self.sendqueue.first().unwrap().has_remaining() {
                self.sendqueue.pop();
                continue;
            }
            let buf = self.sendqueue.first_mut().unwrap();
            match self.socket.try_write(buf.bytes())? {
                None => break,
                Some(n) => {
                    buf.advance(n);
                    continue
                }
            };
        }
        if self.sendqueue.is_empty() {
            self.interest.remove(Ready::writable());
            self.interest.insert(Ready::readable());
        };
        Ok(())
    }

    pub fn shutdown(self) {
        match self.worker.job {
            Some(j) => {
                self.queues.clone().add_job(j)
            },
            None => {},
        }
        self.socket.shutdown(Shutdown::Both).unwrap();
    }
}

impl GearmanServer {
    pub fn new(server_socket: TcpListener, queues: QueueHolder) -> GearmanServer {
        GearmanServer {
            token_counter: 1,
            remotes: HashMap::new(),
            socket: server_socket,
            queues: queues,
        }
    }

}

impl Handler for GearmanServer {
    type Timeout = usize;
    type Message = ();

    fn ready(&mut self, event_loop: &mut EventLoop<GearmanServer>,
             token: Token, events: Ready)
    {
        if events.is_readable() {
            match token {
                SERVER_TOKEN => {
                    let remote_socket = match self.socket.accept() {
                        Err(e) => {
                            error!("Accept error: {}", e);
                            return;
                        },
                        Ok((sock, addr)) => sock,
                    };

                    self.token_counter += 1;
                    let new_token = Token(self.token_counter);

                    self.remotes.insert(new_token, GearmanRemote::new(remote_socket, self.queues.clone()));

                    event_loop.register(&self.remotes[&new_token].socket,
                                        new_token, Ready::readable(),
                                        PollOpt::edge() | PollOpt::oneshot()).unwrap();
                },
                token => {
                    let mut shutdown = false;
                    {
                        let mut remote = self.remotes.get_mut(&token).unwrap();
                        match remote.read() {
                            Ok(_) => event_loop.reregister(&remote.socket, token, remote.interest,
                                                          PollOpt::edge() | PollOpt::oneshot()).unwrap(),
                            Err(e) => {
                                shutdown = true;
                            }
                        }
                    }
                    if shutdown {
                        let remote = self.remotes.remove(&token).unwrap();
                        remote.shutdown();
                    }
                },
            }
        }

        if events.is_writable() {
            match token {
                SERVER_TOKEN => panic!("Received writable event for server socket."),
                token => {
                    let mut shutdown = false;
                    {
                        let mut remote = self.remotes.get_mut(&token).unwrap();
                        match remote.write() {
                            Ok(_) => event_loop.reregister(&remote.socket, token, remote.interest,
                                                          PollOpt::edge() | PollOpt::oneshot()).unwrap(),
                            Err(e) => {
                                shutdown = true;
                            }
                        }
                    }
                    if shutdown {
                        let remote = self.remotes.remove(&token).unwrap();
                        remote.shutdown();
                    }
                },
            }
        }
    }
}
<commit_msg>fix shutdown import<commit_after>use std::collections::HashMap;
use std::net::{SocketAddr};
use std::io;

use byteorder::{ByteOrder, BigEndian};
use mio::deprecated::{EventLoop, Handler, TryWrite};
use mio::*;
use mio::tcp::*;
use bytes::buf::{Buf, ByteBuf};

use constants::*;
use packet::{Packet, PacketMagic, PTYPES, EofError};
use queues::QueueHolder;
use worker::Worker;

struct GearmanRemote {
    socket: TcpStream,
    packet: Packet,
    queues: QueueHolder,
    worker: Worker,
    sendqueue: Vec<ByteBuf>,
    interest: Ready,
}

pub struct GearmanServer {
    pub socket: TcpListener,
    remotes: HashMap<Token, GearmanRemote>,
    token_counter: usize,
    queues: QueueHolder,
}

const SERVER_TOKEN: Token = Token(0);

impl GearmanRemote {
    pub fn new(socket: TcpStream, queues: QueueHolder) -> GearmanRemote {
        GearmanRemote {
            socket: socket,
            packet: Packet::new(),
            queues: queues,
            worker: Worker::new(),
            sendqueue: Vec::with_capacity(2), // Don't need many bufs
            interest: Ready::readable(),
        }
    }

    pub fn read(&mut self) -> Result<(), EofError> {
        match self.packet.from_socket(&mut self.socket, &mut self.worker, self.queues.clone())? {
            None => debug!("Done reading"),
            Some(p) => {
                info!("Queueing {:?}", &p);
                self.sendqueue.push(ByteBuf::from_slice(&p.to_byteslice()));
                self.interest.remove(Ready::readable());
                self.interest.insert(Ready::writable());
            }
        }
        Ok(())
    }

    pub fn write(&mut self) -> Result<(), io::Error> {
        while !self.sendqueue.is_empty() {
            if !self.sendqueue.first().unwrap().has_remaining() {
                self.sendqueue.pop();
                continue;
            }
            let buf = self.sendqueue.first_mut().unwrap();
            match self.socket.try_write(buf.bytes())? {
                None => break,
                Some(n) => {
                    buf.advance(n);
                    continue
                }
            };
        }
        if self.sendqueue.is_empty() {
            self.interest.remove(Ready::writable());
            self.interest.insert(Ready::readable());
        };
        Ok(())
    }

    pub fn shutdown(self) {
        match self.worker.job {
            Some(j) => {
                self.queues.clone().add_job(j)
            },
            None => {},
        }
        self.socket.shutdown(Shutdown::Both).unwrap();
    }
}

impl GearmanServer {
    pub fn new(server_socket: TcpListener, queues: QueueHolder) -> GearmanServer {
        GearmanServer {
            token_counter: 1,
            remotes: HashMap::new(),
            socket: server_socket,
            queues: queues,
        }
    }

}

impl Handler for GearmanServer {
    type Timeout = usize;
    type Message = ();

    fn ready(&mut self, event_loop: &mut EventLoop<GearmanServer>,
             token: Token, events: Ready)
    {
        if events.is_readable() {
            match token {
                SERVER_TOKEN => {
                    let remote_socket = match self.socket.accept() {
                        Err(e) => {
                            error!("Accept error: {}", e);
                            return;
                        },
                        Ok((sock, addr)) => sock,
                    };

                    self.token_counter += 1;
                    let new_token = Token(self.token_counter);

                    self.remotes.insert(new_token, GearmanRemote::new(remote_socket, self.queues.clone()));

                    event_loop.register(&self.remotes[&new_token].socket,
                                        new_token, Ready::readable(),
                                        PollOpt::edge() | PollOpt::oneshot()).unwrap();
                },
                token => {
                    let mut shutdown = false;
                    {
                        let mut remote = self.remotes.get_mut(&token).unwrap();
                        match remote.read() {
                            Ok(_) => event_loop.reregister(&remote.socket, token, remote.interest,
                                                          PollOpt::edge() | PollOpt::oneshot()).unwrap(),
                            Err(e) => {
                                shutdown = true;
                            }
                        }
                    }
                    if shutdown {
                        let remote = self.remotes.remove(&token).unwrap();
                        remote.shutdown();
                    }
                },
            }
        }

        if events.is_writable() {
            match token {
                SERVER_TOKEN => panic!("Received writable event for server socket."),
                token => {
                    let mut shutdown = false;
                    {
                        let mut remote = self.remotes.get_mut(&token).unwrap();
                        match remote.write() {
                            Ok(_) => event_loop.reregister(&remote.socket, token, remote.interest,
                                                          PollOpt::edge() | PollOpt::oneshot()).unwrap(),
                            Err(e) => {
                                shutdown = true;
                            }
                        }
                    }
                    if shutdown {
                        let remote = self.remotes.remove(&token).unwrap();
                        remote.shutdown();
                    }
                },
            }
        }
    }
}
<|endoftext|>


<commit_before>use std::comm::{Receiver, Sender};
use uuid::Uuid;
use std::io::timer::{sleep, Timer};
use std::time::duration::Duration;
use std::sync::{Mutex, Arc};
use sync::atomic::{Relaxed, SeqCst, AtomicBool, AtomicUint};
use std::task::TaskBuilder;

use addr::Addr;
use broadcast::Broadcast;
use result::{GossipResult, GossipError, UnknownError};

use node::Node;
use health::{Health};
use state::State;

/// A message channel only works for broadcasts, not other server communication
/// items. Thus, these are external messages while others are most likely internal.
type MessageChan = (Sender<Broadcast>, Receiver<Broadcast>);

type ItemChan = (Sender<ProcessMessage>, Receiver<ProcessMessage>);

#[deriving(Show, PartialEq, Clone)]
pub enum ProcessMessage {
    Shutdown
}

/// Current status of a particular server.
#[deriving(Show, PartialEq, Clone)]
pub enum Status {
    Initializing,
    Running,
    ShuttingDown,
    Failing
}

/// All the possible types of messages we can send to most of the channels that communicate with
/// the server. The most common one is the `Message` variant, which initiates/receives a new
/// broadcast/gossip to the appropriate nodes.
#[deriving(Show, PartialEq)]
pub enum ServerMessage {
    /// A broadcast/gossip that can be sent and/or received!
    Message(Broadcast)
}

pub struct Process {
    cluster: Mutex<Cluster>,
    status: Mutex<Status>,
    running: AtomicBool,
    server_sender: Sender<ProcessMessage>,
    local_tx: Sender<ProcessMessage>,
    local_rx: Receiver<ProcessMessage>
}

impl Process {
    pub fn new(sender: Sender<ProcessMessage>) -> Process {
        let (tx, rx) = channel();

        Process {
            cluster: Mutex::new(Cluster::new(Uuid::new_v4().to_hyphenated_str())),
            status: Mutex::new(Initializing),
            running: AtomicBool::new(false),
            server_sender: sender,
            local_tx: tx,
            local_rx: rx
        }
    }

    pub fn bind(&self) {
        for msg in self.local_rx.iter() {
            match msg {
                Shutdown => break
            }
        }
    }
}

pub struct Cluster {
    name: String,
    members: Vec<Node>
}

impl Cluster {
    pub fn new(name: String) -> Cluster {
        Cluster {
            name: name,
            members: Vec::new()
        }
    }
}

/// A single node/member that belongs to a cluster. A server is the central piece
/// in the gossip system. Each server contains some basic metadata about itself and about
/// the cluster membership.
pub struct Server {
    id: Uuid,
    addr: Option<Addr>,
    local_tx: Sender<ProcessMessage>,
    local_rx: Receiver<ProcessMessage>,
    process: Option<Sender<ProcessMessage>>
}

impl Server {
    /// ```rust
    /// use gossip::{Server};
    ///
    /// let mut server = Server::new();
    /// ```
    pub fn new() -> Server {
        let (tx, rx) = channel();

        let server = Server {
            id: Uuid::new_v4(),
            addr: None,
            local_tx: tx,
            local_rx: rx,
            process: None
        };

        server
    }

    pub fn bind(&mut self, ip: &str, port: u16) {
        self.addr = Some(Addr::new(ip, port));

        let local_tx = self.local_tx.clone();
        let (sender, receiver) = channel();

        TaskBuilder::new().named("ProcessTask").spawn(proc() {
            let process = Process::new(local_tx);
            let tx = process.local_tx.clone();
            sender.send(tx);
            process.bind();
        });

        self.process = Some(receiver.recv());
    }

    fn send(&self, msg: ProcessMessage) -> GossipResult<()> {
        match self.process {
            Some(ref p) => p.send(msg),
            None => return Err(GossipError::new("Failed to send message. Process is not online.", UnknownError))
        }

        Ok(())
    }

    pub fn shutdown(&self) -> GossipResult<()> {
        try!(self.send(Shutdown));
        Ok(())
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::sync::{Arc, Mutex};

    #[test]
    fn should_be_initializing() {
        let server = Server::new();
    }

    #[test]
    fn should_have_no_addr_or_port() {
        let server = Server::new();
        assert!(server.addr.is_none());
    }

    #[test]
    fn should_have_empty_server_list() {
        let server = Server::new();
    }

    #[test]
    fn should_not_hang_when_calling_listen() {
        let mut server = Server::new();
        server.bind("localhost", 5666);
        server.shutdown();
    }
}
<commit_msg>Added docs and stuff<commit_after>use std::comm::{Receiver, Sender};
use uuid::Uuid;
use std::io::timer::{sleep, Timer};
use std::time::duration::Duration;
use std::sync::{Mutex, Arc};
use sync::atomic::{Relaxed, SeqCst, AtomicBool, AtomicUint};
use std::task::TaskBuilder;

use addr::Addr;
use broadcast::Broadcast;
use result::{GossipResult, GossipError, UnknownError};

use node::Node;
use health::{Health};
use state::State;

/// A message channel only works for broadcasts, not other server communication
/// items. Thus, these are external messages while others are most likely internal.
type MessageChan = (Sender<Broadcast>, Receiver<Broadcast>);

type ItemChan = (Sender<ProcessMessage>, Receiver<ProcessMessage>);

#[deriving(Show, PartialEq, Clone)]
pub enum ProcessMessage {
    Shutdown
}

/// Current status of a particular server.
#[deriving(Show, PartialEq, Clone)]
pub enum Status {
    Initializing,
    Running,
    ShuttingDown,
    Failing
}

/// All the possible types of messages we can send to most of the channels that communicate with
/// the server. The most common one is the `Message` variant, which initiates/receives a new
/// broadcast/gossip to the appropriate nodes.
#[deriving(Show, PartialEq)]
pub enum ServerMessage {
    /// A broadcast/gossip that can be sent and/or received!
    Message(Broadcast)
}

pub struct Process {
    cluster: Mutex<Cluster>,
    status: Mutex<Status>,
    running: AtomicBool,
    server_sender: Sender<ProcessMessage>,
    local_tx: Sender<ProcessMessage>,
    local_rx: Receiver<ProcessMessage>
}

impl Process {
    pub fn new(sender: Sender<ProcessMessage>) -> Process {
        let (tx, rx) = channel();

        Process {
            cluster: Mutex::new(Cluster::new(Uuid::new_v4().to_hyphenated_str())),
            status: Mutex::new(Initializing),
            running: AtomicBool::new(false),
            server_sender: sender,
            local_tx: tx,
            local_rx: rx
        }
    }

    pub fn bind(&self) {
        for msg in self.local_rx.iter() {
            match msg {
                Shutdown => break
            }
        }
    }
}

pub struct Cluster {
    name: String,
    members: Vec<Node>
}

impl Cluster {
    pub fn new(name: String) -> Cluster {
        Cluster {
            name: name,
            members: Vec::new()
        }
    }
}

/// A single node/member that belongs to a cluster. A server is the central piece
/// in the gossip system. Each server contains some basic metadata about itself and about
/// the cluster membership.
pub struct Server {
    id: Uuid,
    addr: Option<Addr>,
    local_tx: Sender<ProcessMessage>,
    local_rx: Receiver<ProcessMessage>,
    process: Option<Sender<ProcessMessage>>
}

impl Server {

    /// Create a new unbound local server. This will create
    /// an isolated server that has not yet connected to any
    /// network system and is currently offline.
    ///
    /// Starting the server is a manual process.
    ///
    /// ```rust
    /// use gossip::{Server};
    ///
    /// let mut server = Server::new();
    /// ```
    pub fn new() -> Server {
        let (tx, rx) = channel();

        let server = Server {
            id: Uuid::new_v4(),
            addr: None,
            local_tx: tx,
            local_rx: rx,
            process: None
        };

        server
    }

    /// Bind the server/node to the local specified
    /// interface ports. This will bind to the TCP
    /// stack, as it's the only supported mechanism.
    ///
    /// This will go and spawn the required tasks and start
    /// running everything!
    ///
    /// ```rust
    /// #![allow(unused_must_use)]
    /// use gossip::{Server};
    ///
    /// let mut server = Server::new();
    /// server.bind("localhost", 8777);
    ///
    /// // We need to shutdown otherwise it'll hang.
    /// server.shutdown();
    /// ```
    pub fn bind(&mut self, ip: &str, port: u16) {
        self.addr = Some(Addr::new(ip, port));

        let local_tx = self.local_tx.clone();
        let (sender, receiver) = channel();

        TaskBuilder::new().named("ProcessTask").spawn(proc() {
            let process = Process::new(local_tx);
            let tx = process.local_tx.clone();
            sender.send(tx);
            process.bind();
        });

        self.process = Some(receiver.recv());
    }

    /// A wrapper to the `process` Sender. Because it's an
    /// option, we'll just want to work with results because
    /// it's more composable.
    fn process_send(&self, msg: ProcessMessage) -> GossipResult<()> {
        match self.process {
            Some(ref p) => p.send(msg),
            None => return Err(GossipError::new("Failed to send message. Process is not online.", UnknownError))
        }

        Ok(())
    }

    /// Allow the ability to shutdown a running server from
    /// another task.
    ///
    /// ```rust
    /// use gossip::{Server};
    ///
    /// let mut server = Server::new();
    ///
    /// server.bind("localhost", 4555);
    /// match server.shutdown() {
    ///     Ok(_) => println!("Server has successfully been terminated!"),
    ///     Err(err) => fail!("Error while shutting down the server: {}", err)
    /// }
    /// ```
    pub fn shutdown(&self) -> GossipResult<()> {
        try!(self.process_send(Shutdown));
        Ok(())
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::sync::{Arc, Mutex};

    #[test]
    fn should_be_initializing() {
        let server = Server::new();
    }

    #[test]
    fn should_have_no_addr_or_port() {
        let server = Server::new();
        assert!(server.addr.is_none());
    }

    #[test]
    fn should_have_empty_server_list() {
        let server = Server::new();
    }

    #[test]
    fn should_not_hang_when_calling_listen() {
        let mut server = Server::new();
        server.bind("localhost", 5666);
        server.shutdown();
    }
}
<|endoftext|>


<commit_before>//! # Socket code around `libc`
//!
//! ## Notes
//!
//! This module provides a low level one-to-one mapping between `libc` system call wrappers
//! with defaults specific to netlink sockets as well as a higher level API for simplification
//! of netlink code.

use std::io;
use std::os::unix::io::{AsRawFd,IntoRawFd,RawFd};
use std::marker::PhantomData;
use std::mem::{zeroed,size_of};

use buffering::copy::{StreamReadBuffer,StreamWriteBuffer};
use libc::{self,c_int,c_void};

use {Nl,MAX_NL_LENGTH};
use err::{NlError,Nlmsgerr};
use consts::{self,AddrFamily,CtrlCmd,CtrlAttr,CtrlAttrMcastGrp,GenlId,NlmF,NlFamily,NlType};
use genl::Genlmsghdr;
use nlattr::Nlattr;
use nl::Nlmsghdr;

/// Handle for the socket file descriptor
pub struct NlSocket<T, P> {
    fd: c_int,
    data_type: PhantomData<T>,
    data_payload: PhantomData<P>,
}

impl<T, P> NlSocket<T, P> where T: NlType {
    /// Wrapper around `socket()` syscall filling in the netlink-specific information
    pub fn new(proto: NlFamily) -> Result<Self, io::Error> {
        let fd = match unsafe {
            libc::socket(AddrFamily::Netlink.into(), libc::SOCK_RAW, proto.into())
        } {
            i if i >= 0 => Ok(i),
            _ => Err(io::Error::last_os_error()),
        }?;
        Ok(NlSocket { fd, data_type: PhantomData, data_payload: PhantomData })
    }

    /// Set underlying socket file descriptor to be blocking
    pub fn block(&mut self) -> Result<&mut Self, io::Error> {
        match unsafe { libc::fcntl(self.fd, libc::F_SETFL,
                                   libc::fcntl(self.fd, libc::F_GETFL, 0) & !libc::O_NONBLOCK) } {
            i if i < 0 => return Err(io::Error::last_os_error()),
            _ => Ok(self),
        }
    }

    /// Set underlying socket file descriptor to be non blocking
    pub fn nonblock(&mut self) -> Result<&mut Self, io::Error> {
        match unsafe { libc::fcntl(self.fd, libc::F_SETFL,
                                   libc::fcntl(self.fd, libc::F_GETFL, 0) | libc::O_NONBLOCK) } {
            i if i < 0 => return Err(io::Error::last_os_error()),
            _ => Ok(self),
        }
    }

    /// Determines if underlying file descriptor is blocking - `Stream` feature will throw an
    /// error if this function returns false
    pub fn is_blocking(&self) -> Result<bool, io::Error> {
        let is_blocking = match unsafe { libc::fcntl(self.fd, libc::F_GETFL, 0) } {
            i if i >= 0 => i & libc::O_NONBLOCK == 0,
            _ => return Err(io::Error::last_os_error()),
        };
        Ok(is_blocking)
    }

    /// Use this function to bind to a netlink ID and subscribe to groups. See netlink(7)
    /// man pages for more information on netlink IDs and groups.
    pub fn bind(&mut self, pid: Option<u32>, groups: Vec<u32>) -> Result<(), io::Error> {
        let mut nladdr = unsafe { zeroed::<libc::sockaddr_nl>() };
        nladdr.nl_family = libc::c_int::from(AddrFamily::Netlink) as u16;
        nladdr.nl_pid = pid.unwrap_or(0);
        nladdr.nl_groups = groups.into_iter().fold(0, |acc, next| {
            acc | (1 << (next - 1))
        });
        match unsafe {
            libc::bind(self.fd, &nladdr as *const _ as *const libc::sockaddr,
                       size_of::<libc::sockaddr_nl>() as u32)
        } {
            i if i >= 0 => Ok(()),
            _ => Err(io::Error::last_os_error()),
        }
    }

    /// Send message encoded as byte slice to the netlink ID specified in the netlink header
    /// (`neli::nl::Nlmsghdr`)
    pub fn send<B>(&mut self, buf: B, flags: i32) -> Result<libc::ssize_t, io::Error> where B: AsRef<[u8]> {
        match unsafe {
            libc::send(self.fd, buf.as_ref() as *const _ as *const c_void, buf.as_ref().len(), flags)
        } {
            i if i >= 0 => Ok(i),
            _ => Err(io::Error::last_os_error()),
        }
    }

    /// Receive message encoded as byte slice from the netlink socket
    pub fn recv<'a, B>(&mut self, mut buf: B, flags: i32) -> Result<libc::ssize_t, io::Error> where B: AsMut<[u8]> {
        match unsafe {
            libc::recv(self.fd, buf.as_mut() as *mut _ as *mut c_void, buf.as_mut().len(), flags)
        } {
            i if i >= 0 => Ok(i),
            _ => Err(io::Error::last_os_error()),
        }
    }

    /// Equivalent of `socket` and `bind` calls.
    pub fn connect(proto: NlFamily, pid: Option<u32>, groups: Vec<u32>)
                   -> Result<Self, io::Error> {
        let mut s = try!(NlSocket::new(proto));
        try!(s.bind(pid, groups));
        Ok(s)
    }
}

impl<T, P> NlSocket<T, P> where T: Nl + NlType, P: Nl {
    /// Convenience function to send an `Nlmsghdr` struct
    pub fn send_nl(&mut self, msg: Nlmsghdr<T, P>) -> Result<(), NlError> {
        let mut mem = StreamWriteBuffer::new_growable(Some(msg.asize()));
        msg.serialize(&mut mem)?;
        self.send(mem, 0)?;
        Ok(())
    }

    /// Convenience function to receive an `Nlmsghdr` struct
    pub fn recv_nl(&mut self, buf_sz: Option<usize>) -> Result<Nlmsghdr<T, P>, NlError> {
        let mut mem = vec![0; buf_sz.unwrap_or(MAX_NL_LENGTH)];
        let mem_read = self.recv(&mut mem, 0)?;
        mem.truncate(mem_read as usize);
        Ok(Nlmsghdr::<T, P>::deserialize(&mut StreamReadBuffer::new(mem))?)
    }

    /// Convenience function to receive an `Nlmsghdr` struct with function type parameters
    /// that determine deserialization type
    pub fn recv_nl_typed<TT, PP>(&mut self, buf_sz: Option<usize>)
            -> Result<Nlmsghdr<TT, PP>, NlError> where TT: NlType, PP: Nl {
        let mut mem = vec![0; buf_sz.unwrap_or(MAX_NL_LENGTH)];
        let mem_read = self.recv(&mut mem, 0)?;
        mem.truncate(mem_read as usize);
        Ok(Nlmsghdr::<TT, PP>::deserialize(&mut StreamReadBuffer::new(mem))?)
    }

    /// Consume an ACK and return an error if an ACK is not found
    pub fn recv_ack(&mut self, buf_sz: Option<usize>) -> Result<(), NlError> {
        let ack = self.recv_nl_typed::<consts::Nlmsg, Nlmsgerr<consts::Nlmsg>>(buf_sz)?;
        if ack.nl_type == consts::Nlmsg::Error && ack.nl_payload.error == 0 {
            Ok(())
        } else {
            Err(NlError::NoAck)
        }
    }
}

impl NlSocket<GenlId, Genlmsghdr<CtrlCmd>> {
    /// Create generic netlink resolution socket
    pub fn new_genl() -> Result<NlSocket<GenlId, Genlmsghdr<CtrlCmd>>, io::Error> {
        Self::connect(NlFamily::Generic, None, Vec::new())
    }

    fn get_genl_family(&mut self, family_name: &str)
            -> Result<Nlmsghdr<GenlId, Genlmsghdr<CtrlCmd>>, NlError> {
        let attrs = vec![Nlattr::new_str_payload(None, CtrlAttr::FamilyName, family_name)?];
        let genlhdr = Genlmsghdr::new(CtrlCmd::Getfamily, 2, attrs)?;
        let nlhdr = Nlmsghdr::new(None, GenlId::Ctrl,
                               vec![NlmF::Request], None, None, genlhdr);
        self.send_nl(nlhdr)?;

        Ok(self.recv_nl(Some(4096))?)
    }

    /// Convenience function for resolving a `&str` containing the multicast group name to a
    /// numeric netlink ID
    pub fn resolve_genl_family(&mut self, family_name: &str) -> Result<u16, NlError> {
        let nlhdr = self.get_genl_family(family_name)?;
        let mut handle = nlhdr.nl_payload.get_attr_handle::<CtrlAttr>();
        Ok(handle.get_payload_with::<u16>(CtrlAttr::FamilyId, None)?)
    }

    /// Convenience function for resolving a `&str` containing the multicast group name to a
    /// numeric netlink ID
    pub fn resolve_nl_mcast_group(&mut self, family_name: &str, mcast_name: &str)
            -> Result<u32, NlError> {
        let nlhdr = self.get_genl_family(family_name)?;
        let mut handle = nlhdr.nl_payload.get_attr_handle::<CtrlAttr>();
        let mut mcast_groups = handle.get_nested_attributes::<u16>(CtrlAttr::McastGroups)?;
        mcast_groups.parse_nested_attributes()?;
        let mut id = None;
        if let Some(iter) = mcast_groups.iter() {
            for attribute in iter {
                let attribute_len = attribute.nla_len;
                let mut handle = attribute.get_attr_handle();
                let string = handle.get_payload_with::<String>(CtrlAttrMcastGrp::Name,
                    Some(attribute_len as usize))?;
                if string.as_str() == mcast_name {
                    id = handle.get_payload_with::<u32>(CtrlAttrMcastGrp::Id, None).ok();
                }
            }
        }
        id.ok_or(NlError::new("Failed to resolve multicast group ID"))
    }
}

impl<T, P> AsRawFd for NlSocket<T, P> {
    fn as_raw_fd(&self) -> RawFd {
        self.fd
    }
}

impl<T, P> IntoRawFd for NlSocket<T, P> {
    fn into_raw_fd(self) -> RawFd {
        self.fd
    }
}

impl<T, P> io::Read for NlSocket<T, P> {
    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
        match unsafe {
            libc::recv(self.fd, buf as *mut _ as *mut c_void, buf.len(), 0)
        } {
            i if i >= 0 => Ok(i as usize),
            _ => Err(io::Error::last_os_error()),
        }
    }
}

/// Tokio-specific features for neli
#[cfg(feature = "stream")]
pub mod tokio {
    use super::*;

    use std::io::{ErrorKind,Read};

    use mio::{self,Evented,Ready};
    use tokio::prelude::{Async,AsyncRead,Stream};
    use tokio::reactor::PollEvented2;

    /// Tokio-enabled Netlink socket struct
    pub struct NlSocket<T, P>(PollEvented2<super::NlSocket<T, P>>);

    impl<T, P> NlSocket<T, P> where T: NlType {
        /// Setup NlSocket for use with tokio - set to nonblocking state and wrap in polling mechanism
        pub fn new(mut sock: super::NlSocket<T, P>) -> io::Result<Self> {
            if sock.is_blocking()? {
                sock.nonblock()?;
            }
            Ok(NlSocket(PollEvented2::new(sock)))
        }
    }

    impl<T, P> Read for NlSocket<T, P> {
        fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
            self.0.get_mut().read(buf)
        }
    }

    impl<T, P> AsyncRead for NlSocket<T, P> {}

    impl<T, P> Stream for NlSocket<T, P> where T: NlType, P: Nl {
        type Item = Nlmsghdr<T, P>;
        type Error = io::Error;

        fn poll(&mut self) -> Result<Async<Option<Self::Item>>, Self::Error> {
            let readiness = self.0.poll_read_ready(Ready::readable())?;
            match readiness {
                Async::NotReady => return Ok(Async::NotReady),
                Async::Ready(_) => (),
            }

            let mut mem = vec![0; MAX_NL_LENGTH];
            let bytes_written = match self.read(mem.as_mut_slice()) {
                Ok(0) => return Ok(Async::Ready(None)),
                Ok(i) => i,
                Err(e) => {
                    if e.kind() == ErrorKind::WouldBlock {
                        self.0.clear_read_ready(Ready::readable())?;
                        return Ok(Async::NotReady);
                    } else {
                        return Err(e);
                    }
                }
            };
            mem.truncate(bytes_written);
            let mut mem_read = StreamReadBuffer::new(mem);
            let hdr = match Nlmsghdr::<T, P>::deserialize(&mut mem_read) {
                Ok(h) => h,
                Err(_) => return Err(io::Error::from(io::ErrorKind::Other)),
            };
            Ok(Async::Ready(Some(hdr)))
        }
    }

    impl<T, P> Evented for super::NlSocket<T, P> {
        fn register(&self, poll: &mio::Poll, token: mio::Token, interest: mio::Ready,
                    opts: mio::PollOpt) -> io::Result<()> {
            poll.register(&mio::unix::EventedFd(&self.as_raw_fd()), token, interest, opts)
        }

        fn reregister(&self, poll: &mio::Poll, token: mio::Token, interest: mio::Ready,
                      opts: mio::PollOpt) -> io::Result<()> {
            poll.reregister(&mio::unix::EventedFd(&self.as_raw_fd()), token, interest, opts)
        }

        fn deregister(&self, poll: &mio::Poll) -> io::Result<()> {
            poll.deregister(&mio::unix::EventedFd(&self.as_raw_fd()))
        }
    }
}

impl<T, P> Drop for NlSocket<T, P> {
    /// Closes underlying file descriptor to avoid file descriptor leaks.
    fn drop(&mut self) {
        unsafe { libc::close(self.fd); }
    }
}

#[cfg(test)]
mod test {
    use super::*;
    use consts::{CtrlCmd,Nlmsg};
    use genl::Genlmsghdr;

    #[test]
    fn test_socket_creation() {
       NlSocket::<Nlmsg, Genlmsghdr<CtrlCmd>>::connect(NlFamily::Generic, None, Vec::new()).unwrap();
    }
}
<commit_msg>Add Iterator implementation for NlSocket and rework message receive functions to be friendlier for parsing multiple messages from a single buffer one at a time<commit_after>//! # Socket code around `libc`
//!
//! ## Notes
//!
//! This module provides a low level one-to-one mapping between `libc` system call wrappers
//! with defaults specific to netlink sockets as well as a higher level API for simplification
//! of netlink code.

use std::io;
use std::os::unix::io::{AsRawFd,IntoRawFd,RawFd};
use std::marker::PhantomData;
use std::mem::{zeroed,size_of};

use buffering::copy::{StreamReadBuffer,StreamWriteBuffer};
use libc::{self,c_int,c_void};

use {Nl,MAX_NL_LENGTH};
use err::{NlError,Nlmsgerr};
use consts::{self,AddrFamily,CtrlCmd,CtrlAttr,CtrlAttrMcastGrp,GenlId,NlmF,NlFamily,NlType};
use genl::Genlmsghdr;
use nlattr::Nlattr;
use nl::Nlmsghdr;

/// Handle for the socket file descriptor
pub struct NlSocket<T, P> {
    fd: c_int,
    data_type: PhantomData<T>,
    data_payload: PhantomData<P>,
    recv_buffer: Option<StreamReadBuffer<Vec<u8>>>,
}

impl<T, P> NlSocket<T, P> where T: NlType {
    /// Wrapper around `socket()` syscall filling in the netlink-specific information
    pub fn new(proto: NlFamily) -> Result<Self, io::Error> {
        let fd = match unsafe {
            libc::socket(AddrFamily::Netlink.into(), libc::SOCK_RAW, proto.into())
        } {
            i if i >= 0 => Ok(i),
            _ => Err(io::Error::last_os_error()),
        }?;
        Ok(NlSocket { fd, data_type: PhantomData, data_payload: PhantomData, recv_buffer: None, })
    }

    /// Set underlying socket file descriptor to be blocking
    pub fn block(&mut self) -> Result<&mut Self, io::Error> {
        match unsafe { libc::fcntl(self.fd, libc::F_SETFL,
                                   libc::fcntl(self.fd, libc::F_GETFL, 0) & !libc::O_NONBLOCK) } {
            i if i < 0 => return Err(io::Error::last_os_error()),
            _ => Ok(self),
        }
    }

    /// Set underlying socket file descriptor to be non blocking
    pub fn nonblock(&mut self) -> Result<&mut Self, io::Error> {
        match unsafe { libc::fcntl(self.fd, libc::F_SETFL,
                                   libc::fcntl(self.fd, libc::F_GETFL, 0) | libc::O_NONBLOCK) } {
            i if i < 0 => return Err(io::Error::last_os_error()),
            _ => Ok(self),
        }
    }

    /// Determines if underlying file descriptor is blocking - `Stream` feature will throw an
    /// error if this function returns false
    pub fn is_blocking(&self) -> Result<bool, io::Error> {
        let is_blocking = match unsafe { libc::fcntl(self.fd, libc::F_GETFL, 0) } {
            i if i >= 0 => i & libc::O_NONBLOCK == 0,
            _ => return Err(io::Error::last_os_error()),
        };
        Ok(is_blocking)
    }

    /// Use this function to bind to a netlink ID and subscribe to groups. See netlink(7)
    /// man pages for more information on netlink IDs and groups.
    pub fn bind(&mut self, pid: Option<u32>, groups: Vec<u32>) -> Result<(), io::Error> {
        let mut nladdr = unsafe { zeroed::<libc::sockaddr_nl>() };
        nladdr.nl_family = libc::c_int::from(AddrFamily::Netlink) as u16;
        nladdr.nl_pid = pid.unwrap_or(0);
        nladdr.nl_groups = groups.into_iter().fold(0, |acc, next| {
            acc | (1 << (next - 1))
        });
        match unsafe {
            libc::bind(self.fd, &nladdr as *const _ as *const libc::sockaddr,
                       size_of::<libc::sockaddr_nl>() as u32)
        } {
            i if i >= 0 => Ok(()),
            _ => Err(io::Error::last_os_error()),
        }
    }

    /// Send message encoded as byte slice to the netlink ID specified in the netlink header
    /// (`neli::nl::Nlmsghdr`)
    pub fn send<B>(&mut self, buf: B, flags: i32) -> Result<libc::ssize_t, io::Error> where B: AsRef<[u8]> {
        match unsafe {
            libc::send(self.fd, buf.as_ref() as *const _ as *const c_void, buf.as_ref().len(), flags)
        } {
            i if i >= 0 => Ok(i),
            _ => Err(io::Error::last_os_error()),
        }
    }

    /// Receive message encoded as byte slice from the netlink socket
    pub fn recv<'a, B>(&mut self, mut buf: B, flags: i32) -> Result<libc::ssize_t, io::Error> where B: AsMut<[u8]> {
        match unsafe {
            libc::recv(self.fd, buf.as_mut() as *mut _ as *mut c_void, buf.as_mut().len(), flags)
        } {
            i if i >= 0 => Ok(i),
            _ => Err(io::Error::last_os_error()),
        }
    }

    /// Equivalent of `socket` and `bind` calls.
    pub fn connect(proto: NlFamily, pid: Option<u32>, groups: Vec<u32>)
                   -> Result<Self, io::Error> {
        let mut s = try!(NlSocket::new(proto));
        try!(s.bind(pid, groups));
        Ok(s)
    }
}

impl<T, P> Iterator for NlSocket<T, P> where T: Nl + NlType, P: Nl {
    type Item = Nlmsghdr<T, P>;

    fn next(&mut self) -> Option<Self::Item> {
        if let Ok(next) = self.recv_nl_next() {
            next
        } else {
            self.recv_buffer = None;
            None
        }
    }
}

impl<T, P> NlSocket<T, P> where T: Nl + NlType, P: Nl {
    /// Convenience function to send an `Nlmsghdr` struct
    pub fn send_nl(&mut self, msg: Nlmsghdr<T, P>) -> Result<(), NlError> {
        let mut mem = StreamWriteBuffer::new_growable(Some(msg.asize()));
        msg.serialize(&mut mem)?;
        self.send(mem, 0)?;
        Ok(())
    }

    /// Convenience function to begin receiving a stream of `Nlmsghdr` structs
    pub fn recv_nl(&mut self, buf_sz: Option<usize>) -> Result<(), NlError> {
        let mut mem = vec![0; buf_sz.unwrap_or(MAX_NL_LENGTH)];
        let mem_read = self.recv(&mut mem, 0)?;
        mem.truncate(mem_read as usize);
        self.recv_buffer = Some(StreamReadBuffer::new(mem));
        Ok(())
    }

    /// Convenience function to yield the next `Nlmsghdr` struct
    pub fn recv_nl_next(&mut self) -> Result<Option<Nlmsghdr<T, P>>, NlError> {
        let (msg, at_end) = if let Some(ref mut reader) = self.recv_buffer {
            let msg = Nlmsghdr::<T, P>::deserialize(reader)?;
            (Ok(Some(msg)), reader.at_end())
        } else {
            (Ok(None), false)
        };
        if at_end {
            self.recv_buffer = None;
        }
        msg
    }

    /// Convenience function to yield the next `Nlmsghdr` struct with differing types from the
    /// `NlSocket` struct types specified
    pub fn recv_nl_next_override<TT, PP>(&mut self) -> Result<Option<Nlmsghdr<TT, PP>>, NlError>
            where TT: NlType + Nl, PP: Nl {
        let (msg, at_end) = if let Some(ref mut reader) = self.recv_buffer {
            let msg = Nlmsghdr::<TT, PP>::deserialize(reader)?;
            (Ok(Some(msg)), reader.at_end())
        } else {
            (Ok(None), false)
        };
        if at_end {
            self.recv_buffer = None;
        }
        msg
    }

    /// Consume an ACK and return an error if an ACK is not found
    pub fn recv_ack(&mut self, buf_sz: Option<usize>) -> Result<(), NlError> {
        self.recv_nl(buf_sz)?;
        if let Some(ack) = self.recv_nl_next_override::<consts::Nlmsg, Nlmsgerr<consts::Nlmsg>>()? {
            if ack.nl_type == consts::Nlmsg::Error && ack.nl_payload.error == 0 {
                Ok(())
            } else {
                Err(NlError::NoAck)
            }
        } else {
            Err(NlError::NoAck)
        }
    }
}

impl NlSocket<GenlId, Genlmsghdr<CtrlCmd>> {
    /// Create generic netlink resolution socket
    pub fn new_genl() -> Result<NlSocket<GenlId, Genlmsghdr<CtrlCmd>>, io::Error> {
        Self::connect(NlFamily::Generic, None, Vec::new())
    }

    fn get_genl_family(&mut self, family_name: &str)
            -> Result<Nlmsghdr<GenlId, Genlmsghdr<CtrlCmd>>, NlError> {
        let attrs = vec![Nlattr::new_str_payload(None, CtrlAttr::FamilyName, family_name)?];
        let genlhdr = Genlmsghdr::new(CtrlCmd::Getfamily, 2, attrs)?;
        let nlhdr = Nlmsghdr::new(None, GenlId::Ctrl,
                               vec![NlmF::Request], None, None, genlhdr);
        self.send_nl(nlhdr)?;

        self.recv_nl(Some(4096))?;
        if let Some(msg) = self.recv_nl_next()? {
            Ok(msg)
        } else {
            Err(NlError::new("No genetlink message received in response to request"))
        }
    }

    /// Convenience function for resolving a `&str` containing the multicast group name to a
    /// numeric netlink ID
    pub fn resolve_genl_family(&mut self, family_name: &str) -> Result<u16, NlError> {
        let nlhdr = self.get_genl_family(family_name)?;
        let mut handle = nlhdr.nl_payload.get_attr_handle::<CtrlAttr>();
        Ok(handle.get_payload_with::<u16>(CtrlAttr::FamilyId, None)?)
    }

    /// Convenience function for resolving a `&str` containing the multicast group name to a
    /// numeric netlink ID
    pub fn resolve_nl_mcast_group(&mut self, family_name: &str, mcast_name: &str)
            -> Result<u32, NlError> {
        let nlhdr = self.get_genl_family(family_name)?;
        let mut handle = nlhdr.nl_payload.get_attr_handle::<CtrlAttr>();
        let mut mcast_groups = handle.get_nested_attributes::<u16>(CtrlAttr::McastGroups)?;
        mcast_groups.parse_nested_attributes()?;
        let mut id = None;
        if let Some(iter) = mcast_groups.iter() {
            for attribute in iter {
                let attribute_len = attribute.nla_len;
                let mut handle = attribute.get_attr_handle();
                let string = handle.get_payload_with::<String>(CtrlAttrMcastGrp::Name,
                    Some(attribute_len as usize))?;
                if string.as_str() == mcast_name {
                    id = handle.get_payload_with::<u32>(CtrlAttrMcastGrp::Id, None).ok();
                }
            }
        }
        id.ok_or(NlError::new("Failed to resolve multicast group ID"))
    }
}

impl<T, P> AsRawFd for NlSocket<T, P> {
    fn as_raw_fd(&self) -> RawFd {
        self.fd
    }
}

impl<T, P> IntoRawFd for NlSocket<T, P> {
    fn into_raw_fd(self) -> RawFd {
        self.fd
    }
}

impl<T, P> io::Read for NlSocket<T, P> {
    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
        match unsafe {
            libc::recv(self.fd, buf as *mut _ as *mut c_void, buf.len(), 0)
        } {
            i if i >= 0 => Ok(i as usize),
            _ => Err(io::Error::last_os_error()),
        }
    }
}

/// Tokio-specific features for neli
#[cfg(feature = "stream")]
pub mod tokio {
    use super::*;

    use std::io::{ErrorKind,Read};

    use mio::{self,Evented,Ready};
    use tokio::prelude::{Async,AsyncRead,Stream};
    use tokio::reactor::PollEvented2;

    /// Tokio-enabled Netlink socket struct
    pub struct NlSocket<T, P>(PollEvented2<super::NlSocket<T, P>>);

    impl<T, P> NlSocket<T, P> where T: NlType {
        /// Setup NlSocket for use with tokio - set to nonblocking state and wrap in polling mechanism
        pub fn new(mut sock: super::NlSocket<T, P>) -> io::Result<Self> {
            if sock.is_blocking()? {
                sock.nonblock()?;
            }
            Ok(NlSocket(PollEvented2::new(sock)))
        }
    }

    impl<T, P> Read for NlSocket<T, P> {
        fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
            self.0.get_mut().read(buf)
        }
    }

    impl<T, P> AsyncRead for NlSocket<T, P> {}

    impl<T, P> Stream for NlSocket<T, P> where T: NlType, P: Nl {
        type Item = Nlmsghdr<T, P>;
        type Error = io::Error;

        fn poll(&mut self) -> Result<Async<Option<Self::Item>>, Self::Error> {
            let readiness = self.0.poll_read_ready(Ready::readable())?;
            match readiness {
                Async::NotReady => return Ok(Async::NotReady),
                Async::Ready(_) => (),
            }

            let mut mem = vec![0; MAX_NL_LENGTH];
            let bytes_written = match self.read(mem.as_mut_slice()) {
                Ok(0) => return Ok(Async::Ready(None)),
                Ok(i) => i,
                Err(e) => {
                    if e.kind() == ErrorKind::WouldBlock {
                        self.0.clear_read_ready(Ready::readable())?;
                        return Ok(Async::NotReady);
                    } else {
                        return Err(e);
                    }
                }
            };
            mem.truncate(bytes_written);
            let mut mem_read = StreamReadBuffer::new(mem);
            let hdr = match Nlmsghdr::<T, P>::deserialize(&mut mem_read) {
                Ok(h) => h,
                Err(_) => return Err(io::Error::from(io::ErrorKind::Other)),
            };
            Ok(Async::Ready(Some(hdr)))
        }
    }

    impl<T, P> Evented for super::NlSocket<T, P> {
        fn register(&self, poll: &mio::Poll, token: mio::Token, interest: mio::Ready,
                    opts: mio::PollOpt) -> io::Result<()> {
            poll.register(&mio::unix::EventedFd(&self.as_raw_fd()), token, interest, opts)
        }

        fn reregister(&self, poll: &mio::Poll, token: mio::Token, interest: mio::Ready,
                      opts: mio::PollOpt) -> io::Result<()> {
            poll.reregister(&mio::unix::EventedFd(&self.as_raw_fd()), token, interest, opts)
        }

        fn deregister(&self, poll: &mio::Poll) -> io::Result<()> {
            poll.deregister(&mio::unix::EventedFd(&self.as_raw_fd()))
        }
    }
}

impl<T, P> Drop for NlSocket<T, P> {
    /// Closes underlying file descriptor to avoid file descriptor leaks.
    fn drop(&mut self) {
        unsafe { libc::close(self.fd); }
    }
}

#[cfg(test)]
mod test {
    use super::*;
    use consts::{CtrlCmd,Nlmsg};
    use genl::Genlmsghdr;

    #[test]
    fn test_socket_creation() {
       NlSocket::<Nlmsg, Genlmsghdr<CtrlCmd>>::connect(NlFamily::Generic, None, Vec::new()).unwrap();
    }
}
<|endoftext|>


<commit_before>use std::io::{IoResult, IoError, IoErrorKind, TcpStream, ConnectionRefused,
              ConnectionFailed, OtherIoError};

enum AuthMethod<'s> {
    NoAuth,
    UPass(&'s str, &'s str)
}

pub struct Socks5<'a> {
    socks_host: &'a str,
    socks_port: u16,
    socks_auth: AuthMethod<'a>,
}

impl<'a> Socks5<'a> {
    pub fn new(host: &'a str, port: u16) -> Socks5 {
        Socks5 {
            socks_host: host,
            socks_port: port,
            socks_auth: NoAuth
        }
    }

    pub fn login(&mut self, uname: &'a str, passwd: &'a str) {
        self.socks_auth = UPass(uname, passwd);
    }

    pub fn connect(&mut self, host: &str, port: u16) -> IoResult<TcpStream> {
        let mut stream = try!(TcpStream::connect(self.socks_host, self.socks_port));
        try!(stream.write([0x05u8]));
        match self.socks_auth {
            NoAuth => { try!(stream.write([0x01u8, 0x00])); },
            UPass(uname, passwd) => { try!(stream.write([0x01u8, 0x02])); }
        }

        if try!(stream.read_u8()) != 0x05 {
            return io_err(OtherIoError, "Unexpected SOCKS version number");
        }

        match try!(stream.read_u8()) {
            0x00 => {
                match self.socks_auth {
                    NoAuth => { /* Continue */ },
                    _ => return io_err(OtherIoError,
                        "Wrong authentication method from server")
                }
            }
            0x02 => {
                match self.socks_auth {
                    UPass(uname, passwd) => {
                        try!(stream.write([0x01u8, uname.len() as u8]));
                        try!(stream.write_str(uname));
                        try!(stream.write([passwd.len() as u8]));
                        try!(stream.write_str(passwd));

                        if try!(stream.read_u8()) != 0x01 {
                            return io_err(OtherIoError,
                                "Invalid authentication version");
                        }

                        if try!(stream.read_u8()) != 0x00 {
                            return io_err(OtherIoError, "Authentication failed");
                        }
                    }
                    _ => { return io_err(OtherIoError,
                            "Wrong authentication method from server");
                    }
                }
            }
            _ => { return io_err(OtherIoError,
                    "Wrong authentication method from server"); }
        }

        try!(stream.write([0x05u8, 0x01, 0x00, 0x03, host.len() as u8]));
        try!(stream.write_str(host));
        try!(stream.write_be_u16(port));

        if try!(stream.read_u8()) != 0x05 {
            return io_err(OtherIoError, "Invalid SOCKS version number");
        }

        match try!(stream.read_u8()) {
            0x00 => {
                let _null = try!(stream.read_u8());
                let _addrtype = try!(stream.read_u8());
                let _addrlen = try!(stream.read_u8());
                let mut _addr: Vec<u8> = vec![];
                for i in range(0, _addrlen) {
                    _addr.push(try!(stream.read_u8()));
                }
                let _port = try!(stream.read_be_u16());
                Ok(stream)
            }
            0x01 => io_err(OtherIoError, "General failure"),
            0x02 => io_err(OtherIoError, "Connection not allowed by ruleset"),
            0x03 => io_err(OtherIoError, "Network unreachable"),
            0x04 => io_err(OtherIoError, "Host unreachable"),
            0x05 => io_err(OtherIoError, "Connection refused by destination"),
            0x06 => io_err(OtherIoError, "TTL expired"),
            0x07 => io_err(OtherIoError, "Protocol Error"),
            0x08 => io_err(OtherIoError, "Address type not supported"),
            _ => io_err(OtherIoError, "Unknown error")
        }
    }
}

fn io_err<T>(kind: IoErrorKind, desc: &'static str) -> IoResult<T> {
    Err(IoError { kind: kind, desc: desc, detail: None })
}

<commit_msg>Fail for unhandled cases.<commit_after>use std::io::{IoResult, IoError, IoErrorKind, TcpStream, ConnectionRefused,
              ConnectionFailed, OtherIoError};

enum AuthMethod<'s> {
    NoAuth,
    UPass(&'s str, &'s str)
}

pub struct Socks5<'a> {
    socks_host: &'a str,
    socks_port: u16,
    socks_auth: AuthMethod<'a>,
}

impl<'a> Socks5<'a> {
    pub fn new(host: &'a str, port: u16) -> Socks5 {
        Socks5 {
            socks_host: host,
            socks_port: port,
            socks_auth: NoAuth
        }
    }

    pub fn login(&mut self, uname: &'a str, passwd: &'a str) {
        self.socks_auth = UPass(uname, passwd);
    }

    pub fn connect(&mut self, host: &str, port: u16) -> IoResult<TcpStream> {
        let mut stream = try!(TcpStream::connect(self.socks_host, self.socks_port));
        try!(stream.write([0x05u8]));
        match self.socks_auth {
            NoAuth => { try!(stream.write([0x01u8, 0x00])); },
            UPass(uname, passwd) => { try!(stream.write([0x01u8, 0x02])); }
        }

        if try!(stream.read_u8()) != 0x05 {
            return io_err(OtherIoError, "Unexpected SOCKS version number");
        }

        match try!(stream.read_u8()) {
            0x00 => {
                match self.socks_auth {
                    NoAuth => { /* Continue */ },
                    _ => return io_err(OtherIoError,
                        "Wrong authentication method from server")
                }
            }
            0x02 => {
                match self.socks_auth {
                    UPass(uname, passwd) => {
                        try!(stream.write([0x01u8, uname.len() as u8]));
                        try!(stream.write_str(uname));
                        try!(stream.write([passwd.len() as u8]));
                        try!(stream.write_str(passwd));

                        if try!(stream.read_u8()) != 0x01 {
                            return io_err(OtherIoError,
                                "Invalid authentication version");
                        }

                        if try!(stream.read_u8()) != 0x00 {
                            return io_err(ConnectionRefused, "Authentication failed");
                        }
                    }
                    _ => { return io_err(OtherIoError,
                            "Wrong authentication method from server");
                    }
                }
            }
            0xFF => { return io_err(ConnectionRefused,
                      "Server refused authentication methods"); }
            _ => { return io_err(OtherIoError,
                    "Wrong authentication method from server"); }
        }

        try!(stream.write([0x05u8, 0x01, 0x00, 0x03, host.len() as u8]));
        try!(stream.write_str(host));
        try!(stream.write_be_u16(port));

        if try!(stream.read_u8()) != 0x05 {
            return io_err(OtherIoError, "Invalid SOCKS version number");
        }

        match try!(stream.read_u8()) {
            0x00 => {
                let _null = try!(stream.read_u8());
                let addrtype = try!(stream.read_u8());
                if addrtype != 0x03 {
                    return io_err(OtherIoError, "Unimplemented");
                }
                let _addrlen = try!(stream.read_u8());
                let mut _addr: Vec<u8> = vec![];
                for i in range(0, _addrlen) {
                    _addr.push(try!(stream.read_u8()));
                }
                let _port = try!(stream.read_be_u16());
                Ok(stream)
            }
            0x01 => io_err(OtherIoError, "General failure"),
            0x02 => io_err(OtherIoError, "Connection not allowed by ruleset"),
            0x03 => io_err(OtherIoError, "Network unreachable"),
            0x04 => io_err(OtherIoError, "Host unreachable"),
            0x05 => io_err(OtherIoError, "Connection refused by destination"),
            0x06 => io_err(OtherIoError, "TTL expired"),
            0x07 => io_err(OtherIoError, "Protocol Error"),
            0x08 => io_err(OtherIoError, "Address type not supported"),
            _ => io_err(OtherIoError, "Unknown error")
        }
    }
}

fn io_err<T>(kind: IoErrorKind, desc: &'static str) -> IoResult<T> {
    Err(IoError { kind: kind, desc: desc, detail: None })
}

<|endoftext|>


<commit_before>use orbclient::{self, Renderer, WindowFlag};
use orbclient::color::Color;
use std::cell::{Cell, RefCell};
use std::collections::VecDeque;
use std::sync::Arc;

use super::{Event, Point, Rect, Widget};
use theme::WINDOW_BACKGROUND;
use traits::Resize;

extern crate orbfont;

pub struct WindowRenderer<'a> {
    inner: &'a mut orbclient::Window,
    font: &'a Option<orbfont::Font>
}

impl<'a> WindowRenderer<'a> {
    pub fn new(inner: &'a mut orbclient::Window, font: &'a Option<orbfont::Font>) -> WindowRenderer<'a> {
        WindowRenderer { inner: inner, font: font }
    }
}

impl<'a> Renderer for WindowRenderer<'a> {
    fn width(&self) -> u32 {
        self.inner.width()
    }

    fn height(&self) -> u32 {
        self.inner.height()
    }

    fn data(&self) -> &[Color] {
        self.inner.data()
    }

    fn data_mut(&mut self) -> &mut [Color] {
        self.inner.data_mut()
    }

    fn sync(&mut self) -> bool {
        self.inner.sync()
    }

    fn char(&mut self, x: i32, y: i32, c: char, color: Color) {
        if let Some(ref font) = *self.font {
            let mut buf = [0; 4];
            font.render(&c.encode_utf8(&mut buf), 16.0).draw(self.inner, x, y, color)
        }else{
            self.inner.char(x, y, c, color);
        }
    }
}

impl<'a> Drop for WindowRenderer<'a> {
    fn drop(&mut self) {
        self.inner.sync();
    }
}

pub struct Window {
    inner: RefCell<orbclient::Window>,
    font: Option<orbfont::Font>,
    pub widgets: RefCell<Vec<Arc<Widget>>>,
    pub widget_focus: Cell<usize>,
    pub bg: Cell<Color>,
    pub running: Cell<bool>,
    resize_callback: RefCell<Option<Arc<Fn(&Window, u32, u32)>>>,
    mouse_point: Point,
    mouse_left: bool,
    mouse_middle: bool,
    mouse_right: bool,
    events: VecDeque<Event>,
    redraw: bool,
}

impl Resize for Window {
    fn emit_resize(&self, width: u32, height: u32) {
        if let Some(ref resize_callback) = *self.resize_callback.borrow() {
            resize_callback(self, width, height);
        }
    }

    fn on_resize<T: Fn(&Self, u32, u32) + 'static>(&self, func: T) -> &Self {
        *self.resize_callback.borrow_mut() = Some(Arc::new(func));
        self
    }
}

impl Window {
    pub fn new(rect: Rect, title: &str) -> Self {
        Window::new_flags(rect, title, &[])
    }

    pub fn new_flags(rect: Rect, title: &str, flags: &[WindowFlag]) -> Self {
        let mut events = VecDeque::new();
        events.push_back(Event::Init);
        Window {
            inner: RefCell::new(orbclient::Window::new_flags(rect.x, rect.y, rect.width, rect.height, title, flags).unwrap()),
            font: orbfont::Font::find(None, None, None).ok(),
            widgets: RefCell::new(Vec::new()),
            widget_focus: Cell::new(0),
            bg: Cell::new(WINDOW_BACKGROUND),
            running: Cell::new(true),
            resize_callback: RefCell::new(None),
            mouse_point: Point::new(0, 0),
            mouse_left: false,
            mouse_right: false,
            mouse_middle: false,
            events: events,
            redraw: true,
        }
    }

    pub fn set_size(&self, width: u32, height: u32) {
        let mut inner = self.inner.borrow_mut();
        (*inner).set_size(width, height);
    }

    pub fn set_title(&self, title: &str) {
        let mut inner = self.inner.borrow_mut();
        (*inner).set_title(title);
    }

    pub fn close(&self) {
        self.running.set(false);
    }

    pub fn add<T: Widget>(&self, widget: &Arc<T>) -> usize {
        let mut widgets = self.widgets.borrow_mut();
        let id = widgets.len();
        widgets.push(widget.clone());
        id
    }

    pub fn draw(&self) {
        let mut inner = self.inner.borrow_mut();
        inner.set(self.bg.get());

        let mut renderer = WindowRenderer::new(&mut *inner, &self.font);
        for i in 0..self.widgets.borrow().len() {
            if let Some(widget) = self.widgets.borrow().get(i) {
                widget.draw(&mut renderer, self.widget_focus.get() == i);
            }
        }
    }

    pub fn step(&mut self) {
        self.drain_orbital_events();
        self.drain_events();
    }

    pub fn drain_events(&mut self) {
        while let Some(event) = self.events.pop_front() {
            match event {
                Event::Resize { width, height } => {
                    self.emit_resize(width, height);
                },
                _ => ()
            }

            for i in 0..self.widgets.borrow().len() {
                if let Some(widget) = self.widgets.borrow().get(i) {
                    if widget.event(event, self.widget_focus.get() == i, &mut self.redraw) {
                        if self.widget_focus.get() != i {
                            self.widget_focus.set(i);
                            self.redraw = true;
                        }
                    }
                }
            }
        }
    }

    pub fn drain_orbital_events(&mut self) {
        for orbital_event in self.inner.borrow_mut().events() {
            match orbital_event.to_option() {
                orbclient::EventOption::Mouse(mouse_event) => {
                    self.mouse_point.x = mouse_event.x;
                    self.mouse_point.y = mouse_event.y;

                    self.events.push_back(Event::Mouse {
                        point: self.mouse_point,
                        left_button: self.mouse_left,
                        middle_button: self.mouse_middle,
                        right_button: self.mouse_right,
                    })
                },
                orbclient::EventOption::Button(button_event) => {
                    self.mouse_left = button_event.left;
                    self.mouse_middle = button_event.middle;
                    self.mouse_right = button_event.right;

                    self.events.push_back(Event::Mouse {
                        point: self.mouse_point,
                        left_button: self.mouse_left,
                        middle_button: self.mouse_middle,
                        right_button: self.mouse_right,
                    })
                },
                orbclient::EventOption::Scroll(scroll_event) => {
                    self.events.push_back(Event::Scroll {
                        x: scroll_event.x,
                        y: scroll_event.y,
                    })
                },
                orbclient::EventOption::Key(key_event) => {
                    if key_event.pressed {
                        match key_event.scancode {
                            orbclient::K_BKSP => self.events.push_back(Event::Backspace),
                            orbclient::K_DEL => self.events.push_back(Event::Delete),
                            orbclient::K_HOME => self.events.push_back(Event::Home),
                            orbclient::K_END => self.events.push_back(Event::End),
                            orbclient::K_UP => self.events.push_back(Event::UpArrow),
                            orbclient::K_DOWN => self.events.push_back(Event::DownArrow),
                            orbclient::K_LEFT => self.events.push_back(Event::LeftArrow),
                            orbclient::K_RIGHT => self.events.push_back(Event::RightArrow),
                            _ => {
                                match key_event.character {
                                    '\0' => (),
                                    '\x1B' => (),
                                    '\n' => self.events.push_back(Event::Enter),
                                    _ => self.events.push_back(Event::Text { c: key_event.character }),
                                }
                            }
                        }
                    }
                },
                orbclient::EventOption::Resize(resize_event) => {
                    self.redraw = true;
                    self.events.push_back(Event::Resize {
                        width: resize_event.width,
                        height: resize_event.height,
                    });
                },
                orbclient::EventOption::Quit(_quit_event) => {
                    self.running.set(false);
                },
                _ => (),
            };
        }

    }

    pub fn exec(&mut self) {
        'event: while self.running.get() {
            self.drain_events();
            self.draw_if_needed();
            self.drain_orbital_events();
        }
    }

    pub fn needs_redraw(&mut self) {
        self.redraw = true;
    }

    pub fn draw_if_needed(&mut self) {
        if self.redraw {
            self.draw();
            self.redraw = false;
        }
    }
}
<commit_msg>Add functions to inspect inner window<commit_after>use orbclient::{self, Renderer, WindowFlag};
use orbclient::color::Color;
use std::cell::{Cell, RefCell};
use std::collections::VecDeque;
use std::sync::Arc;

use super::{Event, Point, Rect, Widget};
use theme::WINDOW_BACKGROUND;
use traits::Resize;

extern crate orbfont;

pub struct WindowRenderer<'a> {
    inner: &'a mut orbclient::Window,
    font: &'a Option<orbfont::Font>
}

impl<'a> WindowRenderer<'a> {
    pub fn new(inner: &'a mut orbclient::Window, font: &'a Option<orbfont::Font>) -> WindowRenderer<'a> {
        WindowRenderer { inner: inner, font: font }
    }
}

impl<'a> Renderer for WindowRenderer<'a> {
    fn width(&self) -> u32 {
        self.inner.width()
    }

    fn height(&self) -> u32 {
        self.inner.height()
    }

    fn data(&self) -> &[Color] {
        self.inner.data()
    }

    fn data_mut(&mut self) -> &mut [Color] {
        self.inner.data_mut()
    }

    fn sync(&mut self) -> bool {
        self.inner.sync()
    }

    fn char(&mut self, x: i32, y: i32, c: char, color: Color) {
        if let Some(ref font) = *self.font {
            let mut buf = [0; 4];
            font.render(&c.encode_utf8(&mut buf), 16.0).draw(self.inner, x, y, color)
        }else{
            self.inner.char(x, y, c, color);
        }
    }
}

impl<'a> Drop for WindowRenderer<'a> {
    fn drop(&mut self) {
        self.inner.sync();
    }
}

pub struct Window {
    inner: RefCell<orbclient::Window>,
    font: Option<orbfont::Font>,
    pub widgets: RefCell<Vec<Arc<Widget>>>,
    pub widget_focus: Cell<usize>,
    pub bg: Cell<Color>,
    pub running: Cell<bool>,
    resize_callback: RefCell<Option<Arc<Fn(&Window, u32, u32)>>>,
    mouse_point: Point,
    mouse_left: bool,
    mouse_middle: bool,
    mouse_right: bool,
    events: VecDeque<Event>,
    redraw: bool,
}

impl Resize for Window {
    fn emit_resize(&self, width: u32, height: u32) {
        if let Some(ref resize_callback) = *self.resize_callback.borrow() {
            resize_callback(self, width, height);
        }
    }

    fn on_resize<T: Fn(&Self, u32, u32) + 'static>(&self, func: T) -> &Self {
        *self.resize_callback.borrow_mut() = Some(Arc::new(func));
        self
    }
}

impl Window {
    pub fn new(rect: Rect, title: &str) -> Self {
        Window::new_flags(rect, title, &[])
    }

    pub fn new_flags(rect: Rect, title: &str, flags: &[WindowFlag]) -> Self {
        let mut events = VecDeque::new();
        events.push_back(Event::Init);
        Window {
            inner: RefCell::new(orbclient::Window::new_flags(rect.x, rect.y, rect.width, rect.height, title, flags).unwrap()),
            font: orbfont::Font::find(None, None, None).ok(),
            widgets: RefCell::new(Vec::new()),
            widget_focus: Cell::new(0),
            bg: Cell::new(WINDOW_BACKGROUND),
            running: Cell::new(true),
            resize_callback: RefCell::new(None),
            mouse_point: Point::new(0, 0),
            mouse_left: false,
            mouse_right: false,
            mouse_middle: false,
            events: events,
            redraw: true,
        }
    }

    pub fn x(&self) -> i32 {
        let inner = self.inner.borrow();
        (*inner).x()
    }

    pub fn y(&self) -> i32 {
        let inner = self.inner.borrow();
        (*inner).y()
    }

    pub fn width(&self) -> u32 {
        let inner = self.inner.borrow();
        (*inner).width()
    }

    pub fn height(&self) -> u32 {
        let inner = self.inner.borrow();
        (*inner).height()
    }

    pub fn title(&self) -> String {
        let inner = self.inner.borrow();
        (*inner).title()
    }

    pub fn set_pos(&self, x: i32, y: i32) {
        let mut inner = self.inner.borrow_mut();
        (*inner).set_pos(x, y);
    }

    pub fn set_size(&self, width: u32, height: u32) {
        let mut inner = self.inner.borrow_mut();
        (*inner).set_size(width, height);
    }

    pub fn set_title(&self, title: &str) {
        let mut inner = self.inner.borrow_mut();
        (*inner).set_title(title);
    }

    pub fn close(&self) {
        self.running.set(false);
    }

    pub fn add<T: Widget>(&self, widget: &Arc<T>) -> usize {
        let mut widgets = self.widgets.borrow_mut();
        let id = widgets.len();
        widgets.push(widget.clone());
        id
    }

    pub fn draw(&self) {
        let mut inner = self.inner.borrow_mut();
        inner.set(self.bg.get());

        let mut renderer = WindowRenderer::new(&mut *inner, &self.font);
        for i in 0..self.widgets.borrow().len() {
            if let Some(widget) = self.widgets.borrow().get(i) {
                widget.draw(&mut renderer, self.widget_focus.get() == i);
            }
        }
    }

    pub fn step(&mut self) {
        self.drain_orbital_events();
        self.drain_events();
    }

    pub fn drain_events(&mut self) {
        while let Some(event) = self.events.pop_front() {
            match event {
                Event::Resize { width, height } => {
                    self.emit_resize(width, height);
                },
                _ => ()
            }

            for i in 0..self.widgets.borrow().len() {
                if let Some(widget) = self.widgets.borrow().get(i) {
                    if widget.event(event, self.widget_focus.get() == i, &mut self.redraw) {
                        if self.widget_focus.get() != i {
                            self.widget_focus.set(i);
                            self.redraw = true;
                        }
                    }
                }
            }
        }
    }

    pub fn drain_orbital_events(&mut self) {
        for orbital_event in self.inner.borrow_mut().events() {
            match orbital_event.to_option() {
                orbclient::EventOption::Mouse(mouse_event) => {
                    self.mouse_point.x = mouse_event.x;
                    self.mouse_point.y = mouse_event.y;

                    self.events.push_back(Event::Mouse {
                        point: self.mouse_point,
                        left_button: self.mouse_left,
                        middle_button: self.mouse_middle,
                        right_button: self.mouse_right,
                    })
                },
                orbclient::EventOption::Button(button_event) => {
                    self.mouse_left = button_event.left;
                    self.mouse_middle = button_event.middle;
                    self.mouse_right = button_event.right;

                    self.events.push_back(Event::Mouse {
                        point: self.mouse_point,
                        left_button: self.mouse_left,
                        middle_button: self.mouse_middle,
                        right_button: self.mouse_right,
                    })
                },
                orbclient::EventOption::Scroll(scroll_event) => {
                    self.events.push_back(Event::Scroll {
                        x: scroll_event.x,
                        y: scroll_event.y,
                    })
                },
                orbclient::EventOption::Key(key_event) => {
                    if key_event.pressed {
                        match key_event.scancode {
                            orbclient::K_BKSP => self.events.push_back(Event::Backspace),
                            orbclient::K_DEL => self.events.push_back(Event::Delete),
                            orbclient::K_HOME => self.events.push_back(Event::Home),
                            orbclient::K_END => self.events.push_back(Event::End),
                            orbclient::K_UP => self.events.push_back(Event::UpArrow),
                            orbclient::K_DOWN => self.events.push_back(Event::DownArrow),
                            orbclient::K_LEFT => self.events.push_back(Event::LeftArrow),
                            orbclient::K_RIGHT => self.events.push_back(Event::RightArrow),
                            _ => {
                                match key_event.character {
                                    '\0' => (),
                                    '\x1B' => (),
                                    '\n' => self.events.push_back(Event::Enter),
                                    _ => self.events.push_back(Event::Text { c: key_event.character }),
                                }
                            }
                        }
                    }
                },
                orbclient::EventOption::Resize(resize_event) => {
                    self.redraw = true;
                    self.events.push_back(Event::Resize {
                        width: resize_event.width,
                        height: resize_event.height,
                    });
                },
                orbclient::EventOption::Quit(_quit_event) => {
                    self.running.set(false);
                },
                _ => (),
            };
        }

    }

    pub fn exec(&mut self) {
        'event: while self.running.get() {
            self.drain_events();
            self.draw_if_needed();
            self.drain_orbital_events();
        }
    }

    pub fn needs_redraw(&mut self) {
        self.redraw = true;
    }

    pub fn draw_if_needed(&mut self) {
        if self.redraw {
            self.draw();
            self.redraw = false;
        }
    }
}
<|endoftext|>


<commit_before>use std::io::{Write, Error, ErrorKind};
use std::io;
use InflateStream;

/// A DEFLATE decoder.
///
/// A struct implementing the `std::io::Write` trait that decompresses DEFLATE
/// encoded data into the given writer `w`.
///
/// # Example
///
/// ```
/// use inflate::writer::InflateWriter;
/// use std::io::Write;
///
/// let encoded = [243, 72, 205, 201, 201, 215, 81, 40, 207, 47, 202, 73, 1, 0];
/// let mut decoder = InflateWriter::new(Vec::new());
/// decoder.write(&encoded).unwrap();
/// let decoded = decoder.finish().unwrap();
/// println!("{}", std::str::from_utf8(&decoded).unwrap()); // prints "Hello, world"
/// ```
pub struct InflateWriter<W: Write> {
    inflater: InflateStream,
    writer: W
}

impl<W: Write> InflateWriter<W> {
    pub fn new(w: W) -> InflateWriter<W> {
        InflateWriter { inflater: InflateStream::new(), writer: w }
    }

    pub fn finish(mut self) -> io::Result<W> {
        try!(self.flush());
        Ok(self.writer)
    }
}

fn update<'a>(inflater: &'a mut InflateStream, buf: &[u8]) -> io::Result<(usize, &'a [u8])>  {
    match inflater.update(buf) {
        Ok(res) => Ok(res),
        Err(m) => return Err(Error::new(ErrorKind::Other, m)),
    }
}
impl<W: Write> Write for InflateWriter<W> {
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
        let mut n = 0;
        while n < buf.len() {
            let (num_bytes_read, result) = try!(update(&mut self.inflater, &buf[n..]));
            n += num_bytes_read;
            self.writer.write(result)?;
        }
        Ok(n)
    }

    fn flush(&mut self) -> io::Result<()> {
        let (_, result) = try!(update(&mut self.inflater, &[]));
        try!(self.writer.write(result));
        Ok(())
    }
}

#[cfg(test)]
mod test {
    use super::InflateWriter;
    use std::io::Write;

    #[test]
    fn inflate_writer() {
       let encoded = [243, 72, 205, 201, 201, 215, 81, 40, 207, 47, 202, 73, 1, 0];
       let mut decoder = InflateWriter::new(Vec::new());
       decoder.write(&encoded).unwrap();
       let decoded = decoder.finish().unwrap();
       assert!(String::from_utf8(decoded).unwrap() == "Hello, world");
    }
}
<commit_msg>Fixes inflate for 1.0<commit_after>use std::io::{Write, Error, ErrorKind};
use std::io;
use InflateStream;

/// A DEFLATE decoder.
///
/// A struct implementing the `std::io::Write` trait that decompresses DEFLATE
/// encoded data into the given writer `w`.
///
/// # Example
///
/// ```
/// use inflate::writer::InflateWriter;
/// use std::io::Write;
///
/// let encoded = [243, 72, 205, 201, 201, 215, 81, 40, 207, 47, 202, 73, 1, 0];
/// let mut decoder = InflateWriter::new(Vec::new());
/// decoder.write(&encoded).unwrap();
/// let decoded = decoder.finish().unwrap();
/// println!("{}", std::str::from_utf8(&decoded).unwrap()); // prints "Hello, world"
/// ```
pub struct InflateWriter<W: Write> {
    inflater: InflateStream,
    writer: W
}

impl<W: Write> InflateWriter<W> {
    pub fn new(w: W) -> InflateWriter<W> {
        InflateWriter { inflater: InflateStream::new(), writer: w }
    }

    pub fn finish(mut self) -> io::Result<W> {
        try!(self.flush());
        Ok(self.writer)
    }
}

fn update<'a>(inflater: &'a mut InflateStream, buf: &[u8]) -> io::Result<(usize, &'a [u8])>  {
    match inflater.update(buf) {
        Ok(res) => Ok(res),
        Err(m) => return Err(Error::new(ErrorKind::Other, m)),
    }
}
impl<W: Write> Write for InflateWriter<W> {
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
        let mut n = 0;
        while n < buf.len() {
            let (num_bytes_read, result) = try!(update(&mut self.inflater, &buf[n..]));
            n += num_bytes_read;
            try!(self.writer.write(result));
        }
        Ok(n)
    }

    fn flush(&mut self) -> io::Result<()> {
        let (_, result) = try!(update(&mut self.inflater, &[]));
        try!(self.writer.write(result));
        Ok(())
    }
}

#[cfg(test)]
mod test {
    use super::InflateWriter;
    use std::io::Write;

    #[test]
    fn inflate_writer() {
       let encoded = [243, 72, 205, 201, 201, 215, 81, 40, 207, 47, 202, 73, 1, 0];
       let mut decoder = InflateWriter::new(Vec::new());
       decoder.write(&encoded).unwrap();
       let decoded = decoder.finish().unwrap();
       assert!(String::from_utf8(decoded).unwrap() == "Hello, world");
    }
}
<|endoftext|>


<commit_before>extern crate ispc;

fn main() {
    // This build script shows how to target a specific of vector ISAs
    // using the target_isas function. We can also compile for all ISAs,
    // in which case ISPC will internally dispatch the function calls to
    // the correct ISA for the host system
    ispc::Config::new()
        .file("src/ao.ispc")
        .target_isas(vec![ispc::TargetISA::SSE2i32x4,
                     ispc::TargetISA::SSE4i32x4,
                     ispc::TargetISA::AVX11i32x8,
                     ispc::TargetISA::AVX2i32x8,
                     ispc::TargetISA::AVX512KNLi32x16,
                     ispc::TargetISA::AVX512SKXi32x16])
        .compile("ao");
}

<commit_msg>aobenhc build failing?<commit_after>extern crate ispc;

fn main() {
    // This build script shows how to target a specific of vector ISAs
    // using the target_isas function. We can also compile for all ISAs,
    // in which case ISPC will internally dispatch the function calls to
    // the correct ISA for the host system
    ispc::Config::new()
        .file("src/ao.ispc")
        .compile("ao");
}

<|endoftext|>