<commit_before>use std::{
    cmp,
    error::Error as StdError,
    future::Future,
    marker::PhantomData,
    net,
    pin::Pin,
    rc::Rc,
    task::{Context, Poll},
};

use actix_codec::{AsyncRead, AsyncWrite};
use actix_rt::time::{sleep, Sleep};
use actix_service::Service;
use actix_utils::future::poll_fn;
use bytes::{Bytes, BytesMut};
use futures_core::ready;
use h2::{
    server::{Connection, SendResponse},
    Ping, PingPong,
};
use http::header::{HeaderValue, CONNECTION, CONTENT_LENGTH, DATE, TRANSFER_ENCODING};
use log::{error, trace};
use pin_project_lite::pin_project;

use crate::{
    body::{BodySize, BoxBody, MessageBody},
    config::ServiceConfig,
    service::HttpFlow,
    Extensions, OnConnectData, Payload, Request, Response, ResponseHead,
};

const CHUNK_SIZE: usize = 16_384;

pin_project! {
    /// Dispatcher for HTTP/2 protocol.
    pub struct Dispatcher<T, S, B, X, U> {
        flow: Rc<HttpFlow<S, X, U>>,
        connection: Connection<T, Bytes>,
        conn_data: Option<Rc<Extensions>>,
        config: ServiceConfig,
        peer_addr: Option<net::SocketAddr>,
        ping_pong: Option<H2PingPong>,
        _phantom: PhantomData<B>
    }
}

impl<T, S, B, X, U> Dispatcher<T, S, B, X, U>
where
    T: AsyncRead + AsyncWrite + Unpin,
{
    pub(crate) fn new(
        mut conn: Connection<T, Bytes>,
        flow: Rc<HttpFlow<S, X, U>>,
        config: ServiceConfig,
        peer_addr: Option<net::SocketAddr>,
        conn_data: OnConnectData,
        timer: Option<Pin<Box<Sleep>>>,
    ) -> Self {
        let ping_pong = config.keep_alive().map(|dur| H2PingPong {
            timer: timer
                .map(|mut timer| {
                    // reset timer if it's received from new function.
                    timer.as_mut().reset(config.now() + dur);
                    timer
                })
                .unwrap_or_else(|| Box::pin(sleep(dur))),
            on_flight: false,
            ping_pong: conn.ping_pong().unwrap(),
        });

        Self {
            flow,
            config,
            peer_addr,
            connection: conn,
            conn_data: conn_data.0.map(Rc::new),
            ping_pong,
            _phantom: PhantomData,
        }
    }
}

struct H2PingPong {
    timer: Pin<Box<Sleep>>,
    on_flight: bool,
    ping_pong: PingPong,
}

impl<T, S, B, X, U> Future for Dispatcher<T, S, B, X, U>
where
    T: AsyncRead + AsyncWrite + Unpin,

    S: Service<Request>,
    S::Error: Into<Response<BoxBody>>,
    S::Future: 'static,
    S::Response: Into<Response<B>>,

    B: MessageBody,
{
    type Output = Result<(), crate::error::DispatchError>;

    #[inline]
    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
        let this = self.get_mut();

        loop {
            match Pin::new(&mut this.connection).poll_accept(cx)? {
                Poll::Ready(Some((req, tx))) => {
                    let (parts, body) = req.into_parts();
                    let pl = crate::h2::Payload::new(body);
                    let pl = Payload::H2(pl);
                    let mut req = Request::with_payload(pl);

                    let head = req.head_mut();
                    head.uri = parts.uri;
                    head.method = parts.method;
                    head.version = parts.version;
                    head.headers = parts.headers.into();
                    head.peer_addr = this.peer_addr;

                    req.conn_data = this.conn_data.as_ref().map(Rc::clone);

                    let fut = this.flow.service.call(req);
                    let config = this.config.clone();

                    // multiplex request handling with spawn task
                    actix_rt::spawn(async move {
                        // resolve service call and send response.
                        let res = match fut.await {
                            Ok(res) => handle_response(res.into(), tx, config).await,
                            Err(err) => {
                                let res: Response<BoxBody> = err.into();
                                handle_response(res, tx, config).await
                            }
                        };

                        // log error.
                        if let Err(err) = res {
                            match err {
                                DispatchError::SendResponse(err) => {
                                    trace!("Error sending HTTP/2 response: {:?}", err)
                                }
                                DispatchError::SendData(err) => warn!("{:?}", err),
                                DispatchError::ResponseBody(err) => {
                                    error!("Response payload stream error: {:?}", err)
                                }
                            }
                        }
                    });
                }
                Poll::Ready(None) => return Poll::Ready(Ok(())),
                Poll::Pending => match this.ping_pong.as_mut() {
                    Some(ping_pong) => loop {
                        if ping_pong.on_flight {
                            // When have on flight ping pong. poll pong and and keep alive timer.
                            // on success pong received update keep alive timer to determine the next timing of
                            // ping pong.
                            match ping_pong.ping_pong.poll_pong(cx)? {
                                Poll::Ready(_) => {
                                    ping_pong.on_flight = false;

                                    let dead_line = this.config.keep_alive_expire().unwrap();
                                    ping_pong.timer.as_mut().reset(dead_line);
                                }
                                Poll::Pending => {
                                    return ping_pong.timer.as_mut().poll(cx).map(|_| Ok(()))
                                }
                            }
                        } else {
                            // When there is no on flight ping pong. keep alive timer is used to wait for next
                            // timing of ping pong. Therefore at this point it serves as an interval instead.
                            ready!(ping_pong.timer.as_mut().poll(cx));

                            ping_pong.ping_pong.send_ping(Ping::opaque())?;

                            let dead_line = this.config.keep_alive_expire().unwrap();
                            ping_pong.timer.as_mut().reset(dead_line);

                            ping_pong.on_flight = true;
                        }
                    },
                    None => return Poll::Pending,
                },
            }
        }
    }
}

enum DispatchError {
    SendResponse(h2::Error),
    SendData(h2::Error),
    ResponseBody(Box<dyn StdError>),
}

async fn handle_response<B>(
    res: Response<B>,
    mut tx: SendResponse<Bytes>,
    config: ServiceConfig,
) -> Result<(), DispatchError>
where
    B: MessageBody,
{
    let (res, body) = res.replace_body(());

    // prepare response.
    let mut size = body.size();
    let res = prepare_response(config, res.head(), &mut size);
    let eof = size.is_eof();

    // send response head and return on eof.
    let mut stream = tx
        .send_response(res, eof)
        .map_err(DispatchError::SendResponse)?;

    if eof {
        return Ok(());
    }

    // poll response body and send chunks to client.
    actix_rt::pin!(body);

    while let Some(res) = poll_fn(|cx| body.as_mut().poll_next(cx)).await {
        let mut chunk = res.map_err(|err| DispatchError::ResponseBody(err.into()))?;

        'send: loop {
            // reserve enough space and wait for stream ready.
            stream.reserve_capacity(cmp::min(chunk.len(), CHUNK_SIZE));

            match poll_fn(|cx| stream.poll_capacity(cx)).await {
                // No capacity left. drop body and return.
                None => return Ok(()),
                Some(res) => {
                    // Split chuck to writeable size and send to client.
                    let cap = res.map_err(DispatchError::SendData)?;

                    let len = chunk.len();
                    let bytes = chunk.split_to(cmp::min(cap, len));

                    stream
                        .send_data(bytes, false)
                        .map_err(DispatchError::SendData)?;

                    // Current chuck completely sent. break send loop and poll next one.
                    if chunk.is_empty() {
                        break 'send;
                    }
                }
            }
        }
    }

    // response body streaming finished. send end of stream and return.
    stream
        .send_data(Bytes::new(), true)
        .map_err(DispatchError::SendData)?;

    Ok(())
}

fn prepare_response(
    config: ServiceConfig,
    head: &ResponseHead,
    size: &mut BodySize,
) -> http::Response<()> {
    let mut has_date = false;
    let mut skip_len = size != &BodySize::Stream;

    let mut res = http::Response::new(());
    *res.status_mut() = head.status;
    *res.version_mut() = http::Version::HTTP_2;

    // Content length
    match head.status {
        http::StatusCode::NO_CONTENT
        | http::StatusCode::CONTINUE
        | http::StatusCode::PROCESSING => *size = BodySize::None,
        http::StatusCode::SWITCHING_PROTOCOLS => {
            skip_len = true;
            *size = BodySize::Stream;
        }
        _ => {}
    }

    let _ = match size {
        BodySize::None | BodySize::Stream => None,

        BodySize::Sized(0) => res
            .headers_mut()
            .insert(CONTENT_LENGTH, HeaderValue::from_static("0")),

        BodySize::Sized(len) => {
            let mut buf = itoa::Buffer::new();

            res.headers_mut().insert(
                CONTENT_LENGTH,
                HeaderValue::from_str(buf.format(*len)).unwrap(),
            )
        }
    };

    // copy headers
    for (key, value) in head.headers.iter() {
        match *key {
            // TODO: consider skipping other headers according to:
            //       https://datatracker.ietf.org/doc/html/rfc7540#section-8.1.2.2
            // omit HTTP/1.x only headers
            CONNECTION | TRANSFER_ENCODING => continue,
            CONTENT_LENGTH if skip_len => continue,
            DATE => has_date = true,
            _ => {}
        }

        res.headers_mut().append(key, value.clone());
    }

    // set date header
    if !has_date {
        let mut bytes = BytesMut::with_capacity(29);
        config.set_date_header(&mut bytes);
        res.headers_mut().insert(
            DATE,
            // SAFETY: serialized date-times are known ASCII strings
            unsafe { HeaderValue::from_maybe_shared_unchecked(bytes.freeze()) },
        );
    }

    res
}
<commit_msg>clean future h2 dispatcher<commit_after>use std::{
    cmp,
    error::Error as StdError,
    future::Future,
    marker::PhantomData,
    net,
    pin::Pin,
    rc::Rc,
    task::{Context, Poll},
};

use actix_codec::{AsyncRead, AsyncWrite};
use actix_rt::time::{sleep, Sleep};
use actix_service::Service;
use actix_utils::future::poll_fn;
use bytes::{Bytes, BytesMut};
use futures_core::ready;
use h2::{
    server::{Connection, SendResponse},
    Ping, PingPong,
};
use log::{error, trace};
use pin_project_lite::pin_project;

use crate::{
    body::{BodySize, BoxBody, MessageBody},
    config::ServiceConfig,
    header::{HeaderValue, CONNECTION, CONTENT_LENGTH, DATE, TRANSFER_ENCODING},
    service::HttpFlow,
    Extensions, OnConnectData, Payload, Request, Response, ResponseHead,
};

const CHUNK_SIZE: usize = 16_384;

pin_project! {
    /// Dispatcher for HTTP/2 protocol.
    pub struct Dispatcher<T, S, B, X, U> {
        flow: Rc<HttpFlow<S, X, U>>,
        connection: Connection<T, Bytes>,
        conn_data: Option<Rc<Extensions>>,
        config: ServiceConfig,
        peer_addr: Option<net::SocketAddr>,
        ping_pong: Option<H2PingPong>,
        _phantom: PhantomData<B>
    }
}

