<commit_before>use super::{sockaddr_un, SocketAddr, UnixStream};
use crate::os::unix::io::{AsRawFd, FromRawFd, IntoRawFd, OwnedFd, RawFd};
use crate::path::Path;
use crate::sys::cvt;
use crate::sys::net::Socket;
use crate::sys_common::{AsInner, FromInner, IntoInner};
use crate::{fmt, io, mem};

/// A structure representing a Unix domain socket server.
///
/// # Examples
///
/// ```no_run
/// use std::thread;
/// use std::os::unix::net::{UnixStream, UnixListener};
///
/// fn handle_client(stream: UnixStream) {
///     // ...
/// }
///
/// fn main() -> std::io::Result<()> {
///     let listener = UnixListener::bind("/path/to/the/socket")?;
///
///     // accept connections and process them, spawning a new thread for each one
///     for stream in listener.incoming() {
///         match stream {
///             Ok(stream) => {
///                 /* connection succeeded */
///                 thread::spawn(|| handle_client(stream));
///             }
///             Err(err) => {
///                 /* connection failed */
///                 break;
///             }
///         }
///     }
///     Ok(())
/// }
/// ```
#[stable(feature = "unix_socket", since = "1.10.0")]
pub struct UnixListener(Socket);

#[stable(feature = "unix_socket", since = "1.10.0")]
impl fmt::Debug for UnixListener {
    fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result {
        let mut builder = fmt.debug_struct("UnixListener");
        builder.field("fd", self.0.as_inner());
        if let Ok(addr) = self.local_addr() {
            builder.field("local", &addr);
        }
        builder.finish()
    }
}

impl UnixListener {
    /// Creates a new `UnixListener` bound to the specified socket.
    ///
    /// # Examples
    ///
    /// ```no_run
    /// use std::os::unix::net::UnixListener;
    ///
    /// let listener = match UnixListener::bind("/path/to/the/socket") {
    ///     Ok(sock) => sock,
    ///     Err(e) => {
    ///         println!("Couldn't connect: {:?}", e);
    ///         return
    ///     }
    /// };
    /// ```
    #[stable(feature = "unix_socket", since = "1.10.0")]
    pub fn bind<P: AsRef<Path>>(path: P) -> io::Result<UnixListener> {
        unsafe {
            let inner = Socket::new_raw(libc::AF_UNIX, libc::SOCK_STREAM)?;
            let (addr, len) = sockaddr_un(path.as_ref())?;

            cvt(libc::bind(inner.as_inner().as_raw_fd(), &addr as *const _ as *const _, len as _))?;
            cvt(libc::listen(inner.as_inner().as_raw_fd(), 128))?;

            Ok(UnixListener(inner))
        }
    }

    /// Accepts a new incoming connection to this listener.
    ///
    /// This function will block the calling thread until a new Unix connection
    /// is established. When established, the corresponding [`UnixStream`] and
    /// the remote peer's address will be returned.
    ///
    /// [`UnixStream`]: crate::os::unix::net::UnixStream
    ///
    /// # Examples
    ///
    /// ```no_run
    /// use std::os::unix::net::UnixListener;
    ///
    /// fn main() -> std::io::Result<()> {
    ///     let listener = UnixListener::bind("/path/to/the/socket")?;
    ///
    ///     match listener.accept() {
    ///         Ok((socket, addr)) => println!("Got a client: {:?}", addr),
    ///         Err(e) => println!("accept function failed: {:?}", e),
    ///     }
    ///     Ok(())
    /// }
    /// ```
    #[stable(feature = "unix_socket", since = "1.10.0")]
    pub fn accept(&self) -> io::Result<(UnixStream, SocketAddr)> {
        let mut storage: libc::sockaddr_un = unsafe { mem::zeroed() };
        let mut len = mem::size_of_val(&storage) as libc::socklen_t;
        let sock = self.0.accept(&mut storage as *mut _ as *mut _, &mut len)?;
        let addr = SocketAddr::from_parts(storage, len)?;
        Ok((UnixStream(sock), addr))
    }

    /// Creates a new independently owned handle to the underlying socket.
    ///
    /// The returned `UnixListener` is a reference to the same socket that this
    /// object references. Both handles can be used to accept incoming
    /// connections and options set on one listener will affect the other.
    ///
    /// # Examples
    ///
    /// ```no_run
    /// use std::os::unix::net::UnixListener;
    ///
    /// fn main() -> std::io::Result<()> {
    ///     let listener = UnixListener::bind("/path/to/the/socket")?;
    ///     let listener_copy = listener.try_clone().expect("try_clone failed");
    ///     Ok(())
    /// }
    /// ```
    #[stable(feature = "unix_socket", since = "1.10.0")]
    pub fn try_clone(&self) -> io::Result<UnixListener> {
        self.0.duplicate().map(UnixListener)
    }

    /// Returns the local socket address of this listener.
    ///
    /// # Examples
    ///
    /// ```no_run
    /// use std::os::unix::net::UnixListener;
    ///
    /// fn main() -> std::io::Result<()> {
    ///     let listener = UnixListener::bind("/path/to/the/socket")?;
    ///     let addr = listener.local_addr().expect("Couldn't get local address");
    ///     Ok(())
    /// }
    /// ```
    #[stable(feature = "unix_socket", since = "1.10.0")]
    pub fn local_addr(&self) -> io::Result<SocketAddr> {
        SocketAddr::new(|addr, len| unsafe { libc::getsockname(self.as_raw_fd(), addr, len) })
    }

    /// Moves the socket into or out of nonblocking mode.
    ///
    /// This will result in the `accept` operation becoming nonblocking,
    /// i.e., immediately returning from their calls. If the IO operation is
    /// successful, `Ok` is returned and no further action is required. If the
    /// IO operation could not be completed and needs to be retried, an error
    /// with kind [`io::ErrorKind::WouldBlock`] is returned.
    ///
    /// # Examples
    ///
    /// ```no_run
    /// use std::os::unix::net::UnixListener;
    ///
    /// fn main() -> std::io::Result<()> {
    ///     let listener = UnixListener::bind("/path/to/the/socket")?;
    ///     listener.set_nonblocking(true).expect("Couldn't set non blocking");
    ///     Ok(())
    /// }
    /// ```
    #[stable(feature = "unix_socket", since = "1.10.0")]
    pub fn set_nonblocking(&self, nonblocking: bool) -> io::Result<()> {
        self.0.set_nonblocking(nonblocking)
    }

    /// Returns the value of the `SO_ERROR` option.
    ///
    /// # Examples
    ///
    /// ```no_run
    /// use std::os::unix::net::UnixListener;
    ///
    /// fn main() -> std::io::Result<()> {
    ///     let listener = UnixListener::bind("/tmp/sock")?;
    ///
    ///     if let Ok(Some(err)) = listener.take_error() {
    ///         println!("Got error: {:?}", err);
    ///     }
    ///     Ok(())
    /// }
    /// ```
    ///
    /// # Platform specific
    /// On Redox this always returns `None`.
    #[stable(feature = "unix_socket", since = "1.10.0")]
    pub fn take_error(&self) -> io::Result<Option<io::Error>> {
        self.0.take_error()
    }

