<commit_before>// Copyright:: Copyright (c) 2015-2016 The Habitat Maintainers
//
// The terms of the Evaluation Agreement (Habitat) between Chef Software Inc.
// and the party accessing this file ("Licensee") apply to Licensee's use of
// the Software until such time that the Software is made available under an
// open source license such as the Apache 2.0 License.

extern crate libc;

use std::ffi::{CString, OsString};
use std::os::unix::ffi::OsStringExt;
use std::path::{Path, PathBuf};
use std::ptr;

use ansi_term::Colour::Cyan;
use common;
use hcore;
use hcore::fs::cache_artifact_path;
use hcore::package::{PackageIdent, PackageInstall};
use hcore::url::default_depot_url;

use error::{Error, Result};

const MAX_RETRIES: u8 = 4;

/// Makes an `execv(3)` system call to become a new program.
///
/// Note that if successful, this function will not return.
///
/// # Failures
///
/// * Command and/or command arguments cannot be converted into `CString`
pub fn exec_command(command: PathBuf, args: Vec<OsString>) -> Result<()> {
    let prog = try!(CString::new(command.into_os_string().into_vec()));
    let mut argv: Vec<*const i8> = Vec::with_capacity(args.len() + 2);
    argv.push(prog.as_ptr());
    for arg in args {
        argv.push(try!(CString::new(arg.into_vec())).as_ptr());
    }
    argv.push(ptr::null());

    // Calls `execv(3)` so this will not return, but rather become the program with the given
    // arguments.
    unsafe {
        libc::execv(prog.as_ptr(), argv.as_mut_ptr());
    }
    Ok(())
}

/// Returns the absolute path to the given command from the given package identifier.
///
/// If the package is not locally installed, the package will be installed before recomputing.
/// There are a maximum number of times a re-installation will be attempted before returning an
/// error.
///
/// # Failures
///
/// * If the package is installed but the command cannot be found in the package
/// * If an error occurs when loading the local package from disk
/// * If the maximum number of installation retries has been exceeded
pub fn command_from_pkg(command: &str,
                        ident: &PackageIdent,
                        cache_key_path: &Path,
                        retry: u8)
                        -> Result<PathBuf> {
    if retry > MAX_RETRIES {
        return Err(Error::ExecCommandNotFound(command.to_string()));
    }

    let fs_root_path = Path::new("/");
    match PackageInstall::load(ident, None) {
        Ok(pi) => {
            match try!(find_command_in_pkg(&command, &pi, fs_root_path)) {
                Some(cmd) => Ok(cmd),
                None => return Err(Error::ExecCommandNotFound(command.to_string())),
            }
        }
        Err(hcore::Error::PackageNotFound(_)) => {
            println!("{}",
                     Cyan.bold()
                         .paint(format!("∵ Package for {} not found, installing", &ident)));
            try!(common::command::package::install::from_url(&default_depot_url(),
                                                             ident,
                                                             fs_root_path,
                                                             &cache_artifact_path(None),
                                                             cache_key_path));
            command_from_pkg(&command, &ident, &cache_key_path, retry + 1)
        }
        Err(e) => return Err(Error::from(e)),
    }
}

/// Returns the absolute path to the given command from a given package installation.
///
/// If the command is not found, then `None` is returned.
///
/// # Failures
///
/// * The path entries metadata cannot be loaded
pub fn find_command_in_pkg(command: &str,
                           pkg_install: &PackageInstall,
                           fs_root_path: &Path)
                           -> Result<Option<PathBuf>> {
    for path in try!(pkg_install.paths()) {
        let candidate = fs_root_path.join(try!(path.strip_prefix("/"))).join(command);
        if candidate.is_file() {
            return Ok(Some(path.join(command)));
        }
    }
    Ok(None)
}
<commit_msg>[hab] Fix exec argument logic which corrupted the arg list.<commit_after>// Copyright:: Copyright (c) 2015-2016 The Habitat Maintainers
//
// The terms of the Evaluation Agreement (Habitat) between Chef Software Inc.
// and the party accessing this file ("Licensee") apply to Licensee's use of
// the Software until such time that the Software is made available under an
// open source license such as the Apache 2.0 License.

extern crate libc;

use std;
use std::ffi::{CString, OsString};
use std::os::unix::ffi::OsStrExt;
use std::path::{Path, PathBuf};
use std::ptr;

use ansi_term::Colour::Cyan;
use common;
use hcore;
use hcore::fs::cache_artifact_path;
use hcore::package::{PackageIdent, PackageInstall};
use hcore::url::default_depot_url;

use error::{Error, Result};

const MAX_RETRIES: u8 = 4;

/// Makes an `execv(3)` system call to become a new program.
///
/// Note that if successful, this function will not return.
///
/// # Failures
///
/// * Command and/or command arguments cannot be converted into `CString`
pub fn exec_command(command: PathBuf, args: Vec<OsString>) -> Result<()> {
    // A massive thanks to the `exec` crate which pointed to the correct invocation
    // behavior--namely to pass null-terminated string pointers.
    //
    // Source: https://github.com/faradayio/exec-rs/blob/master/src/lib.rs

    debug!("Calling execv: ({:?}) {:?}", command.display(), &args);
    let prog_cstring = try!(CString::new(command.as_os_str().as_bytes()));
    let arg_cstrings = try!(args.into_iter()
                                .map(|arg| CString::new(arg.as_os_str().as_bytes()))
                                .collect::<std::result::Result<Vec<_>, _>>());
    let mut arg_charptrs: Vec<_> = arg_cstrings.iter()
                                               .map(|arg| {
                                                   arg.as_bytes_with_nul().as_ptr() as *const i8
                                               })
                                               .collect();
    arg_charptrs.insert(0,
                        prog_cstring.clone().as_bytes_with_nul().as_ptr() as *const i8);
    arg_charptrs.push(ptr::null());

    unsafe {
        libc::execv(prog_cstring.as_bytes_with_nul().as_ptr() as *const i8,
                    arg_charptrs.as_mut_ptr());
    }
    Ok(())
}

/// Returns the absolute path to the given command from the given package identifier.
///
/// If the package is not locally installed, the package will be installed before recomputing.
/// There are a maximum number of times a re-installation will be attempted before returning an
/// error.
///
/// # Failures
///
/// * If the package is installed but the command cannot be found in the package
/// * If an error occurs when loading the local package from disk
/// * If the maximum number of installation retries has been exceeded
pub fn command_from_pkg(command: &str,
                        ident: &PackageIdent,
                        cache_key_path: &Path,
                        retry: u8)
                        -> Result<PathBuf> {
    if retry > MAX_RETRIES {
        return Err(Error::ExecCommandNotFound(command.to_string()));
    }

    let fs_root_path = Path::new("/");
    match PackageInstall::load(ident, None) {
        Ok(pi) => {
            match try!(find_command_in_pkg(&command, &pi, fs_root_path)) {
                Some(cmd) => Ok(cmd),
                None => return Err(Error::ExecCommandNotFound(command.to_string())),
            }
        }
        Err(hcore::Error::PackageNotFound(_)) => {
            println!("{}",
                     Cyan.bold()
                         .paint(format!("∵ Package for {} not found, installing", &ident)));
            try!(common::command::package::install::from_url(&default_depot_url(),
                                                             ident,
                                                             fs_root_path,
                                                             &cache_artifact_path(None),
                                                             cache_key_path));
            command_from_pkg(&command, &ident, &cache_key_path, retry + 1)
        }
        Err(e) => return Err(Error::from(e)),
    }
}

/// Returns the absolute path to the given command from a given package installation.
///
/// If the command is not found, then `None` is returned.
///
/// # Failures
///
/// * The path entries metadata cannot be loaded
pub fn find_command_in_pkg(command: &str,
                           pkg_install: &PackageInstall,
                           fs_root_path: &Path)
                           -> Result<Option<PathBuf>> {
    for path in try!(pkg_install.paths()) {
        let candidate = fs_root_path.join(try!(path.strip_prefix("/"))).join(command);
        if candidate.is_file() {
            return Ok(Some(path.join(command)));
        }
    }
    Ok(None)
}
<|endoftext|>


<commit_before>use std::path:: { Path };
use std::env:: { current_dir };
use clap:: { App, Arg };
use nickel::Nickel;
use configuration:: { Configuration };
use routing:: { router };

pub struct Application<'a> {
    cli: App<'a, 'a>
}