impl<T, S, B, X, U> Dispatcher<T, S, B, X, U>
where
    T: AsyncRead + AsyncWrite + Unpin,
{
    pub(crate) fn new(
        mut conn: Connection<T, Bytes>,
        flow: Rc<HttpFlow<S, X, U>>,
        config: ServiceConfig,
        peer_addr: Option<net::SocketAddr>,
        conn_data: OnConnectData,
        timer: Option<Pin<Box<Sleep>>>,
    ) -> Self {
        let ping_pong = config.keep_alive().map(|dur| H2PingPong {
            timer: timer
                .map(|mut timer| {
                    // reset timer if it's received from new function.
                    timer.as_mut().reset(config.now() + dur);
                    timer
                })
                .unwrap_or_else(|| Box::pin(sleep(dur))),
            on_flight: false,
            ping_pong: conn.ping_pong().unwrap(),
        });

        Self {
            flow,
            config,
            peer_addr,
            connection: conn,
            conn_data: conn_data.0.map(Rc::new),
            ping_pong,
            _phantom: PhantomData,
        }
    }
}

struct H2PingPong {
    timer: Pin<Box<Sleep>>,
    on_flight: bool,
    ping_pong: PingPong,
}

impl<T, S, B, X, U> Future for Dispatcher<T, S, B, X, U>
where
    T: AsyncRead + AsyncWrite + Unpin,

    S: Service<Request>,
    S::Error: Into<Response<BoxBody>>,
    S::Future: 'static,
    S::Response: Into<Response<B>>,

    B: MessageBody,
{
    type Output = Result<(), crate::error::DispatchError>;

    #[inline]
    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
        let this = self.get_mut();

        loop {
            match Pin::new(&mut this.connection).poll_accept(cx)? {
                Poll::Ready(Some((req, tx))) => {
                    let (parts, body) = req.into_parts();
                    let pl = crate::h2::Payload::new(body);
                    let pl = Payload::H2(pl);
                    let mut req = Request::with_payload(pl);

                    let head = req.head_mut();
                    head.uri = parts.uri;
                    head.method = parts.method;
                    head.version = parts.version;
                    head.headers = parts.headers.into();
                    head.peer_addr = this.peer_addr;

                    req.conn_data = this.conn_data.as_ref().map(Rc::clone);

                    let fut = this.flow.service.call(req);
                    let config = this.config.clone();

                    // multiplex request handling with spawn task
                    actix_rt::spawn(async move {
                        // resolve service call and send response.
                        let res = match fut.await {
                            Ok(res) => handle_response(res.into(), tx, config).await,
                            Err(err) => {
                                let res: Response<BoxBody> = err.into();
                                handle_response(res, tx, config).await
                            }
                        };

                        // log error.
                        if let Err(err) = res {
                            match err {
                                DispatchError::SendResponse(err) => {
                                    trace!("Error sending HTTP/2 response: {:?}", err)
                                }
                                DispatchError::SendData(err) => warn!("{:?}", err),
                                DispatchError::ResponseBody(err) => {
                                    error!("Response payload stream error: {:?}", err)
                                }
                            }
                        }
                    });
                }
                Poll::Ready(None) => return Poll::Ready(Ok(())),
                Poll::Pending => match this.ping_pong.as_mut() {
                    Some(ping_pong) => loop {
                        if ping_pong.on_flight {
                            // When have on flight ping pong. poll pong and and keep alive timer.
                            // on success pong received update keep alive timer to determine the next timing of
                            // ping pong.
                            match ping_pong.ping_pong.poll_pong(cx)? {
                                Poll::Ready(_) => {
                                    ping_pong.on_flight = false;

                                    let dead_line = this.config.keep_alive_expire().unwrap();
                                    ping_pong.timer.as_mut().reset(dead_line);
                                }
                                Poll::Pending => {
                                    return ping_pong.timer.as_mut().poll(cx).map(|_| Ok(()))
                                }
                            }
                        } else {
                            // When there is no on flight ping pong. keep alive timer is used to wait for next
                            // timing of ping pong. Therefore at this point it serves as an interval instead.
                            ready!(ping_pong.timer.as_mut().poll(cx));

                            ping_pong.ping_pong.send_ping(Ping::opaque())?;

                            let dead_line = this.config.keep_alive_expire().unwrap();
                            ping_pong.timer.as_mut().reset(dead_line);

                            ping_pong.on_flight = true;
                        }
                    },
                    None => return Poll::Pending,
                },
            }
        }
    }
}

enum DispatchError {
    SendResponse(h2::Error),
    SendData(h2::Error),
    ResponseBody(Box<dyn StdError>),
}

async fn handle_response<B>(
    res: Response<B>,
    mut tx: SendResponse<Bytes>,
    config: ServiceConfig,
) -> Result<(), DispatchError>
where
    B: MessageBody,
{
    let (res, body) = res.replace_body(());

    // prepare response.
    let mut size = body.size();
    let res = prepare_response(config, res.head(), &mut size);
    let eof = size.is_eof();

    // send response head and return on eof.
    let mut stream = tx
        .send_response(res, eof)
        .map_err(DispatchError::SendResponse)?;

    if eof {
        return Ok(());
    }

    // poll response body and send chunks to client
    actix_rt::pin!(body);

    while let Some(res) = poll_fn(|cx| body.as_mut().poll_next(cx)).await {
        let mut chunk = res.map_err(|err| DispatchError::ResponseBody(err.into()))?;

        'send: loop {
            let chunk_size = cmp::min(chunk.len(), CHUNK_SIZE);

            // reserve enough space and wait for stream ready.
            stream.reserve_capacity(chunk_size);

            match poll_fn(|cx| stream.poll_capacity(cx)).await {
                // No capacity left. drop body and return.
                None => return Ok(()),

                Some(Err(err)) => return Err(DispatchError::SendData(err)),

                Some(Ok(cap)) => {
                    // split chunk to writeable size and send to client
                    let len = chunk.len();
                    let bytes = chunk.split_to(cmp::min(len, cap));

                    stream
                        .send_data(bytes, false)
                        .map_err(DispatchError::SendData)?;

                    // Current chuck completely sent. break send loop and poll next one.
                    if chunk.is_empty() {
                        break 'send;
                    }
                }
            }
        }
    }

    // response body streaming finished. send end of stream and return.
    stream
        .send_data(Bytes::new(), true)
        .map_err(DispatchError::SendData)?;

    Ok(())
}

fn prepare_response(
    config: ServiceConfig,
    head: &ResponseHead,
    size: &mut BodySize,
) -> http::Response<()> {
    let mut has_date = false;
    let mut skip_len = size != &BodySize::Stream;

    let mut res = http::Response::new(());
    *res.status_mut() = head.status;
    *res.version_mut() = http::Version::HTTP_2;

    // Content length
    match head.status {
        http::StatusCode::NO_CONTENT
        | http::StatusCode::CONTINUE
        | http::StatusCode::PROCESSING => *size = BodySize::None,
        http::StatusCode::SWITCHING_PROTOCOLS => {
            skip_len = true;
            *size = BodySize::Stream;
        }
        _ => {}
    }

    let _ = match size {
        BodySize::None | BodySize::Stream => None,

        BodySize::Sized(0) => res
            .headers_mut()
            .insert(CONTENT_LENGTH, HeaderValue::from_static("0")),

        BodySize::Sized(len) => {
            let mut buf = itoa::Buffer::new();

            res.headers_mut().insert(
                CONTENT_LENGTH,
                HeaderValue::from_str(buf.format(*len)).unwrap(),
            )
        }
    };

    // copy headers
    for (key, value) in head.headers.iter() {
        match *key {
            // TODO: consider skipping other headers according to:
            //       https://datatracker.ietf.org/doc/html/rfc7540#section-8.1.2.2
            // omit HTTP/1.x only headers
            CONNECTION | TRANSFER_ENCODING => continue,
            CONTENT_LENGTH if skip_len => continue,
            DATE => has_date = true,
            _ => {}
        }

        res.headers_mut().append(key, value.clone());
    }

    // set date header
    if !has_date {
        let mut bytes = BytesMut::with_capacity(29);
        config.set_date_header(&mut bytes);
        res.headers_mut().insert(
            DATE,
            // SAFETY: serialized date-times are known ASCII strings
            unsafe { HeaderValue::from_maybe_shared_unchecked(bytes.freeze()) },
        );
    }

    res
}
<|endoftext|>


<commit_before>use Addr;
use std::mem::transmute;
use system::memory::{read, write};

pub fn random() -> f64 {
    let random = unsafe { transmute::<Addr, extern "C" fn() -> f64>(0x80243b40) };
    random()
}

pub fn cdyl_init_async() {
    let cdyl_init_async = unsafe { transmute::<Addr, extern "C" fn()>(0x80022A88) };
    cdyl_init_async();
}

pub fn get_frame_count() -> u32 {
    read(0x80396218)
}

pub fn is_pause_menu_up() -> bool {
    read(0x803EA537) // alternative: 0x80396228
}

pub fn set_wind(direction: u8) {
    write(0x803D894A, direction << 5);
}<commit_msg>Add wrapper function for Rupee Init method<commit_after>use Addr;
use std::mem::transmute;
use system::memory::{read, write};

pub fn random() -> f64 {
    let random = unsafe { transmute::<Addr, extern "C" fn() -> f64>(0x80243b40) };
    random()
}

pub fn cdyl_init_async() {
    let cdyl_init_async = unsafe { transmute::<Addr, extern "C" fn()>(0x80022A88) };
    cdyl_init_async();
}

pub fn dmeter_rupy_init(addr: Addr) {
    let dmeter_rupy_init = unsafe { transmute::<Addr, extern "C" fn(Addr)>(0x801F7868) };
    dmeter_rupy_init(addr);
}

pub fn get_frame_count() -> u32 {
    read(0x80396218)
}

pub fn is_pause_menu_up() -> bool {
    read(0x803EA537) // alternative: 0x80396228
}

pub fn set_wind(direction: u8) {
    write(0x803D894A, direction << 5);
}<|endoftext|>


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

use ::cache::disk::DiskCache;
use ::client::{
    connect_to_server,
};
use ::commands::{
    do_compile,
    request_shutdown,
    request_stats,
};
use env_logger;
use futures::sync::oneshot::{self, Sender};
use futures_cpupool::CpuPool;
use jobserver::Client;
use ::mock_command::*;
use ::server::{
    ServerMessage,
    SccacheServer,
};
use std::fs::File;
use std::io::{
    Cursor,
    Write,
};
use std::net::TcpListener;
use std::path::Path;
use std::process::Command;
use std::sync::{Arc,Mutex,mpsc};
use std::thread;
use std::time::Duration;
use std::u64;
use test::utils::*;
use tokio_core::reactor::Core;

/// Options for running the server in tests.
#[derive(Default)]
struct ServerOptions {
    /// The server's idle shutdown timeout.
    idle_timeout: Option<u64>,
    /// The maximum size of the disk cache.
    cache_size: Option<u64>,
}

