<commit_before>// TODO: Implement a reader and writer from the client / server relationship.
// Then try and get clients to talk to everyone -> specific clients.
// Once that's done, move on to formatting messages via tokio::Encode / Decode.
// Example TCP

#[allow(unused_imports)]
use std::net::SocketAddr;
use std::collections::{HashMap};
use std::rc::Rc;
use std::cell::RefCell;

use std::io;

use futures::AsyncSink;
use futures::sink::{Sink};
use futures::sync::*;

use futures::future::{Future};
use futures::stream::{Stream};

use tokio_core::io::Io;
use tokio_core::net::{TcpListener};
use tokio_core::reactor::{Core};

use client::TctClient;
use stanza::{Stanza, StanzaCodec, UserID};

pub struct TctServer {
    clients: Rc<RefCell<HashMap<UserID, mpsc::UnboundedSender<Stanza>>>>,
    #[allow(dead_code)]
    groups: Rc<RefCell<HashMap<UserID, mpsc::UnboundedSender<Stanza>>>>,
    channel: (mpsc::UnboundedSender<Stanza>, mpsc::UnboundedReceiver<Stanza>),
    core: Core,
    addr: SocketAddr,
}

// Ixnay on the private client, defeats the purpose and actually breaks the design
// in the first place lulz.

impl TctServer {
    /// Creates a new TctServer to be run.
    pub fn new(addr: SocketAddr) -> TctServer {
        let core = Core::new().unwrap();
        TctServer {
            // Odd way of doing this but yeah, core needs to be defined after
            // channel.
            channel: mpsc::unbounded(),
            core: core,
            addr: addr,
            groups: Rc::new(RefCell::new(HashMap::new())),
            clients: Rc::new(RefCell::new(HashMap::new()))
        }
    }

    /// Execute the server, runs in foreground for any application...
    /// This is a temporary design to get something off the ground.
    pub fn run(&mut self) {
        let socket =
            TcpListener::bind(&self.addr, &self.core.handle().clone()).unwrap();
        let handle = self.core.handle();
        let clients = self.clients.clone();
        let sender = self.channel.0.clone();

        println!("Chat service running at port : {}", self.addr.port());

        // For each incoming client connection at address 'addr'
        let server = socket.incoming().for_each(|(stream, addr)| {

            println!("Connected to client : {}", addr);

            let mut server_sender = sender.clone();

            // (writer, reader) == (sink, stream)
            // reader   -> sender
            // receiver -> writer
            let (mut writer, reader) = TctClient::new(stream, addr)
                .framed(StanzaCodec).split();
            let (sender, receiver) = mpsc::unbounded();

            let clients_inner = clients.clone();

            // Every message received over the stream, from client
            let reader = reader.into_future().map_err(|(err, _)| err).and_then(
                |(creds, stream)| {
                    if let Some(Stanza::LoginCredentials{ user, psw: _ }) = creds {
                        println!("User {} logged in!", user);
                        clients_inner.borrow_mut().insert(user, sender);
                        // TODO: System of verification that a user is in the db.
                    } else if let Some(Stanza::Register{ user, psw: _ }) = creds {
                        println!("New user {} logged in!", user);
                        clients_inner.borrow_mut().insert(user, sender);
                    } else {
                        println!("No login credentials supplied from client. \
                                  Closing stream...");
                        // By not giving resources for sender into HashMap,
                        // data does not persist for this client, connection
                        // will then time-out.
                    }
                    stream.for_each(move |msg: Stanza| {
                        println!("Read made for"); // TODO: Give addr here
                        if let Some(to) = msg.to() {
                            clients_inner.borrow_mut().get_mut(&to)
                                .unwrap_or(&mut server_sender) // TODO:
                                .send(msg)
                                .or_else(
                                    |err| Err(io::Error::new(
                                            io::ErrorKind::Other, err)))

                        } else { panic!("Client reported error") }
                    })
                }).map_err(|_| ());

            let receiver = receiver.for_each(move |msg| {
                println!("Writing message to {}", addr);
                let response = match writer.start_send(msg) { // handle it like 'send'
                    Ok(AsyncSink::Ready) => Ok(()),
                    Ok(AsyncSink::NotReady(_)) => panic!("failed to send"),
                    Err(_) => Err(())
                };
                writer.poll_complete().unwrap();
                response
            });

            //let clients = self.clients.clone();
            // TODO: 'select' combinator
            handle.spawn(receiver);
            handle.spawn(reader);

            Ok(())
        });
        self.core.run(server).unwrap();
    }
}

<commit_msg>Doesn't connect to no names!<commit_after>// TODO: Implement a reader and writer from the client / server relationship.
// Then try and get clients to talk to everyone -> specific clients.
// Once that's done, move on to formatting messages via tokio::Encode / Decode.
// Example TCP

#[allow(unused_imports)]
use std::net::SocketAddr;
use std::collections::{HashMap};
use std::rc::Rc;
use std::cell::RefCell;

use std::io;

use futures::AsyncSink;
use futures::sink::{Sink};
use futures::sync::*;

use futures::future::{Future};
use futures::stream::{Stream};

use tokio_core::io::Io;
use tokio_core::net::{TcpListener};
use tokio_core::reactor::{Core};

use client::TctClient;
use stanza::{Stanza, StanzaCodec, UserID};

pub struct TctServer {
    clients: Rc<RefCell<HashMap<UserID, mpsc::UnboundedSender<Stanza>>>>,
    #[allow(dead_code)]
    groups: Rc<RefCell<HashMap<UserID, mpsc::UnboundedSender<Stanza>>>>,
    channel: (mpsc::UnboundedSender<Stanza>, mpsc::UnboundedReceiver<Stanza>),
    core: Core,
    addr: SocketAddr,
}

// Ixnay on the private client, defeats the purpose and actually breaks the design
// in the first place lulz.

impl TctServer {
    /// Creates a new TctServer to be run.
    pub fn new(addr: SocketAddr) -> TctServer {
        let core = Core::new().unwrap();
        TctServer {
            // Odd way of doing this but yeah, core needs to be defined after
            // channel.
            channel: mpsc::unbounded(),
            core: core,
            addr: addr,
            groups: Rc::new(RefCell::new(HashMap::new())),
            clients: Rc::new(RefCell::new(HashMap::new()))
        }
    }

    /// Execute the server, runs in foreground for any application...
    /// This is a temporary design to get something off the ground.
    pub fn run(&mut self) {
        let socket =
            TcpListener::bind(&self.addr, &self.core.handle().clone()).unwrap();
        let handle = self.core.handle();
        let clients = self.clients.clone();
        let sender = self.channel.0.clone();

        println!("Chat service running at port : {}", self.addr.port());

        // For each incoming client connection at address 'addr'
        let server = socket.incoming().for_each(|(stream, addr)| {

            println!("Connected to client : {}", addr);

            let mut server_sender = sender.clone();

            // (writer, reader) == (sink, stream)
            // reader   -> sender
            // receiver -> writer
            let (mut writer, reader) = TctClient::new(stream, addr)
                .framed(StanzaCodec).split();
            let (sender, receiver) = mpsc::unbounded();

            let clients_inner = clients.clone();

            // Every message received over the stream, from client
            let reader = reader.into_future().map_err(|(err, _)| err).and_then(
                |(creds, stream)| {
                    if let Some(Stanza::LoginCredentials{ user, psw: _ }) = creds {
                        if user.len() == 0 {
                            println!("No username supplied. \
                                      Closing stream...");
                        } else {
                            println!("User {} logged in!", user);
                            clients_inner.borrow_mut().insert(user, sender);
                            // TODO: System of verification that a user is in the db.
                        }
                    } else if let Some(Stanza::Register{ user, psw: _ }) = creds {
                        println!("New user {} logged in!", user);
                        clients_inner.borrow_mut().insert(user, sender);
                    } else {
                        println!("No login credentials supplied from client. \
                                  Closing stream...");
                        // By not giving resources for sender into HashMap,
                        // data does not persist for this client, connection
                        // will then time-out.
                    }
                    stream.for_each(move |msg: Stanza| {
                        println!("Read made for"); // TODO: Give addr here
                        if let Some(to) = msg.to() {
                            clients_inner.borrow_mut().get_mut(&to)
                                .unwrap_or(&mut server_sender) // TODO:
                                .send(msg)
                                .or_else(
                                    |err| Err(io::Error::new(
                                            io::ErrorKind::Other, err)))

                        } else { panic!("Client reported error") }
                    })
                }).map_err(|_| ());

            let receiver = receiver.for_each(move |msg| {
                println!("Writing message to {}", addr);
                let response = match writer.start_send(msg) { // handle it like 'send'
                    Ok(AsyncSink::Ready) => Ok(()),
                    Ok(AsyncSink::NotReady(_)) => panic!("failed to send"),
                    Err(_) => Err(())
                };
                writer.poll_complete().unwrap();
                response
            });

            //let clients = self.clients.clone();
            // TODO: 'select' combinator
            handle.spawn(receiver);
            handle.spawn(reader);

            Ok(())
        });
        self.core.run(server).unwrap();
    }
}

<|endoftext|>


<commit_before>use std::marker::PhantomData;
use std::io::{Write, Read};
use std::sync::Mutex;
use std::sync::mpsc::{channel, Receiver, Sender};
use std::thread;
use std::net::SocketAddr;

use hyper::header::{ContentLength, ContentType};
use hyper::method::Method;
use hyper::server::{Server, Handler, Request, Response};
use hyper::uri::RequestUri::AbsolutePath;
use hyper::status::StatusCode;

use command::{WebDriverMessage, WebDriverCommand};
use error::{WebDriverResult, WebDriverError, ErrorStatus};
use httpapi::{WebDriverHttpApi, WebDriverExtensionRoute, VoidWebDriverExtensionRoute};
use response::WebDriverResponse;

enum DispatchMessage<U: WebDriverExtensionRoute> {
    HandleWebDriver(WebDriverMessage<U>, Sender<WebDriverResult<WebDriverResponse>>),
    Quit
}

#[derive(PartialEq, Clone)]
pub struct Session {
    id: String
}

impl Session {
    fn new(id: String) -> Session {
        Session {
            id: id
        }
    }
}

pub trait WebDriverHandler<U: WebDriverExtensionRoute=VoidWebDriverExtensionRoute> : Send {
    fn handle_command(&mut self, session: &Option<Session>, msg: &WebDriverMessage<U>) -> WebDriverResult<WebDriverResponse>;
    fn delete_session(&mut self, session: &Option<Session>);
}

struct Dispatcher<T: WebDriverHandler<U>,
                  U: WebDriverExtensionRoute> {
    handler: T,
    session: Option<Session>,
    extension_type: PhantomData<U>,
}

