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

use std::io::Error;
use std::ops::DerefMut;
use std::sync::{Arc, Mutex};
use std::{mem, thread};
use std::net::{TcpStream, TcpListener};
use std::os::unix::io::{RawFd, AsRawFd, IntoRawFd};
use std::collections::LinkedList;

use libc;
use libc::{c_int, c_void};
use epoll;
use epoll::util::*;
use epoll::EpollEvent;
use errno::errno;
use threadpool::ThreadPool;
use openssl::ssl::{SslStream, SslContext};
use ss::nonblocking::plain::Plain;
use ss::nonblocking::secure::Secure;
use ss::{Socket, Stream, SRecv, SSend, TcpOptions, SocketOptions};

use stats;
use types::*;
use config::Config;

// We need to be able to access our resource pool from several methods
static mut thread_pool: *mut ThreadPool = 0 as *mut ThreadPool;

// Global SslContext
static mut ssl_context: *mut SslContext = 0 as *mut SslContext;

// When added to epoll, these will be the conditions of kernel notification:
//
// EPOLLET  - Fd is in EdgeTriggered mode (notification on state changes)
// EPOLLIN  - Data is available in kerndl buffer
const EVENTS: u32 = event_type::EPOLLET | event_type::EPOLLIN;

/// Starts the epoll wait and incoming connection listener threads.
pub fn begin(config: Config, handler: Box<EventHandler>) {
    // Master socket list
    let sockets = Arc::new(Mutex::new(LinkedList::<Stream>::new()));

    // Resource pool
    let mut tp = ThreadPool::new(config.workers);
    unsafe {
        thread_pool = &mut tp;
    }

    // Wrap our event handler into something that can be safely shared
    // between threads.
    let e_handler = Handler(Box::into_raw(handler));

    // Epoll instance
    let result = epoll::create1(0);
    if result.is_err() {
        let err = result.unwrap_err();
        error!("Unable to create epoll instance: {}", err);
        panic!()
    }
    let epfd = result.unwrap();

    // Epoll wait thread
    let epfd2 = epfd.clone();
    let streams2 = sockets.clone();
    thread::Builder::new()
        .name("Epoll Wait".to_string())
        .spawn(move || {
            event_loop(epfd2, streams2, e_handler);
        })
        .unwrap();

    // New connection thread
    let epfd3 = epfd.clone();
    let streams3 = sockets.clone();
    let prox = thread::Builder::new()
        .name("TCP Incoming Listener".to_string())
        .spawn(move || {
           listen(config, epfd3, streams3);
        })
        .unwrap();

    // Stay alive forever, or at least we hope
    let _ = prox.join();
}

/// Incoming connection listening thread
fn listen(config: Config, epfd: RawFd, streams: StreamList) {
    // Setup server and listening port
    let listener_result = try_setup_tcp_listener(&config);
    if listener_result.is_err() {
        error!("Setting up server: {}", listener_result.unwrap_err());
        return;
    }

    // If we're using SSL, setup our context reference
    if config.ssl.is_some() {
        setup_ssl_context(&config);
    }

    // Begin listening for new connections
    let listener = listener_result.unwrap();
    for accept_result in listener.incoming() {
        match accept_result {
            Ok(tcp_stream) => handle_new_connection(tcp_stream, &config, epfd, streams.clone()),
            Err(e) => error!("Accepting connection: {}", e)
        }
    }

    drop(listener);
}

fn try_setup_tcp_listener(config: &Config) -> Result<TcpListener, Error> {
    let create_result = TcpListener::bind((&config.addr[..], config.port));
    if create_result.is_err() {
        return create_result;
    }

    let listener = create_result.unwrap();
    let server_fd = listener.as_raw_fd();

    // Set the SO_REUSEADDR so that restarts after crashes do not take 5min to unbind
    // the initial port
    unsafe {
        let optval: c_int = 1;
        let opt_result = libc::setsockopt(server_fd,
                                          libc::SOL_SOCKET,
                                          libc::SO_REUSEADDR,
                                          &optval as *const _ as *const c_void,
                                          mem::size_of::<c_int>() as u32);
        if opt_result < 0 {
            return Err(Error::from_raw_os_error(errno().0 as i32));
        }
    }

    Ok(listener)
}

fn setup_ssl_context(config: &Config) {
    unsafe {
        ssl_context = Box::into_raw(Box::new(config.ssl.clone().unwrap()));
    }
}

fn handle_new_connection(tcp_stream: TcpStream, config: &Config, epfd: RawFd, streams: StreamList) {
    // Update our total opened file descriptors
    stats::fd_opened();

    // Create and configure a new socket
    let mut socket = Socket::new(tcp_stream.into_raw_fd());
    let result = setup_new_socket(&mut socket);
    if result.is_err() {
        close_fd(socket.as_raw_fd());
        return;
    }

    // Setup our stream
    let stream = match config.ssl {
        Some(_) => {
            let sock_fd = socket.as_raw_fd();
            let ssl_result = unsafe { SslStream::accept(&(*ssl_context), socket) };
            match ssl_result {
                Ok(ssl_stream) => {
                    let secure_stream = Secure::new(ssl_stream);
                    Stream::new(Box::new(secure_stream))
                }
                Err(ssl_error) => {
                    error!("Creating SslStream: {}", ssl_error);
                    close_fd(sock_fd);
                    return;
                }
            }
        }
        None => {
            let plain_text = Plain::new(socket);
            Stream::new(Box::new(plain_text))
        }
    };

    // Add stream to our server
    let fd = stream.as_raw_fd();
    add_stream_to_master_list(stream, streams.clone());
    add_fd_to_epoll(epfd, fd, streams.clone());
}

fn setup_new_socket(socket: &mut Socket) -> Result<(), ()> {
    let result = socket.set_nonblocking();
    if result.is_err() {
        error!("Setting fd to nonblocking: {}", result.unwrap_err());
        return Err(());
    }

    let result = socket.set_tcp_nodelay(true);
    if result.is_err() {
        error!("Setting tcp_nodelay: {}", result.unwrap_err());
        return Err(());
    }

    let result = socket.set_tcp_keepalive(true);
    if result.is_err() {
        error!("Setting tcp_keepalive: {}", result.unwrap_err());
        return Err(());
    }

    Ok(())
}

/// Event loop for handling all epoll events
fn event_loop(epfd: RawFd, streams: StreamList, handler: Handler) {
    let mut events = Vec::<EpollEvent>::with_capacity(100);
    unsafe {
        events.set_len(100);
    }

    loop {
        match epoll::wait(epfd, &mut events[..], -1) {
            Ok(num_events) => {
                for x in 0..num_events as usize {
                    handle_epoll_event(epfd, &events[x], streams.clone(), handler.clone());
                }
            }
            Err(e) => {
                error!("Error on epoll::wait(): {}", e);
                panic!()
            }
        };
    }
}

/// Finds the stream the epoll event is associated with and parses the event type
/// to hand off to specific handlers
fn handle_epoll_event(epfd: RawFd, event: &EpollEvent, streams: StreamList, handler: Handler) {
    const READ_EVENT: u32 = event_type::EPOLLIN;

    let fd = event.data as RawFd;
    let find_result = try_find_stream_from_fd(streams.clone(), fd);
    if find_result.is_err() {
        remove_fd_from_epoll(epfd, fd);
        close_fd(fd);
    }

    let stream = find_result.unwrap();
    let read_event = if (event.events & READ_EVENT) > 0 {
        true
    } else {
        false
    };

    if read_event {
        handle_read_event(epfd, stream, streams, handler);
        // TODO - Ensure added back to stream list
    } else {
        remove_fd_from_epoll(epfd, fd);
        close_fd(fd);

        unsafe {
            (*thread_pool).execute(move || {
                let Handler(ptr) = handler;
                (*ptr).on_stream_closed(fd);
            });
        }
    }
}

fn try_find_stream_from_fd(streams: StreamList, fd: RawFd) -> Result<Stream, ()> {
    let mut guard = match streams.lock() {
        Ok(g) => g,
        Err(p) => p.into_inner()
    };
    let list = guard.deref_mut();

    let mut found = false;
    let mut index = 1usize;
    for stream in list.iter() {
        if stream.as_raw_fd() == fd {
            found = true;
            break;
        }
        index += 1;
    }

    let result = if found {
        if index == 1 {
            Ok(list.pop_front().unwrap())
        } else {
            let mut split = list.split_off(index - 1);
            let stream = split.pop_front().unwrap();
            list.append(&mut split);
            Ok(stream)
        }
    } else {
        Err(())
    };

    result
}

/// Reads all available data on the stream.
///
/// If a complete message(s) is available, each message will be routed through the
/// resource pool.
///
/// If an error occurs during the read, the stream is dropped from the server.
fn handle_read_event(epfd: RawFd, stream: Stream, streams: StreamList, handler: Handler) {
    let fd = stream.as_raw_fd();

    let mut stream = stream;
    let recv_result = stream.recv();
    if recv_result.is_err() {
        remove_fd_from_epoll(epfd, fd);
        close_fd(fd);

        unsafe {
            (*thread_pool).execute(move || {
                let Handler(ptr) = handler;
                (*ptr).on_stream_closed(fd);
            });
        }
        return;
    }

    let mut msg_queue = stream.drain_rx_queue();
    for msg in msg_queue.drain(..) {
        let s = stream.clone();
        let sl = streams.clone();
        let h = handler.clone();

        if msg_is_stats_request(&msg) {
            handle_stats_request(&msg, epfd, s, sl);
        } else {
            unsafe {
                (*thread_pool).execute(move || {
                    let Handler(ptr) = h;
                    (*ptr).on_data_received(s, msg);
                });
            }
        }
    }

    add_stream_to_master_list(stream, streams);
}

#[inline]
fn msg_is_stats_request(msg: &[u8]) -> bool {
    if msg.len() == 6 && msg[0] == 0x04 && msg[1] == 0x04 {
        true
    } else {
        false
    }
}

fn handle_stats_request(buf: &[u8], epfd: RawFd, stream: Stream, streams: StreamList) {
    let stream_clone = stream.clone();
    let u8ptr: *const u8 = &buf[2] as *const _;
    let f32ptr: *const f32 = u8ptr as *const _;
    unsafe {
        let sec = *f32ptr;
        (*thread_pool).execute(move || {
            let stats_result = stats::as_serialized_buffer(sec);
            if stats_result.is_err() {
                error!("Retrieving stats");
                return;
            }

            let stats_buffer = stats_result.unwrap();

            let mut stream = stream;
            let send_result = stream.send(&stats_buffer[..]);
            if send_result.is_err() {
                error!("Writing to stream: {}", send_result.unwrap_err());
                let fd = stream.as_raw_fd();
                remove_fd_from_epoll(epfd, fd);
                close_fd(fd);
            }
        });
    }

    add_stream_to_master_list(stream_clone, streams);
}

/// Inserts the stream back into the master list of streams
fn add_stream_to_master_list(stream: Stream, streams: StreamList) {
    let mut guard = match streams.lock() {
        Ok(g) => g,
        Err(p) => p.into_inner()
    };

    let stream_list = guard.deref_mut();
    stream_list.push_back(stream);

    stats::conn_recv();
}

/// Adds a new fd to the epoll instance
fn add_fd_to_epoll(epfd: RawFd, fd: RawFd, streams: StreamList) {
    let result = epoll::ctl(epfd,
                            ctl_op::ADD,
                            fd,
                            &mut EpollEvent {
                                data: fd as u64,
                                events: EVENTS,
                            });

    if result.is_err() {
        let e = result.unwrap_err();
        error!("epoll::CtrlError during add: {}", e);
        remove_fd_from_list(fd, streams.clone());
        close_fd(fd);
    }
}