/// Run a server on a background thread, and return a tuple of useful things.
///
/// * The port on which the server is listening.
/// * A `Sender` which can be used to send messages to the server.
///   (Most usefully, ServerMessage::Shutdown.)
/// * An `Arc`-and-`Mutex`-wrapped `MockCommandCreator` which the server will
///   use for all process creation.
/// * The `JoinHandle` for the server thread.
fn run_server_thread<T>(cache_dir: &Path, options: T)
                        -> (u16, Sender<ServerMessage>, Arc<Mutex<MockCommandCreator>>, thread::JoinHandle<()>)
    where T: Into<Option<ServerOptions>> + Send + 'static
{
    let options = options.into();
    let cache_dir = cache_dir.to_path_buf();

    let cache_size = options.as_ref()
                            .and_then(|o| o.cache_size.as_ref())
                            .map(|s| *s)
                            .unwrap_or(u64::MAX);
    let pool = CpuPool::new(1);
    let storage = Arc::new(DiskCache::new(&cache_dir, cache_size, &pool));

    // Create a server on a background thread, get some useful bits from it.
    let (tx, rx) = mpsc::channel();
    let (shutdown_tx, shutdown_rx) = oneshot::channel();
    let handle = thread::spawn(move || {
        let core = Core::new().unwrap();
        let client = unsafe { Client::new() };
        let srv = SccacheServer::new(0, pool, core, client, storage).unwrap();
        let mut srv: SccacheServer<Arc<Mutex<MockCommandCreator>>> = srv;
        assert!(srv.port() > 0);
        if let Some(options) = options {
            if let Some(timeout) = options.idle_timeout {
                 srv.set_idle_timeout(Duration::from_millis(timeout));
            }
        }
        let port = srv.port();
        let creator = srv.command_creator().clone();
        tx.send((port, creator)).unwrap();
        srv.run(shutdown_rx).unwrap();
    });
    let (port, creator) = rx.recv().unwrap();
    (port, shutdown_tx, creator, handle)
}

#[test]
fn test_server_shutdown() {
    let f = TestFixture::new();
    let (port, _sender, _storage, child) = run_server_thread(&f.tempdir.path(), None);
    // Connect to the server.
    let conn = connect_to_server(port).unwrap();
    // Ask it to shut down
    request_shutdown(conn).unwrap();
    // Ensure that it shuts down.
    child.join().unwrap();
}

#[test]
fn test_server_idle_timeout() {
    let f = TestFixture::new();
    // Set a ridiculously low idle timeout.
    let (_port, _sender, _storage, child) = run_server_thread(&f.tempdir.path(), ServerOptions { idle_timeout: Some(1), .. Default::default() });
    // Don't connect to it.
    // Ensure that it shuts down.
    // It would be nice to have an explicit timeout here so we don't hang
    // if something breaks...
    child.join().unwrap();
}

#[test]
fn test_server_stats() {
    let f = TestFixture::new();
    let (port, sender, _storage, child) = run_server_thread(&f.tempdir.path(), None);
    // Connect to the server.
    let conn = connect_to_server(port).unwrap();
    // Ask it for stats.
    let info = request_stats(conn).unwrap();
    assert_eq!(0, info.stats.compile_requests);
    // Now signal it to shut down.
    sender.send(ServerMessage::Shutdown).ok().unwrap();
    // Ensure that it shuts down.
    child.join().unwrap();
}

#[test]
fn test_server_unsupported_compiler() {
    let f = TestFixture::new();
    let (port, sender, server_creator, child) = run_server_thread(&f.tempdir.path(), None);
    // Connect to the server.
    let conn = connect_to_server(port).unwrap();
    {
        let mut c = server_creator.lock().unwrap();
        // The server will check the compiler, so pretend to be an unsupported
        // compiler.
        c.next_command_spawns(Ok(MockChild::new(exit_status(0), "hello", "error")));
    }
    // Ask the server to compile something.
    //TODO: MockCommand should validate these!
    let exe = &f.bins[0];
    let cmdline = vec!["-c".into(), "file.c".into(), "-o".into(), "file.o".into()];
    let cwd = f.tempdir.path();
    let client_creator = new_creator();
    const COMPILER_STDOUT: &'static [u8] = b"some stdout";
    const COMPILER_STDERR: &'static [u8] = b"some stderr";
    {
        let mut c = client_creator.lock().unwrap();
        // Actual client output.
        c.next_command_spawns(Ok(MockChild::new(exit_status(0), COMPILER_STDOUT, COMPILER_STDERR)));
    }
    let mut stdout = Cursor::new(Vec::new());
    let mut stderr = Cursor::new(Vec::new());
    let path = Some(f.paths);
    let mut core = Core::new().unwrap();
    assert_eq!(0, do_compile(client_creator.clone(), &mut core, conn, exe, cmdline, cwd, path, vec![], &mut stdout, &mut stderr).unwrap());
    // Make sure we ran the mock processes.
    assert_eq!(0, server_creator.lock().unwrap().children.len());
    assert_eq!(0, client_creator.lock().unwrap().children.len());
    assert_eq!(COMPILER_STDOUT, &stdout.into_inner()[..]);
    assert_eq!(COMPILER_STDERR, &stderr.into_inner()[..]);
    // Shut down the server.
    sender.send(ServerMessage::Shutdown).ok().unwrap();
    // Ensure that it shuts down.
    child.join().unwrap();
}

#[test]
fn test_server_compile() {
    match env_logger::init() {
        Ok(_) => {},
        Err(_) => {},
    }
    let f = TestFixture::new();
    let (port, sender, server_creator, child) = run_server_thread(&f.tempdir.path(), None);
    // Connect to the server.
    const PREPROCESSOR_STDOUT : &'static [u8] = b"preprocessor stdout";
    const PREPROCESSOR_STDERR : &'static [u8] = b"preprocessor stderr";
    const STDOUT : &'static [u8] = b"some stdout";
    const STDERR : &'static [u8] = b"some stderr";
    let conn = connect_to_server(port).unwrap();
    {
        let mut c = server_creator.lock().unwrap();
        // The server will check the compiler. Pretend it's GCC.
        c.next_command_spawns(Ok(MockChild::new(exit_status(0), "gcc", "")));
        // Preprocessor invocation.
        c.next_command_spawns(Ok(MockChild::new(exit_status(0), PREPROCESSOR_STDOUT, PREPROCESSOR_STDERR)));
        // Compiler invocation.
        //TODO: wire up a way to get data written to stdin.
        let obj = f.tempdir.path().join("file.o");
        c.next_command_calls(move |_| {
            // Pretend to compile something.
            let mut f = File::create(&obj)?;
            f.write_all(b"file contents")?;
            Ok(MockChild::new(exit_status(0), STDOUT, STDERR))
        });
    }
    // Ask the server to compile something.
    //TODO: MockCommand should validate these!
    let exe = &f.bins[0];
    let cmdline = vec!["-c".into(), "file.c".into(), "-o".into(), "file.o".into()];
    let cwd = f.tempdir.path();
    // This creator shouldn't create any processes. It will assert if
    // it tries to.
    let client_creator = new_creator();
    let mut stdout = Cursor::new(Vec::new());
    let mut stderr = Cursor::new(Vec::new());
    let path = Some(f.paths);
    let mut core = Core::new().unwrap();
    assert_eq!(0, do_compile(client_creator.clone(), &mut core, conn, exe, cmdline, cwd, path, vec![], &mut stdout, &mut stderr).unwrap());
    // Make sure we ran the mock processes.
    assert_eq!(0, server_creator.lock().unwrap().children.len());
    assert_eq!(STDOUT, stdout.into_inner().as_slice());
    assert_eq!(STDERR, stderr.into_inner().as_slice());
    // Shut down the server.
    sender.send(ServerMessage::Shutdown).ok().unwrap();
    // Ensure that it shuts down.
    child.join().unwrap();
}

#[test]
fn test_server_port_in_use() {
    // Bind an arbitrary free port.
    let listener = TcpListener::bind("127.0.0.1:0").unwrap();
    let sccache = find_sccache_binary();
    let output = Command::new(&sccache)
        .arg("--start-server")
        .env("SCCACHE_SERVER_PORT", listener.local_addr().unwrap().port().to_string())
        .output()
        .unwrap();
    assert!(!output.status.success());
    let s = String::from_utf8_lossy(&output.stderr);
    assert!(s.contains("Server startup failed:"),
            "Output did not contain 'Failed to start server:':\n========\n{}\n========",
            s);
}
<commit_msg>Add a test demonstrating that no idle timeout breaks shutdown.<commit_after>// Copyright 2016 Mozilla Foundation
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

use ::cache::disk::DiskCache;
use ::client::{
    connect_to_server,
};
use ::commands::{
    do_compile,
    request_shutdown,
    request_stats,
};
use env_logger;
use futures::sync::oneshot::{self, Sender};
use futures_cpupool::CpuPool;
use jobserver::Client;
use ::mock_command::*;
use ::server::{
    ServerMessage,
    SccacheServer,
};
use std::fs::File;
use std::io::{
    Cursor,
    Write,
};
use std::net::TcpListener;
use std::path::Path;
use std::process::Command;
use std::sync::{Arc,Mutex,mpsc};
use std::thread;
use std::time::Duration;
use std::u64;
use test::utils::*;
use tokio_core::reactor::Core;

/// Options for running the server in tests.
#[derive(Default)]
struct ServerOptions {
    /// The server's idle shutdown timeout.
    idle_timeout: Option<u64>,
    /// The maximum size of the disk cache.
    cache_size: Option<u64>,
}

/// Run a server on a background thread, and return a tuple of useful things.
///
/// * The port on which the server is listening.
/// * A `Sender` which can be used to send messages to the server.
///   (Most usefully, ServerMessage::Shutdown.)
/// * An `Arc`-and-`Mutex`-wrapped `MockCommandCreator` which the server will
///   use for all process creation.
/// * The `JoinHandle` for the server thread.
fn run_server_thread<T>(cache_dir: &Path, options: T)
                        -> (u16, Sender<ServerMessage>, Arc<Mutex<MockCommandCreator>>, thread::JoinHandle<()>)
    where T: Into<Option<ServerOptions>> + Send + 'static
{
    let options = options.into();
    let cache_dir = cache_dir.to_path_buf();

    let cache_size = options.as_ref()
                            .and_then(|o| o.cache_size.as_ref())
                            .map(|s| *s)
                            .unwrap_or(u64::MAX);
    let pool = CpuPool::new(1);
    let storage = Arc::new(DiskCache::new(&cache_dir, cache_size, &pool));

    // Create a server on a background thread, get some useful bits from it.
    let (tx, rx) = mpsc::channel();
    let (shutdown_tx, shutdown_rx) = oneshot::channel();
    let handle = thread::spawn(move || {
        let core = Core::new().unwrap();
        let client = unsafe { Client::new() };
        let srv = SccacheServer::new(0, pool, core, client, storage).unwrap();
        let mut srv: SccacheServer<Arc<Mutex<MockCommandCreator>>> = srv;
        assert!(srv.port() > 0);
        if let Some(options) = options {
            if let Some(timeout) = options.idle_timeout {
                 srv.set_idle_timeout(Duration::from_millis(timeout));
            }
        }
        let port = srv.port();
        let creator = srv.command_creator().clone();
        tx.send((port, creator)).unwrap();
        srv.run(shutdown_rx).unwrap();
    });
    let (port, creator) = rx.recv().unwrap();
    (port, shutdown_tx, creator, handle)
}