impl <T: WebDriverHandler<U>,
      U: WebDriverExtensionRoute> Dispatcher<T,U> {
    fn new(handler: T) -> Dispatcher<T,U> {
        Dispatcher {
            handler: handler,
            session: None,
            extension_type: PhantomData
        }
    }

    fn run(&mut self, msg_chan: Receiver<DispatchMessage<U>>) {
        loop {
            match msg_chan.recv() {
                Ok(DispatchMessage::HandleWebDriver(msg, resp_chan)) => {
                    let resp = match self.check_session(&msg) {
                        Ok(_) => self.handler.handle_command(&self.session, &msg),
                        Err(e) => Err(e)
                    };

                    match resp {
                        Ok(WebDriverResponse::NewSession(ref new_session)) => {
                            self.session = Some(Session::new(new_session.sessionId.clone()));
                        },
                        Ok(WebDriverResponse::DeleteSession) => {
                            self.delete_session();
                        },
                        Err(ref x) if x.delete_session() => {
                            self.delete_session();
                        }
                        _ => {}
                    }

                    if resp_chan.send(resp).is_err() {
                        error!("Sending response to the main thread failed");
                    };
                },
                Ok(DispatchMessage::Quit) => {
                    break;
                },
                Err(_) => panic!("Error receiving message in handler")
            }
        }
    }

    fn delete_session(&mut self) {
        debug!("Deleting session");
        self.handler.delete_session(&self.session);
        self.session = None;
    }

    fn check_session(&self, msg: &WebDriverMessage<U>) -> WebDriverResult<()> {
        match msg.session_id {
            Some(ref msg_session_id) => {
                match self.session {
                    Some(ref existing_session) => {
                        if existing_session.id != *msg_session_id {
                            Err(WebDriverError::new(
                                ErrorStatus::InvalidSessionId,
                                format!("Got unexpected session id {} expected {}",
                                        msg_session_id,
                                        existing_session.id)))
                        } else {
                            Ok(())
                        }
                    },
                    None => Ok(())
                }
            },
            None => {
                match self.session {
                    Some(_) => {
                        match msg.command {
                            WebDriverCommand::NewSession => {
                                Err(WebDriverError::new(
                                    ErrorStatus::UnsupportedOperation,
                                    "Session is already started"))
                            },
                            _ => {
                                //This should be impossible
                                error!("Got a message with no session id");
                                Err(WebDriverError::new(
                                    ErrorStatus::UnknownError,
                                    "Got a command with no session?!"))
                            }
                        }
                    },
                    None => {
                        match msg.command {
                            WebDriverCommand::NewSession => Ok(()),

                            _ => Err(WebDriverError::new(
                                ErrorStatus::InvalidSessionId,
                                "Tried to run a command before creating a session"))
                        }
                    }
                }
            }
        }
    }
}

struct HttpHandler<U: WebDriverExtensionRoute> {
    chan: Mutex<Sender<DispatchMessage<U>>>,
    api: Mutex<WebDriverHttpApi<U>>
}

impl <U: WebDriverExtensionRoute> HttpHandler<U> {
    fn new(api: WebDriverHttpApi<U>, chan: Sender<DispatchMessage<U>>) -> HttpHandler<U> {
        HttpHandler {
            chan: Mutex::new(chan),
            api: Mutex::new(api)
        }
    }
}

impl <U: WebDriverExtensionRoute> Handler for HttpHandler<U> {
    fn handle(&self, req: Request, res: Response) {
        let mut req = req;
        let mut res = res;

        let mut body = String::new();
        if let Method::Post = req.method {
            req.read_to_string(&mut body).unwrap();
        }
        debug!("Got request {} {:?}", req.method, req.uri);
        match req.uri {
            AbsolutePath(path) => {
                let msg_result = {
                    // The fact that this locks for basically the whole request doesn't
                    // matter as long as we are only handling one request at a time.
                    match self.api.lock() {
                        Ok(ref api) => {
                            api.decode_request(req.method, &path[..], &body[..])
                        },
                        Err(_) => return
                    }
                };
                let (status, resp_body) = match msg_result {
                    Ok(message) => {
                        let (send_res, recv_res) = channel();
                        match self.chan.lock() {
                            Ok(ref c) => {
                                let res = c.send(DispatchMessage::HandleWebDriver(message,
                                                                                send_res));
                                match res {
                                    Ok(x) => x,
                                    Err(_) => {
                                        error!("Something terrible happened");
                                        return
                                    }
                                }
                            },
                            Err(_) => {
                                error!("Something terrible happened");
                                return
                            }
                        }
                        match recv_res.recv() {
                            Ok(data) => match data {
                                Ok(response) => (StatusCode::Ok, response.to_json_string()),
                                Err(err) => (err.http_status(), err.to_json_string()),
                            },
                            Err(e) => panic!("Error reading response: {:?}", e)
                        }
                    },
                    Err(err) => {
                        (err.http_status(), err.to_json_string())
                    }
                };
                debug!("Returning status {:?}", status);
                debug!("Returning body {}", resp_body);
                {
                    let resp_status = res.status_mut();
                    *resp_status = status;
                }
                res.headers_mut().set(ContentLength(resp_body.len() as u64));
                res.headers_mut().set(ContentType::json());
                let mut stream = res.start().unwrap();
                stream.write_all(&resp_body.as_bytes()).unwrap();
                stream.end().unwrap();
            },
            _ => {}
        }
    }
}

pub fn start<T: 'static+WebDriverHandler<U>,
             U: 'static+WebDriverExtensionRoute>(address: SocketAddr,
                                                 handler: T,
                                                 extension_routes:Vec<(Method, &str, U)>) {
    let (msg_send, msg_recv) = channel();

    let api = WebDriverHttpApi::new(extension_routes);
    let http_handler = HttpHandler::new(api, msg_send);
    let server = Server::http(address).unwrap();

    let builder = thread::Builder::new().name("webdriver dispatcher".to_string());
    builder.spawn(move || {
        let mut dispatcher = Dispatcher::new(handler);
        dispatcher.run(msg_recv)
    }).unwrap();
    server.handle(http_handler).unwrap();
}
<commit_msg>Don't close stream after writing response.<commit_after>use std::marker::PhantomData;
use std::io::{Write, Read};
use std::sync::Mutex;
use std::sync::mpsc::{channel, Receiver, Sender};
use std::thread;
use std::net::SocketAddr;

use hyper::header::{ContentLength, ContentType};
use hyper::method::Method;
use hyper::server::{Server, Handler, Request, Response};
use hyper::uri::RequestUri::AbsolutePath;
use hyper::status::StatusCode;

use command::{WebDriverMessage, WebDriverCommand};
use error::{WebDriverResult, WebDriverError, ErrorStatus};
use httpapi::{WebDriverHttpApi, WebDriverExtensionRoute, VoidWebDriverExtensionRoute};
use response::WebDriverResponse;

enum DispatchMessage<U: WebDriverExtensionRoute> {
    HandleWebDriver(WebDriverMessage<U>, Sender<WebDriverResult<WebDriverResponse>>),
    Quit
}

#[derive(PartialEq, Clone)]
pub struct Session {
    id: String
}

impl Session {
    fn new(id: String) -> Session {
        Session {
            id: id
        }
    }
}

pub trait WebDriverHandler<U: WebDriverExtensionRoute=VoidWebDriverExtensionRoute> : Send {
    fn handle_command(&mut self, session: &Option<Session>, msg: &WebDriverMessage<U>) -> WebDriverResult<WebDriverResponse>;
    fn delete_session(&mut self, session: &Option<Session>);
}

struct Dispatcher<T: WebDriverHandler<U>,
                  U: WebDriverExtensionRoute> {
    handler: T,
    session: Option<Session>,
    extension_type: PhantomData<U>,
}

impl <T: WebDriverHandler<U>,
      U: WebDriverExtensionRoute> Dispatcher<T,U> {
    fn new(handler: T) -> Dispatcher<T,U> {
        Dispatcher {
            handler: handler,
            session: None,
            extension_type: PhantomData
        }
    }

    fn run(&mut self, msg_chan: Receiver<DispatchMessage<U>>) {
        loop {
            match msg_chan.recv() {
                Ok(DispatchMessage::HandleWebDriver(msg, resp_chan)) => {
                    let resp = match self.check_session(&msg) {
                        Ok(_) => self.handler.handle_command(&self.session, &msg),
                        Err(e) => Err(e)
                    };

                    match resp {
                        Ok(WebDriverResponse::NewSession(ref new_session)) => {
                            self.session = Some(Session::new(new_session.sessionId.clone()));
                        },
                        Ok(WebDriverResponse::DeleteSession) => {
                            self.delete_session();
                        },
                        Err(ref x) if x.delete_session() => {
                            self.delete_session();
                        }
                        _ => {}
                    }

                    if resp_chan.send(resp).is_err() {
                        error!("Sending response to the main thread failed");
                    };
                },
                Ok(DispatchMessage::Quit) => {
                    break;
                },
                Err(_) => panic!("Error receiving message in handler")
            }
        }
    }

    fn delete_session(&mut self) {
        debug!("Deleting session");
        self.handler.delete_session(&self.session);
        self.session = None;
    }

    fn check_session(&self, msg: &WebDriverMessage<U>) -> WebDriverResult<()> {
        match msg.session_id {
            Some(ref msg_session_id) => {
                match self.session {
                    Some(ref existing_session) => {
                        if existing_session.id != *msg_session_id {
                            Err(WebDriverError::new(
                                ErrorStatus::InvalidSessionId,
                                format!("Got unexpected session id {} expected {}",
                                        msg_session_id,
                                        existing_session.id)))
                        } else {
                            Ok(())
                        }
                    },
                    None => Ok(())
                }
            },
            None => {
                match self.session {
                    Some(_) => {
                        match msg.command {
                            WebDriverCommand::NewSession => {
                                Err(WebDriverError::new(
                                    ErrorStatus::UnsupportedOperation,
                                    "Session is already started"))
                            },
                            _ => {
                                //This should be impossible
                                error!("Got a message with no session id");
                                Err(WebDriverError::new(
                                    ErrorStatus::UnknownError,
                                    "Got a command with no session?!"))
                            }
                        }
                    },
                    None => {
                        match msg.command {
                            WebDriverCommand::NewSession => Ok(()),

                            _ => Err(WebDriverError::new(
                                ErrorStatus::InvalidSessionId,
                                "Tried to run a command before creating a session"))
                        }
                    }
                }
            }
        }
    }
}

struct HttpHandler<U: WebDriverExtensionRoute> {
    chan: Mutex<Sender<DispatchMessage<U>>>,
    api: Mutex<WebDriverHttpApi<U>>
}

impl <U: WebDriverExtensionRoute> HttpHandler<U> {
    fn new(api: WebDriverHttpApi<U>, chan: Sender<DispatchMessage<U>>) -> HttpHandler<U> {
        HttpHandler {
            chan: Mutex::new(chan),
            api: Mutex::new(api)
        }
    }
}

impl <U: WebDriverExtensionRoute> Handler for HttpHandler<U> {
    fn handle(&self, req: Request, res: Response) {
        let mut req = req;
        let mut res = res;

        let mut body = String::new();
        if let Method::Post = req.method {
            req.read_to_string(&mut body).unwrap();
        }
        debug!("Got request {} {:?}", req.method, req.uri);
        match req.uri {
            AbsolutePath(path) => {
                let msg_result = {
                    // The fact that this locks for basically the whole request doesn't
                    // matter as long as we are only handling one request at a time.
                    match self.api.lock() {
                        Ok(ref api) => {
                            api.decode_request(req.method, &path[..], &body[..])
                        },
                        Err(_) => return
                    }
                };
                let (status, resp_body) = match msg_result {
                    Ok(message) => {
                        let (send_res, recv_res) = channel();
                        match self.chan.lock() {
                            Ok(ref c) => {
                                let res = c.send(DispatchMessage::HandleWebDriver(message,
                                                                                send_res));
                                match res {
                                    Ok(x) => x,
                                    Err(_) => {
                                        error!("Something terrible happened");
                                        return
                                    }
                                }
                            },
                            Err(_) => {
                                error!("Something terrible happened");
                                return
                            }
                        }
                        match recv_res.recv() {
                            Ok(data) => match data {
                                Ok(response) => (StatusCode::Ok, response.to_json_string()),
                                Err(err) => (err.http_status(), err.to_json_string()),
                            },
                            Err(e) => panic!("Error reading response: {:?}", e)
                        }
                    },
                    Err(err) => {
                        (err.http_status(), err.to_json_string())
                    }
                };
                debug!("Returning status {:?}", status);
                debug!("Returning body {}", resp_body);
                {
                    let resp_status = res.status_mut();
                    *resp_status = status;
                }
                res.headers_mut().set(ContentLength(resp_body.len() as u64));
                res.headers_mut().set(ContentType::json());
                let mut stream = res.start().unwrap();
                stream.write_all(&resp_body.as_bytes()).unwrap();
            },
            _ => {}
        }
    }
}