impl<'a> Application<'a> {
    pub fn new() -> Self {
        let app = App::new("Nickel application server")
            .version("1.0")
            .author("Noritaka Horio <holy.shared.design@gmail.com>")
            .arg(Arg::with_name("config")
                    .short("c")
                    .long("config")
                    .value_name("FILE")
                    .help("configuration file")
                    .takes_value(true));

        Application {
            cli: app
        }
    }
    pub fn run(self) {
        let matches = self.cli.get_matches();
        let config = match matches.value_of("config") {
            Some(config_file) => {
                let cwd = current_dir().unwrap();
                let config_path = Path::new(&cwd).join(config_file);
                Configuration::from(config_path.as_path())
            },
            None => Configuration::default()
        };

        let mut server = Nickel::new();
        server.utilize(router());
        server.listen(config.server());
    }
}
<commit_msg>Refactoring<commit_after>use std::path:: { Path };
use std::env:: { current_dir };
use clap:: { App, Arg, ArgMatches };
use nickel::Nickel;
use configuration:: { Configuration };
use routing:: { router };

pub struct Application;

impl Application {
    pub fn new() -> Self {
        Application
    }
    pub fn run(&self) {
        let matches = self.parse_args();
        let config = match matches.value_of("config") {
            Some(config_file) => self.load_configuration_from_file(config_file),
            None => Configuration::default()
        };
        self.start(&config);
    }
    fn parse_args<'a>(&self) -> ArgMatches<'a> {
        App::new("Nickel application server")
            .version("1.0")
            .author("Noritaka Horio <holy.shared.design@gmail.com>")
            .arg(Arg::from_usage("-c --config=[CONFIG] 'configuration file'"))
            .get_matches()
    }
    fn load_configuration_from_file(&self, config_file: &str) -> Configuration {
        let cwd = current_dir().unwrap();
        let config_path = Path::new(&cwd).join(config_file);
        Configuration::from(config_path.as_path())
    }
    fn start(&self, config: &Configuration) {
        let mut server = Nickel::new();
        server.utilize(router());
        server.listen(config.server());
    }
}
<|endoftext|>


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

use hyper::header::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::SessionNotCreated,
                                    "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(ContentType::json());
                res.send(&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 mut server = Server::http(address).unwrap();
    server.keep_alive(None);

    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>server: return hyper::server::Listening so user can access socket addr (#41)<commit_after>use std::io::Read;
use std::marker::PhantomData;
use std::net::SocketAddr;
use std::sync::mpsc::{channel, Receiver, Sender};
use std::sync::Mutex;
use std::thread;

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

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::SessionNotCreated,
                                    "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(ContentType::json());
                res.send(&resp_body.as_bytes()).unwrap();
            },
            _ => {}
        }
    }
}

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

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

    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)
}
<|endoftext|>


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

//!
//! Implements the Roughenough server functionality.
//!

use hex;
use std::io::ErrorKind;
use std::net::{SocketAddr, IpAddr};
use std::process;
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::Arc;
use std::time::{Duration, SystemTime};
use std::io::Write;

use byteorder::{LittleEndian, WriteBytesExt};

use humansize::{FileSize, file_size_opts as fsopts};

use mio::net::{TcpListener, UdpSocket};
use mio::{Events, Poll, PollOpt, Ready, Token};
use mio::tcp::Shutdown;
use mio_extras::timer::Timer;

use crate::{Error, RtMessage, Tag, MIN_REQUEST_LENGTH};
use crate::config::ServerConfig;
use crate::grease::Grease;
use crate::key::{LongTermKey, OnlineKey};
use crate::kms;
use crate::merkle::MerkleTree;
use crate::stats::{AggregatedStats, PerClientStats, ServerStats, ClientStatEntry};

macro_rules! check_ctrlc {
    ($keep_running:expr) => {
        if !$keep_running.load(Ordering::Acquire) {
            warn!("Ctrl-C caught, exiting...");
            return true;
        }
    };
}

// mio event registrations
const EVT_MESSAGE: Token = Token(0);
const EVT_STATUS_UPDATE: Token = Token(1);
const EVT_HEALTH_CHECK: Token = Token(2);

// Canned response to health check request
const HTTP_RESPONSE: &str = "HTTP/1.1 200 OK\nContent-Length: 0\nConnection: close\n\n";

/// The main Roughenough server instance.
///
/// The [ServerConfig](../config/trait.ServerConfig.html) trait specifies the required and optional
/// parameters available for configuring a Roughenoguh server instance.
///
/// Implementations of `ServerConfig` obtain configurations from different back-end sources
/// such as files or environment variables.
///
/// See [the config module](../config/index.html) for more information.
///
pub struct Server {
    config: Box<dyn ServerConfig>,
    online_key: OnlineKey,
    cert_bytes: Vec<u8>,

    socket: UdpSocket,
    health_listener: Option<TcpListener>,
    keep_running: Arc<AtomicBool>,
    poll_duration: Option<Duration>,
    grease: Grease,
    timer: Timer<()>,
    poll: Poll,
    merkle: MerkleTree,
    requests: Vec<(Vec<u8>, SocketAddr)>,
    buf: [u8; 65_536],

    public_key: String,

    stats: Box<dyn ServerStats>,

    // Used to send requests to ourselves in fuzzing mode
    #[cfg(fuzzing)]
    fake_client_socket: UdpSocket,
}

impl Server {
    ///
    /// Create a new server instance from the provided
    /// [`ServerConfig`](../config/trait.ServerConfig.html) trait object instance.
    ///
    pub fn new(config: Box<dyn ServerConfig>) -> Server {
        let online_key = OnlineKey::new();
        let public_key: String;

        let cert_bytes = {
            let seed = match kms::load_seed(config.as_ref()) {
                Ok(seed) => seed,
                Err(e) => {
                    error!("Failed to load seed: {:#?}", e);
                    process::exit(1);
                }
            };
            let mut long_term_key = LongTermKey::new(&seed);
            public_key = hex::encode(long_term_key.public_key());

            long_term_key.make_cert(&online_key).encode().unwrap()
        };

        let keep_running = Arc::new(AtomicBool::new(true));

        let sock_addr = config.udp_socket_addr().expect("udp sock addr");
        let socket = UdpSocket::bind(&sock_addr).expect("failed to bind to socket");

        let poll_duration = Some(Duration::from_millis(100));

        let mut timer: Timer<()> = Timer::default();
        timer.set_timeout(config.status_interval(), ());

        let poll = Poll::new().unwrap();
        poll.register(&socket, EVT_MESSAGE, Ready::readable(), PollOpt::edge())
            .unwrap();
        poll.register(&timer, EVT_STATUS_UPDATE, Ready::readable(), PollOpt::edge())
            .unwrap();

        let health_listener = if let Some(hc_port) = config.health_check_port() {
            let hc_sock_addr: SocketAddr = format!("{}:{}", config.interface(), hc_port)
                .parse()
                .unwrap();

            let tcp_listener = TcpListener::bind(&hc_sock_addr)
                .expect("failed to bind TCP listener for health check");

            poll.register(&tcp_listener, EVT_HEALTH_CHECK, Ready::readable(), PollOpt::edge())
                .unwrap();

            Some(tcp_listener)
        } else {
            None
        };

        let stats: Box<dyn ServerStats> = if config.client_stats_enabled() {
            Box::new(PerClientStats::new())
        } else {
            Box::new(AggregatedStats::new())
        };

        let grease = Grease::new(config.fault_percentage());
        let merkle = MerkleTree::new();
        let requests = Vec::with_capacity(config.batch_size() as usize);

        Server {
            config,
            online_key,
            cert_bytes,

            socket,
            health_listener,

            keep_running,
            poll_duration,
            grease,
            timer,
            poll,
            merkle,
            requests,
            buf: [0u8; 65_536],

            public_key,

            stats,

            #[cfg(fuzzing)]
            fake_client_socket: UdpSocket::bind(&"127.0.0.1:0".parse().unwrap()).unwrap(),
        }
    }

    /// Returns a reference to the server's long-term public key
    pub fn get_public_key(&self) -> &str {
        &self.public_key
    }

    /// Returns a reference to the server's on-line (delegated) key
    pub fn get_online_key(&self) -> &OnlineKey {
        &self.online_key
    }

    /// Returns a reference to the `ServerConfig` this server was configured with
    pub fn get_config(&self) -> &dyn ServerConfig {
        self.config.as_ref()
    }

    /// Returns a reference counted pointer the this server's `keep_running` value.
    pub fn get_keep_running(&self) -> Arc<AtomicBool> {
        self.keep_running.clone()
    }