/// Removes a fd from the epoll instance
fn remove_fd_from_epoll(epfd: RawFd, fd: RawFd) {
    // In kernel versions before 2.6.9, the EPOLL_CTL_DEL operation required
    // a non-null pointer in event, even though this argument is ignored.
    // Since Linux 2.6.9, event can be specified as NULL when using
    // EPOLL_CTL_DEL. We'll be as backwards compatible as possible.
    let _ = epoll::ctl(epfd,
                       ctl_op::DEL,
                       fd,
                       &mut EpollEvent {
                           data: 0 as u64,
                           events: 0 as u32,
                       })
                .map_err(|e| warn!("Epoll CtrlError during del: {}", e));
}

/// Removes stream with fd from master list
fn remove_fd_from_list(fd: RawFd, streams: StreamList) {
    let mut guard = match streams.lock() {
        Ok(g) => g,
        Err(p) => p.into_inner()
    };
    let list = guard.deref_mut();

    let mut found = false;
    let mut index = 1usize;
    for stream in list.iter() {
        if stream.as_raw_fd() == fd {
            found = true;
            break;
        }
        index += 1;
    }

    if !found {
        return;
    }

    if index == 1 {
        list.pop_front();
    } else {
        let mut split = list.split_off(index - 1);
        split.pop_front();
        list.append(&mut split);
    }

    stats::conn_lost();
}

/// Closes a fd with the kernel
fn close_fd(fd: RawFd) {
    unsafe {
        let result = libc::close(fd);
        if result < 0 {
            error!("Error closing fd: {}",
                   Error::from_raw_os_error(result as i32));
            return;
        }
    }
    stats::fd_closed();
}
<commit_msg>Extra debugs for Chris<commit_after>// Copyright 2015 Nathan Sizemore <nathanrsizemore@gmail.com>
//
// This Source Code Form is subject to the terms of the
// Mozilla Public License, v. 2.0. If a copy of the MPL was not
// distributed with this file, You can obtain one at
// http://mozilla.org/MPL/2.0/.

use std::io::Error;
use std::ops::DerefMut;
use std::sync::{Arc, Mutex};
use std::{mem, thread};
use std::net::{TcpStream, TcpListener};
use std::os::unix::io::{RawFd, AsRawFd, IntoRawFd};
use std::collections::LinkedList;

use libc;
use libc::{c_int, c_void};
use epoll;
use epoll::util::*;
use epoll::EpollEvent;
use errno::errno;
use threadpool::ThreadPool;
use openssl::ssl::{SslStream, SslContext};
use ss::nonblocking::plain::Plain;
use ss::nonblocking::secure::Secure;
use ss::{Socket, Stream, SRecv, SSend, TcpOptions, SocketOptions};

use stats;
use types::*;
use config::Config;

// We need to be able to access our resource pool from several methods
static mut thread_pool: *mut ThreadPool = 0 as *mut ThreadPool;

// Global SslContext
static mut ssl_context: *mut SslContext = 0 as *mut SslContext;

// When added to epoll, these will be the conditions of kernel notification:
//
// EPOLLET  - Fd is in EdgeTriggered mode (notification on state changes)
// EPOLLIN  - Data is available in kerndl buffer
const EVENTS: u32 = event_type::EPOLLET | event_type::EPOLLIN;

/// Starts the epoll wait and incoming connection listener threads.
pub fn begin(config: Config, handler: Box<EventHandler>) {
    // Master socket list
    let sockets = Arc::new(Mutex::new(LinkedList::<Stream>::new()));

    // Resource pool
    let mut tp = ThreadPool::new(config.workers);
    unsafe {
        thread_pool = &mut tp;
    }

    // Wrap our event handler into something that can be safely shared
    // between threads.
    let e_handler = Handler(Box::into_raw(handler));

    // Epoll instance
    let result = epoll::create1(0);
    if result.is_err() {
        let err = result.unwrap_err();
        error!("Unable to create epoll instance: {}", err);
        panic!()
    }
    let epfd = result.unwrap();

    // Epoll wait thread
    let epfd2 = epfd.clone();
    let streams2 = sockets.clone();
    thread::Builder::new()
        .name("Epoll Wait".to_string())
        .spawn(move || {
            event_loop(epfd2, streams2, e_handler);
        })
        .unwrap();

    // New connection thread
    let epfd3 = epfd.clone();
    let streams3 = sockets.clone();
    let prox = thread::Builder::new()
        .name("TCP Incoming Listener".to_string())
        .spawn(move || {
           listen(config, epfd3, streams3);
        })
        .unwrap();

    // Stay alive forever, or at least we hope
    let _ = prox.join();
}

/// Incoming connection listening thread
fn listen(config: Config, epfd: RawFd, streams: StreamList) {
    // Setup server and listening port
    let listener_result = try_setup_tcp_listener(&config);
    if listener_result.is_err() {
        error!("Setting up server: {}", listener_result.unwrap_err());
        return;
    }

    // If we're using SSL, setup our context reference
    if config.ssl.is_some() {
        setup_ssl_context(&config);
    }

    // Begin listening for new connections
    let listener = listener_result.unwrap();
    for accept_result in listener.incoming() {
        match accept_result {
            Ok(tcp_stream) => handle_new_connection(tcp_stream, &config, epfd, streams.clone()),
            Err(e) => error!("Accepting connection: {}", e)
        }
    }

    drop(listener);
}

fn try_setup_tcp_listener(config: &Config) -> Result<TcpListener, Error> {
    let create_result = TcpListener::bind((&config.addr[..], config.port));
    if create_result.is_err() {
        return create_result;
    }

    let listener = create_result.unwrap();
    let server_fd = listener.as_raw_fd();

    // Set the SO_REUSEADDR so that restarts after crashes do not take 5min to unbind
    // the initial port
    unsafe {
        let optval: c_int = 1;
        let opt_result = libc::setsockopt(server_fd,
                                          libc::SOL_SOCKET,
                                          libc::SO_REUSEADDR,
                                          &optval as *const _ as *const c_void,
                                          mem::size_of::<c_int>() as u32);
        if opt_result < 0 {
            return Err(Error::from_raw_os_error(errno().0 as i32));
        }
    }

    Ok(listener)
}

fn setup_ssl_context(config: &Config) {
    unsafe {
        ssl_context = Box::into_raw(Box::new(config.ssl.clone().unwrap()));
    }
}

fn handle_new_connection(tcp_stream: TcpStream, config: &Config, epfd: RawFd, streams: StreamList) {
    // Update our total opened file descriptors
    stats::fd_opened();

    // Create and configure a new socket
    let mut socket = Socket::new(tcp_stream.into_raw_fd());
    let result = setup_new_socket(&mut socket);
    if result.is_err() {
        close_fd(socket.as_raw_fd());
        return;
    }

    // Setup our stream
    let stream = match config.ssl {
        Some(_) => {
            let sock_fd = socket.as_raw_fd();
            let ssl_result = unsafe { SslStream::accept(&(*ssl_context), socket) };
            match ssl_result {
                Ok(ssl_stream) => {
                    let secure_stream = Secure::new(ssl_stream);
                    Stream::new(Box::new(secure_stream))
                }
                Err(ssl_error) => {
                    error!("Creating SslStream: {}", ssl_error);
                    close_fd(sock_fd);
                    return;
                }
            }
        }
        None => {
            let plain_text = Plain::new(socket);
            Stream::new(Box::new(plain_text))
        }
    };

    // Add stream to our server
    let fd = stream.as_raw_fd();
    add_stream_to_master_list(stream, streams.clone());
    add_fd_to_epoll(epfd, fd, streams.clone());
}

fn setup_new_socket(socket: &mut Socket) -> Result<(), ()> {
    let result = socket.set_nonblocking();
    if result.is_err() {
        error!("Setting fd to nonblocking: {}", result.unwrap_err());
        return Err(());
    }

    let result = socket.set_tcp_nodelay(true);
    if result.is_err() {
        error!("Setting tcp_nodelay: {}", result.unwrap_err());
        return Err(());
    }

    let result = socket.set_tcp_keepalive(true);
    if result.is_err() {
        error!("Setting tcp_keepalive: {}", result.unwrap_err());
        return Err(());
    }

    Ok(())
}

/// Event loop for handling all epoll events
fn event_loop(epfd: RawFd, streams: StreamList, handler: Handler) {
    let mut events = Vec::<EpollEvent>::with_capacity(100);
    unsafe {
        events.set_len(100);
    }

    loop {
        match epoll::wait(epfd, &mut events[..], -1) {
            Ok(num_events) => {
                for x in 0..num_events as usize {
                    handle_epoll_event(epfd, &events[x], streams.clone(), handler.clone());
                }
            }
            Err(e) => {
                error!("Error on epoll::wait(): {}", e);
                panic!()
            }
        };
    }
}

/// Finds the stream the epoll event is associated with and parses the event type
/// to hand off to specific handlers
fn handle_epoll_event(epfd: RawFd, event: &EpollEvent, streams: StreamList, handler: Handler) {
    const READ_EVENT: u32 = event_type::EPOLLIN;

    let fd = event.data as RawFd;
    let find_result = try_find_stream_from_fd(streams.clone(), fd);
    if find_result.is_err() {
        remove_fd_from_epoll(epfd, fd);
        close_fd(fd);
    }

    let stream = find_result.unwrap();
    let read_event = if (event.events & READ_EVENT) > 0 {
        true
    } else {
        false
    };

    if read_event {
        handle_read_event(epfd, stream, streams, handler);
        // TODO - Ensure added back to stream list
    } else {
        remove_fd_from_epoll(epfd, fd);
        close_fd(fd);

        unsafe {
            (*thread_pool).execute(move || {
                let Handler(ptr) = handler;
                (*ptr).on_stream_closed(fd);
            });
        }
    }
}

fn try_find_stream_from_fd(streams: StreamList, fd: RawFd) -> Result<Stream, ()> {
    let mut guard = match streams.lock() {
        Ok(g) => g,
        Err(p) => p.into_inner()
    };
    let list = guard.deref_mut();

    let mut found = false;
    let mut index = 1usize;
    for stream in list.iter() {
        if stream.as_raw_fd() == fd {
            found = true;
            break;
        }
        index += 1;
    }

    let result = if found {
        if index == 1 {
            Ok(list.pop_front().unwrap())
        } else {
            let mut split = list.split_off(index - 1);
            let stream = split.pop_front().unwrap();
            list.append(&mut split);
            Ok(stream)
        }
    } else {
        Err(())
    };

    result
}

/// Reads all available data on the stream.
///
/// If a complete message(s) is available, each message will be routed through the
/// resource pool.
///
/// If an error occurs during the read, the stream is dropped from the server.
fn handle_read_event(epfd: RawFd, stream: Stream, streams: StreamList, handler: Handler) {
    trace!("handle_read_event");
    let fd = stream.as_raw_fd();

    let mut stream = stream;
    let recv_result = stream.recv();
    if recv_result.is_err() {
        remove_fd_from_epoll(epfd, fd);
        close_fd(fd);

        unsafe {
            (*thread_pool).execute(move || {
                let Handler(ptr) = handler;
                (*ptr).on_stream_closed(fd);
            });
        }
        return;
    }

    let mut msg_queue = stream.drain_rx_queue();
    for msg in msg_queue.drain(..) {
        let s = stream.clone();
        let sl = streams.clone();
        let h = handler.clone();

        if msg_is_stats_request(&msg) {
            handle_stats_request(&msg, epfd, s, sl);
        } else {
            unsafe {
                (*thread_pool).execute(move || {
                    let Handler(ptr) = h;
                    (*ptr).on_data_received(s, msg);
                });
            }
        }
    }

    add_stream_to_master_list(stream, streams);
}