#[test]
fn test_server_shutdown() {
    let f = TestFixture::new();
    let (port, _sender, _storage, child) = run_server_thread(&f.tempdir.path(), None);
    // Connect to the server.
    let conn = connect_to_server(port).unwrap();
    // Ask it to shut down
    request_shutdown(conn).unwrap();
    // Ensure that it shuts down.
    child.join().unwrap();
}

/// The server will shutdown when requested when the idle timeout is disabled.
#[test]
fn test_server_shutdown_no_idle() {
    let f = TestFixture::new();
    // Set a ridiculously low idle timeout.
    let (port, _sender, _storage, child) = run_server_thread(
        &f.tempdir.path(),
        ServerOptions {
            idle_timeout: Some(0),
            ..Default::default()
        },
    );

    let conn = connect_to_server(port).unwrap();
    request_shutdown(conn).unwrap();
    child.join().unwrap();
}

#[test]
fn test_server_idle_timeout() {
    let f = TestFixture::new();
    // Set a ridiculously low idle timeout.
    let (_port, _sender, _storage, child) = run_server_thread(&f.tempdir.path(), ServerOptions { idle_timeout: Some(1), .. Default::default() });
    // Don't connect to it.
    // Ensure that it shuts down.
    // It would be nice to have an explicit timeout here so we don't hang
    // if something breaks...
    child.join().unwrap();
}

#[test]
fn test_server_stats() {
    let f = TestFixture::new();
    let (port, sender, _storage, child) = run_server_thread(&f.tempdir.path(), None);
    // Connect to the server.
    let conn = connect_to_server(port).unwrap();
    // Ask it for stats.
    let info = request_stats(conn).unwrap();
    assert_eq!(0, info.stats.compile_requests);
    // Now signal it to shut down.
    sender.send(ServerMessage::Shutdown).ok().unwrap();
    // Ensure that it shuts down.
    child.join().unwrap();
}

#[test]
fn test_server_unsupported_compiler() {
    let f = TestFixture::new();
    let (port, sender, server_creator, child) = run_server_thread(&f.tempdir.path(), None);
    // Connect to the server.
    let conn = connect_to_server(port).unwrap();
    {
        let mut c = server_creator.lock().unwrap();
        // The server will check the compiler, so pretend to be an unsupported
        // compiler.
        c.next_command_spawns(Ok(MockChild::new(exit_status(0), "hello", "error")));
    }
    // Ask the server to compile something.
    //TODO: MockCommand should validate these!
    let exe = &f.bins[0];
    let cmdline = vec!["-c".into(), "file.c".into(), "-o".into(), "file.o".into()];
    let cwd = f.tempdir.path();
    let client_creator = new_creator();
    const COMPILER_STDOUT: &'static [u8] = b"some stdout";
    const COMPILER_STDERR: &'static [u8] = b"some stderr";
    {
        let mut c = client_creator.lock().unwrap();
        // Actual client output.
        c.next_command_spawns(Ok(MockChild::new(exit_status(0), COMPILER_STDOUT, COMPILER_STDERR)));
    }
    let mut stdout = Cursor::new(Vec::new());
    let mut stderr = Cursor::new(Vec::new());
    let path = Some(f.paths);
    let mut core = Core::new().unwrap();
    assert_eq!(0, do_compile(client_creator.clone(), &mut core, conn, exe, cmdline, cwd, path, vec![], &mut stdout, &mut stderr).unwrap());
    // Make sure we ran the mock processes.
    assert_eq!(0, server_creator.lock().unwrap().children.len());
    assert_eq!(0, client_creator.lock().unwrap().children.len());
    assert_eq!(COMPILER_STDOUT, &stdout.into_inner()[..]);
    assert_eq!(COMPILER_STDERR, &stderr.into_inner()[..]);
    // Shut down the server.
    sender.send(ServerMessage::Shutdown).ok().unwrap();
    // Ensure that it shuts down.
    child.join().unwrap();
}

#[test]
fn test_server_compile() {
    match env_logger::init() {
        Ok(_) => {},
        Err(_) => {},
    }
    let f = TestFixture::new();
    let (port, sender, server_creator, child) = run_server_thread(&f.tempdir.path(), None);
    // Connect to the server.
    const PREPROCESSOR_STDOUT : &'static [u8] = b"preprocessor stdout";
    const PREPROCESSOR_STDERR : &'static [u8] = b"preprocessor stderr";
    const STDOUT : &'static [u8] = b"some stdout";
    const STDERR : &'static [u8] = b"some stderr";
    let conn = connect_to_server(port).unwrap();
    {
        let mut c = server_creator.lock().unwrap();
        // The server will check the compiler. Pretend it's GCC.
        c.next_command_spawns(Ok(MockChild::new(exit_status(0), "gcc", "")));
        // Preprocessor invocation.
        c.next_command_spawns(Ok(MockChild::new(exit_status(0), PREPROCESSOR_STDOUT, PREPROCESSOR_STDERR)));
        // Compiler invocation.
        //TODO: wire up a way to get data written to stdin.
        let obj = f.tempdir.path().join("file.o");
        c.next_command_calls(move |_| {
            // Pretend to compile something.
            let mut f = File::create(&obj)?;
            f.write_all(b"file contents")?;
            Ok(MockChild::new(exit_status(0), STDOUT, STDERR))
        });
    }
    // Ask the server to compile something.
    //TODO: MockCommand should validate these!
    let exe = &f.bins[0];
    let cmdline = vec!["-c".into(), "file.c".into(), "-o".into(), "file.o".into()];
    let cwd = f.tempdir.path();
    // This creator shouldn't create any processes. It will assert if
    // it tries to.
    let client_creator = new_creator();
    let mut stdout = Cursor::new(Vec::new());
    let mut stderr = Cursor::new(Vec::new());
    let path = Some(f.paths);
    let mut core = Core::new().unwrap();
    assert_eq!(0, do_compile(client_creator.clone(), &mut core, conn, exe, cmdline, cwd, path, vec![], &mut stdout, &mut stderr).unwrap());
    // Make sure we ran the mock processes.
    assert_eq!(0, server_creator.lock().unwrap().children.len());
    assert_eq!(STDOUT, stdout.into_inner().as_slice());
    assert_eq!(STDERR, stderr.into_inner().as_slice());
    // Shut down the server.
    sender.send(ServerMessage::Shutdown).ok().unwrap();
    // Ensure that it shuts down.
    child.join().unwrap();
}

#[test]
fn test_server_port_in_use() {
    // Bind an arbitrary free port.
    let listener = TcpListener::bind("127.0.0.1:0").unwrap();
    let sccache = find_sccache_binary();
    let output = Command::new(&sccache)
        .arg("--start-server")
        .env("SCCACHE_SERVER_PORT", listener.local_addr().unwrap().port().to_string())
        .output()
        .unwrap();
    assert!(!output.status.success());
    let s = String::from_utf8_lossy(&output.stderr);
    assert!(s.contains("Server startup failed:"),
            "Output did not contain 'Failed to start server:':\n========\n{}\n========",
            s);
}
<|endoftext|>


<commit_before>
struct tuple_struct(~str, ~str, ~str);

pub fn url_parse(url: &str) -> ~tuple_struct
{
    let mut out = ~"";
    out.push_str(url.slice(0, url.len()));
    let s = tuple_struct(out, ~"", ~"");
    return ~s;
}

mod test_url_parse;
<commit_msg>fix import<commit_after>
pub struct tuple_struct(~str, ~str, ~str);

pub fn url_parse(url: &str) -> ~tuple_struct
{
    let mut out = ~"";
    out.push_str(url.slice(0, url.len()));
    let s = tuple_struct(out, ~"", ~"");
    return ~s;
}

mod test_url_parse;
<|endoftext|>


<commit_before>use std::sync::{Arc, RwLock, RwLockWriteGuard};
use std::collections::HashMap;

use minihttp::client::{Codec, Config as HConfig, Proto, Error, EncoderDone};
use tokio_core::net::TcpStream;
use tokio_core::reactor::Handle;
use tk_pool::Pool;
use tk_pool::uniform::{UniformMx, Config as PConfig};
use abstract_ns::{Router, Resolver, union_stream};
use futures::Stream;
use futures::future::FutureResult;

use intern::Upstream;
use config::http_destinations::Destination;

/// Future that is used for sending a client request
///
/// While we only support fully buffered requests it's fine to use
/// FutureResult, but we will probably change it to something
pub type HttpFuture<S> = FutureResult<EncoderDone<S>, Error>;
pub type HttpPool = Pool<
        Box<Codec<TcpStream, Future=HttpFuture<TcpStream>>+Send>
    >;

pub struct UpstreamRef<'a> {
    pools: &'a HttpPools,
    upstream: &'a Upstream,
}

pub struct UpstreamGuard<'a> {
    guard: RwLockWriteGuard<'a, HashMap<Upstream, HttpPool>>,
    upstream: &'a Upstream,
}

#[derive(Clone)]
pub struct HttpPools {
    plain: Arc<RwLock<HashMap<Upstream, HttpPool>>>,
    // TODO(tailhook) https pools
}

impl HttpPools {
    pub fn new() -> HttpPools {
        HttpPools {
            plain: Arc::new(RwLock::new(HashMap::new())),
        }
    }
    pub fn upstream<'x>(&'x self, dest: &'x Upstream) -> UpstreamRef<'x> {
        UpstreamRef {
            pools: self,
            upstream: &dest,
        }
    }
    pub fn update(&self, cfg: &HashMap<Upstream, Destination>,
        resolver: &Router, handle: &Handle)
    {
        let mut plain = self.plain.write().expect("pools not poisoned");
        let mut to_delete = Vec::new();
        for k in plain.keys() {
            if !cfg.contains_key(k) {
                to_delete.push(k.clone());
            }
        }
        for k in to_delete {
            plain.remove(&k);
        }
        for (k, dest) in cfg {
            // TODO(tailhook) compare destinations
            if !plain.contains_key(k) {
                let h2 = handle.clone();
                let conn_config = HConfig::new()
                    .inflight_request_limit(
                        dest.in_flight_requests_per_backend_connection)
                    .done();
                let pool_config = PConfig::new()
                    .connections_per_address(
                        dest.backend_connections_per_ip_port)
                    .done();
                let stream = union_stream(
                    dest.addresses.iter()
                    .map(|x| resolver.subscribe(x)
                        as Box<Stream<Item=_, Error=_>>)
                    .collect());
                let mx = UniformMx::new(handle, &pool_config, Box::new(stream),
                    move |addr| Proto::connect_tcp(addr, &conn_config, &h2));
                let pool = Pool::create(handle, dest.queue_size_for_503, mx);
                plain.insert(k.clone(), pool);
            }
        }
     }
}

impl<'a> UpstreamRef<'a> {
    pub fn get_mut(&mut self) -> UpstreamGuard<'a> {
        UpstreamGuard {
            guard: self.pools.plain.write().expect("pools not poisoned"),
            upstream: self.upstream,
        }
    }
}