    /// Returns an iterator over incoming connections.
    ///
    /// The iterator will never return [`None`] and will also not yield the
    /// peer's [`SocketAddr`] structure.
    ///
    /// # Examples
    ///
    /// ```no_run
    /// use std::thread;
    /// use std::os::unix::net::{UnixStream, UnixListener};
    ///
    /// fn handle_client(stream: UnixStream) {
    ///     // ...
    /// }
    ///
    /// fn main() -> std::io::Result<()> {
    ///     let listener = UnixListener::bind("/path/to/the/socket")?;
    ///
    ///     for stream in listener.incoming() {
    ///         match stream {
    ///             Ok(stream) => {
    ///                 thread::spawn(|| handle_client(stream));
    ///             }
    ///             Err(err) => {
    ///                 break;
    ///             }
    ///         }
    ///     }
    ///     Ok(())
    /// }
    /// ```
    #[stable(feature = "unix_socket", since = "1.10.0")]
    pub fn incoming(&self) -> Incoming<'_> {
        Incoming { listener: self }
    }
}

#[stable(feature = "unix_socket", since = "1.10.0")]
impl AsRawFd for UnixListener {
    #[inline]
    fn as_raw_fd(&self) -> RawFd {
        self.0.as_inner().as_raw_fd()
    }
}

#[stable(feature = "unix_socket", since = "1.10.0")]
impl FromRawFd for UnixListener {
    #[inline]
    unsafe fn from_raw_fd(fd: RawFd) -> UnixListener {
        UnixListener(Socket::from_inner(FromInner::from_inner(OwnedFd::from_raw_fd(fd))))
    }
}

#[stable(feature = "unix_socket", since = "1.10.0")]
impl IntoRawFd for UnixListener {
    #[inline]
    fn into_raw_fd(self) -> RawFd {
        self.0.into_inner().into_inner().into_raw_fd()
    }
}

#[stable(feature = "unix_socket", since = "1.10.0")]
impl<'a> IntoIterator for &'a UnixListener {
    type Item = io::Result<UnixStream>;
    type IntoIter = Incoming<'a>;

    fn into_iter(self) -> Incoming<'a> {
        self.incoming()
    }
}

/// An iterator over incoming connections to a [`UnixListener`].
///
/// It will never return [`None`].
///
/// # Examples
///
/// ```no_run
/// use std::thread;
/// use std::os::unix::net::{UnixStream, UnixListener};
///
/// fn handle_client(stream: UnixStream) {
///     // ...
/// }
///
/// fn main() -> std::io::Result<()> {
///     let listener = UnixListener::bind("/path/to/the/socket")?;
///
///     for stream in listener.incoming() {
///         match stream {
///             Ok(stream) => {
///                 thread::spawn(|| handle_client(stream));
///             }
///             Err(err) => {
///                 break;
///             }
///         }
///     }
///     Ok(())
/// }
/// ```
#[derive(Debug)]
#[stable(feature = "unix_socket", since = "1.10.0")]
pub struct Incoming<'a> {
    listener: &'a UnixListener,
}

#[stable(feature = "unix_socket", since = "1.10.0")]
impl<'a> Iterator for Incoming<'a> {
    type Item = io::Result<UnixStream>;

    fn next(&mut self) -> Option<io::Result<UnixStream>> {
        Some(self.listener.accept().map(|s| s.0))
    }

    fn size_hint(&self) -> (usize, Option<usize>) {
        (usize::MAX, None)
    }
}
<commit_msg>Implement `AsFd` etc. for `UnixListener`.<commit_after>use super::{sockaddr_un, SocketAddr, UnixStream};
use crate::os::unix::io::{AsFd, AsRawFd, BorrowedFd, FromRawFd, IntoRawFd, OwnedFd, RawFd};
use crate::path::Path;
use crate::sys::cvt;
use crate::sys::net::Socket;
use crate::sys_common::{AsInner, FromInner, IntoInner};
use crate::{fmt, io, mem};

/// A structure representing a Unix domain socket server.
///
/// # Examples
///
/// ```no_run
/// use std::thread;
/// use std::os::unix::net::{UnixStream, UnixListener};
///
/// fn handle_client(stream: UnixStream) {
///     // ...
/// }
///
/// fn main() -> std::io::Result<()> {
///     let listener = UnixListener::bind("/path/to/the/socket")?;
///
///     // accept connections and process them, spawning a new thread for each one
///     for stream in listener.incoming() {
///         match stream {
///             Ok(stream) => {
///                 /* connection succeeded */
///                 thread::spawn(|| handle_client(stream));
///             }
///             Err(err) => {
///                 /* connection failed */
///                 break;
///             }
///         }
///     }
///     Ok(())
/// }
/// ```
#[stable(feature = "unix_socket", since = "1.10.0")]
pub struct UnixListener(Socket);

#[stable(feature = "unix_socket", since = "1.10.0")]
impl fmt::Debug for UnixListener {
    fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result {
        let mut builder = fmt.debug_struct("UnixListener");
        builder.field("fd", self.0.as_inner());
        if let Ok(addr) = self.local_addr() {
            builder.field("local", &addr);
        }
        builder.finish()
    }
}

impl UnixListener {
    /// Creates a new `UnixListener` bound to the specified socket.
    ///
    /// # Examples
    ///
    /// ```no_run
    /// use std::os::unix::net::UnixListener;
    ///
    /// let listener = match UnixListener::bind("/path/to/the/socket") {
    ///     Ok(sock) => sock,
    ///     Err(e) => {
    ///         println!("Couldn't connect: {:?}", e);
    ///         return
    ///     }
    /// };
    /// ```
    #[stable(feature = "unix_socket", since = "1.10.0")]
    pub fn bind<P: AsRef<Path>>(path: P) -> io::Result<UnixListener> {
        unsafe {
            let inner = Socket::new_raw(libc::AF_UNIX, libc::SOCK_STREAM)?;
            let (addr, len) = sockaddr_un(path.as_ref())?;

            cvt(libc::bind(inner.as_inner().as_raw_fd(), &addr as *const _ as *const _, len as _))?;
            cvt(libc::listen(inner.as_inner().as_raw_fd(), 128))?;

            Ok(UnixListener(inner))
        }
    }

    /// Accepts a new incoming connection to this listener.
    ///
    /// This function will block the calling thread until a new Unix connection
    /// is established. When established, the corresponding [`UnixStream`] and
    /// the remote peer's address will be returned.
    ///
    /// [`UnixStream`]: crate::os::unix::net::UnixStream
    ///
    /// # Examples
    ///
    /// ```no_run
    /// use std::os::unix::net::UnixListener;
    ///
    /// fn main() -> std::io::Result<()> {
    ///     let listener = UnixListener::bind("/path/to/the/socket")?;
    ///
    ///     match listener.accept() {
    ///         Ok((socket, addr)) => println!("Got a client: {:?}", addr),
    ///         Err(e) => println!("accept function failed: {:?}", e),
    ///     }
    ///     Ok(())
    /// }
    /// ```
    #[stable(feature = "unix_socket", since = "1.10.0")]
    pub fn accept(&self) -> io::Result<(UnixStream, SocketAddr)> {
        let mut storage: libc::sockaddr_un = unsafe { mem::zeroed() };
        let mut len = mem::size_of_val(&storage) as libc::socklen_t;
        let sock = self.0.accept(&mut storage as *mut _ as *mut _, &mut len)?;
        let addr = SocketAddr::from_parts(storage, len)?;
        Ok((UnixStream(sock), addr))
    }