#[inline]
fn msg_is_stats_request(msg: &[u8]) -> bool {
    if msg.len() == 6 && msg[0] == 0x04 && msg[1] == 0x04 {
        true
    } else {
        false
    }
}

fn handle_stats_request(buf: &[u8], epfd: RawFd, stream: Stream, streams: StreamList) {
    trace!("handle_stats_request");
    let stream_clone = stream.clone();
    let u8ptr: *const u8 = &buf[2] as *const _;
    let f32ptr: *const f32 = u8ptr as *const _;
    unsafe {
        let sec = *f32ptr;
        (*thread_pool).execute(move || {
            let stats_result = stats::as_serialized_buffer(sec);
            if stats_result.is_err() {
                error!("Retrieving stats");
                return;
            }

            let stats_buffer = stats_result.unwrap();

            let mut stream = stream;
            let send_result = stream.send(&stats_buffer[..]);
            if send_result.is_err() {
                error!("Writing to stream: {}", send_result.unwrap_err());
                let fd = stream.as_raw_fd();
                remove_fd_from_epoll(epfd, fd);
                close_fd(fd);
            }
        });
    }

    add_stream_to_master_list(stream_clone, streams);
}

/// Inserts the stream back into the master list of streams
fn add_stream_to_master_list(stream: Stream, streams: StreamList) {
    let mut guard = match streams.lock() {
        Ok(g) => g,
        Err(p) => p.into_inner()
    };

    let stream_list = guard.deref_mut();
    stream_list.push_back(stream);

    stats::conn_recv();
}

/// Adds a new fd to the epoll instance
fn add_fd_to_epoll(epfd: RawFd, fd: RawFd, streams: StreamList) {
    let result = epoll::ctl(epfd,
                            ctl_op::ADD,
                            fd,
                            &mut EpollEvent {
                                data: fd as u64,
                                events: EVENTS,
                            });

    if result.is_err() {
        let e = result.unwrap_err();
        error!("epoll::CtrlError during add: {}", e);
        remove_fd_from_list(fd, streams.clone());
        close_fd(fd);
    }
}

/// Removes a fd from the epoll instance
fn remove_fd_from_epoll(epfd: RawFd, fd: RawFd) {
    // In kernel versions before 2.6.9, the EPOLL_CTL_DEL operation required
    // a non-null pointer in event, even though this argument is ignored.
    // Since Linux 2.6.9, event can be specified as NULL when using
    // EPOLL_CTL_DEL. We'll be as backwards compatible as possible.
    let _ = epoll::ctl(epfd,
                       ctl_op::DEL,
                       fd,
                       &mut EpollEvent {
                           data: 0 as u64,
                           events: 0 as u32,
                       })
                .map_err(|e| warn!("Epoll CtrlError during del: {}", e));
}

/// Removes stream with fd from master list
fn remove_fd_from_list(fd: RawFd, streams: StreamList) {
    let mut guard = match streams.lock() {
        Ok(g) => g,
        Err(p) => p.into_inner()
    };
    let list = guard.deref_mut();

    let mut found = false;
    let mut index = 1usize;
    for stream in list.iter() {
        if stream.as_raw_fd() == fd {
            found = true;
            break;
        }
        index += 1;
    }

    if !found {
        return;
    }

    if index == 1 {
        list.pop_front();
    } else {
        let mut split = list.split_off(index - 1);
        split.pop_front();
        list.append(&mut split);
    }

    stats::conn_lost();
}

/// Closes a fd with the kernel
fn close_fd(fd: RawFd) {
    unsafe {
        let result = libc::close(fd);
        if result < 0 {
            error!("Error closing fd: {}",
                   Error::from_raw_os_error(result as i32));
            return;
        }
    }
    stats::fd_closed();
}
<|endoftext|>


<commit_before>use futures::{future, Async, Future, Poll, Stream};
use hyper::server::conn::Http;
// use hyper::{self, Body, Method, Request, Response, Server, StatusCode};
use failure::Error;
use futures::future::Either;
use rustls;
use rustls::internal::pemfile;
use tokio::net as tcp;
use tokio_rustls::{self, ServerConfigExt};

use tokio::runtime::current_thread;
use tokio::runtime::Builder as RuntimeBuilder;
use tokio::runtime::TaskExecutor;
use tokio_threadpool::Builder as ThreadPoolBuilder;

use std::net::SocketAddr;
use std::{fs, io, mem, str, sync, thread, time};
use StdError;

use base::BaseService;
use config::{Config, Route};
use connect;
use reuse::reuse_address;
use stat::print as stat_print;

pub fn run(config: Config) -> Result<(), Error> {
    let mut threadpool_builder = ThreadPoolBuilder::new();
    threadpool_builder.name_prefix("worker-");

    let runtime = RuntimeBuilder::new().threadpool_builder(threadpool_builder).build()?;

    let executor = runtime.executor();

    let mut tcp_listener: Option<Box<dyn Future<Item = (), Error = ()> + 'static>> = None;
    let mut error = None;

    for addr in config.addrs.clone() {
        match tcp_listener_module(&addr, config.clone(), executor.clone()) {
            Ok(tcp) => {
                tcp_listener = Some(tcp);
                stat_print(
                    &addr,
                    if config.cert.is_some() { "https" } else { "http" },
                    config.routes.values().cloned().collect(),
                );
                break;
            }
            Err(e) => {
                warn!("listen {}: {}", addr, e);
                error = Some(e);
            }
        }
    }

    if tcp_listener.is_none() {
        if let Some(e) = error {
            return Err(e);
        } else {
            unreachable!("tcp-listener and error is None");
        }
    }

    let mut current_thread_runtime = current_thread::Runtime::new()?;
    current_thread_runtime.spawn(tcp_listener.unwrap());
    current_thread_runtime.run()?;
    Ok(())
}

macro_rules! sockets_stream_handle {
    ($socket: expr, $http: expr, $executor: expr) => {
        $socket
            .and_then(|(addr, socket)| Box::new(MethodDetection::new(true, addr, socket).map(|addrs| Some(addrs))))
            .or_else(|e| {
                tcp_listener_sleep_ms(e, 1);
                future::ok(None)
            }).filter_map(|s: Option<Either<(SocketAddr, _), (SocketAddr, _)>>| s)
            // already or_else..
            // .map_err(|e: io::Error| unreachable!(e))
            .for_each(move |addrs| {
                match addrs {
                    Either::A((addr, socket)) => {
                        let connection = $http.serve_connection(socket, BaseService::new(addr));

                        $executor.spawn(connection.then(move |connection_res| {
                            if let Err(e) = connection_res {
                                error!("client[{}]: {}", addr, e.description());
                            }
                            Ok(())
                        }));
                    }
                    Either::B((addr, socket)) => {
                        info!("{} Use CONNECT Method ~", addr);
                        $executor.spawn(connect::process_socket(addr, socket));
                    }
                }
                Ok(())
            })
    };
}

pub fn tcp_listener_module(
    addr: &SocketAddr,
    config: Config,
    executor: TaskExecutor,
) -> Result<Box<dyn Future<Item = (), Error = ()> + 'static>, Error> {
    let tcp = reuse_address(&addr)?;
    let http = Http::new();

    if let Some(cert) = &config.cert {
        let tls_cfg = {
            let certs = load_certs(&cert.pub_)?;
            let key = load_private_key(&cert.key)?;
            let mut cfg = rustls::ServerConfig::new(rustls::NoClientAuth::new());
            cfg.set_single_cert(certs, key)
                .map_err(|e| format_err!("set single cert failed: {:?}", e))?;
            sync::Arc::new(cfg)
        };

        let tcp_listener_module = tcp
            .incoming()
            .and_then(|socket| socket.peer_addr().map(|addr| (addr, socket)))
            .and_then(move |(addr, socket)| tls_cfg.accept_async(socket).map(move |socket| (addr, socket)));
        let tcp_listener_module = sockets_stream_handle!(tcp_listener_module, http, executor);

        Ok(Box::new(tcp_listener_module) as _)
    } else {
        let tcp_listener_module = tcp
            .incoming()
            .and_then(|socket| socket.peer_addr().map(|addr| (addr, socket)));

        let tcp_listener_module = sockets_stream_handle!(tcp_listener_module, http, executor);
        Ok(Box::new(tcp_listener_module) as _)
    }
}

fn tcp_listener_sleep_ms<T: StdError>(error: T, ms: u64) {
    error!(
        "error in accepting connection: {}, tcp_listener will sleep {} ms",
        error.description(),
        ms
    );
    thread::sleep(time::Duration::from_millis(ms));
}

pub struct MethodDetection<S> {
    pub socket: Option<S>,
    pub addr: Option<SocketAddr>,
    pub proxy: bool,
}

impl<S> MethodDetection<S>
where
    S: AsyncPeek + Send + 'static,
{
    pub fn new(proxy: bool, addr: SocketAddr, socket: S) -> Self {
        let addr = Some(addr);
        let socket = Some(socket);
        Self { proxy, addr, socket }
    }
}

impl<S> Future for MethodDetection<S>
where
    S: AsyncPeek + Send + 'static,
{
    type Item = Either<(SocketAddr, S), (SocketAddr, S)>;
    type Error = io::Error;
    fn poll(&mut self) -> Poll<Self::Item, Self::Error> {
        if self.proxy {
            let mut buf = [0u8; 16];

            match self.socket.as_mut().unwrap().async_peek(&mut buf[..]) {
                Ok(Async::NotReady) => Ok(Async::NotReady),
                Ok(Async::Ready(len)) => {
                    debug!("{:?}", str::from_utf8(&buf[..]));

                    let addr = mem::replace(&mut self.addr, None).unwrap();
                    let socket = mem::replace(&mut self.socket, None).unwrap();

                    // 如果太短可能需要忽视这一次，然后等下一次?
                    if len < 7 {
                        if b"CONNECT".starts_with(&buf[..len]) {
                            Ok(Async::Ready(Either::B((addr, socket))))
                        } else {
                            Ok(Async::Ready(Either::A((addr, socket))))
                        }
                    } else {
                        if buf.starts_with(b"CONNECT") {
                            Ok(Async::Ready(Either::B((addr, socket))))
                        } else {
                            Ok(Async::Ready(Either::A((addr, socket))))
                        }
                    }
                }
                Err(e) => Err(e),
            }
        } else {
            let addr = mem::replace(&mut self.addr, None).unwrap();
            let socket = mem::replace(&mut self.socket, None).unwrap();

            Ok(Async::Ready(Either::A((addr, socket))))
        }
    }
}

pub trait AsyncPeek {
    fn async_peek(&mut self, buf: &mut [u8]) -> Result<Async<usize>, io::Error>;
}

impl AsyncPeek for tcp::TcpStream {
    fn async_peek(&mut self, buf: &mut [u8]) -> Result<Async<usize>, io::Error> {
        self.poll_peek(buf)
    }
}

impl AsyncPeek for tokio_rustls::TlsStream<tcp::TcpStream, rustls::ServerSession> {
    fn async_peek(&mut self, buf: &mut [u8]) -> Result<Async<usize>, io::Error> {
        self.get_mut().0.poll_peek(buf)
    }
}

pub fn load_certs(filename: &str) -> Result<Vec<rustls::Certificate>, Error> {
    let certfile = fs::File::open(filename).map_err(|e| format_err!("open certificate file({}) failed: {:?}", filename, e))?;
    let mut reader = io::BufReader::new(certfile);
    pemfile::certs(&mut reader).map_err(|e| format_err!("load certificate({}) failed: {:?}", filename, e))
}