    #[cfg(fuzzing)]
    pub fn send_to_self(&mut self, data: &[u8]) {
        let res = self
            .fake_client_socket
            .send_to(data, &self.socket.local_addr().unwrap());
        info!("Sent to self: {:?}", res);
    }

    /// The main processing function for incoming connections. This method should be
    /// called repeatedly in a loop to process requests. It returns 'true' when the
    /// server has shutdown (due to keep_running being set to 'false').
    ///
    pub fn process_events(&mut self, events: &mut Events) -> bool {
        self.poll
            .poll(events, self.poll_duration)
            .expect("server event poll failed; cannot recover");

        for msg in events.iter() {
            match msg.token() {
                EVT_MESSAGE => {
                    loop {
                        check_ctrlc!(self.keep_running);

                        self.merkle.reset();
                        self.requests.clear();

                        let socket_now_empty = self.collect_requests();

                        if self.requests.is_empty() {
                            break;
                        }

                        self.send_responses();

                        if socket_now_empty {
                            break;
                        }
                    }
                },
                EVT_HEALTH_CHECK => self.handle_health_check(),
                EVT_STATUS_UPDATE => self.handle_status_update(),
                _ => unreachable!(),
            }
        }
        false
    }

    // Read and process client requests from socket until empty or 'batch_size' number of
    // requests have been read.
    fn collect_requests(&mut self) -> bool {
        for i in 0..self.config.batch_size() {
            match self.socket.recv_from(&mut self.buf) {
                Ok((num_bytes, src_addr)) => {
                    match self.nonce_from_request(&self.buf, num_bytes) {
                        Ok(nonce) => {
                            self.stats.add_valid_request(&src_addr.ip());
                            self.requests.push((Vec::from(nonce), src_addr));
                            self.merkle.push_leaf(nonce);
                        }
                        Err(e) => {
                            self.stats.add_invalid_request(&src_addr.ip());

                            info!(
                                "Invalid request: '{:?}' ({} bytes) from {} (#{} in batch, resp #{})",
                                e, num_bytes, src_addr, i,
                                self.stats.total_responses_sent() + u64::from(i)
                            );
                        }
                    }
                }
                Err(e) => match e.kind() {
                    ErrorKind::WouldBlock => {
                        return true;
                    }
                    _ => {
                        error!("Error receiving from socket: {:?}: {:?}", e.kind(), e);
                        return false;
                    }
                },
            };
        }

        false
    }

    // extract the client's nonce from its request
    fn nonce_from_request<'a>(&self, buf: &'a [u8], num_bytes: usize) -> Result<&'a [u8], Error> {
        if num_bytes < MIN_REQUEST_LENGTH as usize {
            return Err(Error::RequestTooShort);
        }

        let tag_count = &buf[..4];
        let expected_nonc = &buf[8..12];
        let expected_pad = &buf[12..16];

        let tag_count_is_2 = tag_count == [0x02, 0x00, 0x00, 0x00];
        let tag1_is_nonc = expected_nonc == Tag::NONC.wire_value();
        let tag2_is_pad = expected_pad == Tag::PAD.wire_value();

        if tag_count_is_2 && tag1_is_nonc && tag2_is_pad {
            Ok(&buf[0x10..0x50])
        } else {
            Err(Error::InvalidRequest)
        }
    }

    fn send_responses(&mut self) {
        let merkle_root = self.merkle.compute_root();

        // The SREP tag is identical for each response
        let srep = self.online_key.make_srep(SystemTime::now(), &merkle_root);

        for (i, &(ref nonce, ref src_addr)) in self.requests.iter().enumerate() {
            let paths = self.merkle.get_paths(i);
            let resp_msg = {
                let r = self.make_response(&srep, &self.cert_bytes, &paths, i as u32);
                if self.grease.should_add_error() { self.grease.add_errors(&r) } else { r }
            };
            let resp_bytes = resp_msg.encode().unwrap();

            let bytes_sent = self
                .socket
                .send_to(&resp_bytes, &src_addr)
                .expect("send_to failed");

            self.stats.add_response(&src_addr.ip(), bytes_sent);

            info!(
                "Responded {} bytes to {} for '{}..' (#{} in batch, resp #{})",
                bytes_sent,
                src_addr,
                hex::encode(&nonce[0..4]),
                i + 1,
                self.stats.total_responses_sent()
            );
        }
    }

    fn make_response(&self, srep: &RtMessage, cert_bytes: &[u8], path: &[u8], idx: u32) -> RtMessage {
        let mut index = [0; 4];
        (&mut index as &mut [u8])
            .write_u32::<LittleEndian>(idx)
            .unwrap();

        let sig_bytes = srep.get_field(Tag::SIG).unwrap();
        let srep_bytes = srep.get_field(Tag::SREP).unwrap();

        let mut response = RtMessage::new(5);
        response.add_field(Tag::SIG, sig_bytes).unwrap();
        response.add_field(Tag::PATH, path).unwrap();
        response.add_field(Tag::SREP, srep_bytes).unwrap();
        response.add_field(Tag::CERT, cert_bytes).unwrap();
        response.add_field(Tag::INDX, &index).unwrap();

        response
    }

    fn handle_health_check(&mut self) {
        let listener = self.health_listener.as_ref().unwrap();
        match listener.accept() {
            Ok((ref mut stream, src_addr)) => {
                info!("health check from {}", src_addr);
                self.stats.add_health_check(&src_addr.ip());

                match stream.write(HTTP_RESPONSE.as_bytes()) {
                    Ok(_) => (),
                    Err(e) => warn!("error writing health check {}", e),
                };

                match stream.shutdown(Shutdown::Both) {
                    Ok(_) => (),
                    Err(e) => warn!("error in health check socket shutdown {}", e),
                }
            }
            Err(ref e) if e.kind() == ErrorKind::WouldBlock => {
                debug!("blocking in TCP health check");
            }
            Err(e) => {
                warn!("unexpected health check error {}", e);
            }
        }
    }

    fn handle_status_update(&mut self) {
        let mut vec: Vec<(&IpAddr, &ClientStatEntry)> = self.stats.iter().collect();
        vec.sort_by(|lhs, rhs| lhs.1.valid_requests.cmp(&rhs.1.valid_requests));

        for (addr, counts) in vec {
            info!(
                "{:16}: {} valid, {} invalid requests; {} responses ({} sent)",
                format!("{}", addr), counts.valid_requests, counts.invalid_requests,
                counts.responses_sent,
                counts.bytes_sent.file_size(fsopts::BINARY).unwrap()
            );
        }

        info!(
            "Totals: {} unique clients; {} valid, {} invalid requests; {} responses ({} sent)",
            self.stats.total_unique_clients(),
            self.stats.total_valid_requests(), self.stats.total_invalid_requests(),
            self.stats.total_responses_sent(),
            self.stats.total_bytes_sent().file_size(fsopts::BINARY).unwrap()
        );

        self.timer.set_timeout(self.config.status_interval(), ());
    }
}
<commit_msg>stats: sort status update in actual descending order<commit_after>// Copyright 2017-2019 int08h LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

//!
//! Implements the Roughenough server functionality.
//!

use hex;
use std::io::ErrorKind;
use std::net::{SocketAddr, IpAddr};
use std::process;
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::Arc;
use std::time::{Duration, SystemTime};
use std::io::Write;

use byteorder::{LittleEndian, WriteBytesExt};

use humansize::{FileSize, file_size_opts as fsopts};

use mio::net::{TcpListener, UdpSocket};
use mio::{Events, Poll, PollOpt, Ready, Token};
use mio::tcp::Shutdown;
use mio_extras::timer::Timer;

use crate::{Error, RtMessage, Tag, MIN_REQUEST_LENGTH};
use crate::config::ServerConfig;
use crate::grease::Grease;
use crate::key::{LongTermKey, OnlineKey};
use crate::kms;
use crate::merkle::MerkleTree;
use crate::stats::{AggregatedStats, PerClientStats, ServerStats, ClientStatEntry};

macro_rules! check_ctrlc {
    ($keep_running:expr) => {
        if !$keep_running.load(Ordering::Acquire) {
            warn!("Ctrl-C caught, exiting...");
            return true;
        }
    };
}

// mio event registrations
const EVT_MESSAGE: Token = Token(0);
const EVT_STATUS_UPDATE: Token = Token(1);
const EVT_HEALTH_CHECK: Token = Token(2);

// Canned response to health check request
const HTTP_RESPONSE: &str = "HTTP/1.1 200 OK\nContent-Length: 0\nConnection: close\n\n";