impl<'a> UpstreamGuard<'a> {
    pub fn get_mut(&mut self) -> Option<&mut HttpPool> {
        self.guard.get_mut(self.upstream)
    }
}
<commit_msg>Applies client timeout configuration<commit_after>use std::sync::{Arc, RwLock, RwLockWriteGuard};
use std::collections::HashMap;

use minihttp::client::{Codec, Config as HConfig, Proto, Error, EncoderDone};
use tokio_core::net::TcpStream;
use tokio_core::reactor::Handle;
use tk_pool::Pool;
use tk_pool::uniform::{UniformMx, Config as PConfig};
use abstract_ns::{Router, Resolver, union_stream};
use futures::Stream;
use futures::future::FutureResult;

use intern::Upstream;
use config::http_destinations::Destination;

/// Future that is used for sending a client request
///
/// While we only support fully buffered requests it's fine to use
/// FutureResult, but we will probably change it to something
pub type HttpFuture<S> = FutureResult<EncoderDone<S>, Error>;
pub type HttpPool = Pool<
        Box<Codec<TcpStream, Future=HttpFuture<TcpStream>>+Send>
    >;

pub struct UpstreamRef<'a> {
    pools: &'a HttpPools,
    upstream: &'a Upstream,
}

pub struct UpstreamGuard<'a> {
    guard: RwLockWriteGuard<'a, HashMap<Upstream, HttpPool>>,
    upstream: &'a Upstream,
}

#[derive(Clone)]
pub struct HttpPools {
    plain: Arc<RwLock<HashMap<Upstream, HttpPool>>>,
    // TODO(tailhook) https pools
}

impl HttpPools {
    pub fn new() -> HttpPools {
        HttpPools {
            plain: Arc::new(RwLock::new(HashMap::new())),
        }
    }
    pub fn upstream<'x>(&'x self, dest: &'x Upstream) -> UpstreamRef<'x> {
        UpstreamRef {
            pools: self,
            upstream: &dest,
        }
    }
    pub fn update(&self, cfg: &HashMap<Upstream, Destination>,
        resolver: &Router, handle: &Handle)
    {
        let mut plain = self.plain.write().expect("pools not poisoned");
        let mut to_delete = Vec::new();
        for k in plain.keys() {
            if !cfg.contains_key(k) {
                to_delete.push(k.clone());
            }
        }
        for k in to_delete {
            plain.remove(&k);
        }
        for (k, dest) in cfg {
            // TODO(tailhook) compare destinations
            if !plain.contains_key(k) {
                let h2 = handle.clone();
                let conn_config = HConfig::new()
                    .inflight_request_limit(
                        dest.in_flight_requests_per_backend_connection)
                    .keep_alive_timeout(*dest.keep_alive_timeout)
                    .safe_pipeline_timeout(*dest.safe_pipeline_timeout)
                    .max_request_timeout(*dest.max_request_timeout)
                    .done();
                let pool_config = PConfig::new()
                    .connections_per_address(
                        dest.backend_connections_per_ip_port)
                    .done();
                let stream = union_stream(
                    dest.addresses.iter()
                    .map(|x| resolver.subscribe(x)
                        as Box<Stream<Item=_, Error=_>>)
                    .collect());
                let mx = UniformMx::new(handle, &pool_config, Box::new(stream),
                    move |addr| Proto::connect_tcp(addr, &conn_config, &h2));
                let pool = Pool::create(handle, dest.queue_size_for_503, mx);
                plain.insert(k.clone(), pool);
            }
        }
     }
}

impl<'a> UpstreamRef<'a> {
    pub fn get_mut(&mut self) -> UpstreamGuard<'a> {
        UpstreamGuard {
            guard: self.pools.plain.write().expect("pools not poisoned"),
            upstream: self.upstream,
        }
    }
}

impl<'a> UpstreamGuard<'a> {
    pub fn get_mut(&mut self) -> Option<&mut HttpPool> {
        self.guard.get_mut(self.upstream)
    }
}
<|endoftext|>


<commit_before>//! GDAL Vector Data
//!
//! ## Reading
//!
//! ```
//! use std::path::Path;
//! use gdal::vector::Dataset;
//!
//! let ds = Dataset::open(Path::new("fixtures/roads.geojson")).unwrap();
//! let layer = ds.layer(0).unwrap();
//! for f in layer.features() {
//!     println!("{} {}", f.field("highway").unwrap().as_string(), f.wkt());
//! }
//! ```

pub use vector::dataset::Dataset;
pub use vector::layer::{Layer, FeatureIterator};
pub use vector::feature::Feature;
pub use vector::geometry::{Geometry, ToGdal};

mod ogr;
pub mod dataset;
pub mod layer;
pub mod feature;
pub mod geometry;
pub mod geom;

#[cfg(test)]
mod tests;
<commit_msg>update example<commit_after>//! GDAL Vector Data
//!
//! ## Reading
//!
//! ```
//! use std::path::Path;
//! use gdal::vector::Dataset;
//!
//! let dataset = Dataset::open(Path::new("fixtures/roads.geojson")).unwrap();
//! let layer = dataset.layer(0).unwrap();
//! for feature in layer.features() {
//!     let highway_field = feature.field("highway").unwrap();
//!     let geometry = feature.geometry();
//!     println!("{} {}", highway_field.as_string(), geometry.wkt());
//! }
//! ```

pub use vector::dataset::Dataset;
pub use vector::layer::{Layer, FeatureIterator};
pub use vector::feature::Feature;
pub use vector::geometry::{Geometry, ToGdal};

mod ogr;
pub mod dataset;
pub mod layer;
pub mod feature;
pub mod geometry;
pub mod geom;

#[cfg(test)]
mod tests;
<|endoftext|>


<commit_before>use utils::bigint::M256;
use utils::gas::Gas;
use utils::address::Address;

use super::{Memory, VectorMemory, Stack, VectorStack, PC, VectorPC, Result, Error};
use blockchain::{Block, FakeVectorBlock};
use transaction::{Transaction, VectorTransaction};

use std::borrow::BorrowMut;
use std::marker::PhantomData;

pub trait Machine {
    type P: PC;
    type M: Memory;
    type Sta: Stack;
    type T: Transaction;
    type B: Block;
    type Sub: Machine;

    fn pc(&self) -> &Self::P;
    fn pc_mut(&mut self) -> &mut Self::P;
    fn memory(&self) -> &Self::M;
    fn memory_mut(&mut self) -> &mut Self::M;
    fn stack(&self) -> &Self::Sta;
    fn stack_mut(&mut self) -> &mut Self::Sta;

    fn transaction(&self) -> &Self::T;
    fn block(&self) -> &Self::B;
    fn block_mut(&mut self) -> &mut Self::B;

    fn use_gas(&mut self, gas: Gas);
    fn used_gas(&self) -> Gas;
    fn return_values(&self) -> &[u8];
    fn set_return_values(&mut self, data: &[u8]);

    fn fork<F: FnOnce(&mut Self::Sub)>(&mut self, gas: Gas, from:
                                       Address, to: Address, value: M256,
                                       memory_in_start: M256,
                                       memory_in_len: M256,
                                       memory_out_start: M256,
                                       memory_out_len: M256, f: F);

    fn step(&mut self) -> Result<()> where Self: Sized {
        if self.pc().stopped() {
            return Err(Error::Stopped)
        }

        if !self.available_gas().is_valid() {
            return Err(Error::EmptyGas)
        }

        let opcode = self.pc_mut().peek_opcode()?;
        let before = opcode.gas_cost_before(self)?;
        opcode.run(self)?;
        let after = opcode.gas_cost_after(self);

        self.pc_mut().read_opcode().unwrap();
        self.use_gas(before);
        self.use_gas(after);

        Ok(())
    }

    fn fire(&mut self) -> Result<()> where Self: Sized {
        loop {
            let result = self.step();

            if result.is_err() {
                match result.err().unwrap() {
                    Error::Stopped => return Ok(()),
                    err => return Err(err),
                }
            }
        }
    }

    fn available_gas(&self) -> Gas {
        self.transaction().gas_limit() - self.used_gas()
    }
}

pub struct VectorMachine<B0, BR> {
    pc: VectorPC,
    memory: VectorMemory,
    stack: VectorStack,
    transaction: VectorTransaction,
    return_values: Vec<u8>,
    block: Option
,
    used_gas: Gas,
    _block_marker: PhantomData<B0>,
}

impl<B0: Block, BR: AsRef<B0> + AsMut<B0>> VectorMachine<B0, BR> {
    pub fn new(code: &[u8], data: &[u8], gas_limit: Gas,
               transaction: VectorTransaction, block: BR) -> Self {
        VectorMachine {
            pc: VectorPC::new(code),
            memory: VectorMemory::new(),
            stack: VectorStack::new(),
            transaction: transaction,
            return_values: Vec::new(),
            block: Some(block),
            used_gas: Gas::zero(),
            _block_marker: PhantomData,
        }
    }
}

impl<B0: Block, BR: AsRef<B0> + AsMut<B0>> Machine for VectorMachine<B0, BR> {
    type P = VectorPC;
    type M = VectorMemory;
    type Sta = VectorStack;
    type T = VectorTransaction;
    type B = B0;
    type Sub = Self;

    fn return_values(&self) -> &[u8] {
        self.return_values.as_ref()
    }

    fn set_return_values(&mut self, val: &[u8]) {
        self.return_values = val.into();
    }

    fn use_gas(&mut self, gas: Gas) {
        self.used_gas += gas;
    }

    fn used_gas(&self) -> Gas {
        self.used_gas
    }

    fn pc(&self) -> &Self::P {
        &self.pc
    }

    fn pc_mut(&mut self) -> &mut Self::P {
        &mut self.pc
    }

    fn memory(&self) -> &Self::M {
        &self.memory
    }

    fn memory_mut(&mut self) -> &mut Self::M {
        &mut self.memory
    }

    fn stack(&self) -> &Self::Sta {
        &self.stack
    }

    fn stack_mut(&mut self) -> &mut Self::Sta {
        &mut self.stack
    }

    fn transaction(&self) -> &Self::T {
        &self.transaction
    }

    fn block(&self) -> &Self::B {
        self.block.as_ref().unwrap().as_ref()
    }

    fn block_mut(&mut self) -> &mut Self::B {
        self.block.as_mut().unwrap().as_mut()
    }

    fn fork<F: FnOnce(&mut Self::Sub)>(&mut self, gas: Gas, from: Address, to: Address,
                                       value: M256,
                                       memory_in_start: M256, memory_in_len: M256,
                                       memory_out_start: M256, memory_out_len: M256,
                                       f: F) {
        use std::mem::swap;

        let from = from;
        let mem_in_start: usize = memory_in_start.into();
        let mem_in_len: usize = memory_in_len.into();
        let mem_in_end: usize = mem_in_start + mem_in_len;
        let mem_in: Vec<u8> = self.memory().as_ref()[mem_in_start..mem_in_end].into();

        let mut submachine = Self {
            pc: VectorPC::new(if to == from { self.pc().code() }
                              else { self.block().account_code(to) }),
            memory: VectorMemory::new(),
            stack: VectorStack::new(),
            transaction: VectorTransaction::message_call(from, to, value, mem_in.as_ref(), gas),
            return_values: Vec::new(),
            block: None,
            used_gas: Gas::zero(),
            _block_marker: PhantomData,
        };

        // We swap the block into the sub-machine if necessary. The
        // current old block should never be referenced and will be
        // replaced back when the call finishes.
        swap(&mut self.block, &mut submachine.block);
        f(self);
        swap(&mut self.block, &mut submachine.block);

        let mem_out_start: usize = memory_out_start.into();
        let mem_out_len: usize = memory_out_len.into();
        let mem_out_end: usize = mem_out_start + mem_out_len;

        for i in 0..mem_out_end {
            self.memory_mut().write_raw(memory_out_start + i.into(), submachine.return_values[i]);
        }
    }
}