pub fn load_private_key(filename: &str) -> Result<rustls::PrivateKey, Error> {
    let keyfile = fs::File::open(filename).map_err(|e| format_err!("open private key file({}) failed: {:?}", filename, e))?;
    let mut reader = io::BufReader::new(keyfile);
    let keys =
        pemfile::rsa_private_keys(&mut reader).map_err(|e| format_err!("load private key({}) failed: {:?}", filename, e))?;
    assert!(keys.len() == 1);
    Ok(keys[0].clone())
}
<commit_msg>replace a macro by impl trait and generic method.<commit_after>use futures::{future, Async, Future, Poll, Stream};
use hyper::server::conn::Http;
// use hyper::{self, Body, Method, Request, Response, Server, StatusCode};
use failure::Error;
use futures::future::Either;
use rustls;
use rustls::internal::pemfile;
use tokio::net as tcp;
use tokio_rustls::{self, ServerConfigExt};

use tokio::io::{AsyncRead, AsyncWrite};
use tokio::runtime::current_thread;
use tokio::runtime::Builder as RuntimeBuilder;
use tokio::runtime::TaskExecutor;
use tokio_threadpool::Builder as ThreadPoolBuilder;

use std::net::SocketAddr;
use std::{fs, io, mem, str, sync, thread, time};
use StdError;

use base::BaseService;
use config::{Config, Route};
use connect;
use reuse::reuse_address;
use stat::print as stat_print;

pub fn run(config: Config) -> Result<(), Error> {
    let mut threadpool_builder = ThreadPoolBuilder::new();
    threadpool_builder.name_prefix("worker-");

    let runtime = RuntimeBuilder::new().threadpool_builder(threadpool_builder).build()?;

    let executor = runtime.executor();

    let mut tcp_listener: Option<Box<dyn Future<Item = (), Error = ()> + 'static>> = None;
    let mut error = None;

    for addr in config.addrs.clone() {
        match tcp_listener_module(&addr, config.clone(), executor.clone()) {
            Ok(tcp) => {
                tcp_listener = Some(tcp);
                stat_print(
                    &addr,
                    if config.cert.is_some() { "https" } else { "http" },
                    config.routes.values().cloned().collect(),
                );
                break;
            }
            Err(e) => {
                warn!("listen {}: {}", addr, e);
                error = Some(e);
            }
        }
    }

    if tcp_listener.is_none() {
        if let Some(e) = error {
            return Err(e);
        } else {
            unreachable!("tcp-listener and error is None");
        }
    }

    let mut current_thread_runtime = current_thread::Runtime::new()?;
    current_thread_runtime.spawn(tcp_listener.unwrap());
    current_thread_runtime.run()?;
    Ok(())
}

pub fn tcp_listener_module(
    addr: &SocketAddr,
    config: Config,
    executor: TaskExecutor,
) -> Result<Box<dyn Future<Item = (), Error = ()> + 'static>, Error> {
    let tcp = reuse_address(&addr)?;
    let http = Http::new();

    if let Some(cert) = &config.cert {
        let tls_cfg = {
            let certs = load_certs(&cert.pub_)?;
            let key = load_private_key(&cert.key)?;
            let mut cfg = rustls::ServerConfig::new(rustls::NoClientAuth::new());
            cfg.set_single_cert(certs, key)
                .map_err(|e| format_err!("set single cert failed: {:?}", e))?;
            sync::Arc::new(cfg)
        };

        let tcp_listener_module = tcp
            .incoming()
            .and_then(|socket| socket.peer_addr().map(|addr| (addr, socket)))
            .and_then(move |(addr, socket)| tls_cfg.accept_async(socket).map(move |socket| (addr, socket)));

        Ok(Box::new(sockets_stream_handle(tcp_listener_module, http, executor)) as _)
    } else {
        let tcp_listener_module = tcp
            .incoming()
            .and_then(|socket| socket.peer_addr().map(|addr| (addr, socket)));

        Ok(Box::new(sockets_stream_handle(tcp_listener_module, http, executor)) as _)
    }
}

pub fn sockets_stream_handle<RW>(
    socket: impl Stream<Item = (SocketAddr, RW), Error = io::Error> + Send + 'static,
    http: Http,
    executor: TaskExecutor,
) -> impl Future<Item = (), Error = ()> + Send + 'static
where
    RW: AsyncRead + AsyncWrite + AsyncPeek + Send + 'static,
{
    socket
        .and_then(|(addr, socket)| Box::new(MethodDetection::new(true, addr, socket).map(|addrs| Some(addrs))))
        .or_else(|e| {
            tcp_listener_sleep_ms(e, 1);
            future::ok(None)
        }).filter_map(|s: Option<Either<(SocketAddr, _), (SocketAddr, _)>>| s)
        // already or_else..
        // .map_err(|e: io::Error| unreachable!(e))
        .for_each(move |addrs| {
            match addrs {
                Either::A((addr, socket)) => {
                    let connection = http.serve_connection(socket, BaseService::new(addr));

                    executor.spawn(connection.then(move |connection_res| {
                        if let Err(e) = connection_res {
                            error!("client[{}]: {}", addr, e.description());
                        }
                        Ok(())
                    }));
                }
                Either::B((addr, socket)) => {
                    info!("{} Use CONNECT Method ~", addr);
                    executor.spawn(connect::process_socket(addr, socket));
                }
            }
            Ok(())
        })
}

// todo: async if need
pub fn tcp_listener_sleep_ms<T: StdError>(error: T, ms: u64) {
    error!(
        "error in accepting connection: {}, tcp_listener will sleep {} ms",
        error.description(),
        ms
    );
    thread::sleep(time::Duration::from_millis(ms));
}

pub struct MethodDetection<S> {
    pub socket: Option<S>,
    pub addr: Option<SocketAddr>,
    pub proxy: bool,
}

impl<S> MethodDetection<S>
where
    S: AsyncPeek + Send + 'static,
{
    pub fn new(proxy: bool, addr: SocketAddr, socket: S) -> Self {
        let addr = Some(addr);
        let socket = Some(socket);
        Self { proxy, addr, socket }
    }
}

impl<S> Future for MethodDetection<S>
where
    S: AsyncPeek + Send + 'static,
{
    type Item = Either<(SocketAddr, S), (SocketAddr, S)>;
    type Error = io::Error;
    fn poll(&mut self) -> Poll<Self::Item, Self::Error> {
        if self.proxy {
            let mut buf = [0u8; 16];

            match self.socket.as_mut().unwrap().async_peek(&mut buf[..]) {
                Ok(Async::NotReady) => Ok(Async::NotReady),
                Ok(Async::Ready(len)) => {
                    debug!("{:?}", str::from_utf8(&buf[..]));

                    let addr = mem::replace(&mut self.addr, None).unwrap();
                    let socket = mem::replace(&mut self.socket, None).unwrap();

                    // 如果太短可能需要忽视这一次，然后等下一次?
                    if len < 7 {
                        if b"CONNECT".starts_with(&buf[..len]) {
                            Ok(Async::Ready(Either::B((addr, socket))))
                        } else {
                            Ok(Async::Ready(Either::A((addr, socket))))
                        }
                    } else {
                        if buf.starts_with(b"CONNECT") {
                            Ok(Async::Ready(Either::B((addr, socket))))
                        } else {
                            Ok(Async::Ready(Either::A((addr, socket))))
                        }
                    }
                }
                Err(e) => Err(e),
            }
        } else {
            let addr = mem::replace(&mut self.addr, None).unwrap();
            let socket = mem::replace(&mut self.socket, None).unwrap();

            Ok(Async::Ready(Either::A((addr, socket))))
        }
    }
}

pub trait AsyncPeek {
    fn async_peek(&mut self, buf: &mut [u8]) -> Result<Async<usize>, io::Error>;
}

impl AsyncPeek for tcp::TcpStream {
    fn async_peek(&mut self, buf: &mut [u8]) -> Result<Async<usize>, io::Error> {
        self.poll_peek(buf)
    }
}

impl AsyncPeek for tokio_rustls::TlsStream<tcp::TcpStream, rustls::ServerSession> {
    fn async_peek(&mut self, buf: &mut [u8]) -> Result<Async<usize>, io::Error> {
        self.get_mut().0.poll_peek(buf)
    }
}

pub fn load_certs(filename: &str) -> Result<Vec<rustls::Certificate>, Error> {
    let certfile = fs::File::open(filename).map_err(|e| format_err!("open certificate file({}) failed: {:?}", filename, e))?;
    let mut reader = io::BufReader::new(certfile);
    pemfile::certs(&mut reader).map_err(|e| format_err!("load certificate({}) failed: {:?}", filename, e))
}

pub fn load_private_key(filename: &str) -> Result<rustls::PrivateKey, Error> {
    let keyfile = fs::File::open(filename).map_err(|e| format_err!("open private key file({}) failed: {:?}", filename, e))?;
    let mut reader = io::BufReader::new(keyfile);
    let keys =
        pemfile::rsa_private_keys(&mut reader).map_err(|e| format_err!("load private key({}) failed: {:?}", filename, e))?;
    assert!(keys.len() == 1);
    Ok(keys[0].clone())
}
<|endoftext|>


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

use std::net::{Ipv4Addr, SocketAddr, SocketAddrV4};
use std::str::from_utf8;
use std::sync::Arc;

use slog::{debug, info, o, Logger};
use tokio::io::Result;
use tokio::net::UdpSocket;

use crate::config::{Config, ConnectionConfig};

/// Server is the UDP server main implementation
pub struct Server {
    log: Logger,
}

impl Server {
    pub fn new(base: Logger) -> Self {
        let log = base.new(o!("source" => "server::Server"));
        return Server { log };
    }

    /// start the async processing of UDP packets
    pub async fn run(self, config: Arc<Config>) -> Result<()> {
        info!(self.log, "Starting on port {}", config.local.port);
        match &config.connections {
            ConnectionConfig::Sender { address, .. } => {
                info!(self.log, "Sender configuration"; "address" => address)
            }
            ConnectionConfig::Receiver { endpoints } => {
                info!(self.log, "Receiver configuration"; "endpoints" => endpoints.len())
            }
        };

        let mut socket = Server::bind(&config).await?;
        let mut buf: Vec<u8> = vec![0; 65535];
        loop {
            debug!(self.log, "awaiting packet");
            let (size, addr) = socket.recv_from(&mut buf).await?;
            // TOXO: make this async
            self.receive_packet(&buf[..size], addr);
        }
    }

    /// bind binds the local configured port
    async fn bind(config: &Config) -> Result<UdpSocket> {
        let addr = SocketAddrV4::new(Ipv4Addr::new(0, 0, 0, 0), config.local.port);
        return UdpSocket::bind(addr).await;
    }

    /// receive_packet provides the logic for what to do when a packet comes in!
    fn receive_packet(&self, buf: &[u8], addr: SocketAddr) {
        let s = from_utf8(buf).unwrap();
        info!(self.log, "Packet Received from {}: {}", addr, s);
    }
}

#[cfg(test)]
mod tests {
    use std::net::{IpAddr, Ipv4Addr, SocketAddr};

    use crate::config::{Config, ConnectionConfig, Local};
    use crate::server::Server;