    /// Creates a new independently owned handle to the underlying socket.
    ///
    /// The returned `UnixListener` is a reference to the same socket that this
    /// object references. Both handles can be used to accept incoming
    /// connections and options set on one listener will affect the other.
    ///
    /// # Examples
    ///
    /// ```no_run
    /// use std::os::unix::net::UnixListener;
    ///
    /// fn main() -> std::io::Result<()> {
    ///     let listener = UnixListener::bind("/path/to/the/socket")?;
    ///     let listener_copy = listener.try_clone().expect("try_clone failed");
    ///     Ok(())
    /// }
    /// ```
    #[stable(feature = "unix_socket", since = "1.10.0")]
    pub fn try_clone(&self) -> io::Result<UnixListener> {
        self.0.duplicate().map(UnixListener)
    }

    /// Returns the local socket address of this listener.
    ///
    /// # Examples
    ///
    /// ```no_run
    /// use std::os::unix::net::UnixListener;
    ///
    /// fn main() -> std::io::Result<()> {
    ///     let listener = UnixListener::bind("/path/to/the/socket")?;
    ///     let addr = listener.local_addr().expect("Couldn't get local address");
    ///     Ok(())
    /// }
    /// ```
    #[stable(feature = "unix_socket", since = "1.10.0")]
    pub fn local_addr(&self) -> io::Result<SocketAddr> {
        SocketAddr::new(|addr, len| unsafe { libc::getsockname(self.as_raw_fd(), addr, len) })
    }

    /// Moves the socket into or out of nonblocking mode.
    ///
    /// This will result in the `accept` operation becoming nonblocking,
    /// i.e., immediately returning from their calls. If the IO operation is
    /// successful, `Ok` is returned and no further action is required. If the
    /// IO operation could not be completed and needs to be retried, an error
    /// with kind [`io::ErrorKind::WouldBlock`] is returned.
    ///
    /// # Examples
    ///
    /// ```no_run
    /// use std::os::unix::net::UnixListener;
    ///
    /// fn main() -> std::io::Result<()> {
    ///     let listener = UnixListener::bind("/path/to/the/socket")?;
    ///     listener.set_nonblocking(true).expect("Couldn't set non blocking");
    ///     Ok(())
    /// }
    /// ```
    #[stable(feature = "unix_socket", since = "1.10.0")]
    pub fn set_nonblocking(&self, nonblocking: bool) -> io::Result<()> {
        self.0.set_nonblocking(nonblocking)
    }

    /// Returns the value of the `SO_ERROR` option.
    ///
    /// # Examples
    ///
    /// ```no_run
    /// use std::os::unix::net::UnixListener;
    ///
    /// fn main() -> std::io::Result<()> {
    ///     let listener = UnixListener::bind("/tmp/sock")?;
    ///
    ///     if let Ok(Some(err)) = listener.take_error() {
    ///         println!("Got error: {:?}", err);
    ///     }
    ///     Ok(())
    /// }
    /// ```
    ///
    /// # Platform specific
    /// On Redox this always returns `None`.
    #[stable(feature = "unix_socket", since = "1.10.0")]
    pub fn take_error(&self) -> io::Result<Option<io::Error>> {
        self.0.take_error()
    }

    /// Returns an iterator over incoming connections.
    ///
    /// The iterator will never return [`None`] and will also not yield the
    /// peer's [`SocketAddr`] structure.
    ///
    /// # Examples
    ///
    /// ```no_run
    /// use std::thread;
    /// use std::os::unix::net::{UnixStream, UnixListener};
    ///
    /// fn handle_client(stream: UnixStream) {
    ///     // ...
    /// }
    ///
    /// fn main() -> std::io::Result<()> {
    ///     let listener = UnixListener::bind("/path/to/the/socket")?;
    ///
    ///     for stream in listener.incoming() {
    ///         match stream {
    ///             Ok(stream) => {
    ///                 thread::spawn(|| handle_client(stream));
    ///             }
    ///             Err(err) => {
    ///                 break;
    ///             }
    ///         }
    ///     }
    ///     Ok(())
    /// }
    /// ```
    #[stable(feature = "unix_socket", since = "1.10.0")]
    pub fn incoming(&self) -> Incoming<'_> {
        Incoming { listener: self }
    }
}

#[stable(feature = "unix_socket", since = "1.10.0")]
impl AsRawFd for UnixListener {
    #[inline]
    fn as_raw_fd(&self) -> RawFd {
        self.0.as_inner().as_raw_fd()
    }
}

#[stable(feature = "unix_socket", since = "1.10.0")]
impl FromRawFd for UnixListener {
    #[inline]
    unsafe fn from_raw_fd(fd: RawFd) -> UnixListener {
        UnixListener(Socket::from_inner(FromInner::from_inner(OwnedFd::from_raw_fd(fd))))
    }
}

#[stable(feature = "unix_socket", since = "1.10.0")]
impl IntoRawFd for UnixListener {
    #[inline]
    fn into_raw_fd(self) -> RawFd {
        self.0.into_inner().into_inner().into_raw_fd()
    }
}

#[unstable(feature = "io_safety", issue = "87074")]
impl AsFd for UnixListener {
    #[inline]
    fn as_fd(&self) -> BorrowedFd<'_> {
        self.0.as_inner().as_fd()
    }
}

#[unstable(feature = "io_safety", issue = "87074")]
impl From<OwnedFd> for UnixListener {
    #[inline]
    fn from(fd: OwnedFd) -> UnixListener {
        UnixListener(Socket::from_inner(FromInner::from_inner(OwnedFd::from(fd))))
    }
}

#[unstable(feature = "io_safety", issue = "87074")]
impl From<UnixListener> for OwnedFd {
    #[inline]
    fn from(listener: UnixListener) -> OwnedFd {
        listener.0.into_inner().into_inner().into()
    }
}

#[stable(feature = "unix_socket", since = "1.10.0")]
impl<'a> IntoIterator for &'a UnixListener {
    type Item = io::Result<UnixStream>;
    type IntoIter = Incoming<'a>;

    fn into_iter(self) -> Incoming<'a> {
        self.incoming()
    }
}

/// An iterator over incoming connections to a [`UnixListener`].
///
/// It will never return [`None`].
///
/// # Examples
///
/// ```no_run
/// use std::thread;
/// use std::os::unix::net::{UnixStream, UnixListener};
///
/// fn handle_client(stream: UnixStream) {
///     // ...
/// }
///
/// fn main() -> std::io::Result<()> {
///     let listener = UnixListener::bind("/path/to/the/socket")?;
///
///     for stream in listener.incoming() {
///         match stream {
///             Ok(stream) => {
///                 thread::spawn(|| handle_client(stream));
///             }
///             Err(err) => {
///                 break;
///             }
///         }
///     }
///     Ok(())
/// }
/// ```
#[derive(Debug)]
#[stable(feature = "unix_socket", since = "1.10.0")]
pub struct Incoming<'a> {
    listener: &'a UnixListener,
}