pub fn start<T: 'static+WebDriverHandler<U>,
             U: 'static+WebDriverExtensionRoute>(address: SocketAddr,
                                                 handler: T,
                                                 extension_routes:Vec<(Method, &str, U)>) {
    let (msg_send, msg_recv) = channel();

    let api = WebDriverHttpApi::new(extension_routes);
    let http_handler = HttpHandler::new(api, msg_send);
    let server = Server::http(address).unwrap();

    let builder = thread::Builder::new().name("webdriver dispatcher".to_string());
    builder.spawn(move || {
        let mut dispatcher = Dispatcher::new(handler);
        dispatcher.run(msg_recv)
    }).unwrap();
    server.handle(http_handler).unwrap();
}
<|endoftext|>


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

use chasher::player_hash;
use common::*;
use game::Game;
use serde::ser::Serialize;
use serde_json;
use std::collections::{HashMap, HashSet};
use std::io::ErrorKind;
use std::net::{TcpListener, TcpStream};
use std::sync::mpsc::{Receiver, Sender, channel};
use std::thread;
use std::time::Duration;
use tungstenite::{Error, Message, WebSocket, accept};
use tungstenite::protocol::Role;

#[derive(Serialize, Deserialize, Debug)]
pub struct ServerConfig
{
    tick_time_ms: u64,
    server_port: u16,
    game_start_ticks: u64,
    game_turns: u64,
}

#[derive(PartialEq, Eq, Hash, Clone)]
enum ConnectionType
{
    Viewer,
    Player(u64),
    Unknown,
}

struct Connection
{
    role: ConnectionType,
    socket: WebSocket<TcpStream>,
}

enum MessageResponse
{
    Mail(u64, Message),
    Disconnected(u64),
}

pub struct Server
{
    connections: HashMap<u64, Connection>,
    incoming_connections: Receiver<WebSocket<TcpStream>>,
    incoming_messages: Receiver<MessageResponse>,
    message_sender: Sender<MessageResponse>,
    queue: HashMap<u64, Player>,
    seed: u64,
    tick_time: u64,
    game_id: u64,
    game_start_ticks: u64,
    game_turns: u64,
    game_start_ticks_left: u64,
    game: Option<Game>,
}

impl Drop for Connection
{
    fn drop(&mut self)
    {
        let _ = self.socket.close();
    }
}

impl Connection
{
    #[cfg_attr(rustfmt, rustfmt_skip)]
    pub fn new(ws: WebSocket<TcpStream>, sender: Sender<MessageResponse>, seed: u64) -> Connection
    {
        let mut ws = ws;
        let conn = Connection {
            socket: WebSocket::from_raw_socket(ws.get_ref().try_clone().unwrap(), Role::Server),
            role: ConnectionType::Unknown,
        };
        thread::spawn(move || loop
        {
            match ws.read_message()
            {
                Ok(x) => sender.send(MessageResponse::Mail(seed, x)).unwrap(),
                Err(_) =>
                {
                    sender.send(MessageResponse::Disconnected(seed)).unwrap();
                    break;
                }
            }
        });
        conn
    }
}

impl Server
{
    pub fn new(config: ServerConfig) -> Server
    {
        let (sen, rec) = channel();
        let server_port = config.server_port;
        thread::spawn(move || {
            let listener = TcpListener::bind(("0.0.0.0", server_port)).unwrap();
            for stream in listener.incoming()
            {
                match stream
                {
                    Ok(s) =>
                    {
                        match accept(s)
                        {
                            Ok(w) => sen.send(w).unwrap(),
                            Err(_) => error!("Error while accepting WebSocket connection"),
                        }
                    }
                    Err(why) => error!("Error while accepting connection: {}", why),
                }
            }
        });
        let (sen_msg, rec_msg) = channel();
        Server {
            connections: Default::default(),
            incoming_connections: rec,
            incoming_messages: rec_msg,
            message_sender: sen_msg,
            seed: 0,
            tick_time: config.tick_time_ms,
            game_id: 0,
            queue: Default::default(),
            game_start_ticks: config.game_start_ticks,
            game_turns: config.game_turns,
            game_start_ticks_left: config.game_turns,
            game: None,
        }
    }

    pub fn send_overview(&mut self)
    {
        if self.game.is_none()
        {
            error!("Trying to generate overview while there is no game");
            return;
        }
        let overview: Overview;
        let mut gaming_ids: HashSet<u64> = Default::default();
        {
            let game = self.game.as_ref().unwrap();
            let mut tokens = game.tokens.clone();
            let mut players = game.players.clone();
            let tokens = tokens.drain().collect();
            let players = players.drain().map(|(_, x)| x).collect();
            overview = Overview {
                game_id: game.game_id,
                turns_left: game.turns_left,
                ms_for_turn: self.tick_time,
                grid: game.grid.clone(),
                tokens: tokens,
                players: players,
            };
            for v in game.players.keys()
            {
                gaming_ids.insert(*v);
            }
        }
        let mut conn_info: HashSet<(u64, ConnectionType)> = Default::default();
        for (id, conn) in &self.connections
        {
            conn_info.insert((*id, conn.role.clone()));
        }
        for (id, role) in conn_info.drain()
        {
            match role
            {
                ConnectionType::Viewer => self.send_data(id, &overview),
                ConnectionType::Player(user_game_id) if gaming_ids.contains(&user_game_id) => self.send_data(id, &overview),
                _ => continue,
            }
        }
    }

    pub fn main(&mut self)
    {
        while let Ok(x) = self.incoming_connections.try_recv()
        {
            self.connections.insert(self.seed, Connection::new(x, self.message_sender.clone(), self.seed));
            info!("Accepted connection #{}", self.seed);
            self.seed += 1;
        }
        while let Ok(x) = self.incoming_messages.try_recv()
        {
            if let MessageResponse::Disconnected(id) = x
            {
                self.queue.remove(&id);
                if self.game.is_some()
                {
                    self.game
                        .as_mut()
                        .unwrap()
                        .players
                        .remove(&id);
                }
                self.connections.remove(&id);
                info!("Closed connection #{}", id);
            }
            else if let MessageResponse::Mail(id, msg) = x
            {
                let msg = match msg.into_text()
                {
                    Ok(s) => s,
                    Err(why) =>
                    {
                        error!("Received garbage: {}", why);
                        continue;
                    }
                };
                let msg: ClientMessage = match serde_json::from_str(msg.as_str())
                {
                    Ok(s) => s,
                    Err(why) =>
                    {
                        error!("Received garbage: {}", why);
                        continue;
                    }
                };
                if let ClientMessage::HandShake(role) = msg
                {
                    self.handle_accept(id, role, false);
                }
                else if let ClientMessage::Command(command) = msg
                {
                    self.handle_command(id, command);
                }
                else
                {
                    unreachable!();
                }
            }
            else
            {
                unreachable!();
            }
        }
        if self.game.is_none() && self.game_start_ticks_left > 0 && self.queue.len() >= 2
        {
            info!("Game staring in {} ticks", self.game_start_ticks_left);
            self.game_start_ticks_left -= 1;
        }
        else if self.game.is_none() && self.game_start_ticks_left == 0 && self.queue.len() >= 2
        {
            info!("Game started with {} players", self.queue.len());
            let mut players: HashMap<u64, Player> = Default::default();
            for (player_id, player_value) in self.queue.drain()
            {
                players.insert(player_id, player_value);
            }
            self.game = Some(Game::new(players, self.game_turns, self.game_id));
            self.game
                .as_mut()
                .unwrap()
                .tick();
            self.send_overview();
            self.game_id += 1;
        }
        else if self.game.is_some() &&
                  !self.game
                       .as_ref()
                       .unwrap()
                       .finished()
        {
            self.game
                .as_mut()
                .unwrap()
                .tick();
            self.send_overview();
        }
        else if self.game.is_some() &&
                  self.game
                      .as_ref()
                      .unwrap()
                      .finished()
        {
            info!("Game ended!");
            let mut to_readd: Vec<(u64, ClientRole)> = Vec::new();
            for (key, player) in &self.game
                                      .as_ref()
                                      .unwrap()
                                      .players
            {
                to_readd.push((*key, ClientRole::Player(PlayerInfo { name: player.name.clone() })));
            }
            self.game = None;
            self.game_start_ticks_left = self.game_start_ticks;
            for player in to_readd
            {
                self.handle_accept(player.0, player.1, true);
            }
        }
        thread::sleep(Duration::from_millis(self.tick_time));
    }

    #[cfg_attr(feature = "cargo-clippy", allow(map_entry))]
    fn handle_accept(&mut self, id: u64, role: ClientRole, not_interactive: bool)
    {
        match role
        {
            ClientRole::Viewer =>
            {
                {
                    let v = self.connections.get_mut(&id);
                    if v.is_none()
                    {
                        return;
                    }
                    let v = v.unwrap();
                    v.role = ConnectionType::Viewer;
                }
                if !not_interactive
                {
                    self.send_data(id, &ServerResponse::Ok);
                    info!("Viewer connected");
                }
            }
            ClientRole::Player(info) =>
            {
                let user_game_id = player_hash(&info);
                if self.queue.contains_key(&user_game_id)
                {
                    if !not_interactive
                    {
                        self.send_data(id, &ServerResponse::Error(String::from("Username already taken")));
                    }
                }
                else
                {
                    let info_name = String::from(info.name.trim());
                    {
                        let conn = self.connections.get_mut(&id);
                        if conn.is_none()
                        {
                            error!("Trying to register a non-existent WebSocket as player!");
                            return;
                        }
                        let conn = conn.unwrap();
                        conn.role = ConnectionType::Player(user_game_id);
                        self.queue.insert(user_game_id,
                                          Player {
                                              name: info_name.clone(),
                                              id: user_game_id,
                                              points: 0,
                                              position: Point { x: 0, y: 0 },
                                          });
                        if self.game.is_none()
                        {
                            self.game_start_ticks_left = self.game_start_ticks;
                        }
                    }
                    if !not_interactive
                    {
                        info!("Player connected: {}", info_name);
                        self.send_data(id, &ServerResponse::Ok);
                    }
                }
            }
        };
    }