    #[tokio::test]
    async fn bind() {
        let config = Config {
            local: Local { port: 12345 },
            connections: ConnectionConfig::Receiver {
                endpoints: Vec::new(),
            },
        };
        let socket = Server::bind(&config).await.unwrap();
        let addr = socket.local_addr().unwrap();

        let expected = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 0, 1)), 12345);
        assert_eq!(expected, addr)
    }
}
<commit_msg>Not failing checks!<commit_after>/*
 * Copyright 2020 Google LLC All Rights Reserved.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

use std::net::{Ipv4Addr, SocketAddr, SocketAddrV4};
use std::str::from_utf8;
use std::sync::Arc;

use slog::{debug, error, info, o, Logger};
use tokio::io::Result;
use tokio::net::UdpSocket;

use crate::config::{Config, ConnectionConfig};

/// Server is the UDP server main implementation
pub struct Server {
    log: Logger,
}

impl Server {
    pub fn new(base: Logger) -> Self {
        let log = base.new(o!("source" => "server::Server"));
        return Server { log };
    }

    /// start the async processing of UDP packets
    pub async fn run(self, config: Arc<Config>) -> Result<()> {
        info!(self.log, "Starting on port {}", config.local.port);
        match &config.connections {
            ConnectionConfig::Sender { address, .. } => {
                info!(self.log, "Sender configuration"; "address" => address)
            }
            ConnectionConfig::Receiver { endpoints } => {
                info!(self.log, "Receiver configuration"; "endpoints" => endpoints.len())
            }
        };

        let mut socket = Server::bind(&config).await?;
        let mut buf: Vec<u8> = vec![0; 65535];
        loop {
            debug!(self.log, "awaiting packet");
            let (size, addr) = socket.recv_from(&mut buf).await?;
            let result = buf.clone();
            let async_log = self.log.clone();
            tokio::spawn(
                async move { Server::receive_packet(async_log, &result[..size], addr).await },
            );
        }
    }

    /// bind binds the local configured port
    async fn bind(config: &Config) -> Result<UdpSocket> {
        let addr = SocketAddrV4::new(Ipv4Addr::new(0, 0, 0, 0), config.local.port);
        return UdpSocket::bind(addr).await;
    }

    /// receive_packet provides the logic for what to do when a packet comes in!
    async fn receive_packet(log: Logger, buf: &[u8], addr: SocketAddr) -> Result<()> {
        let s = from_utf8(buf).unwrap();
        info!(log, "Packet Received from {}: {}", addr, s);
        return Ok(());
    }
}

#[cfg(test)]
mod tests {
    use std::net::{IpAddr, Ipv4Addr, SocketAddr};

    use crate::config::{Config, ConnectionConfig, Local};
    use crate::server::Server;

    #[tokio::test]
    async fn bind() {
        let config = Config {
            local: Local { port: 12345 },
            connections: ConnectionConfig::Receiver {
                endpoints: Vec::new(),
            },
        };
        let socket = Server::bind(&config).await.unwrap();
        let addr = socket.local_addr().unwrap();

        let expected = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 0, 1)), 12345);
        assert_eq!(expected, addr)
    }
}
<|endoftext|>


<commit_before>use std::ptr::null;
use std::num::zero;
use std::time::duration::Duration;
use std::default::Default;
use std::io::process::Process;
use libc::types::os::common::posix01::timespec;
use time::{Timespec, get_time};

pub use libc::consts::os::posix88::{SIGTERM, SIGINT, SIGQUIT, EINTR};
use libc::{c_int, pid_t};

static SIGCHLD: c_int = 17;
static WNOHANG: c_int = 1;

#[deriving(Show)]
pub enum Signal {
    Terminate(int),  // Actual signal for termination: INT, TERM, QUIT...
    Child(pid_t, int),  //  pid and result code
    Reboot,
    Timeout,  // Not actually a OS signal, but it's a signal for our app
}

#[deriving(Default)]
#[repr(C)]
struct CSignalInfo {
    signo: c_int,
    pid: pid_t,
    status: c_int,
}

extern {
    fn block_all_signals();
    fn wait_any_signal(ptr: *mut CSignalInfo, timeout: *const timespec)
        -> c_int;
    fn waitpid(pid: pid_t, status: *mut c_int, options: c_int) -> pid_t;
}

pub fn block_all() {
    unsafe { block_all_signals() };
}

fn _convert_status(status: i32) -> int {
    if status & 0xff == 0 {
        return ((status & 0xff00) >> 8) as int;
    }
    return (128 + (status & 0x7f)) as int;  // signal
}

pub fn wait_next(reboot_supported: bool, timeout: Option<Timespec>) -> Signal {
    let mut status: i32 = 0;
    let pid = unsafe { waitpid(-1, &mut status, WNOHANG) };
    if pid > 0 {
        return Child(pid, _convert_status(status));
    }
    loop {
        let mut ptr = Default::default();
        let res = match timeout {
            Some(tm) => {
                let dur = tm - get_time();
                let c_dur = if dur < zero() {
                    timespec { tv_sec: 0, tv_nsec: 0 }
                } else {
                    timespec {
                        tv_sec: dur.num_seconds(),
                        tv_nsec: (dur - Duration::seconds(dur.num_seconds()))
                            .num_nanoseconds().unwrap(),
                    }
                };
                unsafe { wait_any_signal(&mut ptr, &c_dur) }
            }
            None => {
                unsafe { wait_any_signal(&mut ptr, null()) }
            }
        };
        if res != 0 {
            //  Any error is ok, because application should be always prepared
            //  for spurious timeouts
            //  only EAGAIN and EINTR expected
            return Timeout;
        }
        match ptr.signo {
            SIGCHLD => {
                loop {
                    status = 0;
                    let rc = unsafe { waitpid(ptr.pid, &mut status, WNOHANG) };
                    if rc < 0 {
                        if rc == EINTR {
                            continue;
                        }
                        fail!("Failure {} not expected", rc);
                    }
                    assert_eq!(rc, ptr.pid);
                    break;
                }
                assert_eq!(_convert_status(status), ptr.status as int);
                return Child(ptr.pid, ptr.status as int);
            }
            SIGQUIT if reboot_supported => {
                return Reboot;
            }
            sig@SIGTERM | sig@SIGINT | sig@SIGQUIT => {
                return Terminate(sig as int);
            }
            _ => continue,   // TODO(tailhook) improve logging
        }
    }
}

pub fn send_signal(pid: pid_t, sig: int) {
    Process::kill(pid, sig).ok();
}
<commit_msg>Better reporting of signal crashes<commit_after>use std::io::IoError;
use std::ptr::null;
use std::num::zero;
use std::time::duration::Duration;
use std::default::Default;
use std::io::process::Process;
use libc::types::os::common::posix01::timespec;
use time::{Timespec, get_time};

pub use libc::consts::os::posix88::{SIGTERM, SIGINT, SIGQUIT, EINTR};
use libc::{c_int, pid_t};

static SIGCHLD: c_int = 17;
static WNOHANG: c_int = 1;

#[deriving(Show)]
pub enum Signal {
    Terminate(int),  // Actual signal for termination: INT, TERM, QUIT...
    Child(pid_t, int),  //  pid and result code
    Reboot,
    Timeout,  // Not actually a OS signal, but it's a signal for our app
}

#[deriving(Default)]
#[repr(C)]
struct CSignalInfo {
    signo: c_int,
    pid: pid_t,
    status: c_int,
}

extern {
    fn block_all_signals();
    fn wait_any_signal(ptr: *mut CSignalInfo, timeout: *const timespec)
        -> c_int;
    fn waitpid(pid: pid_t, status: *mut c_int, options: c_int) -> pid_t;
}

pub fn block_all() {
    unsafe { block_all_signals() };
}

fn _convert_status(status: i32) -> int {
    if status & 0xff == 0 {
        return ((status & 0xff00) >> 8) as int;
    }
    return (128 + (status & 0x7f)) as int;  // signal
}

pub fn wait_next(reboot_supported: bool, timeout: Option<Timespec>) -> Signal {
    let mut status: i32 = 0;
    let pid = unsafe { waitpid(-1, &mut status, WNOHANG) };
    if pid > 0 {
        return Child(pid, _convert_status(status));
    }
    loop {
        let mut ptr = Default::default();
        let res = match timeout {
            Some(tm) => {
                let dur = tm - get_time();
                let c_dur = if dur < zero() {
                    timespec { tv_sec: 0, tv_nsec: 0 }
                } else {
                    timespec {
                        tv_sec: dur.num_seconds(),
                        tv_nsec: (dur - Duration::seconds(dur.num_seconds()))
                            .num_nanoseconds().unwrap(),
                    }
                };
                unsafe { wait_any_signal(&mut ptr, &c_dur) }
            }
            None => {
                unsafe { wait_any_signal(&mut ptr, null()) }
            }
        };
        if res != 0 {
            //  Any error is ok, because application should be always prepared
            //  for spurious timeouts
            //  only EAGAIN and EINTR expected
            return Timeout;
        }
        match ptr.signo {
            SIGCHLD => {
                loop {
                    status = 0;
                    let rc = unsafe { waitpid(ptr.pid, &mut status, WNOHANG) };
                    if rc < 0 {
                        if rc == EINTR {
                            continue;
                        }
                        fail!("Failure '{}' not expected, on death of {}",
                            IoError::last_error(), ptr.pid);
                    }
                    assert_eq!(rc, ptr.pid);
                    break;
                }
                assert_eq!(_convert_status(status), ptr.status as int);
                return Child(ptr.pid, ptr.status as int);
            }
            SIGQUIT if reboot_supported => {
                return Reboot;
            }
            sig@SIGTERM | sig@SIGINT | sig@SIGQUIT => {
                return Terminate(sig as int);
            }
            _ => continue,   // TODO(tailhook) improve logging
        }
    }
}

pub fn send_signal(pid: pid_t, sig: int) {
    Process::kill(pid, sig).ok();
}
<|endoftext|>


<commit_before>use std::net::ToSocketAddrs;
use std::sync::{Arc, RwLock};
use std::collections::HashMap;
use hyper::server::{Request, Response, Handler};
use hyper::server::Server as HyperServer;

use middleware::MiddlewareStack;
use request;
use response;

pub struct Server {
    middleware_stack: MiddlewareStack,
    templates: response::TemplateCache
}

// FIXME: Any better coherence solutions?
struct ArcServer(Arc<Server>);

impl Handler for ArcServer {
    fn handle<'a, 'k>(&'a self, req: Request<'a, 'k>, res: Response<'a>) {
        let nickel_req = request::Request::from_internal(req);
        let nickel_res = response::Response::from_internal(res, &self.0.templates);

        self.0.middleware_stack.invoke(nickel_req, nickel_res);
    }
}

impl Server {
    pub fn new(middleware_stack: MiddlewareStack) -> Server {
        Server {
            middleware_stack: middleware_stack,
            templates: RwLock::new(HashMap::new())
        }
    }

    pub fn serve<T: ToSocketAddrs>(self, addr: T) {
        let arc = ArcServer(Arc::new(self));
        let server = HyperServer::http(arc);
        let _ = server.listen(addr);
    }
}

<commit_msg>fix(server): ensure panic on invalid listen addr<commit_after>use std::net::ToSocketAddrs;
use std::sync::{Arc, RwLock};
use std::collections::HashMap;
use hyper::server::{Request, Response, Handler};
use hyper::server::Server as HyperServer;

use middleware::MiddlewareStack;
use request;
use response;

pub struct Server {
    middleware_stack: MiddlewareStack,
    templates: response::TemplateCache
}

// FIXME: Any better coherence solutions?
struct ArcServer(Arc<Server>);

impl Handler for ArcServer {
    fn handle<'a, 'k>(&'a self, req: Request<'a, 'k>, res: Response<'a>) {
        let nickel_req = request::Request::from_internal(req);
        let nickel_res = response::Response::from_internal(res, &self.0.templates);

        self.0.middleware_stack.invoke(nickel_req, nickel_res);
    }
}

impl Server {
    pub fn new(middleware_stack: MiddlewareStack) -> Server {
        Server {
            middleware_stack: middleware_stack,
            templates: RwLock::new(HashMap::new())
        }
    }

    pub fn serve<T: ToSocketAddrs>(self, addr: T) {
        let arc = ArcServer(Arc::new(self));
        let server = HyperServer::http(arc);
        let _ = server.listen(addr).unwrap();
    }
}

#[test]
#[should_panic]
fn invalid_listen_addr() {
    Server::new(MiddlewareStack::new()).serve("invalid_addr");
}
<|endoftext|>


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

use protocol::{self, mojang, packet};
use world;
use world::block::{self, BlockSet};
use rand::{self, Rng};
use std::sync::{Arc, RwLock, Mutex};
use std::sync::mpsc;
use std::thread;
use std::collections::HashMap;
use resources;
use openssl;
use console;
use render;
use auth;
use cgmath::{self, Vector};
use glutin::VirtualKeyCode;

pub struct Server {
    conn: Option<protocol::Conn>,
    read_queue: Option<mpsc::Receiver<Result<packet::Packet, protocol::Error>>>,
    pub world: world::World,

    resources: Arc<RwLock<resources::Manager>>,
    console: Arc<Mutex<console::Console>>,
    version: usize,

    pub position: cgmath::Vector3<f64>,
    pub yaw: f64,
    pub pitch: f64,

    pressed_keys: HashMap<VirtualKeyCode, bool>,

    tick_timer: f64,
}

macro_rules! handle_packet {
    ($s:ident $pck:ident {
        $($packet:ident => $func:ident,)*
    }) => (
        match $pck {
        $(
            protocol::packet::Packet::$packet(val) => $s.$func(val),
        )*
            _ => {},
        }
    )
}

impl Server {

    pub fn connect(resources: Arc<RwLock<resources::Manager>>, console: Arc<Mutex<console::Console>>, address: &str) -> Result<Server, protocol::Error> {
        let mut conn = try!(protocol::Conn::new(address));

        let host = conn.host.clone();
        let port = conn.port;
        try!(conn.write_packet(protocol::packet::handshake::serverbound::Handshake {
             protocol_version: protocol::VarInt(protocol::SUPPORTED_PROTOCOL),
             host: host,
             port: port,
             next: protocol::VarInt(2),
         }));
        conn.state = protocol::State::Login;
        try!(conn.write_packet(protocol::packet::login::serverbound::LoginStart {
            username: "Thinkofdeath".to_owned()
        }));

        let packet = match try!(conn.read_packet()) {
            protocol::packet::Packet::EncryptionRequest(val) => val,
            protocol::packet::Packet::LoginDisconnect(val) => return Err(protocol::Error::Disconnect(val.reason)),
            val => return Err(protocol::Error::Err(format!("Wrong packet: {:?}", val))),
        };

        let mut key = openssl::PublicKey::new(&packet.public_key.data);
        let shared = openssl::gen_random(16);

        let shared_e = key.encrypt(&shared);
        let token_e = key.encrypt(&packet.verify_token.data);

        let profile = {
            let console = console.lock().unwrap();
            mojang::Profile {
                username: console.get(auth::CL_USERNAME).clone(),
                id: console.get(auth::CL_UUID).clone(),
                access_token: console.get(auth::AUTH_TOKEN).clone(),
            }
        };
        try!(profile.join_server(&packet.server_id, &shared, &packet.public_key.data));

        try!(conn.write_packet(protocol::packet::login::serverbound::EncryptionResponse {
            shared_secret: protocol::LenPrefixedBytes::new(shared_e),
            verify_token: protocol::LenPrefixedBytes::new(token_e),
        }));

        let mut read = conn.clone(); // TODO
        let mut write = conn.clone();

        read.enable_encyption(&shared, true);
        write.enable_encyption(&shared, false);

        loop {
           match try!(read.read_packet()) {
               protocol::packet::Packet::SetInitialCompression(val) => {
                   read.set_compresssion(val.threshold.0, true);
                   write.set_compresssion(val.threshold.0, false);
               }
               protocol::packet::Packet::LoginSuccess(val) => {
                   debug!("Login: {} {}", val.username, val.uuid);
                   read.state = protocol::State::Play;
                   write.state = protocol::State::Play;
                   break;
               }
               protocol::packet::Packet::LoginDisconnect(val) => return Err(protocol::Error::Disconnect(val.reason)),
               val => return Err(protocol::Error::Err(format!("Wrong packet: {:?}", val))),
           }
        }

        let (tx, rx) = mpsc::channel();
        thread::spawn(move || {
            loop {
                let pck = read.read_packet();
                let was_error = pck.is_err();
                if let Err(_) = tx.send(pck) {
                    return;
                }
                if was_error {
                    return;
                }
            }
        });

        let version = resources.read().unwrap().version();
        Ok(Server {
            conn: Some(write),
            read_queue: Some(rx),
            world: world::World::new(),
            resources: resources,
            console: console,
            version: version,

            pressed_keys: HashMap::new(),

            position: cgmath::Vector3::zero(),
            yaw: 0.0,
            pitch: 0.0,

            tick_timer: 0.0,
        })
    }

    pub fn dummy_server(resources: Arc<RwLock<resources::Manager>>, console: Arc<Mutex<console::Console>>) -> Server {
        let mut world = world::World::new();
        let mut rng = rand::thread_rng();
        for x in -7*16 .. 7*16 {
            for z in -7*16 .. 7*16 {
                let h = rng.gen_range(3, 10);
                for y in 0 .. h {
                    world.set_block(x, y, z, block::MISSING.base());
                }
            }
        }
        let version = resources.read().unwrap().version();
        Server {
            conn: None,
            read_queue: None,
            world: world,

            version: version,
            resources: resources,
            console: console,

            pressed_keys: HashMap::new(),

            position: cgmath::Vector3::new(0.5, 13.2, 0.5),
            yaw: 0.0,
            pitch: 0.0,

            tick_timer: 0.0,
        }
    }

    pub fn is_connected(&self) -> bool {
        self.conn.is_some()
    }

    pub fn tick(&mut self, renderer: &mut render::Renderer, delta: f64) {
        let version = self.resources.read().unwrap().version();
        if version != self.version {
            self.version = version;
            self.world.flag_dirty_all();
        }

        if let Some(rx) = self.read_queue.take() {
            while let Ok(pck) = rx.try_recv() {
                match pck {
                    Ok(pck) => handle_packet!{
                        self pck {
                            KeepAliveClientbound => on_keep_alive,
                            ChunkData => on_chunk_data,
                            ChunkUnload => on_chunk_unload,
                            TeleportPlayer => on_teleport,
                        }
                    },
                    Err(err) => panic!("Err: {:?}", err),
                }
            }
            self.read_queue = Some(rx);
        }

        let (forward, yaw) = self.calculate_movement();

        if self.world.is_chunk_loaded((self.position.x as i32) >> 4, (self.position.z as i32) >> 4) {
                let mut speed = 4.317 / 60.0;
                if self.is_key_pressed(VirtualKeyCode::LShift) {
                    speed = 5.612 / 60.0;
                }
                // TODO: only do this for flying
                speed *= 2.5;

                if self.is_key_pressed(VirtualKeyCode::Space) {
                    self.position.y += speed * delta;
                }
                if self.is_key_pressed(VirtualKeyCode::LControl) {
                    self.position.y -= speed * delta;
                }
                self.position.x += forward * yaw.cos() * delta * speed;
                self.position.z -= forward * yaw.sin() * delta * speed;
        }

        self.tick_timer += delta;
        while self.tick_timer >= 3.0 && self.is_connected() {
            self.minecraft_tick();
            self.tick_timer -= 3.0;
        }

        // Copy to camera
        renderer.camera.pos = cgmath::Point::from_vec(self.position + cgmath::Vector3::new(0.0, 1.8, 0.0));
        renderer.camera.yaw = self.yaw;
        renderer.camera.pitch = self.pitch;
    }

    fn calculate_movement(&self) -> (f64, f64) {
        use std::f64::consts::PI;
        let mut forward = 0.0f64;
        let mut yaw = self.yaw - (PI/2.0);
        if self.is_key_pressed(VirtualKeyCode::W) || self.is_key_pressed(VirtualKeyCode::S) {
            forward = 1.0;
            if self.is_key_pressed(VirtualKeyCode::S) {
                yaw += PI;
            }
        }
        let mut change = 0.0;
        if self.is_key_pressed(VirtualKeyCode::A) {
            change = (PI / 2.0) / (forward.abs() + 1.0);
        }
        if self.is_key_pressed(VirtualKeyCode::D) {
            change = -(PI / 2.0) / (forward.abs() + 1.0);
        }
        if self.is_key_pressed(VirtualKeyCode::A) || self.is_key_pressed(VirtualKeyCode::D) {
            forward = 1.0;
        }
        if self.is_key_pressed(VirtualKeyCode::S) {
            yaw -= change;
        } else {
            yaw += change;
        }

        (forward, yaw)
    }

    pub fn minecraft_tick(&mut self) {

        // Sync our position to the server
        let packet = packet::play::serverbound::PlayerPositionLook {
            x: self.position.x,
            y: self.position.y,
            z: self.position.z,
            yaw: self.yaw as f32,
            pitch: self.pitch as f32,
            on_ground: false,
        };
        self.write_packet(packet);
    }

    pub fn key_press(&mut self, down: bool, key: VirtualKeyCode) {
        self.pressed_keys.insert(key, down);
    }

    fn is_key_pressed(&self, key: VirtualKeyCode) -> bool {
        self.pressed_keys.get(&key).map(|v| *v).unwrap_or(false)
    }

    pub fn write_packet<T: protocol::PacketType>(&mut self, p: T) {
        self.conn.as_mut().unwrap().write_packet(p).unwrap(); // TODO handle errors
    }

    fn on_keep_alive(&mut self, keep_alive: packet::play::clientbound::KeepAliveClientbound) {
        self.write_packet(packet::play::serverbound::KeepAliveServerbound {
            id: keep_alive.id,
        });
    }

    fn on_teleport(&mut self, teleport: packet::play::clientbound::TeleportPlayer) {
        // TODO: relative teleports
        self.position.x = teleport.x;
        self.position.y = teleport.y;
        self.position.z = teleport.z;
        self.yaw = teleport.yaw as f64;
        self.pitch = teleport.pitch as f64;

        self.write_packet(packet::play::serverbound::PlayerPositionLook {
            x: teleport.x,
            y: teleport.y,
            z: teleport.z,
            yaw: teleport.yaw,
            pitch: teleport.pitch,
            on_ground: false,
        });
    }

    fn on_chunk_data(&mut self, chunk_data: packet::play::clientbound::ChunkData) {
        self.world.load_chunk(
            chunk_data.chunk_x,
            chunk_data.chunk_z,
            chunk_data.new,
            chunk_data.bitmask.0 as u16,
            chunk_data.data.data
        ).unwrap();
    }

    fn on_chunk_unload(&mut self, chunk_unload: packet::play::clientbound::ChunkUnload) {
        self.world.unload_chunk(chunk_unload.x, chunk_unload.z);
    }
}
<commit_msg>Don't hardcode my username into the login process<commit_after>// Copyright 2015 Matthew Collins
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

use protocol::{self, mojang, packet};
use world;
use world::block::{self, BlockSet};
use rand::{self, Rng};
use std::sync::{Arc, RwLock, Mutex};
use std::sync::mpsc;
use std::thread;
use std::collections::HashMap;
use resources;
use openssl;
use console;
use render;
use auth;
use cgmath::{self, Vector};
use glutin::VirtualKeyCode;

pub struct Server {
    conn: Option<protocol::Conn>,
    read_queue: Option<mpsc::Receiver<Result<packet::Packet, protocol::Error>>>,
    pub world: world::World,

    resources: Arc<RwLock<resources::Manager>>,
    console: Arc<Mutex<console::Console>>,
    version: usize,

    pub position: cgmath::Vector3<f64>,
    pub yaw: f64,
    pub pitch: f64,

    pressed_keys: HashMap<VirtualKeyCode, bool>,

    tick_timer: f64,
}

macro_rules! handle_packet {
    ($s:ident $pck:ident {
        $($packet:ident => $func:ident,)*
    }) => (
        match $pck {
        $(
            protocol::packet::Packet::$packet(val) => $s.$func(val),
        )*
            _ => {},
        }
    )
}

impl Server {

    pub fn connect(resources: Arc<RwLock<resources::Manager>>, console: Arc<Mutex<console::Console>>, address: &str) -> Result<Server, protocol::Error> {
        let mut conn = try!(protocol::Conn::new(address));

        let profile = {
            let console = console.lock().unwrap();
            mojang::Profile {
                username: console.get(auth::CL_USERNAME).clone(),
                id: console.get(auth::CL_UUID).clone(),
                access_token: console.get(auth::AUTH_TOKEN).clone(),
            }
        };

        let host = conn.host.clone();
        let port = conn.port;
        try!(conn.write_packet(protocol::packet::handshake::serverbound::Handshake {
             protocol_version: protocol::VarInt(protocol::SUPPORTED_PROTOCOL),
             host: host,
             port: port,
             next: protocol::VarInt(2),
         }));
        conn.state = protocol::State::Login;
        try!(conn.write_packet(protocol::packet::login::serverbound::LoginStart {
            username: profile.username.clone(),
        }));

        let packet = match try!(conn.read_packet()) {
            protocol::packet::Packet::EncryptionRequest(val) => val,
            protocol::packet::Packet::LoginDisconnect(val) => return Err(protocol::Error::Disconnect(val.reason)),
            val => return Err(protocol::Error::Err(format!("Wrong packet: {:?}", val))),
        };

        let mut key = openssl::PublicKey::new(&packet.public_key.data);
        let shared = openssl::gen_random(16);

        let shared_e = key.encrypt(&shared);
        let token_e = key.encrypt(&packet.verify_token.data);

        try!(profile.join_server(&packet.server_id, &shared, &packet.public_key.data));

        try!(conn.write_packet(protocol::packet::login::serverbound::EncryptionResponse {
            shared_secret: protocol::LenPrefixedBytes::new(shared_e),
            verify_token: protocol::LenPrefixedBytes::new(token_e),
        }));

        let mut read = conn.clone(); // TODO
        let mut write = conn.clone();

        read.enable_encyption(&shared, true);
        write.enable_encyption(&shared, false);

        loop {
           match try!(read.read_packet()) {
               protocol::packet::Packet::SetInitialCompression(val) => {
                   read.set_compresssion(val.threshold.0, true);
                   write.set_compresssion(val.threshold.0, false);
               }
               protocol::packet::Packet::LoginSuccess(val) => {
                   debug!("Login: {} {}", val.username, val.uuid);
                   read.state = protocol::State::Play;
                   write.state = protocol::State::Play;
                   break;
               }
               protocol::packet::Packet::LoginDisconnect(val) => return Err(protocol::Error::Disconnect(val.reason)),
               val => return Err(protocol::Error::Err(format!("Wrong packet: {:?}", val))),
           }
        }

        let (tx, rx) = mpsc::channel();
        thread::spawn(move || {
            loop {
                let pck = read.read_packet();
                let was_error = pck.is_err();
                if let Err(_) = tx.send(pck) {
                    return;
                }
                if was_error {
                    return;
                }
            }
        });

        let version = resources.read().unwrap().version();
        Ok(Server {
            conn: Some(write),
            read_queue: Some(rx),
            world: world::World::new(),
            resources: resources,
            console: console,
            version: version,

            pressed_keys: HashMap::new(),

            position: cgmath::Vector3::zero(),
            yaw: 0.0,
            pitch: 0.0,

            tick_timer: 0.0,
        })
    }

    pub fn dummy_server(resources: Arc<RwLock<resources::Manager>>, console: Arc<Mutex<console::Console>>) -> Server {
        let mut world = world::World::new();
        let mut rng = rand::thread_rng();
        for x in -7*16 .. 7*16 {
            for z in -7*16 .. 7*16 {
                let h = rng.gen_range(3, 10);
                for y in 0 .. h {
                    world.set_block(x, y, z, block::MISSING.base());
                }
            }
        }
        let version = resources.read().unwrap().version();
        Server {
            conn: None,
            read_queue: None,
            world: world,

            version: version,
            resources: resources,
            console: console,

            pressed_keys: HashMap::new(),

            position: cgmath::Vector3::new(0.5, 13.2, 0.5),
            yaw: 0.0,
            pitch: 0.0,

            tick_timer: 0.0,
        }
    }

    pub fn is_connected(&self) -> bool {
        self.conn.is_some()
    }

    pub fn tick(&mut self, renderer: &mut render::Renderer, delta: f64) {
        let version = self.resources.read().unwrap().version();
        if version != self.version {
            self.version = version;
            self.world.flag_dirty_all();
        }

        if let Some(rx) = self.read_queue.take() {
            while let Ok(pck) = rx.try_recv() {
                match pck {
                    Ok(pck) => handle_packet!{
                        self pck {
                            KeepAliveClientbound => on_keep_alive,
                            ChunkData => on_chunk_data,
                            ChunkUnload => on_chunk_unload,
                            TeleportPlayer => on_teleport,
                        }
                    },
                    Err(err) => panic!("Err: {:?}", err),
                }
            }
            self.read_queue = Some(rx);
        }

        let (forward, yaw) = self.calculate_movement();

        if self.world.is_chunk_loaded((self.position.x as i32) >> 4, (self.position.z as i32) >> 4) {
                let mut speed = 4.317 / 60.0;
                if self.is_key_pressed(VirtualKeyCode::LShift) {
                    speed = 5.612 / 60.0;
                }
                // TODO: only do this for flying
                speed *= 2.5;

                if self.is_key_pressed(VirtualKeyCode::Space) {
                    self.position.y += speed * delta;
                }
                if self.is_key_pressed(VirtualKeyCode::LControl) {
                    self.position.y -= speed * delta;
                }
                self.position.x += forward * yaw.cos() * delta * speed;
                self.position.z -= forward * yaw.sin() * delta * speed;
        }

        self.tick_timer += delta;
        while self.tick_timer >= 3.0 && self.is_connected() {
            self.minecraft_tick();
            self.tick_timer -= 3.0;
        }

        // Copy to camera
        renderer.camera.pos = cgmath::Point::from_vec(self.position + cgmath::Vector3::new(0.0, 1.8, 0.0));
        renderer.camera.yaw = self.yaw;
        renderer.camera.pitch = self.pitch;
    }

    fn calculate_movement(&self) -> (f64, f64) {
        use std::f64::consts::PI;
        let mut forward = 0.0f64;
        let mut yaw = self.yaw - (PI/2.0);
        if self.is_key_pressed(VirtualKeyCode::W) || self.is_key_pressed(VirtualKeyCode::S) {
            forward = 1.0;
            if self.is_key_pressed(VirtualKeyCode::S) {
                yaw += PI;
            }
        }
        let mut change = 0.0;
        if self.is_key_pressed(VirtualKeyCode::A) {
            change = (PI / 2.0) / (forward.abs() + 1.0);
        }
        if self.is_key_pressed(VirtualKeyCode::D) {
            change = -(PI / 2.0) / (forward.abs() + 1.0);
        }
        if self.is_key_pressed(VirtualKeyCode::A) || self.is_key_pressed(VirtualKeyCode::D) {
            forward = 1.0;
        }
        if self.is_key_pressed(VirtualKeyCode::S) {
            yaw -= change;
        } else {
            yaw += change;
        }

        (forward, yaw)
    }

    pub fn minecraft_tick(&mut self) {

        // Sync our position to the server
        let packet = packet::play::serverbound::PlayerPositionLook {
            x: self.position.x,
            y: self.position.y,
            z: self.position.z,
            yaw: self.yaw as f32,
            pitch: self.pitch as f32,
            on_ground: false,
        };
        self.write_packet(packet);
    }

    pub fn key_press(&mut self, down: bool, key: VirtualKeyCode) {
        self.pressed_keys.insert(key, down);
    }

    fn is_key_pressed(&self, key: VirtualKeyCode) -> bool {
        self.pressed_keys.get(&key).map(|v| *v).unwrap_or(false)
    }

    pub fn write_packet<T: protocol::PacketType>(&mut self, p: T) {
        self.conn.as_mut().unwrap().write_packet(p).unwrap(); // TODO handle errors
    }

    fn on_keep_alive(&mut self, keep_alive: packet::play::clientbound::KeepAliveClientbound) {
        self.write_packet(packet::play::serverbound::KeepAliveServerbound {
            id: keep_alive.id,
        });
    }

    fn on_teleport(&mut self, teleport: packet::play::clientbound::TeleportPlayer) {
        // TODO: relative teleports
        self.position.x = teleport.x;
        self.position.y = teleport.y;
        self.position.z = teleport.z;
        self.yaw = teleport.yaw as f64;
        self.pitch = teleport.pitch as f64;

        self.write_packet(packet::play::serverbound::PlayerPositionLook {
            x: teleport.x,
            y: teleport.y,
            z: teleport.z,
            yaw: teleport.yaw,
            pitch: teleport.pitch,
            on_ground: false,
        });
    }

    fn on_chunk_data(&mut self, chunk_data: packet::play::clientbound::ChunkData) {
        self.world.load_chunk(
            chunk_data.chunk_x,
            chunk_data.chunk_z,
            chunk_data.new,
            chunk_data.bitmask.0 as u16,
            chunk_data.data.data
        ).unwrap();
    }

    fn on_chunk_unload(&mut self, chunk_unload: packet::play::clientbound::ChunkUnload) {
        self.world.unload_chunk(chunk_unload.x, chunk_unload.z);
    }
}
<|endoftext|>


<commit_before>use std;
use std::collections::HashMap;
use std::{fmt, io, result};
use std::net::{Shutdown, TcpListener, TcpStream};
use std::sync::{Arc, Mutex};
use std::thread;
use std::sync::mpsc::{self, channel, Receiver, Sender};

use message;

#[derive(Debug)]
pub enum Error {
    NotificationSend(mpsc::SendError<Option<message::Notification>>),
    Message(message::Error),
    Io(io::Error),
}

impl std::error::Error for Error {
    fn description(&self) -> &str {
        "Server error"
    }

    fn cause(&self) -> Option<&std::error::Error> {
        match *self {
            Error::NotificationSend(ref err) => Some(err),
            Error::Message(ref err) => Some(err),
            Error::Io(ref err) => Some(err),
        }
    }
}

impl fmt::Display for Error {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        std::error::Error::description(self).fmt(f)
    }
}

impl From<message::Error> for Error {
    fn from(err: message::Error) -> Error {
        Error::Message(err)
    }
}

impl From<mpsc::SendError<Option<message::Notification>>> for Error {
    fn from(err: mpsc::SendError<Option<message::Notification>>) -> Error {
        Error::NotificationSend(err)
    }
}

impl From<io::Error> for Error {
    fn from(err: io::Error) -> Error {
        Error::Io(err)
    }
}

pub type Result<T> = result::Result<T, Error>;

pub fn execute(host: String, port: i32) -> Result<()> {
    let listener = try!(TcpListener::bind(&format!("{}:{}", host, port)[..]));
    let mut channels = Channels { channels: HashMap::new() };

    for stream in listener.incoming() {
        let _ = handle_client(stream, &mut channels);
    }

    Ok(())
}

fn handle_client(stream: result::Result<TcpStream, io::Error>, channels: &mut Channels) -> Result<()> {
    let stream = try!(stream);
    let request = try!(message::JoinRequest::receive(&stream));

    match request {
        message::JoinRequest::Broadcast(ch) =>
            BroadcastHandler::spawn(&stream, channels.fetch(&ch)),
        message::JoinRequest::Watch(ch) =>
            WatchHandler::spawn(&stream, channels.fetch(&ch)),
    }

    Ok(())
}

struct Channels {
    channels: HashMap<String, Arc<Mutex<Channel>>>,
}

struct Channel {
    broadcaster: Option<TcpStream>,
    watchers: Vec<TcpStream>,
}

struct BroadcastHandler {
    stream: TcpStream,
    channel: Arc<Mutex<Channel>>,
}

struct WatchHandler {
    stream: TcpStream,
    channel: Arc<Mutex<Channel>>,
}

impl Channels {
    fn fetch<K: AsRef<str>>(&mut self, key: &K) -> &Arc<Mutex<Channel>> {
        let ch = key.as_ref();

        if self.channels.contains_key(&ch[..]) {
            self.channels.get(&ch[..]).unwrap()
        } else {
            let channel = Channel {
                broadcaster: None,
                watchers: Vec::new(),
            };

            self.channels.insert(ch.to_owned(), Arc::new(Mutex::new(channel)));

            self.fetch(&ch)
        }
    }
}

impl Channel {
    fn takeover(&mut self, stream: TcpStream) -> Result<()> {
        match self.broadcaster.as_mut() {
            Some(mut current) => {
                try!(message::Notification::Closed("Broadcaster has changed".to_owned()).send(&mut current));
                try!(current.shutdown(Shutdown::Both));
            }
            None => (),
        }

        self.broadcaster = Some(stream);
        Ok(())
    }
}

impl BroadcastHandler {
    fn spawn(stream: &TcpStream, channel: &Arc<Mutex<Channel>>) {
        let mut handler = BroadcastHandler {
            stream: stream.try_clone().unwrap(),
            channel: channel.clone(),
        };

        let _ = channel.lock().and_then(|mut ch| {
            let _ = ch.takeover(stream.try_clone().unwrap());
            Ok(ch)
        });

        thread::spawn(move || {
            handler.process().unwrap_or_else(|e| {
                println!("{}", e);
            });
        });
    }

    fn process(&mut self) -> Result<()> {
        try!(message::JoinResponse::Success.send(&mut self.stream));

        let (sender, receiver) = channel();

        try!(self.spawn_relay(sender.clone()));
        try!(self.broadcast(receiver));

        try!(self.stream.shutdown(Shutdown::Both));

        Ok(())
    }

    fn spawn_relay(&self, sender: Sender<Option<message::Notification>>) -> Result<()> {
        let stream = try!(self.stream.try_clone());

        thread::spawn(move || -> Result<()> {
            loop {
                match message::Notification::receive(&stream) {
                    Ok(notification) => {
                        match notification {
                            message::Notification::Output(data)   => {
                                sender.send(Some(message::Notification::Output(data))).unwrap_or_else(|e| println!("{}", e));
                            },
                            message::Notification::Closed(reason) => {
                                try!(sender.send(Some(message::Notification::Closed(reason))));
                            },
                            _ => break,
                        }
                    },
                    Err(_) => break,
                }
            };

            try!(sender.send(None));

            Ok(())
        });

        Ok(())
    }

    fn broadcast(&self, receiver: Receiver<Option<message::Notification>>) -> Result<()> {
        for message in receiver {
            match message {
                Some(notification) => {
                    let mut channel = self.channel.lock().unwrap();

                    for mut watcher in channel.watchers.iter_mut() {
                        let _ = notification.send(watcher);
                    }
                },
                None => break
            }
        }

        Ok(())
    }
}

impl WatchHandler {
    fn spawn(stream: &TcpStream, channel: &Arc<Mutex<Channel>>) {
        let mut handler = WatchHandler {
            stream: stream.try_clone().unwrap(),
            channel: channel.clone(),
        };

        thread::spawn(move || {
            handler.process().unwrap_or_else(|e| {
                println!("{}", e);
            });
        });
    }

    fn process(&mut self) -> Result<()> {
        try!(message::JoinResponse::Success.send(&mut self.stream));

        let mut channel = self.channel.lock().unwrap();

        channel.watchers.push(try!(self.stream.try_clone()));

        match channel.broadcaster.as_mut() {
            Some(mut broadcaster) => message::Notification::WatcherJoined("".to_owned())
                                         .send(&mut broadcaster)
                                         .unwrap(),
            None => (),
        }

        Ok(())
    }
}
<commit_msg>Remove closed stream from channel<commit_after>use std;
use std::collections::HashMap;
use std::{fmt, io, result};
use std::net::{Shutdown, TcpListener, TcpStream};
use std::sync::{Arc, Mutex};
use std::thread;
use std::sync::mpsc::{self, channel, Receiver, Sender};

use message;

#[derive(Debug)]
pub enum Error {
    NotificationSend(mpsc::SendError<Option<message::Notification>>),
    Message(message::Error),
    Io(io::Error),
}

impl std::error::Error for Error {
    fn description(&self) -> &str {
        "Server error"
    }

    fn cause(&self) -> Option<&std::error::Error> {
        match *self {
            Error::NotificationSend(ref err) => Some(err),
            Error::Message(ref err) => Some(err),
            Error::Io(ref err) => Some(err),
        }
    }
}

impl fmt::Display for Error {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        std::error::Error::description(self).fmt(f)
    }
}

impl From<message::Error> for Error {
    fn from(err: message::Error) -> Error {
        Error::Message(err)
    }
}

impl From<mpsc::SendError<Option<message::Notification>>> for Error {
    fn from(err: mpsc::SendError<Option<message::Notification>>) -> Error {
        Error::NotificationSend(err)
    }
}

impl From<io::Error> for Error {
    fn from(err: io::Error) -> Error {
        Error::Io(err)
    }
}

pub type Result<T> = result::Result<T, Error>;

pub fn execute(host: String, port: i32) -> Result<()> {
    let listener = try!(TcpListener::bind(&format!("{}:{}", host, port)[..]));
    let mut channels = Channels { channels: HashMap::new() };

    for stream in listener.incoming() {
        let _ = handle_client(stream, &mut channels);
    }

    Ok(())
}

fn handle_client(stream: result::Result<TcpStream, io::Error>, channels: &mut Channels) -> Result<()> {
    let stream = try!(stream);
    let request = try!(message::JoinRequest::receive(&stream));

    match request {
        message::JoinRequest::Broadcast(ch) =>
            BroadcastHandler::spawn(&stream, channels.fetch(&ch)),
        message::JoinRequest::Watch(ch) =>
            WatchHandler::spawn(&stream, channels.fetch(&ch)),
    }

    Ok(())
}

struct Channels {
    channels: HashMap<String, Arc<Mutex<Channel>>>,
}

struct Channel {
    broadcaster: Option<TcpStream>,
    watchers: Vec<TcpStream>,
}

struct BroadcastHandler {
    stream: TcpStream,
    channel: Arc<Mutex<Channel>>,
}

struct WatchHandler {
    stream: TcpStream,
    channel: Arc<Mutex<Channel>>,
}

impl Channels {
    fn fetch<K: AsRef<str>>(&mut self, key: &K) -> &Arc<Mutex<Channel>> {
        let ch = key.as_ref();

        if self.channels.contains_key(&ch[..]) {
            self.channels.get(&ch[..]).unwrap()
        } else {
            let channel = Channel {
                broadcaster: None,
                watchers: Vec::new(),
            };

            self.channels.insert(ch.to_owned(), Arc::new(Mutex::new(channel)));

            self.fetch(&ch)
        }
    }
}

impl Channel {
    fn takeover(&mut self, stream: TcpStream) -> Result<()> {
        match self.broadcaster.as_mut() {
            Some(mut current) => {
                try!(message::Notification::Closed("Broadcaster has changed".to_owned()).send(&mut current));
                try!(current.shutdown(Shutdown::Both));
            }
            None => (),
        }

        self.broadcaster = Some(stream);
        Ok(())
    }
}

impl BroadcastHandler {
    fn spawn(stream: &TcpStream, channel: &Arc<Mutex<Channel>>) {
        let mut handler = BroadcastHandler {
            stream: stream.try_clone().unwrap(),
            channel: channel.clone(),
        };

        let _ = channel.lock().and_then(|mut ch| {
            let _ = ch.takeover(stream.try_clone().unwrap());
            Ok(ch)
        });

        thread::spawn(move || {
            handler.process().unwrap_or_else(|e| {
                println!("{}", e);
            });

            handler.shutdown();
        });
    }

    fn process(&mut self) -> Result<()> {
        try!(message::JoinResponse::Success.send(&mut self.stream));

        let (sender, receiver) = channel();

        try!(self.spawn_relay(sender.clone()));
        try!(self.broadcast(receiver));

        Ok(())
    }

    fn spawn_relay(&self, sender: Sender<Option<message::Notification>>) -> Result<()> {
        let stream = try!(self.stream.try_clone());

        thread::spawn(move || -> Result<()> {
            loop {
                match message::Notification::receive(&stream) {
                    Ok(notification) => {
                        match notification {
                            message::Notification::Output(data)   => {
                                sender.send(Some(message::Notification::Output(data))).unwrap_or_else(|e| println!("{}", e));
                            },
                            message::Notification::Closed(reason) => {
                                try!(sender.send(Some(message::Notification::Closed(reason))));
                            },
                            _ => break,
                        }
                    },
                    Err(_) => break,
                }
            };

            try!(sender.send(None));

            Ok(())
        });

        Ok(())
    }

    fn broadcast(&self, receiver: Receiver<Option<message::Notification>>) -> Result<()> {
        for message in receiver {
            match message {
                Some(notification) => {
                    let mut channel = self.channel.lock().unwrap();

                    for mut watcher in channel.watchers.iter_mut() {
                        let _ = notification.send(watcher);
                    }
                },
                None => break
            }
        }

        Ok(())
    }

    fn shutdown(&mut self) {
        let mut channel = self.channel.lock().unwrap();

        let _ = self.stream.shutdown(Shutdown::Both);
        channel.broadcaster = None;
    }
}

impl WatchHandler {
    fn spawn(stream: &TcpStream, channel: &Arc<Mutex<Channel>>) {
        let mut handler = WatchHandler {
            stream: stream.try_clone().unwrap(),
            channel: channel.clone(),
        };

        thread::spawn(move || {
            handler.process().unwrap_or_else(|e| {
                println!("{}", e);
            });
        });
    }

    fn process(&mut self) -> Result<()> {
        try!(message::JoinResponse::Success.send(&mut self.stream));

        let mut channel = self.channel.lock().unwrap();

        channel.watchers.push(try!(self.stream.try_clone()));

        match channel.broadcaster.as_mut() {
            Some(mut broadcaster) => message::Notification::WatcherJoined("".to_owned())
                                         .send(&mut broadcaster)
                                         .unwrap(),
            None => (),
        }

        Ok(())
    }
}
<|endoftext|>