#[stable(feature = "unix_socket", since = "1.10.0")]
impl<'a> Iterator for Incoming<'a> {
    type Item = io::Result<UnixStream>;

    fn next(&mut self) -> Option<io::Result<UnixStream>> {
        Some(self.listener.accept().map(|s| s.0))
    }

    fn size_hint(&self) -> (usize, Option<usize>) {
        (usize::MAX, None)
    }
}
<|endoftext|>


<commit_before>//
// base64.rs - base64 implementation
//
// The Base64 Alphabet
//
//  Value Encoding  Value Encoding  Value Encoding  Value Encoding
//   0 A            17 R            34 i            51 z
//   1 B            18 S            35 j            52 0
//   2 C            19 T            36 k            53 1
//   3 D            20 U            37 l            54 2
//   4 E            21 V            38 m            55 3
//   5 F            22 W            39 n            56 4
//   6 G            23 X            40 o            57 5
//   7 H            24 Y            41 p            58 6
//   8 I            25 Z            42 q            59 7
//   9 J            26 a            43 r            60 8
//  10 K            27 b            44 s            61 9
//  11 L            28 c            45 t            62 +
//  12 M            29 d            46 u            63 /
//  13 N            30 e            47 v
//  14 O            31 f            48 w         (pad) =
//  15 P            32 g            49 x
//  16 Q            33 h            50 y
//

use std;

import std::map;

export base64, mk_base64;

iface base64 {
    fn encode(src: [u8]) -> [u8];
    fn decode(src: [u8]) -> [u8];
    fn urlsafe_encode(src: [u8]) -> [u8];
    fn urlsafe_decode(src: [u8]) -> [u8];
}

fn mk_base64() -> base64 {
    type _base64 = {padd: u8, table: [u8]};
    impl of base64 for _base64 {
        fn encode(src: [u8]) -> [u8] {
            let srclen = vec::len(src);
            let targ = if srclen % 3u == 0u {
                vec::init_elt_mut((srclen / 3u) * 4u, 0u8)
            } else {
                vec::init_elt_mut((srclen / 3u + 1u) * 4u, 0u8)
            };
            let input = vec::init_elt_mut(3u, 0u8);
            let output = vec::init_elt_mut(4u, 0u8);
            let curr = 0u, src_curr = 0u;
            let table = self.table;

            while srclen > 2u {
                input[0] = src[src_curr];
                input[1] = src[src_curr + 1u];
                input[2] = src[src_curr + 2u];
                srclen -= 3u; src_curr += 3u;

                output[0] = input[0] >> 2u8;
                output[1] = ((input[0] &  3u8) << 4u8) | (input[1] >> 4u8);
                output[2] = ((input[1] & 15u8) << 2u8) | (input[2] >> 6u8);
                output[3] = input[2] & 63u8;

                targ[curr] = table[output[0]]; curr += 1u;
                targ[curr] = table[output[1]]; curr += 1u;
                targ[curr] = table[output[2]]; curr += 1u;
                targ[curr] = table[output[3]]; curr += 1u;
            }

            if srclen != 0u {
                input[0] = 0u8; input[1] = 0u8; input[2] = 0u8;

                alt srclen {
                  1u { input[0] = src[src_curr]; }
                  2u { input[0] = src[src_curr];
                       input[1] = src[src_curr + 1u]; }
                  _  { }
                }

                output[0] = input[0] >> 2u8;
                output[1] = ((input[0] &  3u8) << 4u8) | (input[1] >> 4u8);
                output[2] = ((input[1] & 15u8) << 2u8) | (input[2] >> 6u8);

                targ[curr] = table[output[0]]; curr += 1u;
                targ[curr] = table[output[1]]; curr += 1u;
                targ[curr] = if srclen == 1u {
                    self.padd
                } else {
                    table[output[2]]
                };
                curr += 1u; targ[curr] = self.padd; curr += 1u;
            }

            vec::from_mut(targ)
        }
        fn decode(src: [u8]) -> [u8] {
            let src = src;
            let srclen = vec::len(src);
            let input  = vec::init_elt_mut(4u, 0u8);
            let output = vec::init_elt_mut(4u, 0u8);
            let targ = if src[srclen - 2u] == self.padd {
                vec::init_elt_mut(srclen / 4u * 3u - 2u, 0u8)
            } else if src[srclen - 1u] == self.padd {
                vec::init_elt_mut(srclen / 4u * 3u - 1u, 0u8)
            } else {
                vec::init_elt_mut(srclen / 4u * 3u, 0u8)
            };
            let c = -1;
            let curr = 0u, src_curr = 0u;
            let table = self.table;

            while srclen > 4u {
                input[0] = src[src_curr];
                input[1] = src[src_curr + 1u];
                input[2] = src[src_curr + 2u];
                input[3] = src[src_curr + 3u];
                srclen -= 4u; src_curr += 4u;

                c = idx(table, input[0]);
                if c == -1 { fail "malformed base64 string"; }
                output[0] = c as u8;

                c = idx(table, input[1]);
                if c == -1 { fail "malformed base64 string"; }
                output[1] = c as u8;

                c = idx(table, input[2]);
                if c == -1 { fail "malformed base64 string"; }
                output[2] = c as u8;

                c = idx(table, input[3]);
                if c == -1 { fail "malformed base64 string"; }
                output[3] = c as u8;

                targ[curr] = (output[0] << 2u8) | (output[1] >> 4u8);
                curr += 1u;

                targ[curr] =
                    ((output[1] & 15u8) << 4u8) | (output[2] >> 2u8);
                curr += 1u;

                targ[curr] = ((output[2] &  3u8) << 6u8) | output[3];
                curr += 1u;
            }

            if srclen == 4u {
                input[0] = src[src_curr];
                input[1] = src[src_curr + 1u];
                input[2] = src[src_curr + 2u];
                input[3] = src[src_curr + 3u];

                c = idx(table, input[0]);
                if c == -1 { fail "malformed base64 string"; }
                output[0] = c as u8;

                c = idx(table, input[1]);
                if c == -1 { fail "malformed base64 string"; }
                output[1] = c as u8;

                targ[curr] = (output[0] << 2u8) | (output[1] >> 4u8);
                curr += 1u;

                if input[2] == self.padd { ret vec::from_mut(targ); }

                c = idx(table, input[2]);
                if c == -1 { fail "malformed base64 string"; }
                output[2] = c as u8;

                targ[curr] = ((output[1] & 15u8) << 4u8) | (output[2] >> 2u8);
                curr += 1u;

                if input[3] == self.padd { ret vec::from_mut(targ); }

                c = idx(table, input[3]);
                if c == -1 { fail "malformed base64 string"; }
                output[3] = c as u8;

                targ[curr] = ((output[2] & 3u8) << 6u8) | output[3];
            } else {
                fail "malformed base64 string";
            }

            vec::from_mut(targ)
        }
        fn urlsafe_encode(src: [u8]) -> [u8] {
            let res = self.encode(src);
            vec::map(res) { |i|
                if i == 43u8 { 45u8 }
                else if i == 47u8 { 95u8 }
                else { i }
            }
        }
        fn urlsafe_decode(src: [u8]) -> [u8] {
            let src0 = vec::map(src) { |i|
                if i == 45u8 { 43u8 }
                else if i == 95u8 { 47u8 }
                else { i }
            };
            self.decode(src0)
        }
    }

    let table = vec::init_elt_mut(64u, 0u8), i = 0u8;
    u8::range(65u8, 91u8)  { |j| table[i] = j; i += 1u8; }
    u8::range(97u8, 123u8) { |j| table[i] = j; i += 1u8; }
    u8::range(48u8, 58u8)  { |j| table[i] = j; i += 1u8; }
    table[i] = 43u8; i += 1u8;
    table[i] = 47u8; i += 1u8;

    {padd: 61u8, table: vec::from_mut(table)} as base64
}