    fn handle_command(&mut self, id: u64, command: ClientCommand)
    {
        if !self.connections.contains_key(&id)
        {
            error!("Wrong ID is trying to execute command");
            return;
        }
        if self.game.is_none()
        {
            self.send_data(id, &ServerResponse::Error(String::from("No active game")));
            return;
        }
        match self.connections[&id].role
        {
            ConnectionType::Player(user_game_id) =>
            {
                match self.game
                          .as_mut()
                          .unwrap()
                          .action(user_game_id, command)
                {
                    Ok(_) => self.send_data(id, &ServerResponse::Ok),
                    Err(s) => self.send_data(id, &ServerResponse::Error(s)),
                }
            }
            _ =>
            {
                self.send_data(id, &ServerResponse::Error(String::from("Operation not allowed")));
            }
        }
    }

    fn send_data<T: ?Sized + Serialize>(&mut self, id: u64, value: &T)
    {
        match serde_json::to_string(value)
        {
            Ok(s) =>
            {
                match self.connections.get_mut(&id)
                {
                    Some(conn) =>
                    {
                        match conn.socket.write_message(Message::text(s))
                        {
                            Ok(_) =>
                            {}
                            Err(Error::Io(x)) =>
                            {
                                match x.kind()
                                {
                                    ErrorKind::WouldBlock =>
                                    {
                                        match conn.socket.write_pending()
                                        {
                                            Ok(_) =>
                                            {}
                                            Err(why) => error!("Unable to flush WebSocket {}", why),
                                        }
                                    }
                                    _ => error!("I/O error"),
                                }
                            }
                            Err(Error::ConnectionClosed) => self.message_sender.send(MessageResponse::Disconnected(id)).unwrap(),
                            Err(_) => error!("Error while sending data"),
                        }
                    }
                    None => error!("Missing id"),
                }
            }
            Err(_) => unreachable!(),
        };
    }
}

impl ServerConfig
{
    pub fn new() -> ServerConfig
    {
        ServerConfig {
            tick_time_ms: 500,
            server_port: 42000,
            game_start_ticks: 60,
            game_turns: 300,
        }
    }
}
<commit_msg>Sanitaze usernames<commit_after>/*
 * This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at http://mozilla.org/MPL/2.0/.
 *
 * Copyright 2017 - Dario Ostuni <dario.ostuni@gmail.com>
 *
 */

use chasher::player_hash;
use common::*;
use game::Game;
use serde::ser::Serialize;
use serde_json;
use std::collections::{HashMap, HashSet};
use std::io::ErrorKind;
use std::net::{TcpListener, TcpStream};
use std::sync::mpsc::{Receiver, Sender, channel};
use std::thread;
use std::time::Duration;
use tungstenite::{Error, Message, WebSocket, accept};
use tungstenite::protocol::Role;

#[derive(Serialize, Deserialize, Debug)]
pub struct ServerConfig
{
    tick_time_ms: u64,
    server_port: u16,
    game_start_ticks: u64,
    game_turns: u64,
}

#[derive(PartialEq, Eq, Hash, Clone)]
enum ConnectionType
{
    Viewer,
    Player(u64),
    Unknown,
}

struct Connection
{
    role: ConnectionType,
    socket: WebSocket<TcpStream>,
}

enum MessageResponse
{
    Mail(u64, Message),
    Disconnected(u64),
}

pub struct Server
{
    connections: HashMap<u64, Connection>,
    incoming_connections: Receiver<WebSocket<TcpStream>>,
    incoming_messages: Receiver<MessageResponse>,
    message_sender: Sender<MessageResponse>,
    queue: HashMap<u64, Player>,
    seed: u64,
    tick_time: u64,
    game_id: u64,
    game_start_ticks: u64,
    game_turns: u64,
    game_start_ticks_left: u64,
    game: Option<Game>,
}

impl Drop for Connection
{
    fn drop(&mut self)
    {
        let _ = self.socket.close();
    }
}

impl Connection
{
    #[cfg_attr(rustfmt, rustfmt_skip)]
    pub fn new(ws: WebSocket<TcpStream>, sender: Sender<MessageResponse>, seed: u64) -> Connection
    {
        let mut ws = ws;
        let conn = Connection {
            socket: WebSocket::from_raw_socket(ws.get_ref().try_clone().unwrap(), Role::Server),
            role: ConnectionType::Unknown,
        };
        thread::spawn(move || loop
        {
            match ws.read_message()
            {
                Ok(x) => sender.send(MessageResponse::Mail(seed, x)).unwrap(),
                Err(_) =>
                {
                    sender.send(MessageResponse::Disconnected(seed)).unwrap();
                    break;
                }
            }
        });
        conn
    }
}

impl Server
{
    pub fn new(config: ServerConfig) -> Server
    {
        let (sen, rec) = channel();
        let server_port = config.server_port;
        thread::spawn(move || {
            let listener = TcpListener::bind(("0.0.0.0", server_port)).unwrap();
            for stream in listener.incoming()
            {
                match stream
                {
                    Ok(s) =>
                    {
                        match accept(s)
                        {
                            Ok(w) => sen.send(w).unwrap(),
                            Err(_) => error!("Error while accepting WebSocket connection"),
                        }
                    }
                    Err(why) => error!("Error while accepting connection: {}", why),
                }
            }
        });
        let (sen_msg, rec_msg) = channel();
        Server {
            connections: Default::default(),
            incoming_connections: rec,
            incoming_messages: rec_msg,
            message_sender: sen_msg,
            seed: 0,
            tick_time: config.tick_time_ms,
            game_id: 0,
            queue: Default::default(),
            game_start_ticks: config.game_start_ticks,
            game_turns: config.game_turns,
            game_start_ticks_left: config.game_turns,
            game: None,
        }
    }

    pub fn send_overview(&mut self)
    {
        if self.game.is_none()
        {
            error!("Trying to generate overview while there is no game");
            return;
        }
        let overview: Overview;
        let mut gaming_ids: HashSet<u64> = Default::default();
        {
            let game = self.game.as_ref().unwrap();
            let mut tokens = game.tokens.clone();
            let mut players = game.players.clone();
            let tokens = tokens.drain().collect();
            let players = players.drain().map(|(_, x)| x).collect();
            overview = Overview {
                game_id: game.game_id,
                turns_left: game.turns_left,
                ms_for_turn: self.tick_time,
                grid: game.grid.clone(),
                tokens: tokens,
                players: players,
            };
            for v in game.players.keys()
            {
                gaming_ids.insert(*v);
            }
        }
        let mut conn_info: HashSet<(u64, ConnectionType)> = Default::default();
        for (id, conn) in &self.connections
        {
            conn_info.insert((*id, conn.role.clone()));
        }
        for (id, role) in conn_info.drain()
        {
            match role
            {
                ConnectionType::Viewer => self.send_data(id, &overview),
                ConnectionType::Player(user_game_id) if gaming_ids.contains(&user_game_id) => self.send_data(id, &overview),
                _ => continue,
            }
        }
    }

    pub fn main(&mut self)
    {
        while let Ok(x) = self.incoming_connections.try_recv()
        {
            self.connections.insert(self.seed, Connection::new(x, self.message_sender.clone(), self.seed));
            info!("Accepted connection #{}", self.seed);
            self.seed += 1;
        }
        while let Ok(x) = self.incoming_messages.try_recv()
        {
            if let MessageResponse::Disconnected(id) = x
            {
                self.queue.remove(&id);
                if self.game.is_some()
                {
                    self.game
                        .as_mut()
                        .unwrap()
                        .players
                        .remove(&id);
                }
                self.connections.remove(&id);
                info!("Closed connection #{}", id);
            }
            else if let MessageResponse::Mail(id, msg) = x
            {
                let msg = match msg.into_text()
                {
                    Ok(s) => s,
                    Err(why) =>
                    {
                        error!("Received garbage: {}", why);
                        continue;
                    }
                };
                let msg: ClientMessage = match serde_json::from_str(msg.as_str())
                {
                    Ok(s) => s,
                    Err(why) =>
                    {
                        error!("Received garbage: {}", why);
                        continue;
                    }
                };
                if let ClientMessage::HandShake(role) = msg
                {
                    self.handle_accept(id, role, false);
                }
                else if let ClientMessage::Command(command) = msg
                {
                    self.handle_command(id, command);
                }
                else
                {
                    unreachable!();
                }
            }
            else
            {
                unreachable!();
            }
        }
        if self.game.is_none() && self.game_start_ticks_left > 0 && self.queue.len() >= 2
        {
            info!("Game staring in {} ticks", self.game_start_ticks_left);
            self.game_start_ticks_left -= 1;
        }
        else if self.game.is_none() && self.game_start_ticks_left == 0 && self.queue.len() >= 2
        {
            info!("Game started with {} players", self.queue.len());
            let mut players: HashMap<u64, Player> = Default::default();
            for (player_id, player_value) in self.queue.drain()
            {
                players.insert(player_id, player_value);
            }
            self.game = Some(Game::new(players, self.game_turns, self.game_id));
            self.game
                .as_mut()
                .unwrap()
                .tick();
            self.send_overview();
            self.game_id += 1;
        }
        else if self.game.is_some() &&
                  !self.game
                       .as_ref()
                       .unwrap()
                       .finished()
        {
            self.game
                .as_mut()
                .unwrap()
                .tick();
            self.send_overview();
        }
        else if self.game.is_some() &&
                  self.game
                      .as_ref()
                      .unwrap()
                      .finished()
        {
            info!("Game ended!");
            let mut to_readd: Vec<(u64, ClientRole)> = Vec::new();
            for (key, player) in &self.game
                                      .as_ref()
                                      .unwrap()
                                      .players
            {
                to_readd.push((*key, ClientRole::Player(PlayerInfo { name: player.name.clone() })));
            }
            self.game = None;
            self.game_start_ticks_left = self.game_start_ticks;
            for player in to_readd
            {
                self.handle_accept(player.0, player.1, true);
            }
        }
        thread::sleep(Duration::from_millis(self.tick_time));
    }

    #[cfg_attr(feature = "cargo-clippy", allow(map_entry))]
    fn handle_accept(&mut self, id: u64, role: ClientRole, not_interactive: bool)
    {
        match role
        {
            ClientRole::Viewer =>
            {
                {
                    let v = self.connections.get_mut(&id);
                    if v.is_none()
                    {
                        return;
                    }
                    let v = v.unwrap();
                    v.role = ConnectionType::Viewer;
                }
                if !not_interactive
                {
                    self.send_data(id, &ServerResponse::Ok);
                    info!("Viewer connected");
                }
            }
            ClientRole::Player(info) =>
            {
                let mut info = info;
                info.name = String::from(info.name.trim());
                let user_game_id = player_hash(&info);
                if self.queue.contains_key(&user_game_id) || "\n\r\t ".chars().any(|x| info.name.contains(x))
                {
                    if !not_interactive
                    {
                        self.send_data(id, &ServerResponse::Error(String::from("Username already taken")));
                    }
                }
                else
                {
                    {
                        let conn = self.connections.get_mut(&id);
                        if conn.is_none()
                        {
                            error!("Trying to register a non-existent WebSocket as player!");
                            return;
                        }
                        let conn = conn.unwrap();
                        conn.role = ConnectionType::Player(user_game_id);
                        self.queue.insert(user_game_id,
                                          Player {
                                              name: info.name.clone(),
                                              id: user_game_id,
                                              points: 0,
                                              position: Point { x: 0, y: 0 },
                                          });
                        if self.game.is_none()
                        {
                            self.game_start_ticks_left = self.game_start_ticks;
                        }
                    }
                    if !not_interactive
                    {
                        info!("Player connected: {}", info.name);
                        self.send_data(id, &ServerResponse::Ok);
                    }
                }
            }
        };
    }