/// The main Roughenough server instance.
///
/// The [ServerConfig](../config/trait.ServerConfig.html) trait specifies the required and optional
/// parameters available for configuring a Roughenoguh server instance.
///
/// Implementations of `ServerConfig` obtain configurations from different back-end sources
/// such as files or environment variables.
///
/// See [the config module](../config/index.html) for more information.
///
pub struct Server {
    config: Box<dyn ServerConfig>,
    online_key: OnlineKey,
    cert_bytes: Vec<u8>,

    socket: UdpSocket,
    health_listener: Option<TcpListener>,
    keep_running: Arc<AtomicBool>,
    poll_duration: Option<Duration>,
    grease: Grease,
    timer: Timer<()>,
    poll: Poll,
    merkle: MerkleTree,
    requests: Vec<(Vec<u8>, SocketAddr)>,
    buf: [u8; 65_536],

    public_key: String,

    stats: Box<dyn ServerStats>,

    // Used to send requests to ourselves in fuzzing mode
    #[cfg(fuzzing)]
    fake_client_socket: UdpSocket,
}

impl Server {
    ///
    /// Create a new server instance from the provided
    /// [`ServerConfig`](../config/trait.ServerConfig.html) trait object instance.
    ///
    pub fn new(config: Box<dyn ServerConfig>) -> Server {
        let online_key = OnlineKey::new();
        let public_key: String;

        let cert_bytes = {
            let seed = match kms::load_seed(config.as_ref()) {
                Ok(seed) => seed,
                Err(e) => {
                    error!("Failed to load seed: {:#?}", e);
                    process::exit(1);
                }
            };
            let mut long_term_key = LongTermKey::new(&seed);
            public_key = hex::encode(long_term_key.public_key());

            long_term_key.make_cert(&online_key).encode().unwrap()
        };

        let keep_running = Arc::new(AtomicBool::new(true));

        let sock_addr = config.udp_socket_addr().expect("udp sock addr");
        let socket = UdpSocket::bind(&sock_addr).expect("failed to bind to socket");

        let poll_duration = Some(Duration::from_millis(100));

        let mut timer: Timer<()> = Timer::default();
        timer.set_timeout(config.status_interval(), ());

        let poll = Poll::new().unwrap();
        poll.register(&socket, EVT_MESSAGE, Ready::readable(), PollOpt::edge())
            .unwrap();
        poll.register(&timer, EVT_STATUS_UPDATE, Ready::readable(), PollOpt::edge())
            .unwrap();

        let health_listener = if let Some(hc_port) = config.health_check_port() {
            let hc_sock_addr: SocketAddr = format!("{}:{}", config.interface(), hc_port)
                .parse()
                .unwrap();

            let tcp_listener = TcpListener::bind(&hc_sock_addr)
                .expect("failed to bind TCP listener for health check");

            poll.register(&tcp_listener, EVT_HEALTH_CHECK, Ready::readable(), PollOpt::edge())
                .unwrap();

            Some(tcp_listener)
        } else {
            None
        };

        let stats: Box<dyn ServerStats> = if config.client_stats_enabled() {
            Box::new(PerClientStats::new())
        } else {
            Box::new(AggregatedStats::new())
        };

        let grease = Grease::new(config.fault_percentage());
        let merkle = MerkleTree::new();
        let requests = Vec::with_capacity(config.batch_size() as usize);

        Server {
            config,
            online_key,
            cert_bytes,

            socket,
            health_listener,

            keep_running,
            poll_duration,
            grease,
            timer,
            poll,
            merkle,
            requests,
            buf: [0u8; 65_536],

            public_key,

            stats,

            #[cfg(fuzzing)]
            fake_client_socket: UdpSocket::bind(&"127.0.0.1:0".parse().unwrap()).unwrap(),
        }
    }

    /// Returns a reference to the server's long-term public key
    pub fn get_public_key(&self) -> &str {
        &self.public_key
    }

    /// Returns a reference to the server's on-line (delegated) key
    pub fn get_online_key(&self) -> &OnlineKey {
        &self.online_key
    }

    /// Returns a reference to the `ServerConfig` this server was configured with
    pub fn get_config(&self) -> &dyn ServerConfig {
        self.config.as_ref()
    }

    /// Returns a reference counted pointer the this server's `keep_running` value.
    pub fn get_keep_running(&self) -> Arc<AtomicBool> {
        self.keep_running.clone()
    }

    #[cfg(fuzzing)]
    pub fn send_to_self(&mut self, data: &[u8]) {
        let res = self
            .fake_client_socket
            .send_to(data, &self.socket.local_addr().unwrap());
        info!("Sent to self: {:?}", res);
    }

    /// The main processing function for incoming connections. This method should be
    /// called repeatedly in a loop to process requests. It returns 'true' when the
    /// server has shutdown (due to keep_running being set to 'false').
    ///
    pub fn process_events(&mut self, events: &mut Events) -> bool {
        self.poll
            .poll(events, self.poll_duration)
            .expect("server event poll failed; cannot recover");

        for msg in events.iter() {
            match msg.token() {
                EVT_MESSAGE => {
                    loop {
                        check_ctrlc!(self.keep_running);

                        self.merkle.reset();
                        self.requests.clear();

                        let socket_now_empty = self.collect_requests();

                        if self.requests.is_empty() {
                            break;
                        }

                        self.send_responses();

                        if socket_now_empty {
                            break;
                        }
                    }
                },
                EVT_HEALTH_CHECK => self.handle_health_check(),
                EVT_STATUS_UPDATE => self.handle_status_update(),
                _ => unreachable!(),
            }
        }
        false
    }

    // Read and process client requests from socket until empty or 'batch_size' number of
    // requests have been read.
    fn collect_requests(&mut self) -> bool {
        for i in 0..self.config.batch_size() {
            match self.socket.recv_from(&mut self.buf) {
                Ok((num_bytes, src_addr)) => {
                    match self.nonce_from_request(&self.buf, num_bytes) {
                        Ok(nonce) => {
                            self.stats.add_valid_request(&src_addr.ip());
                            self.requests.push((Vec::from(nonce), src_addr));
                            self.merkle.push_leaf(nonce);
                        }
                        Err(e) => {
                            self.stats.add_invalid_request(&src_addr.ip());

                            info!(
                                "Invalid request: '{:?}' ({} bytes) from {} (#{} in batch, resp #{})",
                                e, num_bytes, src_addr, i,
                                self.stats.total_responses_sent() + u64::from(i)
                            );
                        }
                    }
                }
                Err(e) => match e.kind() {
                    ErrorKind::WouldBlock => {
                        return true;
                    }
                    _ => {
                        error!("Error receiving from socket: {:?}: {:?}", e.kind(), e);
                        return false;
                    }
                },
            };
        }

        false
    }