fn idx(elems: [const u8], x: u8) -> int {
    let i = 0;
    for e in elems {
        if e == x { ret i; }
        i += 1;
    }
    ret -1;
}

#[cfg(test)]
mod tests {
    import core::str::{from_bytes, bytes};
    enum mode {
        t_decode,
        t_encode,
        t_urlsafe_encode,
        t_urlsafe_decode
    }
    fn setup(t: mode) -> map::hashmap<str, str> {
        let m = map::new_str_hash::<str>();
        alt t {
          t_decode | t_urlsafe_decode {
            m.insert("cGxlYXN1cmUu", "pleasure.");
            m.insert("bGVhc3VyZS4=",  "leasure.");
            m.insert("ZWFzdXJlLg==",   "easure.");
            m.insert("YXN1cmUu",        "asure.");
            m.insert("c3VyZS4=",         "sure.");
          }
          t_encode | t_urlsafe_encode {
            m.insert("pleasure.", "cGxlYXN1cmUu");
            m.insert("leasure.",  "bGVhc3VyZS4=");
            m.insert("easure.",   "ZWFzdXJlLg==");
            m.insert("asure.",        "YXN1cmUu");
            m.insert("sure.",         "c3VyZS4=");
          }
        }
        m
    }
    fn do_test(t: mode) {
        let b64 = mk_base64();
        let m = setup(t);
        m.keys { |k|
            let expected = m.get(k);
            let actual = alt t {
              t_decode { from_bytes(b64.decode(bytes(k))) }
              t_encode { from_bytes(b64.encode(bytes(k))) }
              t_urlsafe_decode { from_bytes(b64.urlsafe_decode(bytes(k))) }
              t_urlsafe_encode { from_bytes(b64.urlsafe_encode(bytes(k))) }
            };
            #debug("expected: %s", expected);
            #debug("actual:   %s", actual);
            assert expected == actual;
        }
    }
    #[test]
    fn decode() { do_test(t_decode); }
    #[test]
    fn encode() { do_test(t_encode); }
    #[test]
    fn urlsafe_decode() { do_test(t_urlsafe_decode); }
    #[test]
    fn urlsafe_encode() { do_test(t_urlsafe_encode); }
}
<commit_msg>changes<commit_after>//
// base64.rs - base64 implementation
//
// The Base64 Alphabet
//
//  Value Encoding  Value Encoding  Value Encoding  Value Encoding
//   0 A            17 R            34 i            51 z
//   1 B            18 S            35 j            52 0
//   2 C            19 T            36 k            53 1
//   3 D            20 U            37 l            54 2
//   4 E            21 V            38 m            55 3
//   5 F            22 W            39 n            56 4
//   6 G            23 X            40 o            57 5
//   7 H            24 Y            41 p            58 6
//   8 I            25 Z            42 q            59 7
//   9 J            26 a            43 r            60 8
//  10 K            27 b            44 s            61 9
//  11 L            28 c            45 t            62 +
//  12 M            29 d            46 u            63 /
//  13 N            30 e            47 v
//  14 O            31 f            48 w         (pad) =
//  15 P            32 g            49 x
//  16 Q            33 h            50 y
//

use std;

import std::map;

export base64, mk_base64;

iface base64 {
    fn encode(src: [u8]) -> [u8];
    fn decode(src: [u8]) -> [u8];
    fn urlsafe_encode(src: [u8]) -> [u8];
    fn urlsafe_decode(src: [u8]) -> [u8];
}

fn mk_base64() -> base64 {
    type _base64 = {padd: u8, table: [u8]};
    impl of base64 for _base64 {
        fn encode(src: [u8]) -> [u8] {
            let srclen = vec::len(src);
            let targ = if srclen % 3u == 0u {
                vec::init_elt_mut((srclen / 3u) * 4u, 0u8)
            } else {
                vec::init_elt_mut((srclen / 3u + 1u) * 4u, 0u8)
            };
            let input = vec::init_elt_mut(3u, 0u8);
            let output = vec::init_elt_mut(4u, 0u8);
            let curr = 0u, src_curr = 0u;
            let table = self.table;

            while srclen > 2u {
                input[0] = src[src_curr];
                input[1] = src[src_curr + 1u];
                input[2] = src[src_curr + 2u];
                srclen -= 3u; src_curr += 3u;

                output[0] = input[0] >> 2u8;
                output[1] = ((input[0] &  3u8) << 4u8) | (input[1] >> 4u8);
                output[2] = ((input[1] & 15u8) << 2u8) | (input[2] >> 6u8);
                output[3] = input[2] & 63u8;

                targ[curr] = table[output[0]]; curr += 1u;
                targ[curr] = table[output[1]]; curr += 1u;
                targ[curr] = table[output[2]]; curr += 1u;
                targ[curr] = table[output[3]]; curr += 1u;
            }

            if srclen != 0u {
                input[0] = 0u8; input[1] = 0u8; input[2] = 0u8;

                alt srclen {
                  1u { input[0] = src[src_curr]; }
                  2u { input[0] = src[src_curr];
                       input[1] = src[src_curr + 1u]; }
                  _  { }
                }

                output[0] = input[0] >> 2u8;
                output[1] = ((input[0] &  3u8) << 4u8) | (input[1] >> 4u8);
                output[2] = ((input[1] & 15u8) << 2u8) | (input[2] >> 6u8);

                targ[curr] = table[output[0]]; curr += 1u;
                targ[curr] = table[output[1]]; curr += 1u;
                targ[curr] = if srclen == 1u {
                    self.padd
                } else {
                    table[output[2]]
                };
                curr += 1u; targ[curr] = self.padd; curr += 1u;
            }

            vec::from_mut(targ)
        }
        fn decode(src: [u8]) -> [u8] {
            let srclen = vec::len(src);
            let input  = vec::init_elt_mut(4u, 0u8);
            let output = vec::init_elt_mut(4u, 0u8);
            let targ = if src[srclen - 2u] == self.padd {
                vec::init_elt_mut(srclen / 4u * 3u - 2u, 0u8)
            } else if src[srclen - 1u] == self.padd {
                vec::init_elt_mut(srclen / 4u * 3u - 1u, 0u8)
            } else {
                vec::init_elt_mut(srclen / 4u * 3u, 0u8)
            };
            let c = -1;
            let curr = 0u, src_curr = 0u;
            let table = self.table;

            while srclen > 4u {
                input[0] = src[src_curr];
                input[1] = src[src_curr + 1u];
                input[2] = src[src_curr + 2u];
                input[3] = src[src_curr + 3u];
                srclen -= 4u; src_curr += 4u;

                c = idx(table, input[0]);
                if c == -1 { fail "malformed base64 string"; }
                output[0] = c as u8;

                c = idx(table, input[1]);
                if c == -1 { fail "malformed base64 string"; }
                output[1] = c as u8;

                c = idx(table, input[2]);
                if c == -1 { fail "malformed base64 string"; }
                output[2] = c as u8;

                c = idx(table, input[3]);
                if c == -1 { fail "malformed base64 string"; }
                output[3] = c as u8;

                targ[curr] = (output[0] << 2u8) | (output[1] >> 4u8);
                curr += 1u;

                targ[curr] =
                    ((output[1] & 15u8) << 4u8) | (output[2] >> 2u8);
                curr += 1u;

                targ[curr] = ((output[2] &  3u8) << 6u8) | output[3];
                curr += 1u;
            }

            if srclen == 4u {
                input[0] = src[src_curr];
                input[1] = src[src_curr + 1u];
                input[2] = src[src_curr + 2u];
                input[3] = src[src_curr + 3u];

                c = idx(table, input[0]);
                if c == -1 { fail "malformed base64 string"; }
                output[0] = c as u8;

                c = idx(table, input[1]);
                if c == -1 { fail "malformed base64 string"; }
                output[1] = c as u8;

                targ[curr] = (output[0] << 2u8) | (output[1] >> 4u8);
                curr += 1u;

                if input[2] == self.padd { ret vec::from_mut(targ); }

                c = idx(table, input[2]);
                if c == -1 { fail "malformed base64 string"; }
                output[2] = c as u8;

                targ[curr] = ((output[1] & 15u8) << 4u8) | (output[2] >> 2u8);
                curr += 1u;

                if input[3] == self.padd { ret vec::from_mut(targ); }

                c = idx(table, input[3]);
                if c == -1 { fail "malformed base64 string"; }
                output[3] = c as u8;

                targ[curr] = ((output[2] & 3u8) << 6u8) | output[3];
            } else {
                fail "malformed base64 string";
            }

            vec::from_mut(targ)
        }
        fn urlsafe_encode(src: [u8]) -> [u8] {
            let res = self.encode(src);
            vec::map(res) { |i|
                if i == 43u8 { 45u8 }
                else if i == 47u8 { 95u8 }
                else { i }
            }
        }
        fn urlsafe_decode(src: [u8]) -> [u8] {
            let src0 = vec::map(src) { |i|
                if i == 45u8 { 43u8 }
                else if i == 95u8 { 47u8 }
                else { i }
            };
            self.decode(src0)
        }
    }

    let table = vec::init_elt_mut(64u, 0u8), i = 0u8;
    u8::range(65u8, 91u8)  { |j| table[i] = j; i += 1u8; }
    u8::range(97u8, 123u8) { |j| table[i] = j; i += 1u8; }
    u8::range(48u8, 58u8)  { |j| table[i] = j; i += 1u8; }
    table[i] = 43u8; i += 1u8;
    table[i] = 47u8; i += 1u8;

    {padd: 61u8, table: vec::from_mut(table)} as base64
}