    fn handle_command(&mut self, id: u64, command: ClientCommand)
    {
        if !self.connections.contains_key(&id)
        {
            error!("Wrong ID is trying to execute command");
            return;
        }
        if self.game.is_none()
        {
            self.send_data(id, &ServerResponse::Error(String::from("No active game")));
            return;
        }
        match self.connections[&id].role
        {
            ConnectionType::Player(user_game_id) =>
            {
                match self.game
                          .as_mut()
                          .unwrap()
                          .action(user_game_id, command)
                {
                    Ok(_) => self.send_data(id, &ServerResponse::Ok),
                    Err(s) => self.send_data(id, &ServerResponse::Error(s)),
                }
            }
            _ =>
            {
                self.send_data(id, &ServerResponse::Error(String::from("Operation not allowed")));
            }
        }
    }

    fn send_data<T: ?Sized + Serialize>(&mut self, id: u64, value: &T)
    {
        match serde_json::to_string(value)
        {
            Ok(s) =>
            {
                match self.connections.get_mut(&id)
                {
                    Some(conn) =>
                    {
                        match conn.socket.write_message(Message::text(s))
                        {
                            Ok(_) =>
                            {}
                            Err(Error::Io(x)) =>
                            {
                                match x.kind()
                                {
                                    ErrorKind::WouldBlock =>
                                    {
                                        match conn.socket.write_pending()
                                        {
                                            Ok(_) =>
                                            {}
                                            Err(why) => error!("Unable to flush WebSocket {}", why),
                                        }
                                    }
                                    _ => error!("I/O error"),
                                }
                            }
                            Err(Error::ConnectionClosed) => self.message_sender.send(MessageResponse::Disconnected(id)).unwrap(),
                            Err(_) => error!("Error while sending data"),
                        }
                    }
                    None => error!("Missing id"),
                }
            }
            Err(_) => unreachable!(),
        };
    }
}

impl ServerConfig
{
    pub fn new() -> ServerConfig
    {
        ServerConfig {
            tick_time_ms: 500,
            server_port: 42000,
            game_start_ticks: 60,
            game_turns: 300,
        }
    }
}
<|endoftext|>


<commit_before>#![feature(test)]
extern crate test;
extern crate ndarray;

use ndarray::Array;

#[test]
fn char_array()
{
    // test compilation & basics of non-numerical array
    let cc = Array::from_iter("alphabet".chars()).reshape((4, 2));
    assert!(cc.subview(1, 0) == Array::from_iter("apae".chars()));
}

#[bench]
fn time_matmul(b: &mut test::Bencher)
{
    b.iter(|| {
        let mut a: Array<usize, _> = Array::zeros((2, 3));
        for (i, elt) in a.iter_mut().enumerate() {
            *elt = i;
        }

        let mut b: Array<usize, _> = Array::zeros((3, 4));
        for (i, elt) in b.iter_mut().enumerate() {
            *elt = i;
        }

        let c = a.mat_mul(&b);
        unsafe {
            let result = Array::from_vec_dim((2, 4), vec![20, 23, 26, 29, 56, 68, 80, 92]);
            assert!(c == result);
        }
    })
}
<commit_msg>Remove feat test from tests<commit_after>extern crate ndarray;

use ndarray::Array;

#[test]
fn char_array()
{
    // test compilation & basics of non-numerical array
    let cc = Array::from_iter("alphabet".chars()).reshape((4, 2));
    assert!(cc.subview(1, 0) == Array::from_iter("apae".chars()));
}
<|endoftext|>


<commit_before>/* Any copyright is dedicated to the Public Domain.
 * http://creativecommons.org/publicdomain/zero/1.0/ */

#[macro_use]
extern crate autograd;
use autograd::Context;

#[test]
fn negative() {
    let context = new_autograd_context!(f32, 1000);
    let x = context.new_variable(1.5);

    let y = -x;
    assert_eq!(y.value, -1.5);

    context.differentiate(y);
    assert_eq!(context.get_derivative(x), -1.);
    assert_eq!(context.get_derivative(y), 1.);
}

#[test]
fn addition() {
    let context = new_autograd_context!(f32, 1000);
    let x1 = context.new_variable(1.5);
    let x2 = context.new_variable(2.5);

    let y = x1 + x2;
    assert_eq!(y.value, 4.);

    context.differentiate(y);
    assert_eq!(context.get_derivative(x1), 1.);
    assert_eq!(context.get_derivative(x2), 1.);
    assert_eq!(context.get_derivative(y), 1.);
}

#[test]
fn multiplication() {
    let context = new_autograd_context!(f32, 1000);
    let x1 = context.new_variable(1.5);
    let x2 = context.new_variable(2.5);

    let y = x1 * x2;
    assert_eq!(y.value, 3.75);

    context.differentiate(y);
    assert_eq!(context.get_derivative(x1), 2.5);
    assert_eq!(context.get_derivative(x2), 1.5);
    assert_eq!(context.get_derivative(y), 1.);
}

#[test]
fn single_thread_multiple_run() {
    for _ in 0..10 {
        addition();
    }
}

#[test]
fn multi_thread_multiple_run() {
    // TODO Insert purposeful sleep for each task?
    let task_pool = std::sync::TaskPool::new(10);
    let semaphore = std::sync::Arc::new(std::sync::Semaphore::new(-999));
    for _ in 0..1000 {
        let semaphore_task = semaphore.clone();
        task_pool.execute(move || {
            single_thread_multiple_run();
            semaphore_task.release();
        });
    }
    semaphore.access();
}

#[test]
#[should_fail(expected = "This Context instance is in use now. Note that a Context instance is allowed per construction location and per thread. Consequently, it cannot be recursively constructed unless it is destructed. This is a limitation caused by the thread local static variables usages in the current implementation.")]
fn context_lock() {
    fn recursive_context(n: usize) {
        if n > 0 {
            let context = new_autograd_context!(f32, 10);
            recursive_context(n - 1);
        }
    }
    recursive_context(2);
}

#[test]
fn context_capacity_expression_argument() {
    let a = 200;
    let b = 10;
    let context = new_autograd_context!(f32, a/b);
}
<commit_msg>Add unary_operation_test macro.<commit_after>/* Any copyright is dedicated to the Public Domain.
 * http://creativecommons.org/publicdomain/zero/1.0/ */

#[macro_use]
extern crate autograd;
use autograd::Context;
use std::num::Float;
use std::ops::Neg;

macro_rules! unary_operation_test {
    ($name:ident, $x:expr, $y:expr, $dx:expr) => (
        #[test]
        fn $name() {
            let context = new_autograd_context!(f32, 1000);
            let x = context.new_variable($x);

            let y = x.$name();
            assert_eq!(y.value, $y);

            context.differentiate(y);
            assert_eq!(context.get_derivative(x), $dx);
            assert_eq!(context.get_derivative(y), 1.);
        }
    )
}

unary_operation_test!(cos, 0., 1., 0.);
unary_operation_test!(neg, 1.5, -1.5, -1.);

#[test]
fn addition() {
    let context = new_autograd_context!(f32, 1000);
    let x1 = context.new_variable(1.5);
    let x2 = context.new_variable(2.5);

    let y = x1 + x2;
    assert_eq!(y.value, 4.);

    context.differentiate(y);
    assert_eq!(context.get_derivative(x1), 1.);
    assert_eq!(context.get_derivative(x2), 1.);
    assert_eq!(context.get_derivative(y), 1.);
}

#[test]
fn multiplication() {
    let context = new_autograd_context!(f32, 1000);
    let x1 = context.new_variable(1.5);
    let x2 = context.new_variable(2.5);

    let y = x1 * x2;
    assert_eq!(y.value, 3.75);

    context.differentiate(y);
    assert_eq!(context.get_derivative(x1), 2.5);
    assert_eq!(context.get_derivative(x2), 1.5);
    assert_eq!(context.get_derivative(y), 1.);
}

#[test]
fn single_thread_multiple_run() {
    for _ in 0..10 {
        addition();
    }
}

#[test]
fn multi_thread_multiple_run() {
    // TODO Insert purposeful sleep for each task?
    let task_pool = std::sync::TaskPool::new(10);
    let semaphore = std::sync::Arc::new(std::sync::Semaphore::new(-999));
    for _ in 0..1000 {
        let semaphore_task = semaphore.clone();
        task_pool.execute(move || {
            single_thread_multiple_run();
            semaphore_task.release();
        });
    }
    semaphore.access();
}

#[test]
#[should_fail(expected = "This Context instance is in use now. Note that a Context instance is allowed per construction location and per thread. Consequently, it cannot be recursively constructed unless it is destructed. This is a limitation caused by the thread local static variables usages in the current implementation.")]
fn context_lock() {
    fn recursive_context(n: usize) {
        if n > 0 {
            let context = new_autograd_context!(f32, 10);
            recursive_context(n - 1);
        }
    }
    recursive_context(2);
}

#[test]
fn context_capacity_expression_argument() {
    let a = 200;
    let b = 10;
    let context = new_autograd_context!(f32, a/b);
}
<|endoftext|>


<commit_before>/* Any copyright is dedicated to the Public Domain.
 * http://creativecommons.org/publicdomain/zero/1.0/ */

#[macro_use]
extern crate autograd;
use autograd::Context;

#[test]
fn addition() {
    let context = new_autograd_context!(f32, 1000);
    let x1 = context.new_variable(1.5);
    let x2 = context.new_variable(2.5);

    let y = x1 + x2;
    assert_eq!(y.value, 4.);

    context.differentiate(y);
    assert_eq!(context.get_derivative(y), 1.);
    assert_eq!(context.get_derivative(x1), 1.);
    assert_eq!(context.get_derivative(x2), 1.);
}

#[test]
fn multiplication() {
    let context = new_autograd_context!(f32, 1000);
    let x1 = context.new_variable(1.5);
    let x2 = context.new_variable(2.5);

    let y = x1 * x2;
    assert_eq!(y.value, 3.75);

    context.differentiate(y);
    assert_eq!(context.get_derivative(y), 1.);
    assert_eq!(context.get_derivative(x1), 2.5);
    assert_eq!(context.get_derivative(x2), 1.5);
}

#[test]
fn single_thread_multiple_run() {
    for _ in 0..10 {
        addition();
    }
}

#[test]
fn multi_thread_multiple_run() {
    // TODO Insert purposeful sleep for each task?
    let task_pool = std::sync::TaskPool::new(10);
    let semaphore = std::sync::Arc::new(std::sync::Semaphore::new(-999));
    for _ in 0..1000 {
        let semaphore_task = semaphore.clone();
        task_pool.execute(move || {
            single_thread_multiple_run();
            semaphore_task.release();
        });
    }
    semaphore.access();
}

#[test]
#[should_fail(expected = "This Context instance is in use now. Note that a Context instance is allowed per construction location and per thread. Consequently, it cannot be recursively constructed unless it is destructed. This is a limitation caused by the thread local static variables usages in the current implementation.")]
fn context_lock() {
    fn recursive_context(n: usize) {
        if n > 0 {
            let context = new_autograd_context!(f32, 10);
            recursive_context(n - 1);
        }
    }
    recursive_context(2);
}

#[test]
fn context_capacity_expression_argument() {
    let a = 200;
    let b = 10;
    let context = new_autograd_context!(f32, a/b);
}
<commit_msg>Add negative test.<commit_after>/* Any copyright is dedicated to the Public Domain.
 * http://creativecommons.org/publicdomain/zero/1.0/ */