pub type FakeVectorMachine = VectorMachine<FakeVectorBlock, Box<FakeVectorBlock>>;

impl FakeVectorMachine {
    pub fn fake(code: &[u8], data: &[u8], gas_limit: Gas) -> FakeVectorMachine {
        VectorMachine::new(code, data, gas_limit,
                           VectorTransaction::message_call(Address::default(), Address::default(),
                                                           M256::zero(), data, gas_limit),
                           Box::new(FakeVectorBlock::new()))
    }
}
<commit_msg>Problem: the restore logic is incorrect Solution: replace ? with the correct restore logic<commit_after>use utils::bigint::M256;
use utils::gas::Gas;
use utils::address::Address;

use super::{Memory, VectorMemory, Stack, VectorStack, PC, VectorPC, Result, Error};
use blockchain::{Block, FakeVectorBlock};
use transaction::{Transaction, VectorTransaction};

use std::borrow::BorrowMut;
use std::marker::PhantomData;

pub trait Machine {
    type P: PC;
    type M: Memory;
    type Sta: Stack;
    type T: Transaction;
    type B: Block;
    type Sub: Machine;

    fn pc(&self) -> &Self::P;
    fn pc_mut(&mut self) -> &mut Self::P;
    fn memory(&self) -> &Self::M;
    fn memory_mut(&mut self) -> &mut Self::M;
    fn stack(&self) -> &Self::Sta;
    fn stack_mut(&mut self) -> &mut Self::Sta;

    fn transaction(&self) -> &Self::T;
    fn block(&self) -> &Self::B;
    fn block_mut(&mut self) -> &mut Self::B;

    fn use_gas(&mut self, gas: Gas);
    fn used_gas(&self) -> Gas;
    fn return_values(&self) -> &[u8];
    fn set_return_values(&mut self, data: &[u8]);

    fn fork<F: FnOnce(&mut Self::Sub)>(&mut self, gas: Gas, from:
                                       Address, to: Address, value: M256,
                                       memory_in_start: M256,
                                       memory_in_len: M256,
                                       memory_out_start: M256,
                                       memory_out_len: M256, f: F);

    fn step(&mut self) -> Result<()> where Self: Sized {
        if self.pc().stopped() {
            return Err(Error::Stopped)
        }

        if !self.available_gas().is_valid() {
            return Err(Error::EmptyGas)
        }

        let position = self.pc().position();
        let restore = |machine: &mut Self| { machine.pc_mut().jump(position); };

        let opcode = match self.pc_mut().read_opcode() {
            Ok(opcode) => opcode,
            Err(e) => {
                restore(self);
                return Err(e);
            }
        };
        let before = match opcode.gas_cost_before(self) {
            Ok(gas) => gas,
            Err(e) => {
                restore(self);
                return Err(e);
            }
        };
        match opcode.run(self) {
            Ok(()) => (),
            Err(e) => {
                restore(self);
                return Err(e);
            }
        };
        let after = opcode.gas_cost_after(self);

        self.use_gas(before);
        self.use_gas(after);

        Ok(())
    }

    fn fire(&mut self) -> Result<()> where Self: Sized {
        loop {
            let result = self.step();

            if result.is_err() {
                match result.err().unwrap() {
                    Error::Stopped => return Ok(()),
                    err => return Err(err),
                }
            }
        }
    }

    fn available_gas(&self) -> Gas {
        self.transaction().gas_limit() - self.used_gas()
    }
}

pub struct VectorMachine<B0, BR> {
    pc: VectorPC,
    memory: VectorMemory,
    stack: VectorStack,
    transaction: VectorTransaction,
    return_values: Vec<u8>,
    block: Option
,
    used_gas: Gas,
    _block_marker: PhantomData<B0>,
}

impl<B0: Block, BR: AsRef<B0> + AsMut<B0>> VectorMachine<B0, BR> {
    pub fn new(code: &[u8], data: &[u8], gas_limit: Gas,
               transaction: VectorTransaction, block: BR) -> Self {
        VectorMachine {
            pc: VectorPC::new(code),
            memory: VectorMemory::new(),
            stack: VectorStack::new(),
            transaction: transaction,
            return_values: Vec::new(),
            block: Some(block),
            used_gas: Gas::zero(),
            _block_marker: PhantomData,
        }
    }
}

impl<B0: Block, BR: AsRef<B0> + AsMut<B0>> Machine for VectorMachine<B0, BR> {
    type P = VectorPC;
    type M = VectorMemory;
    type Sta = VectorStack;
    type T = VectorTransaction;
    type B = B0;
    type Sub = Self;

    fn return_values(&self) -> &[u8] {
        self.return_values.as_ref()
    }

    fn set_return_values(&mut self, val: &[u8]) {
        self.return_values = val.into();
    }

    fn use_gas(&mut self, gas: Gas) {
        self.used_gas += gas;
    }

    fn used_gas(&self) -> Gas {
        self.used_gas
    }

    fn pc(&self) -> &Self::P {
        &self.pc
    }

    fn pc_mut(&mut self) -> &mut Self::P {
        &mut self.pc
    }

    fn memory(&self) -> &Self::M {
        &self.memory
    }

    fn memory_mut(&mut self) -> &mut Self::M {
        &mut self.memory
    }

    fn stack(&self) -> &Self::Sta {
        &self.stack
    }

    fn stack_mut(&mut self) -> &mut Self::Sta {
        &mut self.stack
    }

    fn transaction(&self) -> &Self::T {
        &self.transaction
    }

    fn block(&self) -> &Self::B {
        self.block.as_ref().unwrap().as_ref()
    }

    fn block_mut(&mut self) -> &mut Self::B {
        self.block.as_mut().unwrap().as_mut()
    }

    fn fork<F: FnOnce(&mut Self::Sub)>(&mut self, gas: Gas, from: Address, to: Address,
                                       value: M256,
                                       memory_in_start: M256, memory_in_len: M256,
                                       memory_out_start: M256, memory_out_len: M256,
                                       f: F) {
        use std::mem::swap;

        let from = from;
        let mem_in_start: usize = memory_in_start.into();
        let mem_in_len: usize = memory_in_len.into();
        let mem_in_end: usize = mem_in_start + mem_in_len;
        let mem_in: Vec<u8> = self.memory().as_ref()[mem_in_start..mem_in_end].into();

        let mut submachine = Self {
            pc: VectorPC::new(if to == from { self.pc().code() }
                              else { self.block().account_code(to) }),
            memory: VectorMemory::new(),
            stack: VectorStack::new(),
            transaction: VectorTransaction::message_call(from, to, value, mem_in.as_ref(), gas),
            return_values: Vec::new(),
            block: None,
            used_gas: Gas::zero(),
            _block_marker: PhantomData,
        };

        // We swap the block into the sub-machine if necessary. The
        // current old block should never be referenced and will be
        // replaced back when the call finishes.
        swap(&mut self.block, &mut submachine.block);
        f(self);
        swap(&mut self.block, &mut submachine.block);

        let mem_out_start: usize = memory_out_start.into();
        let mem_out_len: usize = memory_out_len.into();
        let mem_out_end: usize = mem_out_start + mem_out_len;

        for i in 0..mem_out_end {
            self.memory_mut().write_raw(memory_out_start + i.into(), submachine.return_values[i]);
        }
    }
}

pub type FakeVectorMachine = VectorMachine<FakeVectorBlock, Box<FakeVectorBlock>>;

impl FakeVectorMachine {
    pub fn fake(code: &[u8], data: &[u8], gas_limit: Gas) -> FakeVectorMachine {
        VectorMachine::new(code, data, gas_limit,
                           VectorTransaction::message_call(Address::default(), Address::default(),
                                                           M256::zero(), data, gas_limit),
                           Box::new(FakeVectorBlock::new()))
    }
}
<|endoftext|>


<commit_before>// Copyright (c) 2016 Sandstorm Development Group, Inc.
// Licensed under the MIT License:
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.

use gj::{Promise};
use capnp::Error;
use std::cell::Cell;
use std::rc::Rc;
use sandstorm::web_session_capnp::web_session::web_socket_stream;

#[repr(u8)]
pub enum OpCode {
    Continue = 0,
    Utf8Payload = 1,
    BinaryPayload = 2,
    Terminate = 8,
    Ping = 9,
    Pong = 10,
}

pub fn encode_text_message(params: web_socket_stream::send_bytes_params::Builder,
                           message: &str)
{
    encode_message(params, OpCode::Utf8Payload, message.as_bytes())
}

pub fn encode_message(mut params: web_socket_stream::send_bytes_params::Builder,
                      opcode: OpCode, message: &[u8])
{
    // TODO(perf) avoid this allocation
    let mut bytes: Vec<u8> = Vec::new();
    bytes.push(0x80 | opcode as u8);
    if message.len() < 126 {
        bytes.push(message.len() as u8);
    } else if message.len() < 1 << 16  {
        // 16 bits
        bytes.push(0x7e);
        bytes.push((message.len() >> 8) as u8);
        bytes.push(message.len() as u8);
    } else {
        // 64 bits
        bytes.push(0x7f);
        bytes.push((message.len() >> 56) as u8);
        bytes.push((message.len() >> 48) as u8);
        bytes.push((message.len() >> 40) as u8);
        bytes.push((message.len() >> 32) as u8);
        bytes.push((message.len() >> 16) as u8);
        bytes.push((message.len() >> 24) as u8);
        bytes.push((message.len() >> 8) as u8);
        bytes.push(message.len() as u8);
    }

    bytes.extend_from_slice(message);

    params.set_message(&bytes[..]);
}

pub enum Message {
  Text(String),
  Data(Vec<u8>),
}

pub trait MessageHandler {
    fn handle_message(&mut self, message: Message) -> Promise<(), Error>;
}

fn do_ping_pong(client_stream: web_socket_stream::Client,
                timer: ::gjio::Timer,
                awaiting_pong: Rc<Cell<bool>>) -> Promise<(), Error>
{
    let mut req = client_stream.send_bytes_request();
    req.get().set_message(&[0x89, 0]); // PING
    let promise = req.send().promise;
    awaiting_pong.set(true);
    promise.then(move|_| {
        timer.after_delay(::std::time::Duration::new(10, 0)).lift().then(move |_| {
            if awaiting_pong.get() {
                Promise::err(Error::failed("pong not received within 10 seconds".into()))
            } else {
                do_ping_pong(client_stream, timer, awaiting_pong)
            }
        })
    })
}

enum PreviousFrames {
   None,
   Text(String),
   Data(Vec<u8>)
}