fn idx(elems: [const u8], x: u8) -> int {
    let i = 0;
    for e in elems {
        if e == x { ret i; }
        i += 1;
    }
    ret -1;
}

#[cfg(test)]
mod tests {
    import core::str::{from_bytes, bytes};
    enum mode {
        t_decode,
        t_encode,
        t_urlsafe_encode,
        t_urlsafe_decode
    }
    fn setup(t: mode) -> map::hashmap<str, str> {
        let m = map::new_str_hash::<str>();
        alt t {
          t_decode | t_urlsafe_decode {
            m.insert("cGxlYXN1cmUu", "pleasure.");
            m.insert("bGVhc3VyZS4=",  "leasure.");
            m.insert("ZWFzdXJlLg==",   "easure.");
            m.insert("YXN1cmUu",        "asure.");
            m.insert("c3VyZS4=",         "sure.");
          }
          t_encode | t_urlsafe_encode {
            m.insert("pleasure.", "cGxlYXN1cmUu");
            m.insert("leasure.",  "bGVhc3VyZS4=");
            m.insert("easure.",   "ZWFzdXJlLg==");
            m.insert("asure.",        "YXN1cmUu");
            m.insert("sure.",         "c3VyZS4=");
          }
        }
        m
    }
    fn do_test(t: mode) {
        let b64 = mk_base64();
        let m = setup(t);
        m.keys { |k|
            let expected = m.get(k);
            let actual = alt t {
              t_decode { from_bytes(b64.decode(bytes(k))) }
              t_encode { from_bytes(b64.encode(bytes(k))) }
              t_urlsafe_decode { from_bytes(b64.urlsafe_decode(bytes(k))) }
              t_urlsafe_encode { from_bytes(b64.urlsafe_encode(bytes(k))) }
            };
            #debug("expected: %s", expected);
            #debug("actual:   %s", actual);
            assert expected == actual;
        }
    }
    #[test]
    fn decode() { do_test(t_decode); }
    #[test]
    fn encode() { do_test(t_encode); }
    #[test]
    fn urlsafe_decode() { do_test(t_urlsafe_decode); }
    #[test]
    fn urlsafe_encode() { do_test(t_urlsafe_encode); }
}
<|endoftext|>


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

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

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

pub enum ScheduledTask {
    ContinueBootstrap,
    ContinueHealthCheck,
    ContinueRefresh
}

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

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

enum Status {
    Bootstrapping,
    Idle
}

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