#[macro_use]
extern crate autograd;
use autograd::Context;

#[test]
fn negative() {
    let context = new_autograd_context!(f32, 1000);
    let x = context.new_variable(1.5);

    let y = -x;
    assert_eq!(y.value, -1.5);

    context.differentiate(y);
    assert_eq!(context.get_derivative(x), -1.);
    assert_eq!(context.get_derivative(y), 1.);
}

#[test]
fn addition() {
    let context = new_autograd_context!(f32, 1000);
    let x1 = context.new_variable(1.5);
    let x2 = context.new_variable(2.5);

    let y = x1 + x2;
    assert_eq!(y.value, 4.);

    context.differentiate(y);
    assert_eq!(context.get_derivative(x1), 1.);
    assert_eq!(context.get_derivative(x2), 1.);
    assert_eq!(context.get_derivative(y), 1.);
}

#[test]
fn multiplication() {
    let context = new_autograd_context!(f32, 1000);
    let x1 = context.new_variable(1.5);
    let x2 = context.new_variable(2.5);

    let y = x1 * x2;
    assert_eq!(y.value, 3.75);

    context.differentiate(y);
    assert_eq!(context.get_derivative(x1), 2.5);
    assert_eq!(context.get_derivative(x2), 1.5);
    assert_eq!(context.get_derivative(y), 1.);
}

#[test]
fn single_thread_multiple_run() {
    for _ in 0..10 {
        addition();
    }
}

#[test]
fn multi_thread_multiple_run() {
    // TODO Insert purposeful sleep for each task?
    let task_pool = std::sync::TaskPool::new(10);
    let semaphore = std::sync::Arc::new(std::sync::Semaphore::new(-999));
    for _ in 0..1000 {
        let semaphore_task = semaphore.clone();
        task_pool.execute(move || {
            single_thread_multiple_run();
            semaphore_task.release();
        });
    }
    semaphore.access();
}

#[test]
#[should_fail(expected = "This Context instance is in use now. Note that a Context instance is allowed per construction location and per thread. Consequently, it cannot be recursively constructed unless it is destructed. This is a limitation caused by the thread local static variables usages in the current implementation.")]
fn context_lock() {
    fn recursive_context(n: usize) {
        if n > 0 {
            let context = new_autograd_context!(f32, 10);
            recursive_context(n - 1);
        }
    }
    recursive_context(2);
}

#[test]
fn context_capacity_expression_argument() {
    let a = 200;
    let b = 10;
    let context = new_autograd_context!(f32, a/b);
}
<|endoftext|>


<commit_before>// We add this `extern crate` here to ensure that bindgen is up-to-date and
// rebuilt, even though we aren't using any of its types or functions here, only
// indirectly calling the executable.
#[allow(dead_code)]
extern crate bindgen;

use std::env;
use std::fs;
use std::io::Read;
use std::path::{Path, PathBuf};
use std::process;

const TEST_BATCH_DEFAULT_SIZE: usize = 16;

fn spawn_run_bindgen<P, Q, R>(run_bindgen: P, bindgen: Q, header: R) -> process::Child
    where P: AsRef<Path>,
          Q: AsRef<Path>,
          R: AsRef<Path>
{
    let run_bindgen = run_bindgen.as_ref();
    let bindgen = bindgen.as_ref();
    let header = header.as_ref();

    // Convert from "tests/headers/foo.hpp" to "tests/expectations/foo.rs" by
    // saving the filename, popping off "headers/foo.hpp", pushing
    // "expectations", pushing the saved filename, and finally modifying the
    // extension.

    let mut expected = PathBuf::from(header);
    let file_name = expected.file_name()
        .expect("Should have filename")
        .to_os_string();
    expected.pop();
    expected.pop();
    expected.push("expectations");
    expected.push(file_name);
    expected.set_extension("rs");

    let mut cmd = process::Command::new(run_bindgen);
    cmd.stdout(process::Stdio::piped())
        .stderr(process::Stdio::piped())
        .arg(bindgen)
        .arg(header)
        .arg(expected);

    if cfg!(feature = "llvm_stable") {
        cmd.arg("--feature")
            .arg("llvm_stable");
    }

    cmd.spawn()
        .expect("Should be able to spawn run-bindgen.py child process")
}

#[test]
fn run_bindgen_tests() {
    let crate_root = env::var("CARGO_MANIFEST_DIR")
        .expect("should have CARGO_MANIFEST_DIR environment variable");

    let mut run_bindgen = PathBuf::from(&crate_root);
    run_bindgen.push("tests");
    run_bindgen.push("tools");
    run_bindgen.push("run-bindgen.py");

    let mut bindgen = PathBuf::from(&crate_root);
    bindgen.push("target");
    bindgen.push("debug");
    bindgen.push("bindgen");
    if !bindgen.is_file() {
        panic!("{} is not a file! Build bindgen before running tests.",
               bindgen.display());
    }

    let mut headers_dir = PathBuf::from(&crate_root);
    headers_dir.push("tests");
    headers_dir.push("headers");

    let entries = fs::read_dir(&headers_dir)
        .expect("Should read directory")
        .map(|result| result.expect("Should read directory entry"));

    let tests = entries.filter(|entry| {
        match entry.path().extension().map(|s| s.to_str()) {
            Some(Some("h")) |
            Some(Some("hpp")) => true,
            _ => false,
        }
    }).collect::<Vec<_>>();

    let batch_size = env::var("BINDGEN_TEST_BATCH_SIZE")
        .ok()
        .and_then(|x| x.parse::<usize>().ok())
        .unwrap_or(TEST_BATCH_DEFAULT_SIZE);

    // Spawn batch_size child to run in parallel
    // and wait on all of them before processing the next batch

    let children = tests.chunks(batch_size).map(|x| {
        x.iter().map(|entry| {
            let child = spawn_run_bindgen(run_bindgen.clone(), bindgen.clone(), entry.path());
            (entry.path(), child)
        }).collect::<Vec<_>>()
    });

    let failures: Vec<_> = children.flat_map(|x| {
        x.into_iter().filter_map(|(path, mut child)| {
            let passed = child.wait()
                .expect("Should wait on child process")
                .success();

            if passed { None } else { Some((path, child)) }
        })
    })
    .collect();

    let num_failures = failures.len();

    for (path, child) in failures {
        println!("FAIL: {}", path.display());

        let mut buf = String::new();

        child.stdout
            .expect("should have stdout piped")
            .read_to_string(&mut buf)
            .expect("should read child's stdout");
        for line in buf.lines() {
            println!("child stdout> {}", line);
        }

        child.stderr
            .expect("should have stderr piped")
            .read_to_string(&mut buf)
            .expect("should read child's stderr");
        for line in buf.lines() {
            println!("child stderr> {}", line);
        }
    }

    if num_failures > 0 {
        panic!("{} test failures!", num_failures);
    }
}
<commit_msg>Fix a small typo and expand batching comment<commit_after>// We add this `extern crate` here to ensure that bindgen is up-to-date and
// rebuilt, even though we aren't using any of its types or functions here, only
// indirectly calling the executable.
#[allow(dead_code)]
extern crate bindgen;

use std::env;
use std::fs;
use std::io::Read;
use std::path::{Path, PathBuf};
use std::process;

const TEST_BATCH_DEFAULT_SIZE: usize = 16;

fn spawn_run_bindgen<P, Q, R>(run_bindgen: P, bindgen: Q, header: R) -> process::Child
    where P: AsRef<Path>,
          Q: AsRef<Path>,
          R: AsRef<Path>
{
    let run_bindgen = run_bindgen.as_ref();
    let bindgen = bindgen.as_ref();
    let header = header.as_ref();

    // Convert from "tests/headers/foo.hpp" to "tests/expectations/foo.rs" by
    // saving the filename, popping off "headers/foo.hpp", pushing
    // "expectations", pushing the saved filename, and finally modifying the
    // extension.

    let mut expected = PathBuf::from(header);
    let file_name = expected.file_name()
        .expect("Should have filename")
        .to_os_string();
    expected.pop();
    expected.pop();
    expected.push("expectations");
    expected.push(file_name);
    expected.set_extension("rs");

    let mut cmd = process::Command::new(run_bindgen);
    cmd.stdout(process::Stdio::piped())
        .stderr(process::Stdio::piped())
        .arg(bindgen)
        .arg(header)
        .arg(expected);

    if cfg!(feature = "llvm_stable") {
        cmd.arg("--feature")
            .arg("llvm_stable");
    }

    cmd.spawn()
        .expect("Should be able to spawn run-bindgen.py child process")
}

#[test]
fn run_bindgen_tests() {
    let crate_root = env::var("CARGO_MANIFEST_DIR")
        .expect("should have CARGO_MANIFEST_DIR environment variable");

    let mut run_bindgen = PathBuf::from(&crate_root);
    run_bindgen.push("tests");
    run_bindgen.push("tools");
    run_bindgen.push("run-bindgen.py");

    let mut bindgen = PathBuf::from(&crate_root);
    bindgen.push("target");
    bindgen.push("debug");
    bindgen.push("bindgen");
    if !bindgen.is_file() {
        panic!("{} is not a file! Build bindgen before running tests.",
               bindgen.display());
    }

    let mut headers_dir = PathBuf::from(&crate_root);
    headers_dir.push("tests");
    headers_dir.push("headers");

    let entries = fs::read_dir(&headers_dir)
        .expect("Should read directory")
        .map(|result| result.expect("Should read directory entry"));

    let tests = entries.filter(|entry| {
        match entry.path().extension().map(|s| s.to_str()) {
            Some(Some("h")) |
            Some(Some("hpp")) => true,
            _ => false,
        }
    }).collect::<Vec<_>>();

    let batch_size = env::var("BINDGEN_TEST_BATCH_SIZE")
        .ok()
        .and_then(|x| x.parse::<usize>().ok())
        .unwrap_or(TEST_BATCH_DEFAULT_SIZE);

    // Spawn `batch_size` children to run in parallel and wait on all of them
    // before processing the next batch. This puts a limit on the resources
    // consumed when testing, so that we don't overload the system.

    let children = tests.chunks(batch_size).map(|x| {
        x.iter().map(|entry| {
            let child = spawn_run_bindgen(run_bindgen.clone(), bindgen.clone(), entry.path());
            (entry.path(), child)
        }).collect::<Vec<_>>()
    });

    let failures: Vec<_> = children.flat_map(|x| {
        x.into_iter().filter_map(|(path, mut child)| {
            let passed = child.wait()
                .expect("Should wait on child process")
                .success();

            if passed { None } else { Some((path, child)) }
        })
    })
    .collect();

    let num_failures = failures.len();

    for (path, child) in failures {
        println!("FAIL: {}", path.display());

        let mut buf = String::new();

        child.stdout
            .expect("should have stdout piped")
            .read_to_string(&mut buf)
            .expect("should read child's stdout");
        for line in buf.lines() {
            println!("child stdout> {}", line);
        }

        child.stderr
            .expect("should have stderr piped")
            .read_to_string(&mut buf)
            .expect("should read child's stderr");
        for line in buf.lines() {
            println!("child stderr> {}", line);
        }
    }

    if num_failures > 0 {
        panic!("{} test failures!", num_failures);
    }
}
<|endoftext|>


<commit_before>use std::fmt;