    // extract the client's nonce from its request
    fn nonce_from_request<'a>(&self, buf: &'a [u8], num_bytes: usize) -> Result<&'a [u8], Error> {
        if num_bytes < MIN_REQUEST_LENGTH as usize {
            return Err(Error::RequestTooShort);
        }

        let tag_count = &buf[..4];
        let expected_nonc = &buf[8..12];
        let expected_pad = &buf[12..16];

        let tag_count_is_2 = tag_count == [0x02, 0x00, 0x00, 0x00];
        let tag1_is_nonc = expected_nonc == Tag::NONC.wire_value();
        let tag2_is_pad = expected_pad == Tag::PAD.wire_value();

        if tag_count_is_2 && tag1_is_nonc && tag2_is_pad {
            Ok(&buf[0x10..0x50])
        } else {
            Err(Error::InvalidRequest)
        }
    }

    fn send_responses(&mut self) {
        let merkle_root = self.merkle.compute_root();

        // The SREP tag is identical for each response
        let srep = self.online_key.make_srep(SystemTime::now(), &merkle_root);

        for (i, &(ref nonce, ref src_addr)) in self.requests.iter().enumerate() {
            let paths = self.merkle.get_paths(i);
            let resp_msg = {
                let r = self.make_response(&srep, &self.cert_bytes, &paths, i as u32);
                if self.grease.should_add_error() { self.grease.add_errors(&r) } else { r }
            };
            let resp_bytes = resp_msg.encode().unwrap();

            let bytes_sent = self
                .socket
                .send_to(&resp_bytes, &src_addr)
                .expect("send_to failed");

            self.stats.add_response(&src_addr.ip(), bytes_sent);

            info!(
                "Responded {} bytes to {} for '{}..' (#{} in batch, resp #{})",
                bytes_sent,
                src_addr,
                hex::encode(&nonce[0..4]),
                i + 1,
                self.stats.total_responses_sent()
            );
        }
    }

    fn make_response(&self, srep: &RtMessage, cert_bytes: &[u8], path: &[u8], idx: u32) -> RtMessage {
        let mut index = [0; 4];
        (&mut index as &mut [u8])
            .write_u32::<LittleEndian>(idx)
            .unwrap();

        let sig_bytes = srep.get_field(Tag::SIG).unwrap();
        let srep_bytes = srep.get_field(Tag::SREP).unwrap();

        let mut response = RtMessage::new(5);
        response.add_field(Tag::SIG, sig_bytes).unwrap();
        response.add_field(Tag::PATH, path).unwrap();
        response.add_field(Tag::SREP, srep_bytes).unwrap();
        response.add_field(Tag::CERT, cert_bytes).unwrap();
        response.add_field(Tag::INDX, &index).unwrap();

        response
    }

    fn handle_health_check(&mut self) {
        let listener = self.health_listener.as_ref().unwrap();
        match listener.accept() {
            Ok((ref mut stream, src_addr)) => {
                info!("health check from {}", src_addr);
                self.stats.add_health_check(&src_addr.ip());

                match stream.write(HTTP_RESPONSE.as_bytes()) {
                    Ok(_) => (),
                    Err(e) => warn!("error writing health check {}", e),
                };

                match stream.shutdown(Shutdown::Both) {
                    Ok(_) => (),
                    Err(e) => warn!("error in health check socket shutdown {}", e),
                }
            }
            Err(ref e) if e.kind() == ErrorKind::WouldBlock => {
                debug!("blocking in TCP health check");
            }
            Err(e) => {
                warn!("unexpected health check error {}", e);
            }
        }
    }

    fn handle_status_update(&mut self) {
        let mut vec: Vec<(&IpAddr, &ClientStatEntry)> = self.stats.iter().collect();
        // sort in descending order
        vec.sort_by(|lhs, rhs| rhs.1.valid_requests.cmp(&lhs.1.valid_requests));

        for (addr, counts) in vec {
            info!(
                "{:16}: {} valid, {} invalid requests; {} responses ({} sent)",
                format!("{}", addr), counts.valid_requests, counts.invalid_requests,
                counts.responses_sent,
                counts.bytes_sent.file_size(fsopts::BINARY).unwrap()
            );
        }

        info!(
            "Totals: {} unique clients; {} valid, {} invalid requests; {} responses ({} sent)",
            self.stats.total_unique_clients(),
            self.stats.total_valid_requests(), self.stats.total_invalid_requests(),
            self.stats.total_responses_sent(),
            self.stats.total_bytes_sent().file_size(fsopts::BINARY).unwrap()
        );

        self.timer.set_timeout(self.config.status_interval(), ());
    }
}
<|endoftext|>


<commit_before>extern crate hyper;
extern crate rand;

use std::io;
use std::io::Write;
use std::fs;
use std::fs::File;
use std::borrow::Borrow;
use std::str::FromStr;

use std::sync::Mutex;

use hyper::{Get, Post};
use hyper::server::{Request, Response};
use hyper::uri::RequestUri::AbsolutePath;
use hyper::header::ContentType;
use hyper::server::Handler;
use hyper::net::Fresh;

use hyper::status::StatusCode::{UnsupportedMediaType, InternalServerError, NotFound, BadRequest};

use util::read_whole;

pub use image_url::*;

macro_rules! try_print(
    ($e:expr) => {{
        match $e {
            Ok(v) => v,
            Err(e) => { println!("Error: {}", e); return; }
        }
    }}
);

pub struct Server<'a> {
    pub image_dir : &'a str,
    pub content_type: ContentType,
    pub server_url: &'a str,
    pub run_address: &'a str,
    pub used_urls: Mutex<UsedUrlSet>,
}

impl <'a> Server<'a> {

    pub fn new(image_dir: &'a str, content_type: ContentType, run_address: &'a str, server_url: &'a str) -> Server<'a> {
        let used_urls = match fs::read_dir(image_dir) {
            Ok(paths) => {
                let mut urls = UsedUrlSet::new();
                for path in paths {
                    match path.unwrap().file_name().into_string() {
                        Ok(s) => {
                            match ImageUrl::from_str(s.split(".").nth(0).unwrap()) {
                                Ok(num) => { urls.insert(num); }
                                Err(_) => { continue; }
                            }
                        },
                        Err(_) => { continue; }
                    }
                }

                urls
            }
            Err(_) => UsedUrlSet::new()
        };

        for i in used_urls.iter() {
            println!("url exists: {}", i);
        }

        Server { image_dir: image_dir,
                 content_type: content_type,
                 server_url: server_url,
                 run_address: run_address,
                 used_urls: Mutex::new(used_urls),
                 }
    }

    // retrieves an image path
    // name.png => <image_dir>/name
    // name => <image_dir>/name
    fn get_image_path(&self, path: &str) -> String {
        // remove extension
        let base = path.split(".").nth(0).unwrap();
        self.image_dir.to_string() + base
    }

    fn retrieve_image(&self, path: &str) -> io::Result<Vec<u8>> {
        let full_path = self.get_image_path(path);
        match fs::metadata(&full_path) {
            Ok(_) => {
                let mut file = File::open(&full_path).unwrap();
                return read_whole(&mut file);
            },
            Err(e) => Err(e)
        }
    }

    fn handle_get_image_req(&self, path: &str, mut res: Response) {
        match self.retrieve_image(path) {
            Ok(data) => {
                res.headers_mut().set(self.content_type.clone());
                try_print!(res.send(data.borrow()));
            },
            Err(e) => {
                println!("Error: {}", e);
                let status = match e.kind() {
                    io::ErrorKind::NotFound | io::ErrorKind::Other => NotFound,
                    _ => InternalServerError
                };

                *res.status_mut() = status;
                try_print!(res.send(status.to_string().as_bytes()));
            }
        }
    }

    fn upload_image(&self, data: &[u8]) -> Option<String> {

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

        let mut num = gen_image_url();
        while used_urls.contains(&num) {
            num = gen_image_url();
        }

        let file_name = self.get_image_path(num.to_string().borrow());
        used_urls.insert(ImageUrl(num.0));

        return match File::create(&file_name) {
            Ok(mut file) => {
                match file.write_all(data) {
                    Ok(_) => Some(self.server_url.to_string() + "/" + num.to_string().borrow()),
                    Err(_) => None
                }
            }
            Err(_) => None
        }
    }

    fn on_req(&self, mut req: Request, mut res: Response) {
        // determine what request they are making:
        // if it is a GET method, then retrieve the image
        // if it is a POST method:
        //    - then ensure it is an image
        //    - save it locally with a random 6 character long 62 bit number (a-zA-Z0-9)
        //    - give them the URL

        match req.method {
            Get => {
                match req.uri {
                    AbsolutePath(path) => self.handle_get_image_req(&path[..], res),
                    _ => ()
                }
            }
            Post => {
                println!("received post");
                // check to see if the request has a content type
                // and it is something that we can use. If so: then we will upload it
                if !req.headers.has::<ContentType>() {
                    println!("Error: Post request has no content type");
                    *res.status_mut() = BadRequest;
                    return;
                }

                if *req.headers.get::<ContentType>().unwrap() != self.content_type {
                    println!("UnsupportedMediaType: {}", *req.headers.get::<ContentType>().unwrap());
                    println!("{}", req.headers);
                    *res.status_mut() = UnsupportedMediaType;
                    return;
                }

                // give back the URL to the client
                let data = try_print!(read_whole(&mut req));
                let url = self.upload_image(data.borrow());
                try_print!(res.send(url.unwrap().as_bytes()));
            },
            _ => ()
        }
    }
}

impl <'a> Handler for Server<'a> {
    fn handle<'b, 'k>(&'b self, req: Request<'b, 'k>, res: Response<'b, Fresh>) {
        self.on_req(req, res);
    }
}
<commit_msg>check for dir<commit_after>extern crate hyper;
extern crate rand;

use std::io;
use std::io::Write;
use std::fs;
use std::fs::File;
use std::borrow::Borrow;
use std::str::FromStr;

use std::sync::Mutex;

use hyper::{Get, Post};
use hyper::server::{Request, Response};
use hyper::uri::RequestUri::AbsolutePath;
use hyper::header::ContentType;
use hyper::server::Handler;
use hyper::net::Fresh;

use hyper::status::StatusCode::{UnsupportedMediaType, InternalServerError, NotFound, BadRequest};

use util::read_whole;

pub use image_url::*;

macro_rules! try_print(
    ($e:expr) => {{
        match $e {
            Ok(v) => v,
            Err(e) => { println!("Error: {}", e); return; }
        }
    }}
);

pub struct Server<'a> {
    pub image_dir : &'a str,
    pub content_type: ContentType,
    pub server_url: &'a str,
    pub run_address: &'a str,
    pub used_urls: Mutex<UsedUrlSet>,
}