#[derive(Debug)]
enum ParserState {
    NotStarted,
    DoneFirstByte { fin: bool, opcode: u8},
    ReadingLongPayloadLength {fin: bool, opcode: u8, masked: bool,
                            payload_len_bytes_read: usize, payload_len_so_far: u64 },
    ReadingMask { fin: bool, opcode: u8, mask_bytes_read: usize, payload_len: u64,
                  mask_so_far: [u8; 4] },
    ReadingPayload { fin: bool, opcode: u8, payload_len: u64, mask: [u8; 4], bytes_so_far: Vec<u8> },
}

struct ParseResult {
    frame: Vec<u8>,
    opcode: u8,
    fin: bool,
}

impl ParserState {
    fn done_payload_length(bytes_read: usize,
                           fin: bool, opcode: u8, masked: bool, payload_len: u64)
                           -> (ParserState, (usize, Option<ParseResult>))
    {
        use self::ParserState::*;
        if masked {
            (ReadingMask { fin: fin, opcode: opcode, payload_len: payload_len,
                           mask_bytes_read: 0, mask_so_far: [0; 4] },
             (bytes_read, None))
        } else if payload_len == 0 {
            (NotStarted,
             (bytes_read, Some(ParseResult { frame: Vec::new(), fin: fin, opcode: opcode })))
        } else {
            (ReadingPayload { fin: fin, opcode: opcode,
                              payload_len: payload_len, mask: [0; 4],
                              bytes_so_far: Vec::new() },
              (bytes_read, None))
        }
    }

    /// returns number of bytes consumed and the complete message, if there is one.
    fn advance(&mut self, buf: &[u8]) -> (usize, Option<ParseResult>) {
        use self::ParserState::*;
        let (new_state, result) = match self {
            &mut NotStarted => {
                if buf.len() < 1 {
                    return (0, None)
                }

                (DoneFirstByte { fin: (buf[0] & 0x80) != 0, opcode: buf[0] & 0xf }, (1, None))
            }
            &mut DoneFirstByte { fin, opcode } => {
                if buf.len() < 1 {
                    return (0, None)
                }

                let masked = (buf[0] & 0x80) != 0;

                match buf[0] & 0x7f {
                    126 => {
                        (ReadingLongPayloadLength {
                            fin: fin,
                            opcode: opcode,
                            masked: masked,
                            payload_len_bytes_read: 6,
                            payload_len_so_far: 0,
                        }, (1, None))
                    }
                    127 => {
                        (ReadingLongPayloadLength {
                            fin: fin,
                            opcode: opcode,
                            masked: masked,
                            payload_len_bytes_read: 0,
                            payload_len_so_far: 0,
                        }, (1, None))
                    }
                    n => ParserState::done_payload_length(1, fin, opcode, masked, n as u64)
                }
            }
            &mut ReadingLongPayloadLength { fin, opcode, masked, payload_len_bytes_read,
                                            payload_len_so_far } => {
                let mut idx = 0;
                let mut new_so_far = payload_len_so_far;
                while idx + payload_len_bytes_read < 8 && idx < buf.len() {
                    new_so_far += (buf[idx] as u64) << (8 * (7 - idx - payload_len_bytes_read));
                    idx += 1;
                }

                if buf.len() + payload_len_bytes_read < 8 {
                    (ReadingLongPayloadLength {
                        fin: fin,
                        opcode: opcode,
                        masked: masked,
                        payload_len_bytes_read: idx + payload_len_bytes_read,
                        payload_len_so_far: new_so_far,
                    }, (idx, None))
                } else {
                    ParserState::done_payload_length(idx, fin, opcode, masked, new_so_far)
                }
            }
            &mut ReadingMask { fin, opcode, mask_bytes_read, payload_len, mask_so_far } => {
                let mut idx = 0;
                let mut new_so_far = mask_so_far;
                while idx + mask_bytes_read < 4 && idx < buf.len() {
                    new_so_far[idx] = buf[idx];
                    idx += 1;
                }

                if buf.len() + mask_bytes_read < 4 {
                    (ReadingMask {
                        fin: fin,
                        opcode: opcode,
                        payload_len: payload_len,
                        mask_bytes_read: idx + mask_bytes_read,
                        mask_so_far: new_so_far,
                    }, (idx, None))
                } else if payload_len == 0 {
                    (NotStarted,
                     (idx, Some(ParseResult { frame: Vec::new(), fin: fin, opcode: opcode })))
                } else {
                    (ReadingPayload { fin: fin, opcode: opcode, mask: new_so_far,
                                      bytes_so_far: Vec::new(),
                                      payload_len: payload_len },
                     (idx, None))
                }
            }
            &mut ReadingPayload { fin, opcode, payload_len, mask, ref mut bytes_so_far } => {
                let mut idx = 0;

                while (bytes_so_far.len() as u64) < payload_len && idx < buf.len() {
                    let mask_byte = mask[bytes_so_far.len() % 4];
                    bytes_so_far.push(buf[idx] ^ mask_byte);
                    idx += 1;
                }

                if (bytes_so_far.len() as u64) < payload_len {
                    return (idx, None)
                } else {
                    let frame = ::std::mem::replace(bytes_so_far, Vec::new());
                    (NotStarted,
                     (idx, Some(ParseResult { frame: frame, fin: fin, opcode: opcode })))
                }

            }
        };

        *self = new_state;
        result
    }
}

pub struct Adapter<T> where T: MessageHandler {
    handler: Option<T>,
    awaiting_pong: Rc<Cell<bool>>,
    ping_pong_promise: Promise<(), Error>,
    client_stream: Option<web_socket_stream::Client>,
    parser_state: ParserState,
    previous_frames: PreviousFrames,
}

impl <T> Adapter<T> where T: MessageHandler {
    pub fn new(handler: T,
               client_stream: web_socket_stream::Client,
               timer: ::gjio::Timer)
               -> Adapter<T> {
        let awaiting = Rc::new(Cell::new(false));
        let ping_pong_promise = do_ping_pong(
            client_stream.clone(),
            timer,
            awaiting.clone()
        ).map_else(|r| match r {
            Ok(_) => Ok(()),
            Err(e) => {
                println!("error while pinging client: {}", e);
                Ok(())
            }
        }).eagerly_evaluate();

        Adapter {
            handler: Some(handler),
            awaiting_pong: awaiting,
            ping_pong_promise: ping_pong_promise,
            client_stream: Some(client_stream),
            parser_state: ParserState::NotStarted,
            previous_frames: PreviousFrames::None,
        }
    }

    fn process_message(&mut self) -> Promise<(), Error> {
        let frames = ::std::mem::replace(&mut self.previous_frames,
                                         PreviousFrames::None);
        let message = match frames {
            PreviousFrames::None => {
                return Promise::err(Error::failed(format!("message has no frames")));
            }
            PreviousFrames::Data(d) => {
                Message::Data(d)
            }
            PreviousFrames::Text(t) => {
                Message::Text(t)
            }
        };

        match self.handler {
            Some(ref mut h) => h.handle_message(message),
            None => Promise::ok(()),
        }
    }
}

impl <T> web_socket_stream::Server for Adapter<T> where T: MessageHandler {
    fn send_bytes(&mut self,
                  params: web_socket_stream::SendBytesParams,
                  _results: web_socket_stream::SendBytesResults)
                  -> Promise<(), Error>
    {
        let message = pry!(pry!(params.get()).get_message());
        let mut result_promise = Promise::ok(());
        let mut num_bytes_read = 0;
        while num_bytes_read < message.len() {
            let (n, result) = self.parser_state.advance(&message[num_bytes_read..]);
            num_bytes_read += n;
            match result {
                None => (),
                Some(ParseResult { frame, opcode, fin }) => {
                    match opcode {
                        0x0 => { // CONTINUE
                            match &mut self.previous_frames {
                                &mut PreviousFrames::None => {
                                    return Promise::err(Error::failed(
                                        format!("CONTINUE frame received, but there are no \
                                                 previous frames.")));
                                }
                                &mut PreviousFrames::Data(ref mut data) => {
                                    data.extend_from_slice(&frame[..])
                                }
                                &mut PreviousFrames::Text(ref mut text) => {
                                    text.push_str(&pry!(String::from_utf8(frame)))
                                }
                            }

                            if fin {
                                let promise = self.process_message();
                                result_promise = result_promise.then(|_| promise);
                            }
                        }
                        0x1 => { // UTF-8 PAYLOAD
                            self.previous_frames =
                                PreviousFrames::Text(pry!(String::from_utf8(frame)));

                            if fin {
                                let promise = self.process_message();
                                result_promise = result_promise.then(|_| promise);
                            }
                        }
                        0x2 => { // BINARY PAYLOAD
                            self.previous_frames = PreviousFrames::Data(frame);

                            if fin {
                                let promise = self.process_message();
                                result_promise = result_promise.then(|_| promise);
                            }
                        }
                        0x8 => { // TERMINATE
                            self.handler = None;
                            self.ping_pong_promise = Promise::ok(());
                            self.client_stream = None;
                        }
                        0x9 => { // PING
                            match &self.client_stream {
                                &None => (),
                                &Some(ref client) => {
                                    println!("responding to ping from client");
                                    let req = client.send_bytes_request();
                                    let promise = req.send().promise.map(|_| Ok(()));
                                    result_promise = result_promise.then(|_| promise);
                                }
                            }
                        }
                        0xa => { // PONG
                            self.awaiting_pong.set(false);
                        }
                        _ => { // OTHER
                            println!("unrecognized websocket opcode {}", opcode);
                        }
                    }
                }
            }
        }

        result_promise
    }
}
<commit_msg>Don't allow huge messages.<commit_after>// Copyright (c) 2016 Sandstorm Development Group, Inc.
// Licensed under the MIT License:
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.

use gj::{Promise};
use capnp::Error;
use std::cell::Cell;
use std::rc::Rc;
use sandstorm::web_session_capnp::web_session::web_socket_stream;

#[repr(u8)]
pub enum OpCode {
    Continue = 0,
    Utf8Payload = 1,
    BinaryPayload = 2,
    Terminate = 8,
    Ping = 9,
    Pong = 10,
}

pub fn encode_text_message(params: web_socket_stream::send_bytes_params::Builder,
                           message: &str)
{
    encode_message(params, OpCode::Utf8Payload, message.as_bytes())
}

pub fn encode_message(mut params: web_socket_stream::send_bytes_params::Builder,
                      opcode: OpCode, message: &[u8])
{
    // TODO(perf) avoid this allocation
    let mut bytes: Vec<u8> = Vec::new();
    bytes.push(0x80 | opcode as u8);
    if message.len() < 126 {
        bytes.push(message.len() as u8);
    } else if message.len() < 1 << 16  {
        // 16 bits
        bytes.push(0x7e);
        bytes.push((message.len() >> 8) as u8);
        bytes.push(message.len() as u8);
    } else {
        // 64 bits
        bytes.push(0x7f);
        bytes.push((message.len() >> 56) as u8);
        bytes.push((message.len() >> 48) as u8);
        bytes.push((message.len() >> 40) as u8);
        bytes.push((message.len() >> 32) as u8);
        bytes.push((message.len() >> 16) as u8);
        bytes.push((message.len() >> 24) as u8);
        bytes.push((message.len() >> 8) as u8);
        bytes.push(message.len() as u8);
    }

    bytes.extend_from_slice(message);

    params.set_message(&bytes[..]);
}