pub struct Memory {
    rom: u8,
    ram_0x4000_0x7FFF: usize,
    ram_0x8000_0xBFFF: usize,
    ram_0xC000_0xFFFF: usize,

    rom0: Vec<u8>,
    rom1: Vec<u8>,
    bank: [Vec<u8>; 8],
}

impl Memory {
    pub fn new(_rom0: Vec<u8>, _rom1: Vec<u8>) -> Memory {
        Memory {
            rom: 0,
            ram_0x4000_0x7FFF: 5,
            ram_0x8000_0xBFFF: 2,
            ram_0xC000_0xFFFF: 0,

            rom0: _rom0,
            rom1: _rom1,

            bank: [
                vec![0; 16 * 1024],
                vec![0; 16 * 1024],
                vec![0; 16 * 1024],
                vec![0; 16 * 1024],
                vec![0; 16 * 1024],
                vec![0; 16 * 1024],
                vec![0; 16 * 1024],
                vec![0; 16 * 1024]
            ]
        }
    }

    pub fn read_word(&self, addr: u16) -> u8 {
        if addr >= 0x0000 && addr <= 0x3FFF {
            match self.rom {
                0 => self.rom0[addr as usize],
                1 => self.rom1[addr as usize],
                _ => unreachable!()
            }
        } else if addr >= 0x4000 && addr <= 0x7FFF {
            self.bank[self.ram_0x4000_0x7FFF][(addr - 0x4000) as usize]
        } else if addr >= 0xC000 && addr <= 0xFFFF {
            self.bank[self.ram_0xC000_0xFFFF][(addr - 0xC000) as usize]
        } else {
            panic!("Trying to read from unrecognized address: {:#x}", addr);
        }
    }

    pub fn write_word(&mut self, addr: u16, val: u8) {
        if addr >= 0xC000 && addr <= 0xFFFF {
            self.bank[self.ram_0xC000_0xFFFF][(addr - 0xC000) as usize] = val;
        } else if addr >= 0x4000 && addr <= 0x7FFF {
            self.bank[self.ram_0x4000_0x7FFF][(addr - 0x4000) as usize] = val;
        } else {
            panic!("Trying to write to unrecognized address: {:#x}", addr);
        }
    }

    pub fn change_bank(&mut self, val: u8) {
        self.ram_0xC000_0xFFFF = val as usize;
    }

    pub fn change_rom_bank(&mut self, val: u8) {
        self.rom = val;
    }
}

impl fmt::Debug for Memory {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        writeln!(f, "ram_0xC000_0xFFFF: {}
                     ram_0x8000_0xBFFF: {}
                     ram_0x4000_0x7FFF: {}
                     rom: {}
                    ",
                    self.ram_0xC000_0xFFFF,
                    self.ram_0x8000_0xBFFF,
                    self.ram_0x4000_0x7FFF,
                    self.rom)
    }
}
<commit_msg>Add RAM between 0x8000 and 0xC000<commit_after>use std::fmt;

pub struct Memory {
    rom: u8,
    ram_0x4000_0x7FFF: usize,
    ram_0x8000_0xBFFF: usize,
    ram_0xC000_0xFFFF: usize,

    rom0: Vec<u8>,
    rom1: Vec<u8>,
    bank: [Vec<u8>; 8],
}

impl Memory {
    pub fn new(_rom0: Vec<u8>, _rom1: Vec<u8>) -> Memory {
        Memory {
            rom: 0,
            ram_0x4000_0x7FFF: 5,
            ram_0x8000_0xBFFF: 2,
            ram_0xC000_0xFFFF: 0,

            rom0: _rom0,
            rom1: _rom1,

            bank: [
                vec![0; 16 * 1024],
                vec![0; 16 * 1024],
                vec![0; 16 * 1024],
                vec![0; 16 * 1024],
                vec![0; 16 * 1024],
                vec![0; 16 * 1024],
                vec![0; 16 * 1024],
                vec![0; 16 * 1024]
            ]
        }
    }

    pub fn read_word(&self, addr: u16) -> u8 {
        if addr >= 0x0000 && addr <= 0x3FFF {
            match self.rom {
                0 => self.rom0[addr as usize],
                1 => self.rom1[addr as usize],
                _ => unreachable!()
            }
        } else if addr >= 0x4000 && addr <= 0x7FFF {
            self.bank[self.ram_0x4000_0x7FFF][(addr - 0x4000) as usize]
        } else if addr >= 0x8000 && addr <= 0xBFFF {
            self.bank[self.ram_0x8000_0xBFFF][(addr - 0x8000) as usize]
        } else if addr >= 0xC000 && addr <= 0xFFFF {
            self.bank[self.ram_0xC000_0xFFFF][(addr - 0xC000) as usize]
        } else {
            panic!("Trying to read from unrecognized address: {:#x}", addr);
        }
    }

    pub fn write_word(&mut self, addr: u16, val: u8) {
        if addr >= 0xC000 && addr <= 0xFFFF {
            self.bank[self.ram_0xC000_0xFFFF][(addr - 0xC000) as usize] = val;
        } else if addr >= 0x8000 && addr <= 0xBFFF {
            self.bank[self.ram_0x8000_0xBFFF][(addr - 0x8000) as usize] = val;
        } else if addr >= 0x4000 && addr <= 0x7FFF {
            self.bank[self.ram_0x4000_0x7FFF][(addr - 0x4000) as usize] = val;
        } else {
            panic!("Trying to write to unrecognized address: {:#x}", addr);
        }
    }

    pub fn change_bank(&mut self, val: u8) {
        self.ram_0xC000_0xFFFF = val as usize;
    }

    pub fn change_rom_bank(&mut self, val: u8) {
        self.rom = val;
    }
}

impl fmt::Debug for Memory {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        writeln!(f, "ram_0xC000_0xFFFF: {}
                     ram_0x8000_0xBFFF: {}
                     ram_0x4000_0x7FFF: {}
                     rom: {}
                    ",
                    self.ram_0xC000_0xFFFF,
                    self.ram_0x8000_0xBFFF,
                    self.ram_0x4000_0x7FFF,
                    self.rom)
    }
}
<|endoftext|>


<commit_before>// Copyright 2020 Google LLC
//
// Use of this source code is governed by an MIT-style license that can be found
// in the LICENSE file or at https://opensource.org/licenses/MIT.

use std::io::Cursor;

use byteorder::BigEndian;

pub fn encode<'a, I>(iter: I) -> Encode<I>
where
    I: Iterator<Item = &'a [u8]>,
{
    Encode {
        iter: iter,
        cur: Cursor::new(vec!()),
    }
}

pub fn decode<R>(buf: R) -> Decode<R>
where
    R: std::io::Read,
{
    Decode {
        buf: buf,
    }
}

pub struct Encode<I> {
    iter: I,
    cur: Cursor<Vec<u8>>,
}

impl<'a, I> Encode<I>
where
    I: Iterator<Item = &'a [u8]>,
{
    fn is_empty(&self) -> bool {
        self.cur.position() == self.cur.get_ref().len() as u64
    }

    fn pull(&mut self) -> std::io::Result<()> {
        use std::io::Write as _;
        use byteorder::WriteBytesExt as _;

        let data = match self.iter.next() {
            Some(data) => data,
            None => return Ok(()),
        };

        self.cur.get_mut().clear();
        self.cur.set_position(0);

        self.cur.write_u64::<BigEndian>(data.len() as u64)?;
        self.cur.write_all(data)?;
        self.cur.set_position(0);

        Ok(())
    }
}

impl<'a, I> std::io::Read for Encode<I>
where
    I: Iterator<Item = &'a [u8]>,
{
    fn read(&mut self, buf: &mut [u8]) -> std::io::Result<usize> {
        if self.is_empty() {
            self.pull()?;
        }

        self.cur.read(buf)
    }
}

pub struct Decode<R> {
    buf: R,
}

impl<R: std::io::Read> Decode<R> {

    fn read_len(&mut self) -> std::io::Result<Option<usize>> {
        use byteorder::ReadBytesExt as _;

        let mut buf = [0; 8];
        match self.buf.read(&mut buf[..])? {
            8 => (),
            0 => return Ok(None),
            _ => return Err(SizeTagError.into()),
        }

        let len = (&buf[..]).read_u64::<BigEndian>()? as usize;
        Ok(Some(len))
    }
}

impl<R: std::io::Read> Iterator for Decode<R> {

    type Item = std::io::Result<Vec<u8>>;

    fn next(&mut self) -> Option<std::io::Result<Vec<u8>>> {
        let len = match self.read_len() {
            Ok(Some(len)) => len,
            Ok(None) => return None,
            Err(error) => return Some(Err(error)),
        };

        let mut buf = vec!(0; len);
        match self.buf.read_exact(&mut buf[..]) {
            Ok(()) => (),
            Err(error) => return Some(Err(error)),
        }

        Some(Ok(buf))
    }
}

#[derive(Debug)]
struct SizeTagError;

impl std::fmt::Display for SizeTagError {

    fn fmt(&self, fmt: &mut std::fmt::Formatter) -> std::fmt::Result {
        write!(fmt, "incorrect length of the size tag")
    }
}

impl std::error::Error for SizeTagError {
}

impl From<SizeTagError> for std::io::Error {

    fn from(error: SizeTagError) -> std::io::Error {
        std::io::Error::new(std::io::ErrorKind::InvalidData, error)
    }
}

#[cfg(test)]
pub mod tests {

    use std::io::Read as _;

    use super::*;

    #[test]
    pub fn test_encode_empty() {
        let mut stream = encode(std::iter::empty());

        let mut output = vec!();
        stream.read_to_end(&mut output).unwrap();

        assert!(output.is_empty());
    }

    #[test]
    pub fn test_encode_single_chunk() {
        let mut stream = encode(vec!(&b"foo"[..]).into_iter());

        let mut output = vec!();
        stream.read_to_end(&mut output).unwrap();

        assert_eq!(output, b"\x00\x00\x00\x00\x00\x00\x00\x03foo");
    }

    #[test]
    pub fn test_encode_single_empty_chunk() {
        let mut stream = encode(vec!(&b""[..]).into_iter());

        let mut output = vec!();
        stream.read_to_end(&mut output).unwrap();

        assert_eq!(output, b"\x00\x00\x00\x00\x00\x00\x00\x00");
    }

    #[test]
    pub fn test_encode_multiple_empty_chunks() {
        let mut stream = encode(vec!(&b""[..], &b""[..]).into_iter());

        let mut output = vec!();
        stream.read_to_end(&mut output).unwrap();

        assert_eq!(&output[..8], b"\x00\x00\x00\x00\x00\x00\x00\x00");
        assert_eq!(&output[8..], b"\x00\x00\x00\x00\x00\x00\x00\x00");
    }

    #[test]
    pub fn test_decode_empty() {
        let mut iter = decode(&b""[..]);
        assert!(iter.next().is_none());
    }

    #[test]
    pub fn test_decode_single_chunk() {
        let mut iter = decode(&b"\x00\x00\x00\x00\x00\x00\x00\x03foo"[..])
            .map(Result::unwrap);

        assert_eq!(iter.next(), Some(b"foo".to_vec()));
        assert_eq!(iter.next(), None);
    }

    #[test]
    pub fn test_decode_empty_chunk() {
        let mut iter = decode(&b"\x00\x00\x00\x00\x00\x00\x00\x00"[..])
            .map(Result::unwrap);

        assert_eq!(iter.next(), Some(b"".to_vec()));
        assert_eq!(iter.next(), None);
    }