impl <'a> Server<'a> {

    pub fn new(image_dir: &'a str, content_type: ContentType, run_address: &'a str, server_url: &'a str) -> Server<'a> {
        let used_urls = match fs::read_dir(image_dir) {
            Ok(paths) => {
                let mut urls = UsedUrlSet::new();
                for path in paths {
                    match path.unwrap().file_name().into_string() {
                        Ok(s) => {
                            match ImageUrl::from_str(s.split(".").nth(0).unwrap()) {
                                Ok(num) => { urls.insert(num); }
                                Err(_) => { continue; }
                            }
                        },
                        Err(_) => { continue; }
                    }
                }

                urls
            }
            Err(_) => UsedUrlSet::new()
        };

        for i in used_urls.iter() {
            println!("url exists: {}", i);
        }

        Server { image_dir: image_dir,
                 content_type: content_type,
                 server_url: server_url,
                 run_address: run_address,
                 used_urls: Mutex::new(used_urls),
                 }
    }

    // retrieves an image path
    // name.png => <image_dir>/name
    // name => <image_dir>/name
    fn get_image_path(&self, path: &str) -> String {
        // remove extension
        let base = path.split(".").nth(0).unwrap();
        self.image_dir.to_string() + base
    }

    fn retrieve_image(&self, path: &str) -> io::Result<Vec<u8>> {
        let full_path = self.get_image_path(path);
        match fs::metadata(&full_path) {
            Ok(metadata) => {
                if metadata.is_file() {
                    let mut file = File::open(&full_path).unwrap();
                    return read_whole(&mut file);
                } else {
                    return Err(io::Error::new(io::ErrorKind::NotFound, "Directory"));
                }
            },
            Err(e) => Err(e)
        }
    }

    fn handle_get_image_req(&self, path: &str, mut res: Response) {
        match self.retrieve_image(path) {
            Ok(data) => {
                res.headers_mut().set(self.content_type.clone());
                try_print!(res.send(data.borrow()));
            },
            Err(e) => {
                println!("Error: {}", e);
                let status = match e.kind() {
                    io::ErrorKind::NotFound | io::ErrorKind::Other => NotFound,
                    _ => InternalServerError
                };

                *res.status_mut() = status;
                try_print!(res.send(status.to_string().as_bytes()));
            }
        }
    }

    fn upload_image(&self, data: &[u8]) -> Option<String> {

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

        let mut num = gen_image_url();
        while used_urls.contains(&num) {
            num = gen_image_url();
        }

        let file_name = self.get_image_path(num.to_string().borrow());
        used_urls.insert(ImageUrl(num.0));

        return match File::create(&file_name) {
            Ok(mut file) => {
                match file.write_all(data) {
                    Ok(_) => Some(self.server_url.to_string() + "/" + num.to_string().borrow()),
                    Err(_) => None
                }
            }
            Err(_) => None
        }
    }

    fn on_req(&self, mut req: Request, mut res: Response) {
        // determine what request they are making:
        // if it is a GET method, then retrieve the image
        // if it is a POST method:
        //    - then ensure it is an image
        //    - save it locally with a random 6 character long 62 bit number (a-zA-Z0-9)
        //    - give them the URL

        match req.method {
            Get => {
                match req.uri {
                    AbsolutePath(path) => self.handle_get_image_req(&path[..], res),
                    _ => ()
                }
            }
            Post => {
                println!("received post");
                // check to see if the request has a content type
                // and it is something that we can use. If so: then we will upload it
                if !req.headers.has::<ContentType>() {
                    println!("Error: Post request has no content type");
                    *res.status_mut() = BadRequest;
                    return;
                }

                if *req.headers.get::<ContentType>().unwrap() != self.content_type {
                    println!("UnsupportedMediaType: {}", *req.headers.get::<ContentType>().unwrap());
                    println!("{}", req.headers);
                    *res.status_mut() = UnsupportedMediaType;
                    return;
                }

                // give back the URL to the client
                let data = try_print!(read_whole(&mut req));
                let url = self.upload_image(data.borrow());
                try_print!(res.send(url.unwrap().as_bytes()));
            },
            _ => ()
        }
    }
}

impl <'a> Handler for Server<'a> {
    fn handle<'b, 'k>(&'b self, req: Request<'b, 'k>, res: Response<'b, Fresh>) {
        self.on_req(req, res);
    }
}
<|endoftext|>


<commit_before>use luminance::shader::program::Program as LProgram;
use luminance::shader::stage::{Stage, StageError, Type};
use std::fs::File;
use std::io::{BufRead, BufReader};
use std::ops::Deref;
use std::path::Path;

pub use luminance::shader::program::{ProgramError, Uniform, Uniformable, UniformBuilder,
                                     UniformInterface, UniformWarning};
use luminance::vertex::Vertex;

use resource::{Load, LoadError, ResCache};

#[derive(Debug)]
pub enum ShaderError {
  StageError(StageError),
  ProgramError(ProgramError)
}

pub fn new_program<In, Out, Uni>(tcs_src: &str,
                                 tes_src: &str,
                                 vs_src: &str,
                                 gs_src: &str,
                                 fs_src: &str)
                                 -> Result<(LProgram<In, Out, Uni>, Vec<UniformWarning>), ProgramError>
    where In: Vertex,
          Uni: UniformInterface {
  let stages = compile_stages(tcs_src, tes_src, vs_src, gs_src, fs_src);

  match stages {
    Ok((tess, vs, gs, fs)) => {
      if let Some((tcs, tes)) = tess {
        if let Some(gs) = gs {
          LProgram::new(Some((&tcs, &tes)), &vs, Some(&gs), &fs)
        } else {
          LProgram::new(Some((&tcs, &tes)), &vs, None, &fs)
        }
      } else if let Some(gs) = gs {
        LProgram::new(None, &vs, Some(&gs), &fs)
      } else {
        LProgram::new(None, &vs, None, &fs)
      }
    },
    Err(stage_error) => {
      Err(ProgramError::LinkFailed(format!("{:?}", stage_error)))
    }
  }
}

// Take raw shader sources and turn them into stages.
fn compile_stages(tcs_src: &str, tes_src: &str, vs_src: &str, gs_src: &str, fs_src: &str) -> Result<(Option<(Stage, Stage)>, Stage, Option<Stage>, Stage), StageError> {
  let tess = if !tcs_src.is_empty() && !tes_src.is_empty() {
    Some((Stage::new(Type::TessellationControlShader, tcs_src)?,
          Stage::new(Type::TessellationEvaluationShader, tes_src)?))
  } else {
    None
  };

  let vs = Stage::new(Type::VertexShader, vs_src)?;
  let gs = if !gs_src.is_empty() { Some(Stage::new(Type::GeometryShader, gs_src)?) } else { None };
  let fs = Stage::new(Type::FragmentShader, fs_src)?;

  Ok((tess, vs, gs, fs))
}

/// Shader program.
///
/// If the program is retrieved from the cache, the path must point to a file containing all the
/// stages. A stage begins with a stage pragma indicating which kind of stage the following source
/// is. Here’s a list of all pragma for each kind of stage:
///
/// - `#tcs`: *tessellation control stage*
/// - `#tes`: *tessellation evaluation stage*
/// - `#vs`: *vertex stage*
/// - `#gs`: *geometry stage*
/// - `#fs`: *fragment stage*
///
/// A stage starts at such a pragma listed above and ends at the next, different pragma. You cannot
/// use twice the same pragma in a file.
///
/// At the top of the file, if you don’t put a pragma, you can use `//` to add comments, or die.
pub struct Program<In, Out, Uni> {
  program: LProgram<In, Out, Uni>
}

impl<In, Out, Uni> Deref for Program<In, Out, Uni> {
  type Target = LProgram<In, Out, Uni>;

  fn deref(&self) -> &Self::Target {
    &self.program
  }
}

impl<In, Out, Uni> Load for Program<In, Out, Uni> where In: Vertex, Uni: UniformInterface {
  type Args = ();

  const TY_STR: &'static str = "shaders";