pub enum Message {
  Text(String),
  Data(Vec<u8>),
}

pub trait MessageHandler {
    fn handle_message(&mut self, message: Message) -> Promise<(), Error>;
}

fn do_ping_pong(client_stream: web_socket_stream::Client,
                timer: ::gjio::Timer,
                awaiting_pong: Rc<Cell<bool>>) -> Promise<(), Error>
{
    let mut req = client_stream.send_bytes_request();
    req.get().set_message(&[0x89, 0]); // PING
    let promise = req.send().promise;
    awaiting_pong.set(true);
    promise.then(move|_| {
        timer.after_delay(::std::time::Duration::new(10, 0)).lift().then(move |_| {
            if awaiting_pong.get() {
                Promise::err(Error::failed("pong not received within 10 seconds".into()))
            } else {
                do_ping_pong(client_stream, timer, awaiting_pong)
            }
        })
    })
}

enum PreviousFrames {
   None,
   Text(String),
   Data(Vec<u8>)
}

#[derive(Debug)]
enum ParserState {
    NotStarted,
    DoneFirstByte { fin: bool, opcode: u8},
    ReadingLongPayloadLength {fin: bool, opcode: u8, masked: bool,
                            payload_len_bytes_read: usize, payload_len_so_far: u64 },
    ReadingMask { fin: bool, opcode: u8, mask_bytes_read: usize, payload_len: u64,
                  mask_so_far: [u8; 4] },
    ReadingPayload { fin: bool, opcode: u8, payload_len: u64, mask: [u8; 4], bytes_so_far: Vec<u8> },
}

struct ParseResult {
    frame: Vec<u8>,
    opcode: u8,
    fin: bool,
}

impl ParserState {
    fn done_payload_length(bytes_read: usize,
                           fin: bool, opcode: u8, masked: bool, payload_len: u64)
                           -> (ParserState, (usize, Option<ParseResult>))
    {
        use self::ParserState::*;
        if masked {
            (ReadingMask { fin: fin, opcode: opcode, payload_len: payload_len,
                           mask_bytes_read: 0, mask_so_far: [0; 4] },
             (bytes_read, None))
        } else if payload_len == 0 {
            (NotStarted,
             (bytes_read, Some(ParseResult { frame: Vec::new(), fin: fin, opcode: opcode })))
        } else {
            (ReadingPayload { fin: fin, opcode: opcode,
                              payload_len: payload_len, mask: [0; 4],
                              bytes_so_far: Vec::new() },
              (bytes_read, None))
        }
    }

    /// returns number of bytes consumed and the complete message, if there is one.
    fn advance(&mut self, buf: &[u8]) -> (usize, Option<ParseResult>) {
        use self::ParserState::*;
        let (new_state, result) = match self {
            &mut NotStarted => {
                if buf.len() < 1 {
                    return (0, None)
                }

                (DoneFirstByte { fin: (buf[0] & 0x80) != 0, opcode: buf[0] & 0xf }, (1, None))
            }
            &mut DoneFirstByte { fin, opcode } => {
                if buf.len() < 1 {
                    return (0, None)
                }

                let masked = (buf[0] & 0x80) != 0;

                match buf[0] & 0x7f {
                    126 => {
                        (ReadingLongPayloadLength {
                            fin: fin,
                            opcode: opcode,
                            masked: masked,
                            payload_len_bytes_read: 6,
                            payload_len_so_far: 0,
                        }, (1, None))
                    }
                    127 => {
                        (ReadingLongPayloadLength {
                            fin: fin,
                            opcode: opcode,
                            masked: masked,
                            payload_len_bytes_read: 0,
                            payload_len_so_far: 0,
                        }, (1, None))
                    }
                    n => ParserState::done_payload_length(1, fin, opcode, masked, n as u64)
                }
            }
            &mut ReadingLongPayloadLength { fin, opcode, masked, payload_len_bytes_read,
                                            payload_len_so_far } => {
                let mut idx = 0;
                let mut new_so_far = payload_len_so_far;
                while idx + payload_len_bytes_read < 8 && idx < buf.len() {
                    new_so_far += (buf[idx] as u64) << (8 * (7 - idx - payload_len_bytes_read));
                    idx += 1;
                }

                if buf.len() + payload_len_bytes_read < 8 {
                    (ReadingLongPayloadLength {
                        fin: fin,
                        opcode: opcode,
                        masked: masked,
                        payload_len_bytes_read: idx + payload_len_bytes_read,
                        payload_len_so_far: new_so_far,
                    }, (idx, None))
                } else {
                    ParserState::done_payload_length(idx, fin, opcode, masked, new_so_far)
                }
            }
            &mut ReadingMask { fin, opcode, mask_bytes_read, payload_len, mask_so_far } => {
                let mut idx = 0;
                let mut new_so_far = mask_so_far;
                while idx + mask_bytes_read < 4 && idx < buf.len() {
                    new_so_far[idx] = buf[idx];
                    idx += 1;
                }

                if buf.len() + mask_bytes_read < 4 {
                    (ReadingMask {
                        fin: fin,
                        opcode: opcode,
                        payload_len: payload_len,
                        mask_bytes_read: idx + mask_bytes_read,
                        mask_so_far: new_so_far,
                    }, (idx, None))
                } else if payload_len == 0 {
                    (NotStarted,
                     (idx, Some(ParseResult { frame: Vec::new(), fin: fin, opcode: opcode })))
                } else {
                    (ReadingPayload { fin: fin, opcode: opcode, mask: new_so_far,
                                      bytes_so_far: Vec::new(),
                                      payload_len: payload_len },
                     (idx, None))
                }
            }
            &mut ReadingPayload { fin, opcode, payload_len, mask, ref mut bytes_so_far } => {
                let mut idx = 0;

                while (bytes_so_far.len() as u64) < payload_len && idx < buf.len() {
                    let mask_byte = mask[bytes_so_far.len() % 4];
                    bytes_so_far.push(buf[idx] ^ mask_byte);
                    idx += 1;
                }

                if (bytes_so_far.len() as u64) < payload_len {
                    return (idx, None)
                } else {
                    let frame = ::std::mem::replace(bytes_so_far, Vec::new());
                    (NotStarted,
                     (idx, Some(ParseResult { frame: frame, fin: fin, opcode: opcode })))
                }

            }
        };

        *self = new_state;
        result
    }
}

pub struct Adapter<T> where T: MessageHandler {
    handler: Option<T>,
    awaiting_pong: Rc<Cell<bool>>,
    ping_pong_promise: Promise<(), Error>,
    client_stream: Option<web_socket_stream::Client>,
    parser_state: ParserState,
    previous_frames: PreviousFrames,
}

impl <T> Adapter<T> where T: MessageHandler {
    pub fn new(handler: T,
               client_stream: web_socket_stream::Client,
               timer: ::gjio::Timer)
               -> Adapter<T> {
        let awaiting = Rc::new(Cell::new(false));
        let ping_pong_promise = do_ping_pong(
            client_stream.clone(),
            timer,
            awaiting.clone()
        ).map_else(|r| match r {
            Ok(_) => Ok(()),
            Err(e) => {
                println!("error while pinging client: {}", e);
                Ok(())
            }
        }).eagerly_evaluate();

        Adapter {
            handler: Some(handler),
            awaiting_pong: awaiting,
            ping_pong_promise: ping_pong_promise,
            client_stream: Some(client_stream),
            parser_state: ParserState::NotStarted,
            previous_frames: PreviousFrames::None,
        }
    }

    fn process_message(&mut self) -> Promise<(), Error> {
        let frames = ::std::mem::replace(&mut self.previous_frames,
                                         PreviousFrames::None);
        let message = match frames {
            PreviousFrames::None => {
                return Promise::err(Error::failed(format!("message has no frames")));
            }
            PreviousFrames::Data(d) => {
                Message::Data(d)
            }
            PreviousFrames::Text(t) => {
                Message::Text(t)
            }
        };

        match self.handler {
            Some(ref mut h) => h.handle_message(message),
            None => Promise::ok(()),
        }
    }
}

impl <T> web_socket_stream::Server for Adapter<T> where T: MessageHandler {
    fn send_bytes(&mut self,
                  params: web_socket_stream::SendBytesParams,
                  _results: web_socket_stream::SendBytesResults)
                  -> Promise<(), Error>
    {
        let message = pry!(pry!(params.get()).get_message());
        let mut result_promise = Promise::ok(());
        let mut num_bytes_read = 0;
        while num_bytes_read < message.len() {
            let (n, result) = self.parser_state.advance(&message[num_bytes_read..]);
            num_bytes_read += n;
            match result {
                None => (),
                Some(ParseResult { frame, opcode, fin }) => {
                    match opcode {
                        0x0 => { // CONTINUE
                            match &mut self.previous_frames {
                                &mut PreviousFrames::None => {
                                    return Promise::err(Error::failed(
                                        format!("CONTINUE frame received, but there are no \
                                                 previous frames.")));
                                }
                                &mut PreviousFrames::Data(ref mut data) => {
                                    data.extend_from_slice(&frame[..]);
                                    if data.len() > (1 << 20) { // 1 MB
                                        return Promise::err(Error::failed(
                                            format!("Websocket message is too big. Please split \
                                                     the message into chunks smaller than 1MB.")));
                                    }
                                }
                                &mut PreviousFrames::Text(ref mut text) => {
                                    text.push_str(&pry!(String::from_utf8(frame)));
                                    if text.len() > (1 << 20) { // 1 MB
                                        return Promise::err(Error::failed(
                                            format!("Websocket message is too big. Please split \
                                                     the message into chunks smaller than 1MB.")));
                                    }

                                }
                            }

                            if fin {
                                let promise = self.process_message();
                                result_promise = result_promise.then(|_| promise);
                            }
                        }
                        0x1 => { // UTF-8 PAYLOAD
                            self.previous_frames =
                                PreviousFrames::Text(pry!(String::from_utf8(frame)));

                            if fin {
                                let promise = self.process_message();
                                result_promise = result_promise.then(|_| promise);
                            }
                        }
                        0x2 => { // BINARY PAYLOAD
                            self.previous_frames = PreviousFrames::Data(frame);

                            if fin {
                                let promise = self.process_message();
                                result_promise = result_promise.then(|_| promise);
                            }
                        }
                        0x8 => { // TERMINATE
                            self.handler = None;
                            self.ping_pong_promise = Promise::ok(());
                            self.client_stream = None;
                        }
                        0x9 => { // PING
                            match &self.client_stream {
                                &None => (),
                                &Some(ref client) => {
                                    println!("responding to ping from client");
                                    let req = client.send_bytes_request();
                                    let promise = req.send().promise.map(|_| Ok(()));
                                    result_promise = result_promise.then(|_| promise);
                                }
                            }
                        }
                        0xa => { // PONG
                            self.awaiting_pong.set(false);
                        }
                        _ => { // OTHER
                            println!("unrecognized websocket opcode {}", opcode);
                        }
                    }
                }
            }
        }

        result_promise
    }
}
<|endoftext|>