impl Network {
    pub fn new<T: ToSocketAddrs>(self_address: Address, host: T, routers: Vec<Box<node::Node>>) -> Network {
        let host = host.to_socket_addrs().unwrap().next().unwrap();
        let mapped_host = stun::get_mapped_address(host).unwrap();
        let self_node = node::UdpNode::new(self_address, mapped_host);
        let routing_table = RoutingTable::new(8, self_address, routers);

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

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

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

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

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

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

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

                match query {
                    Query::FindNode(target) => {
                        let response: Vec<u8> = outgoing::create_find_node_response(
                                transaction_id,
                                &self.self_node,
                                self.routing_table.nearest_to(&target, false));
                        origin.send(response);

                        self.insert_node(origin).unwrap();
                        {
                            if let Some(origin) = self.routing_table.find_node(&origin_address) {
                                origin.received_query(transaction_id);
                            }
                        }
                    },
                    Query::Packet(payload) => {
                        for listener in &self.event_listeners {
                            listener.send(Event::ReceivedPacket(origin_address, payload.clone())).unwrap();
                        }
                        let response = outgoing::create_packet_response(
                            transaction_id, &self.self_node);
                        origin.send(response);
                    },
                    Query::Ping => {
                        let response = outgoing::create_ping_response(
                            transaction_id,
                            &self.self_node);
                        origin.send(response);

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

            Message::Response(transaction_id, origin, response) => {
                let origin_address = origin.address();
                match response {
                    Response::FindNode(mut nodes) => {
                        let mut encounted_new_node = false;
                        let mut tail = vec![origin];
                        for node in nodes.drain(..).chain(tail.drain(..)) {
                            match self.insert_node(node) {
                                Ok(InsertOutcome::Inserted) => {
                                    encounted_new_node = true;
                                }
                                Err(error) => { panic!(error) }
                                _ => { }
                            }
                        }

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

                        match self.pending_actions.remove(&transaction_id) {
                            Some(TableAction::Bootstrap(timeout)) => {
                                event_loop.clear_timeout(timeout);
                                if encounted_new_node {
                                    self.continue_bootstrap(event_loop);
                                } else {
                                    self.continue_health_check(event_loop);
                                    self.continue_refresh(event_loop)
                                }
                            }
                            Some(TableAction::HealthCheck(_)) => {
                            },
                            Some(TableAction::RefreshBucket(_)) => {
                            }
                            None => { }
                        }

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

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

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

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

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

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

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

    fn health_check(&mut self) -> TransactionId {
        self.routing_table.remove_bad_nodes();
        let transaction_id = self.transaction_ids.generate();
        if let Some(node) = self.routing_table.questionable_nodes().get_mut(0) {
            let query = outgoing::create_ping_query(
                transaction_id, &self.self_node);
            node.send(query);
            node.sent_query(transaction_id);
        }
        transaction_id
    }

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

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

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

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

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

struct Handler {
    network: Network
}

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

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

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

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

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

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

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

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

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

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

    });
    shutdown_sender
}
<commit_msg>Insert any nodes that send a query into the routing table<commit_after>use address::Address;
use messages::outgoing;
use mio;
use node;
use routing_table::{InsertOutcome, InsertionResult, RoutingTable};
use std::collections::HashMap;
use std::io::Cursor;
use std::net::{UdpSocket, ToSocketAddrs, SocketAddr};
use std::sync::mpsc;
use std::thread;
use stun;
use transaction::{TransactionId, TransactionIdGenerator};

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

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

pub enum ScheduledTask {
    ContinueBootstrap,
    ContinueHealthCheck,
    ContinueRefresh
}

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

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

enum Status {
    Bootstrapping,
    Idle
}

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

impl Network {
    pub fn new<T: ToSocketAddrs>(self_address: Address, host: T, routers: Vec<Box<node::Node>>) -> Network {
        let host = host.to_socket_addrs().unwrap().next().unwrap();
        let mapped_host = stun::get_mapped_address(host).unwrap();
        let self_node = node::UdpNode::new(self_address, mapped_host);
        let routing_table = RoutingTable::new(8, self_address, routers);

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

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

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

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

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

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

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

                match query {
                    Query::FindNode(target) => {
                        let response: Vec<u8> = outgoing::create_find_node_response(
                                transaction_id,
                                &self.self_node,
                                self.routing_table.nearest_to(&target, false));
                        origin.send(response);
                    },
                    Query::Packet(payload) => {
                        for listener in &self.event_listeners {
                            listener.send(Event::ReceivedPacket(origin_address, payload.clone())).unwrap();
                        }
                        let response = outgoing::create_packet_response(
                            transaction_id, &self.self_node);
                        origin.send(response);
                    },
                    Query::Ping => {
                        let response = outgoing::create_ping_response(
                            transaction_id,
                            &self.self_node);
                        origin.send(response);
                    }
                }

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

            Message::Response(transaction_id, origin, response) => {
                let origin_address = origin.address();
                match response {
                    Response::FindNode(mut nodes) => {
                        let mut encounted_new_node = false;
                        let mut tail = vec![origin];
                        for node in nodes.drain(..).chain(tail.drain(..)) {
                            match self.insert_node(node) {
                                Ok(InsertOutcome::Inserted) => {
                                    encounted_new_node = true;
                                }
                                Err(error) => { panic!(error) }
                                _ => { }
                            }
                        }

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

                        match self.pending_actions.remove(&transaction_id) {
                            Some(TableAction::Bootstrap(timeout)) => {
                                event_loop.clear_timeout(timeout);
                                if encounted_new_node {
                                    self.continue_bootstrap(event_loop);
                                } else {
                                    self.continue_health_check(event_loop);
                                    self.continue_refresh(event_loop)
                                }
                            }
                            Some(TableAction::HealthCheck(_)) => {
                            },
                            Some(TableAction::RefreshBucket(_)) => {
                            }
                            None => { }
                        }

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

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

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

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

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

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

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

    fn health_check(&mut self) -> TransactionId {
        self.routing_table.remove_bad_nodes();
        let transaction_id = self.transaction_ids.generate();
        if let Some(node) = self.routing_table.questionable_nodes().get_mut(0) {
            let query = outgoing::create_ping_query(
                transaction_id, &self.self_node);
            node.send(query);
            node.sent_query(transaction_id);
        }
        transaction_id
    }

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

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

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

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

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

struct Handler {
    network: Network
}

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

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

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

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

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

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

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

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

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

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

    });
    shutdown_sender
}
<|endoftext|>


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

pub enum ScheduledTask {
    ContinueBootstrap,
    ContinueHealthCheck,
    ContinueRefresh
}

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

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

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

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

enum Status {
    Bootstrapping,
    Idle
}

pub struct Network {
    host: SocketAddr,
    routing_table: RoutingTable,
    self_node: Box<node::Node + 'static>,
    transaction_ids: TransactionIdGenerator,
    status: Status,
    pending_actions: HashMap<TransactionId, TableAction>,
    event_listeners: Vec<mpsc::Sender<Event>>
}

impl Network {
    pub fn new<N: node::Node + 'static, T: ToSocketAddrs>(self_node: N, host: T, routers: Vec<Box<node::Node>>) -> Network {
        let host = host.to_socket_addrs().unwrap().next().unwrap();
        let self_address = self_node.address().clone();
        let routing_table = RoutingTable::new(8, self_address, routers);

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

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

        create_incoming_udp_channel(self.host, event_loop.channel());
        event_loop.channel().send(OneshotTask::StartBootstrap).unwrap();
        let mut handler = Handler::new(self);
        let task_sender = event_loop.channel();
        thread::spawn(move|| event_loop.run(&mut handler).unwrap());
        task_sender
    }

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

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

        match message {
            Message::Query(transaction_id, origin, query) => {
                let origin_address = origin.address();
                match query {
                    Query::FindNode(target) => {
                        let response: Vec<u8> = outgoing::create_find_node_response(
                                transaction_id,
                                &self.self_node,
                                self.routing_table.nearest_to(&target, false));
                        origin.send(response);

                        self.insert_node(origin).unwrap();
                        {
                            if let Some(origin) = self.routing_table.find_node(&origin_address) {
                                origin.received_query(transaction_id);
                            }
                        }
                    },
                    Query::Packet(payload) => {
                        for listener in &self.event_listeners {
                            listener.send(Event::ReceivedPacket(origin_address, payload.clone())).unwrap();
                        }
                        let response = outgoing::create_packet_response(
                            transaction_id, &self.self_node);
                        origin.send(response);
                    },
                    Query::Ping => {
                        let response = outgoing::create_ping_response(
                            transaction_id,
                            &self.self_node);
                        origin.send(response);

                        self.insert_node(origin).unwrap();
                        if let Some(origin) = self.routing_table.find_node(&origin_address) {
                            origin.received_query(transaction_id);
                        }
                    }
                }
            }
            Message::Response(transaction_id, origin, response) => {
                let origin_address = origin.address();
                match response {
                    Response::FindNode(mut nodes) => {
                        let mut encounted_new_node = false;
                        let mut tail = vec![origin];
                        for node in nodes.drain(..).chain(tail.drain(..)) {
                            match self.insert_node(node) {
                                Ok(InsertOutcome::Inserted) => {
                                    encounted_new_node = true;
                                }
                                Err(error) => { panic!(error) }
                                _ => { }
                            }
                        }

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

                        match self.pending_actions.remove(&transaction_id) {
                            Some(TableAction::Bootstrap(timeout)) => {
                                event_loop.clear_timeout(timeout);
                                if encounted_new_node {
                                    self.continue_bootstrap(event_loop);
                                } else {
                                    self.continue_health_check(event_loop);
                                    self.continue_refresh(event_loop)
                                }
                            }
                            Some(TableAction::HealthCheck(_)) => {
                            },
                            Some(TableAction::RefreshBucket(_)) => {
                            }
                            None => { }
                        }

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

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

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

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

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

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

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

    fn health_check(&mut self) -> TransactionId {
        self.routing_table.remove_bad_nodes();
        let transaction_id = self.transaction_ids.generate();
        if let Some(node) = self.routing_table.questionable_nodes().get_mut(0) {
            let query = outgoing::create_ping_query(
                transaction_id, &self.self_node);
            node.send(query);
            node.sent_query(transaction_id);
        }
        transaction_id
    }

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

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

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

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

struct Handler {
    network: Network
}

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

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

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

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

fn create_incoming_udp_channel(host: SocketAddr, sender: TaskSender) {
    thread::spawn(move || {
        let host = ("0.0.0.0", host.port());
        let socket = UdpSocket::bind(host).unwrap();
        loop {
            let mut buf = [0; 4096];
            match socket.recv_from(&mut buf) {
                Ok((size, _src)) => {
                    sender
                        .send(OneshotTask::Incoming(buf[..size].iter().cloned().collect()))
                        .unwrap_or_else(|err| info!("Couldn't handling incoming: {:?}", err));
                }
                Err(e) => panic!("Error receiving from server: {}", e)
            }
        }
    });
}
<commit_msg>Silence unused arg warning<commit_after>use address::Address;
use messages::outgoing;
use mio;
use node;
use routing_table::{InsertOutcome, InsertionResult, RoutingTable};
use std::collections::HashMap;
use std::io::Cursor;
use std::net::{UdpSocket, ToSocketAddrs, SocketAddr};
use std::sync::mpsc;
use std::thread;
use transaction::{TransactionId, TransactionIdGenerator};

pub enum ScheduledTask {
    ContinueBootstrap,
    ContinueHealthCheck,
    ContinueRefresh
}

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

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

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

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

enum Status {
    Bootstrapping,
    Idle
}

pub struct Network {
    host: SocketAddr,
    routing_table: RoutingTable,
    self_node: Box<node::Node + 'static>,
    transaction_ids: TransactionIdGenerator,
    status: Status,
    pending_actions: HashMap<TransactionId, TableAction>,
    event_listeners: Vec<mpsc::Sender<Event>>
}

impl Network {
    pub fn new<N: node::Node + 'static, T: ToSocketAddrs>(self_node: N, host: T, routers: Vec<Box<node::Node>>) -> Network {
        let host = host.to_socket_addrs().unwrap().next().unwrap();
        let self_address = self_node.address().clone();
        let routing_table = RoutingTable::new(8, self_address, routers);

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

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

        create_incoming_udp_channel(self.host, event_loop.channel());
        event_loop.channel().send(OneshotTask::StartBootstrap).unwrap();
        let mut handler = Handler::new(self);
        let task_sender = event_loop.channel();
        thread::spawn(move|| event_loop.run(&mut handler).unwrap());
        task_sender
    }

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

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

        match message {
            Message::Query(transaction_id, origin, query) => {
                let origin_address = origin.address();
                match query {
                    Query::FindNode(target) => {
                        let response: Vec<u8> = outgoing::create_find_node_response(
                                transaction_id,
                                &self.self_node,
                                self.routing_table.nearest_to(&target, false));
                        origin.send(response);

                        self.insert_node(origin).unwrap();
                        {
                            if let Some(origin) = self.routing_table.find_node(&origin_address) {
                                origin.received_query(transaction_id);
                            }
                        }
                    },
                    Query::Packet(payload) => {
                        for listener in &self.event_listeners {
                            listener.send(Event::ReceivedPacket(origin_address, payload.clone())).unwrap();
                        }
                        let response = outgoing::create_packet_response(
                            transaction_id, &self.self_node);
                        origin.send(response);
                    },
                    Query::Ping => {
                        let response = outgoing::create_ping_response(
                            transaction_id,
                            &self.self_node);
                        origin.send(response);

                        self.insert_node(origin).unwrap();
                        if let Some(origin) = self.routing_table.find_node(&origin_address) {
                            origin.received_query(transaction_id);
                        }
                    }
                }
            }
            Message::Response(transaction_id, origin, response) => {
                let origin_address = origin.address();
                match response {
                    Response::FindNode(mut nodes) => {
                        let mut encounted_new_node = false;
                        let mut tail = vec![origin];
                        for node in nodes.drain(..).chain(tail.drain(..)) {
                            match self.insert_node(node) {
                                Ok(InsertOutcome::Inserted) => {
                                    encounted_new_node = true;
                                }
                                Err(error) => { panic!(error) }
                                _ => { }
                            }
                        }

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

                        match self.pending_actions.remove(&transaction_id) {
                            Some(TableAction::Bootstrap(timeout)) => {
                                event_loop.clear_timeout(timeout);
                                if encounted_new_node {
                                    self.continue_bootstrap(event_loop);
                                } else {
                                    self.continue_health_check(event_loop);
                                    self.continue_refresh(event_loop)
                                }
                            }
                            Some(TableAction::HealthCheck(_)) => {
                            },
                            Some(TableAction::RefreshBucket(_)) => {
                            }
                            None => { }
                        }

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

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

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

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

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

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

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

    fn health_check(&mut self) -> TransactionId {
        self.routing_table.remove_bad_nodes();
        let transaction_id = self.transaction_ids.generate();
        if let Some(node) = self.routing_table.questionable_nodes().get_mut(0) {
            let query = outgoing::create_ping_query(
                transaction_id, &self.self_node);
            node.send(query);
            node.sent_query(transaction_id);
        }
        transaction_id
    }

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

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

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

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

struct Handler {
    network: Network
}

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

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

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

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

fn create_incoming_udp_channel(host: SocketAddr, sender: TaskSender) {
    thread::spawn(move || {
        let host = ("0.0.0.0", host.port());
        let socket = UdpSocket::bind(host).unwrap();
        loop {
            let mut buf = [0; 4096];
            match socket.recv_from(&mut buf) {
                Ok((size, _src)) => {
                    sender
                        .send(OneshotTask::Incoming(buf[..size].iter().cloned().collect()))
                        .unwrap_or_else(|err| info!("Couldn't handling incoming: {:?}", err));
                }
                Err(e) => panic!("Error receiving from server: {}", e)
            }
        }
    });
}
<|endoftext|>