  fn load<P>(path: P, _: &mut ResCache, _: Self::Args) -> Result<Self, LoadError> where P: AsRef<Path> {
    let path = path.as_ref();

    info!("loading shader: {:?}", path);

    enum CurrentStage {
      VS,
      FS,
      GS,
      TCS,
      TES
    }

    fn add_line_to_src(src: &mut String, line: &str, line_nb: usize) {
      *src += &format!("#line {}\n{}\n", line_nb, line);
    }

    match File::open(path) {
      Ok(fh) => {
        let buffered = BufReader::new(fh);
        let mut tcs_src = String::new();
        let mut tes_src = String::new();
        let mut vs_src = String::new();
        let mut gs_src = String::new();
        let mut fs_src = String::new();
        let mut current_stage: Option<CurrentStage> = None;

        for (line_nb, line) in buffered.lines().enumerate() {
          let line_nb = line_nb + 1;
          let line = line.unwrap();
          let trimmed = line.trim();

          if trimmed.starts_with("#vs") {
            if !vs_src.is_empty() {
              return Err(LoadError::ParseFailed(format!("(line {}) several #vs sections", line_nb)));
            }

            info!("  found a vertex shader");

            current_stage = Some(CurrentStage::VS);
            continue;
          } else if trimmed.starts_with("#fs") {
            if !fs_src.is_empty() {
              return Err(LoadError::ParseFailed(format!("(line {}) several #fs sections", line_nb)));
            }

            info!("  found a fragment shader");

            current_stage = Some(CurrentStage::FS);
            continue;
          } else if trimmed.starts_with("#gs") {
            if !gs_src.is_empty() {
              return Err(LoadError::ParseFailed(format!("(line {}) several #gs sections", line_nb)));
            }

            info!("  found a geometry shader");

            current_stage = Some(CurrentStage::GS);
            continue;
          } else if trimmed.starts_with("#tcs") {
            if !tcs_src.is_empty() {
              return Err(LoadError::ParseFailed(format!("(line {}) several #tcs sections", line_nb)));
            }

            info!("  found a tessellation control shader");

            current_stage = Some(CurrentStage::TCS);
            continue;
          } else if trimmed.starts_with("#tes") {
            if !tes_src.is_empty() {
              return Err(LoadError::ParseFailed(format!("(line {}) several #tes sections", line_nb)));
            }

            info!("  found a tessellation evaluation shader");

            current_stage = Some(CurrentStage::TES);
            continue;
          } else if current_stage.is_none() && !trimmed.is_empty() && !trimmed.starts_with("//") && !trimmed.starts_with("\n") {
            return Err(LoadError::ParseFailed(format!("(line {}) not in a shader stage nor a comment", line_nb)));
          }

          match current_stage {
            Some(CurrentStage::VS) => {
              add_line_to_src(&mut vs_src, trimmed, line_nb);
            },
            Some(CurrentStage::FS) => {
              add_line_to_src(&mut fs_src, trimmed, line_nb);
            },
            Some(CurrentStage::GS) => {
              add_line_to_src(&mut gs_src, trimmed, line_nb);
            },
            Some(CurrentStage::TCS) => {
              add_line_to_src(&mut tcs_src, trimmed, line_nb);
            },
            Some(CurrentStage::TES) => {
              add_line_to_src(&mut tes_src, trimmed, line_nb);
            },
            None => {}
          }
        }

        let (program, warnings) = new_program(&tcs_src, &tes_src, &vs_src, &gs_src, &fs_src)
          .map_err(|e| LoadError::ConversionFailed(format!("{:#?}", e)))?;

        // check for semantic errors
        for warning in warnings {
          warn!("uniform warning: {:?}", warning);
        }

        Ok(
          Program {
            program: program
          }
        )
      },
      Err(_) => Err(LoadError::FileNotFound(path.to_owned()))
    }
  }
}

pub trait UnwrapOrUnbound<T> {
  fn unwrap_or_unbound(self, builder: &UniformBuilder, warnings: &mut Vec<UniformWarning>) -> Uniform<T> where T: Uniformable;
}

impl<T> UnwrapOrUnbound<T> for Result<Uniform<T>, UniformWarning> {
  fn unwrap_or_unbound(self, builder: &UniformBuilder, warnings: &mut Vec<UniformWarning>) -> Uniform<T> where T: Uniformable {
    self.unwrap_or_else(|w| {
      warnings.push(w);
      builder.unbound()
    })
  }
}
<commit_msg>Added from_ to Program so that it’s possible to build Programs in an easy way.<commit_after>use luminance::shader::program::Program as LProgram;
use luminance::shader::stage::{Stage, StageError, Type};
use std::fs::File;
use std::io::{BufRead, BufReader};
use std::ops::Deref;
use std::path::Path;

pub use luminance::shader::program::{ProgramError, Uniform, Uniformable, UniformBuilder,
                                     UniformInterface, UniformWarning};
use luminance::vertex::Vertex;

use resource::{Load, LoadError, ResCache};

#[derive(Debug)]
pub enum ShaderError {
  StageError(StageError),
  ProgramError(ProgramError)
}

/// Create a new luminance::Program from a set of shader strings.
fn new_program<In, Out, Uni>(tcs_src: &str,
                             tes_src: &str,
                             vs_src: &str,
                             gs_src: &str,
                             fs_src: &str)
                             -> Result<(LProgram<In, Out, Uni>, Vec<UniformWarning>), ProgramError>
    where In: Vertex,
          Uni: UniformInterface {
  let stages = compile_stages(tcs_src, tes_src, vs_src, gs_src, fs_src);

  match stages {
    Ok((tess, vs, gs, fs)) => {
      if let Some((tcs, tes)) = tess {
        if let Some(gs) = gs {
          LProgram::new(Some((&tcs, &tes)), &vs, Some(&gs), &fs)
        } else {
          LProgram::new(Some((&tcs, &tes)), &vs, None, &fs)
        }
      } else if let Some(gs) = gs {
        LProgram::new(None, &vs, Some(&gs), &fs)
      } else {
        LProgram::new(None, &vs, None, &fs)
      }
    },
    Err(stage_error) => {
      Err(ProgramError::LinkFailed(format!("{:?}", stage_error)))
    }
  }
}

/// Take raw shader sources and turn them into stages.
fn compile_stages(tcs_src: &str, tes_src: &str, vs_src: &str, gs_src: &str, fs_src: &str) -> Result<(Option<(Stage, Stage)>, Stage, Option<Stage>, Stage), StageError> {
  let tess = if !tcs_src.is_empty() && !tes_src.is_empty() {
    Some((Stage::new(Type::TessellationControlShader, tcs_src)?,
          Stage::new(Type::TessellationEvaluationShader, tes_src)?))
  } else {
    None
  };

  let vs = Stage::new(Type::VertexShader, vs_src)?;
  let gs = if !gs_src.is_empty() { Some(Stage::new(Type::GeometryShader, gs_src)?) } else { None };
  let fs = Stage::new(Type::FragmentShader, fs_src)?;

  Ok((tess, vs, gs, fs))
}

/// Shader program.
///
/// If the program is retrieved from the cache, the path must point to a file containing all the
/// stages. A stage begins with a stage pragma indicating which kind of stage the following source
/// is. Here’s a list of all pragmas for each kind of stage:
///
/// - `#tcs`: *tessellation control stage*
/// - `#tes`: *tessellation evaluation stage*
/// - `#vs`: *vertex stage*
/// - `#gs`: *geometry stage*
/// - `#fs`: *fragment stage*
///
/// A stage starts at such a pragma listed above and ends at the next, different pragma. You cannot
/// use twice the same pragma in a file.
///
/// At the top of the file, if you don’t put a pragma, you can use `//` to add comments, or die.
///
/// The end of the file marks the end of the last pragma.
pub struct Program<In, Out, Uni> {
  program: LProgram<In, Out, Uni>
}

impl<In, Out, Uni> Deref for Program<In, Out, Uni> {
  type Target = LProgram<In, Out, Uni>;