    #[test]
    pub fn test_decode_incorrect_size_tag() {
        let mut iter = decode(&b"\x12\x34\x56"[..]);

        let error = iter.next().unwrap().unwrap_err();
        assert_eq!(error.kind(), std::io::ErrorKind::InvalidData);
    }

    #[test]
    pub fn test_decode_short_input() {
        let mut iter = decode(&b"\x00\x00\x00\x00\x00\x00\x00\x42foo"[..]);

        let error = iter.next().unwrap().unwrap_err();
        assert_eq!(error.kind(), std::io::ErrorKind::UnexpectedEof);
    }

    #[test]
    pub fn test_encode_and_decode_multiple_chunks() {
        let chunks = vec! {
            &b"foo"[..],
            &b"bar"[..],
            &b"baz"[..],
        };

        let mut iter = decode(encode(chunks.into_iter()))
            .map(Result::unwrap);

        assert_eq!(iter.next(), Some(b"foo".to_vec()));
        assert_eq!(iter.next(), Some(b"bar".to_vec()));
        assert_eq!(iter.next(), Some(b"baz".to_vec()));
        assert_eq!(iter.next(), None);
    }

    #[test]
    pub fn test_encode_and_decode_multiple_empty_chunks() {
        let chunks = vec!(&b""[..], &b""[..], &b""[..]);

        let mut iter = decode(encode(chunks.into_iter()))
            .map(Result::unwrap);

        assert_eq!(iter.next(), Some(b"".to_vec()));
        assert_eq!(iter.next(), Some(b"".to_vec()));
        assert_eq!(iter.next(), Some(b"".to_vec()));
        assert_eq!(iter.next(), None);
    }
}
<commit_msg>Write documentation for the gzchunked utilities.<commit_after>// Copyright 2020 Google LLC
//
// Use of this source code is governed by an MIT-style license that can be found
// in the LICENSE file or at https://opensource.org/licenses/MIT.

//! Utilities for working with the chunked format.
//!
//! The chunked file format is extremely simple serialization procedure where
//! each chunk of raw bytes is prepended with the size of the chunk (as 64-bit
//! unsigned big-endian integer).
//!
//! Its primary application is streaming encoding and decoding of blobs of data
//! in formats that do not support or are inefficient for these purposes (such
//! as serialized Protocol Buffer messages).

use std::io::Cursor;

use byteorder::BigEndian;

/// Encodes a given iterator over binary blobs into the chunked format.
///
/// This is a streaming encoder and performs the encoding in a lazy way. It
/// should compose well with other streaming encoders (e.g. these offered by
/// the [`flate2`] crate).
///
/// [`flate2`]: https://crates.io/crates/flate2
///
/// # Examples
///
/// ```no_run
/// use std::fs::File;
///
/// let data = [b"foo", b"bar", b"baz"];
///
/// let mut stream = rrg::chunked::encode(data.iter().map(|blob| &blob[..]));
/// let mut file = File::create("output.chunked").unwrap();
/// std::io::copy(&mut stream, &mut file).unwrap();
/// ```
pub fn encode<'a, I>(iter: I) -> Encode<I>
where
    I: Iterator<Item = &'a [u8]>,
{
    Encode {
        iter: iter,
        cur: Cursor::new(vec!()),
    }
}

/// Decodes a buffer in the chunked format into binary blobs.
///
/// This is a streaming decoder and performs the decoding in a lazy way. It
/// should compose well with other streaming decoders (e.g. these offered by the
/// [`flate2`] crate).
///
/// [`flate2`]: https://crates.io/crates/flate2
///
/// # Examples
///
/// ```no_run
/// use std::fs::File;
///
/// let file = File::open("input.chunked").unwrap();
/// for (idx, blob) in rrg::chunked::decode(file).enumerate() {
///     println!("blob #{}: {:?}", idx, blob.unwrap());
/// }
/// ```
pub fn decode<R>(buf: R) -> Decode<R>
where
    R: std::io::Read,
{
    Decode {
        buf: buf,
    }
}

/// Streaming encoder for the chunked format.
///
/// It implements the `Read` trait, lazily polling the underlying chunk iterator
/// as more bytes is needed.
///
/// Instances of this type can be constructed using the [`encode`] function.
///
/// [`encode`]: fn.encode.html
pub struct Encode<I> {
    iter: I,
    cur: Cursor<Vec<u8>>,
}

impl<'a, I> Encode<I>
where
    I: Iterator<Item = &'a [u8]>,
{
    /// Checks whether all the data from the underlying cursor has been read.
    fn is_empty(&self) -> bool {
        self.cur.position() == self.cur.get_ref().len() as u64
    }

    /// Pulls another blob of data from the underlying iterator.
    fn pull(&mut self) -> std::io::Result<()> {
        use std::io::Write as _;
        use byteorder::WriteBytesExt as _;

        let data = match self.iter.next() {
            Some(data) => data,
            None => return Ok(()),
        };

        self.cur.get_mut().clear();
        self.cur.set_position(0);

        self.cur.write_u64::<BigEndian>(data.len() as u64)?;
        self.cur.write_all(data)?;
        self.cur.set_position(0);

        Ok(())
    }
}

impl<'a, I> std::io::Read for Encode<I>
where
    I: Iterator<Item = &'a [u8]>,
{
    fn read(&mut self, buf: &mut [u8]) -> std::io::Result<usize> {
        if self.is_empty() {
            self.pull()?;
        }

        self.cur.read(buf)
    }
}

/// Streaming decoder for the chunked format.
///
/// It implements the `Iterator` trait yielding chunks of decoded blobs, lazily
/// decoding data from the underlying buffer.
///
/// Instances of this type can be constructed using the [`decode`] function.
///
/// [`decode`]: fn.decode.html
pub struct Decode<R> {
    buf: R,
}

impl<R: std::io::Read> Decode<R> {

    /// Reads a size tag from the underlying buffer.
    ///
    /// It will return `None` if the is no more data in the buffer.
    fn read_len(&mut self) -> std::io::Result<Option<usize>> {
        use byteorder::ReadBytesExt as _;

        let mut buf = [0; 8];
        match self.buf.read(&mut buf[..])? {
            8 => (),
            0 => return Ok(None),
            _ => return Err(SizeTagError.into()),
        }

        let len = (&buf[..]).read_u64::<BigEndian>()? as usize;
        Ok(Some(len))
    }
}

impl<R: std::io::Read> Iterator for Decode<R> {

    type Item = std::io::Result<Vec<u8>>;

    fn next(&mut self) -> Option<std::io::Result<Vec<u8>>> {
        let len = match self.read_len() {
            Ok(Some(len)) => len,
            Ok(None) => return None,
            Err(error) => return Some(Err(error)),
        };

        let mut buf = vec!(0; len);
        match self.buf.read_exact(&mut buf[..]) {
            Ok(()) => (),
            Err(error) => return Some(Err(error)),
        }

        Some(Ok(buf))
    }
}

/// An error type for errors encountered when reading the size tag.
#[derive(Debug)]
struct SizeTagError;

impl std::fmt::Display for SizeTagError {

    fn fmt(&self, fmt: &mut std::fmt::Formatter) -> std::fmt::Result {
        write!(fmt, "incorrect length of the size tag")
    }
}

impl std::error::Error for SizeTagError {
}

impl From<SizeTagError> for std::io::Error {

    fn from(error: SizeTagError) -> std::io::Error {
        std::io::Error::new(std::io::ErrorKind::InvalidData, error)
    }
}

#[cfg(test)]
pub mod tests {

    use std::io::Read as _;

    use super::*;

    #[test]
    pub fn test_encode_empty() {
        let mut stream = encode(std::iter::empty());

        let mut output = vec!();
        stream.read_to_end(&mut output).unwrap();

        assert!(output.is_empty());
    }

    #[test]
    pub fn test_encode_single_chunk() {
        let mut stream = encode(vec!(&b"foo"[..]).into_iter());

        let mut output = vec!();
        stream.read_to_end(&mut output).unwrap();

        assert_eq!(output, b"\x00\x00\x00\x00\x00\x00\x00\x03foo");
    }

    #[test]
    pub fn test_encode_single_empty_chunk() {
        let mut stream = encode(vec!(&b""[..]).into_iter());

        let mut output = vec!();
        stream.read_to_end(&mut output).unwrap();

        assert_eq!(output, b"\x00\x00\x00\x00\x00\x00\x00\x00");
    }

    #[test]
    pub fn test_encode_multiple_empty_chunks() {
        let mut stream = encode(vec!(&b""[..], &b""[..]).into_iter());

        let mut output = vec!();
        stream.read_to_end(&mut output).unwrap();

        assert_eq!(&output[..8], b"\x00\x00\x00\x00\x00\x00\x00\x00");
        assert_eq!(&output[8..], b"\x00\x00\x00\x00\x00\x00\x00\x00");
    }

    #[test]
    pub fn test_decode_empty() {
        let mut iter = decode(&b""[..]);
        assert!(iter.next().is_none());
    }

    #[test]
    pub fn test_decode_single_chunk() {
        let mut iter = decode(&b"\x00\x00\x00\x00\x00\x00\x00\x03foo"[..])
            .map(Result::unwrap);

        assert_eq!(iter.next(), Some(b"foo".to_vec()));
        assert_eq!(iter.next(), None);
    }

    #[test]
    pub fn test_decode_empty_chunk() {
        let mut iter = decode(&b"\x00\x00\x00\x00\x00\x00\x00\x00"[..])
            .map(Result::unwrap);

        assert_eq!(iter.next(), Some(b"".to_vec()));
        assert_eq!(iter.next(), None);
    }

    #[test]
    pub fn test_decode_incorrect_size_tag() {
        let mut iter = decode(&b"\x12\x34\x56"[..]);

        let error = iter.next().unwrap().unwrap_err();
        assert_eq!(error.kind(), std::io::ErrorKind::InvalidData);
    }

    #[test]
    pub fn test_decode_short_input() {
        let mut iter = decode(&b"\x00\x00\x00\x00\x00\x00\x00\x42foo"[..]);

        let error = iter.next().unwrap().unwrap_err();
        assert_eq!(error.kind(), std::io::ErrorKind::UnexpectedEof);
    }

    #[test]
    pub fn test_encode_and_decode_multiple_chunks() {
        let chunks = vec! {
            &b"foo"[..],
            &b"bar"[..],
            &b"baz"[..],
        };

        let mut iter = decode(encode(chunks.into_iter()))
            .map(Result::unwrap);

        assert_eq!(iter.next(), Some(b"foo".to_vec()));
        assert_eq!(iter.next(), Some(b"bar".to_vec()));
        assert_eq!(iter.next(), Some(b"baz".to_vec()));
        assert_eq!(iter.next(), None);
    }

    #[test]
    pub fn test_encode_and_decode_multiple_empty_chunks() {
        let chunks = vec!(&b""[..], &b""[..], &b""[..]);

        let mut iter = decode(encode(chunks.into_iter()))
            .map(Result::unwrap);

        assert_eq!(iter.next(), Some(b"".to_vec()));
        assert_eq!(iter.next(), Some(b"".to_vec()));
        assert_eq!(iter.next(), Some(b"".to_vec()));
        assert_eq!(iter.next(), None);
    }
}
<|endoftext|>