  fn deref(&self) -> &Self::Target {
    &self.program
  }
}

// Current stage dispatch.
enum CurrentStage {
  VS,
  FS,
  GS,
  TCS,
  TES
}

struct ShaderSources {
  tcs_src: String,
  tes_src: String,
  vs_src: String,
  gs_src: String,
  fs_src: String
}

// Annotate a line with its original line number.
fn annotate_line_src(src: &mut String, line: &str, line_nb: usize) {
  *src += &format!("#line {}\n{}\n", line_nb, line);
}

// Split a single source into several shader sources.
fn split_shader_stages_sources<R>(buffered: R) -> Result<ShaderSources, LoadError> where R: BufRead {
  let mut current_stage: Option<CurrentStage> = None;
  let mut tcs_src = String::new();
  let mut tes_src = String::new();
  let mut vs_src = String::new();
  let mut gs_src = String::new();
  let mut fs_src = String::new();

  for (line_nb, line) in buffered.lines().enumerate() {
    let line_nb = line_nb + 1;
    let line = line.unwrap();
    let trimmed = line.trim();

    if trimmed.starts_with("#vs") {
      if !vs_src.is_empty() {
        return Err(LoadError::ParseFailed(format!("(line {}) several #vs sections", line_nb)));
      }

      info!("  found a vertex shader");

      current_stage = Some(CurrentStage::VS);
      continue;
    } else if trimmed.starts_with("#fs") {
      if !fs_src.is_empty() {
        return Err(LoadError::ParseFailed(format!("(line {}) several #fs sections", line_nb)));
      }

      info!("  found a fragment shader");

      current_stage = Some(CurrentStage::FS);
      continue;
    } else if trimmed.starts_with("#gs") {
      if !gs_src.is_empty() {
        return Err(LoadError::ParseFailed(format!("(line {}) several #gs sections", line_nb)));
      }

      info!("  found a geometry shader");

      current_stage = Some(CurrentStage::GS);
      continue;
    } else if trimmed.starts_with("#tcs") {
      if !tcs_src.is_empty() {
        return Err(LoadError::ParseFailed(format!("(line {}) several #tcs sections", line_nb)));
      }

      info!("  found a tessellation control shader");

      current_stage = Some(CurrentStage::TCS);
      continue;
    } else if trimmed.starts_with("#tes") {
      if !tes_src.is_empty() {
        return Err(LoadError::ParseFailed(format!("(line {}) several #tes sections", line_nb)));
      }

      info!("  found a tessellation evaluation shader");

      current_stage = Some(CurrentStage::TES);
      continue;
    } else if current_stage.is_none() && !trimmed.is_empty() && !trimmed.starts_with("//") && !trimmed.starts_with("\n") {
      return Err(LoadError::ParseFailed(format!("(line {}) not in a shader stage nor a comment", line_nb)));
    }

    match current_stage {
      Some(CurrentStage::TCS) => {
        annotate_line_src(&mut tcs_src, trimmed, line_nb);
      },
      Some(CurrentStage::TES) => {
        annotate_line_src(&mut tes_src, trimmed, line_nb);
      },
      Some(CurrentStage::VS) => {
        annotate_line_src(&mut vs_src, trimmed, line_nb);
      },
      Some(CurrentStage::GS) => {
        annotate_line_src(&mut gs_src, trimmed, line_nb);
      },
      Some(CurrentStage::FS) => {
        annotate_line_src(&mut fs_src, trimmed, line_nb);
      },
      None => {}
    }
  }

  Ok(ShaderSources {
    tcs_src: tcs_src,
    tes_src: tes_src,
    vs_src: vs_src,
    gs_src: gs_src,
    fs_src: fs_src
  })
}

impl<In, Out, Uni> Program<In, Out, Uni> where In: Vertex, Uni: UniformInterface {
  /// Create a `Program` from a `BufRead`.
  pub fn from_bufread<R>(buffered: R) -> Result<Self, LoadError> where R: BufRead {
    let sources = split_shader_stages_sources(buffered)?;
    let (program, warnings) = new_program(&sources.tcs_src,
                                          &sources.tes_src,
                                          &sources.vs_src,
                                          &sources.gs_src,
                                          &sources.fs_src)
      .map_err(|e| LoadError::ConversionFailed(format!("{:#?}", e)))?;

    // check for semantic errors
    for warning in warnings {
      warn!("uniform warning: {:?}", warning);
    }

    Ok(Program {
      program: program
    })
  }

  /// Create a `Program` from a string – you can for instance use `str` or `String`.
  pub fn from_str<S>(s: S) -> Result<Self, LoadError> where S: for<'a> Into<&'a str> {
    Self::from_bufread(s.into().as_bytes())
  }
}

impl<In, Out, Uni> Load for Program<In, Out, Uni> where In: Vertex, Uni: UniformInterface {
  type Args = ();

  const TY_STR: &'static str = "shaders";

  fn load<P>(path: P, _: &mut ResCache, _: Self::Args) -> Result<Self, LoadError> where P: AsRef<Path> {
    let path = path.as_ref();

    info!("loading shader: {:?}", path);

    enum CurrentStage {
      VS,
      FS,
      GS,
      TCS,
      TES
    }

    fn annotate_line_src(src: &mut String, line: &str, line_nb: usize) {
      *src += &format!("#line {}\n{}\n", line_nb, line);
    }

    match File::open(path) {
      Ok(fh) => {
        let buffered = BufReader::new(fh);
        let mut tcs_src = String::new();
        let mut tes_src = String::new();
        let mut vs_src = String::new();
        let mut gs_src = String::new();
        let mut fs_src = String::new();
        let mut current_stage: Option<CurrentStage> = None;

        for (line_nb, line) in buffered.lines().enumerate() {
          let line_nb = line_nb + 1;
          let line = line.unwrap();
          let trimmed = line.trim();

          if trimmed.starts_with("#vs") {
            if !vs_src.is_empty() {
              return Err(LoadError::ParseFailed(format!("(line {}) several #vs sections", line_nb)));
            }

            info!("  found a vertex shader");

            current_stage = Some(CurrentStage::VS);
            continue;
          } else if trimmed.starts_with("#fs") {
            if !fs_src.is_empty() {
              return Err(LoadError::ParseFailed(format!("(line {}) several #fs sections", line_nb)));
            }

            info!("  found a fragment shader");

            current_stage = Some(CurrentStage::FS);
            continue;
          } else if trimmed.starts_with("#gs") {
            if !gs_src.is_empty() {
              return Err(LoadError::ParseFailed(format!("(line {}) several #gs sections", line_nb)));
            }

            info!("  found a geometry shader");

            current_stage = Some(CurrentStage::GS);
            continue;
          } else if trimmed.starts_with("#tcs") {
            if !tcs_src.is_empty() {
              return Err(LoadError::ParseFailed(format!("(line {}) several #tcs sections", line_nb)));
            }

            info!("  found a tessellation control shader");

            current_stage = Some(CurrentStage::TCS);
            continue;
          } else if trimmed.starts_with("#tes") {
            if !tes_src.is_empty() {
              return Err(LoadError::ParseFailed(format!("(line {}) several #tes sections", line_nb)));
            }

            info!("  found a tessellation evaluation shader");

            current_stage = Some(CurrentStage::TES);
            continue;
          } else if current_stage.is_none() && !trimmed.is_empty() && !trimmed.starts_with("//") && !trimmed.starts_with("\n") {
            return Err(LoadError::ParseFailed(format!("(line {}) not in a shader stage nor a comment", line_nb)));
          }

          match current_stage {
            Some(CurrentStage::VS) => {
              annotate_line_src(&mut vs_src, trimmed, line_nb);
            },
            Some(CurrentStage::FS) => {
              annotate_line_src(&mut fs_src, trimmed, line_nb);
            },
            Some(CurrentStage::GS) => {
              annotate_line_src(&mut gs_src, trimmed, line_nb);
            },
            Some(CurrentStage::TCS) => {
              annotate_line_src(&mut tcs_src, trimmed, line_nb);
            },
            Some(CurrentStage::TES) => {
              annotate_line_src(&mut tes_src, trimmed, line_nb);
            },
            None => {}
          }
        }

        let (program, warnings) = new_program(&tcs_src, &tes_src, &vs_src, &gs_src, &fs_src)
          .map_err(|e| LoadError::ConversionFailed(format!("{:#?}", e)))?;

        // check for semantic errors
        for warning in warnings {
          warn!("uniform warning: {:?}", warning);
        }

        Ok(
          Program {
            program: program
          }
        )
      },
      Err(_) => Err(LoadError::FileNotFound(path.to_owned()))
    }
  }
}

pub trait UnwrapOrUnbound<T> {
  fn unwrap_or_unbound(self, builder: &UniformBuilder, warnings: &mut Vec<UniformWarning>) -> Uniform<T> where T: Uniformable;
}

impl<T> UnwrapOrUnbound<T> for Result<Uniform<T>, UniformWarning> {
  fn unwrap_or_unbound(self, builder: &UniformBuilder, warnings: &mut Vec<UniformWarning>) -> Uniform<T> where T: Uniformable {
    self.unwrap_or_else(|w| {
      warnings.push(w);
      builder.unbound()
    })
  }
}
<|endoftext|>