<commit_before>use std::sync::Arc;
use std::net::SocketAddr;

use bytes::Bytes;

use futures;
use futures::stream::Stream;

use httpbis;
use httpbis::Service as HttpbisService;
use httpbis::HttpScheme;
use httpbis::Header;
use httpbis::Headers;
use httpbis::HttpPartStream;

use tls_api;

use method::MethodDescriptor;

use error::*;
use result;

use httpbis::futures_misc::*;
use futures_grpc::*;

use grpc_frame::*;
use grpc_http_to_response::*;

use req::*;
use resp::*;

#[derive(Default, Debug, Clone)]
pub struct ClientConf {
    pub http: httpbis::ClientConf,
}

impl ClientConf {
    pub fn new() -> ClientConf {
        Default::default()
    }
}

/// gRPC client implementation.
/// Used by generated code.
pub struct Client {
    client: ::std::sync::Arc<httpbis::Client>,
    host: String,
    http_scheme: HttpScheme,
}

impl Client {
    /// Create a client connected to specified host and port.
    pub fn new_plain(host: &str, port: u16, conf: ClientConf)
        -> result::Result<Client>
    {
        let mut conf = conf;
        conf.http.thread_name =
            Some(conf.http.thread_name.unwrap_or_else(|| "grpc-client-loop".to_owned()));

        httpbis::Client::new_plain(host, port, conf.http)
            .map(|client| {
                Client {
                    client: ::std::sync::Arc::new(client),
                    host: host.to_owned(),
                    http_scheme: HttpScheme::Http,
                }
            })
            .map_err(Error::from)
    }

    /// Create a clone of this client but refer to same httpbis::Client.
    pub fn clone(&self)
        -> Client
    {
        Client {
            client: self.client.clone(),
            host: self.host.to_owned(),
            http_scheme: HttpScheme::Http,
        }
    }

    /// Create a client connected to specified host and port.
    pub fn new_tls<C : tls_api::TlsConnector>(host: &str, port: u16, conf: ClientConf)
        -> result::Result<Client>
    {
        let mut conf = conf;
        conf.http.thread_name =
            Some(conf.http.thread_name.unwrap_or_else(|| "grpc-client-loop".to_owned()));

        httpbis::Client::new_tls::<C>(host, port, conf.http)
            .map(|client| {
                Client {
                    client: ::std::sync::Arc::new(client),
                    host: host.to_owned(),
                    http_scheme: HttpScheme::Https,
                }
            })
            .map_err(Error::from)
    }

    pub fn new_expl<C : tls_api::TlsConnector>(addr: &SocketAddr, host: &str, tls: httpbis::ClientTlsOption<C>, conf: ClientConf)
        -> result::Result<Client>
    {
        let mut conf = conf;
        conf.http.thread_name =
            Some(conf.http.thread_name.unwrap_or_else(|| "grpc-client-loop".to_owned()));

        let http_scheme = tls.http_scheme();

        httpbis::Client::new_expl(addr, tls, conf.http)
            .map(|client| {
                Client {
                    client: ::std::sync::Arc::new(client),
                    host: host.to_owned(),
                    http_scheme: http_scheme,
                }
            })
            .map_err(Error::from)
    }

    pub fn new_resp_channel<Resp : Send + 'static>(&self)
        -> futures::Oneshot<(futures::sync::mpsc::UnboundedSender<ResultOrEof<Resp, Error>>, GrpcStream<Resp>)>
    {
        let (one_sender, one_receiver) = futures::oneshot();

        let (sender, receiver) = futures::sync::mpsc::unbounded();
        let receiver: GrpcStream<ResultOrEof<Resp, Error>> =
            Box::new(receiver.map_err(|()| Error::Other("receive from resp channel")));

        let receiver: GrpcStream<Resp> = Box::new(stream_with_eof_and_error(receiver, || Error::Other("unexpected EOF")));

        // TODO: oneshot sender is no longer necessary as we don't use tokio queues
        one_sender.send((sender, receiver)).ok().unwrap();

        one_receiver
    }

    fn call_impl<Req, Resp>(
        &self,
        options: RequestOptions,
        req: StreamingRequest<Req>,
        method: Arc<MethodDescriptor<Req, Resp>>)
        -> StreamingResponse<Resp>
        where
            Req : Send + 'static,
            Resp : Send + 'static,
    {
        info!("start call {}", method.name);

        let mut headers = Headers(vec![
            Header::new(Bytes::from_static(b":method"), Bytes::from_static(b"POST")),
            Header::new(Bytes::from_static(b":path"), method.name.clone()),
            Header::new(Bytes::from_static(b":authority"), self.host.clone()),
            Header::new(Bytes::from_static(b":scheme"), Bytes::from_static(self.http_scheme.as_bytes())),
            Header::new(Bytes::from_static(b"content-type"), Bytes::from_static(b"application/grpc")),
        ]);

        headers.extend(options.metadata.into_headers());

        let request_frames = {
            let method = method.clone();
            req.0
                .and_then(move |req| {
                    let grpc_frame = method.req_marshaller.write(&req)?;
                    Ok(Bytes::from(write_grpc_frame_to_vec(&grpc_frame)))
                })
                .map_err(|_e| httpbis::Error::Other("grpc error")) // TODO: preserve error
        };

        let http_response_stream = self.client
            .start_request(
                headers,
                HttpPartStream::bytes(request_frames));

        let grpc_frames = http_response_to_grpc_frames(http_response_stream);

        grpc_frames.and_then_items(move |frame| method.resp_marshaller.read(frame))
    }

    pub fn call_unary<Req, Resp>(&self, o: RequestOptions, req: Req, method: Arc<MethodDescriptor<Req, Resp>>)
                                 -> SingleResponse<Resp>
            where Req: Send + 'static, Resp: Send + 'static
    {
        self.call_impl(o, StreamingRequest::once(req), method).single()
    }

    pub fn call_server_streaming<Req, Resp>(&self, o: RequestOptions, req: Req, method: Arc<MethodDescriptor<Req, Resp>>)
                                            -> StreamingResponse<Resp>
            where Req: Send + 'static, Resp: Send + 'static
    {
        self.call_impl(o, StreamingRequest::once(req), method)
    }

    pub fn call_client_streaming<Req, Resp>(&self, o: RequestOptions, req: StreamingRequest<Req>, method: Arc<MethodDescriptor<Req, Resp>>)
                                            -> SingleResponse<Resp>
            where Req: Send + 'static, Resp: Send + 'static
    {
        self.call_impl(o, req, method).single()
    }

    pub fn call_bidi<Req, Resp>(&self, o: RequestOptions, req: StreamingRequest<Req>, method: Arc<MethodDescriptor<Req, Resp>>)
                                -> StreamingResponse<Resp>
            where Req: Send + 'static, Resp: Send + 'static
    {
        self.call_impl(o, req, method)
    }
}

fn _assert_types() {
    ::assert_types::assert_send::<Client>();
    ::assert_types::assert_sync::<Client>();
}
<commit_msg>Add the 'te' header to the GRPC requests<commit_after>use std::sync::Arc;
use std::net::SocketAddr;

use bytes::Bytes;

use futures;
use futures::stream::Stream;

use httpbis;
use httpbis::Service as HttpbisService;
use httpbis::HttpScheme;
use httpbis::Header;
use httpbis::Headers;
use httpbis::HttpPartStream;

use tls_api;

use method::MethodDescriptor;

use error::*;
use result;

use httpbis::futures_misc::*;
use futures_grpc::*;

use grpc_frame::*;
use grpc_http_to_response::*;

use req::*;
use resp::*;

#[derive(Default, Debug, Clone)]
pub struct ClientConf {
    pub http: httpbis::ClientConf,
}

impl ClientConf {
    pub fn new() -> ClientConf {
        Default::default()
    }
}

/// gRPC client implementation.
/// Used by generated code.
pub struct Client {
    client: ::std::sync::Arc<httpbis::Client>,
    host: String,
    http_scheme: HttpScheme,
}

impl Client {
    /// Create a client connected to specified host and port.
    pub fn new_plain(host: &str, port: u16, conf: ClientConf)
        -> result::Result<Client>
    {
        let mut conf = conf;
        conf.http.thread_name =
            Some(conf.http.thread_name.unwrap_or_else(|| "grpc-client-loop".to_owned()));

        httpbis::Client::new_plain(host, port, conf.http)
            .map(|client| {
                Client {
                    client: ::std::sync::Arc::new(client),
                    host: host.to_owned(),
                    http_scheme: HttpScheme::Http,
                }
            })
            .map_err(Error::from)
    }

    /// Create a clone of this client but refer to same httpbis::Client.
    pub fn clone(&self)
        -> Client
    {
        Client {
            client: self.client.clone(),
            host: self.host.to_owned(),
            http_scheme: HttpScheme::Http,
        }
    }

    /// Create a client connected to specified host and port.
    pub fn new_tls<C : tls_api::TlsConnector>(host: &str, port: u16, conf: ClientConf)
        -> result::Result<Client>
    {
        let mut conf = conf;
        conf.http.thread_name =
            Some(conf.http.thread_name.unwrap_or_else(|| "grpc-client-loop".to_owned()));

        httpbis::Client::new_tls::<C>(host, port, conf.http)
            .map(|client| {
                Client {
                    client: ::std::sync::Arc::new(client),
                    host: host.to_owned(),
                    http_scheme: HttpScheme::Https,
                }
            })
            .map_err(Error::from)
    }

    pub fn new_expl<C : tls_api::TlsConnector>(addr: &SocketAddr, host: &str, tls: httpbis::ClientTlsOption<C>, conf: ClientConf)
        -> result::Result<Client>
    {
        let mut conf = conf;
        conf.http.thread_name =
            Some(conf.http.thread_name.unwrap_or_else(|| "grpc-client-loop".to_owned()));

        let http_scheme = tls.http_scheme();

        httpbis::Client::new_expl(addr, tls, conf.http)
            .map(|client| {
                Client {
                    client: ::std::sync::Arc::new(client),
                    host: host.to_owned(),
                    http_scheme: http_scheme,
                }
            })
            .map_err(Error::from)
    }

    pub fn new_resp_channel<Resp : Send + 'static>(&self)
        -> futures::Oneshot<(futures::sync::mpsc::UnboundedSender<ResultOrEof<Resp, Error>>, GrpcStream<Resp>)>
    {
        let (one_sender, one_receiver) = futures::oneshot();

        let (sender, receiver) = futures::sync::mpsc::unbounded();
        let receiver: GrpcStream<ResultOrEof<Resp, Error>> =
            Box::new(receiver.map_err(|()| Error::Other("receive from resp channel")));

        let receiver: GrpcStream<Resp> = Box::new(stream_with_eof_and_error(receiver, || Error::Other("unexpected EOF")));

        // TODO: oneshot sender is no longer necessary as we don't use tokio queues
        one_sender.send((sender, receiver)).ok().unwrap();

        one_receiver
    }

    fn call_impl<Req, Resp>(
        &self,
        options: RequestOptions,
        req: StreamingRequest<Req>,
        method: Arc<MethodDescriptor<Req, Resp>>)
        -> StreamingResponse<Resp>
        where
            Req : Send + 'static,
            Resp : Send + 'static,
    {
        info!("start call {}", method.name);

        let mut headers = Headers(vec![
            Header::new(Bytes::from_static(b":method"), Bytes::from_static(b"POST")),
            Header::new(Bytes::from_static(b":path"), method.name.clone()),
            Header::new(Bytes::from_static(b":authority"), self.host.clone()),
            Header::new(Bytes::from_static(b":scheme"), Bytes::from_static(self.http_scheme.as_bytes())),
            Header::new(Bytes::from_static(b"content-type"), Bytes::from_static(b"application/grpc")),
            Header::new(Bytes::from_static(b"te"), Bytes::from_static(b"trailers")),
        ]);

        headers.extend(options.metadata.into_headers());

        let request_frames = {
            let method = method.clone();
            req.0
                .and_then(move |req| {
                    let grpc_frame = method.req_marshaller.write(&req)?;
                    Ok(Bytes::from(write_grpc_frame_to_vec(&grpc_frame)))
                })
                .map_err(|_e| httpbis::Error::Other("grpc error")) // TODO: preserve error
        };

        let http_response_stream = self.client
            .start_request(
                headers,
                HttpPartStream::bytes(request_frames));

        let grpc_frames = http_response_to_grpc_frames(http_response_stream);

        grpc_frames.and_then_items(move |frame| method.resp_marshaller.read(frame))
    }

    pub fn call_unary<Req, Resp>(&self, o: RequestOptions, req: Req, method: Arc<MethodDescriptor<Req, Resp>>)
                                 -> SingleResponse<Resp>
            where Req: Send + 'static, Resp: Send + 'static
    {
        self.call_impl(o, StreamingRequest::once(req), method).single()
    }

    pub fn call_server_streaming<Req, Resp>(&self, o: RequestOptions, req: Req, method: Arc<MethodDescriptor<Req, Resp>>)
                                            -> StreamingResponse<Resp>
            where Req: Send + 'static, Resp: Send + 'static
    {
        self.call_impl(o, StreamingRequest::once(req), method)
    }

    pub fn call_client_streaming<Req, Resp>(&self, o: RequestOptions, req: StreamingRequest<Req>, method: Arc<MethodDescriptor<Req, Resp>>)
                                            -> SingleResponse<Resp>
            where Req: Send + 'static, Resp: Send + 'static
    {
        self.call_impl(o, req, method).single()
    }

    pub fn call_bidi<Req, Resp>(&self, o: RequestOptions, req: StreamingRequest<Req>, method: Arc<MethodDescriptor<Req, Resp>>)
                                -> StreamingResponse<Resp>
            where Req: Send + 'static, Resp: Send + 'static
    {
        self.call_impl(o, req, method)
    }
}

fn _assert_types() {
    ::assert_types::assert_send::<Client>();
    ::assert_types::assert_sync::<Client>();
}
<|endoftext|>


<commit_before>extern crate piston;
extern crate graphics;
extern crate opengl_graphics;
extern crate drag_controller;
extern crate sdl2_window;

use std::path::Path;
use std::rc::Rc;
use std::cell::RefCell;

use graphics::ImageSize;
use drag_controller::{
    DragController,
    Drag
};
use opengl_graphics::{ GlGraphics, OpenGL, Texture };
use sdl2_window::Sdl2Window as Window;

fn main() {
    println!("Click in the red square and drag.");
    println!("Toggle grid with G.");
    println!("Reset grid with R.");

    let opengl = OpenGL::_3_2;
    let window = Window::new(
        opengl,
        piston::window::WindowSettings {
            title: "piston-example-deform".to_string(),
            size: [300, 300],
            fullscreen: false,
            exit_on_esc: true,
            samples: 4,
        }
    );

    let image = Path::new("./bin/assets/rust-logo.png");
    let image = Texture::from_path(&image).unwrap();

    let (width, height) = image.get_size();
    let width = width as f64;
    let height = height as f64;
    let mut grid = graphics::deform::DeformGrid::new(
        [0.0, 0.0, width, height],
        20, 20
    );

    let mut drag = DragController::new();
    let mut draw_grid = true;

    let mut gl = GlGraphics::new(opengl);
    let window = Rc::new(RefCell::new(window));
    for e in piston::events(window) {
        use piston::event::*;

        drag.event(&e, |action| {
            match action {
                Drag::Start(x, y) => {
                    match grid.hit([x, y]) {
                        None => {
                            // Did not hit deformed grid.
                            grid.add_control_point([x, y]);
                        }
                        Some(pos) => {
                            // Add point to deformed grid.
                            grid.add_control_point(pos);
                            let n = grid.ps.len();
                            grid.set_current(n - 1, [x, y]);
                        }
                    }
                    grid.update();
                    true
                }
                Drag::Move(x, y) => {
                    let n = grid.ps.len();
                    grid.set_current(n - 1, [x, y]);
                    grid.update();
                    true
                }
                Drag::End(_, _) => false,
                // Continue dragging when receiving focus.
                Drag::Interrupt => true,
            }
        });
        if let Some(button) = e.press_args() {
            use piston::input::Button::Keyboard;
            use piston::input::keyboard::Key;

            if button == Keyboard(Key::G) {
                draw_grid = !draw_grid;
                println!("Draw grid {}", draw_grid);
            } else if button == Keyboard(Key::R) {
                grid.reset_control_points();
                grid.reset_vertices_and_texture_coords();
                grid.update();
                println!("Reset grid");
            }
        }
        if let Some(args) = e.render_args() {
            gl.draw(
                [0, 0, args.width as i32, args.height as i32],
                |c, g| {

            graphics::clear(graphics::color::WHITE, g);

            // Draw deformed image.
            grid.draw_image(&image, &c.draw_state, c.transform, g);

            if draw_grid {
                // Draw grid.
                grid.draw_vertical_lines(
                    &graphics::Line::new([0.0, 1.0, 0.0, 1.0], 0.5),
                    &c.draw_state,
                    c.transform,
                    g
                );
                grid.draw_horizontal_lines(
                    &graphics::Line::new([0.0, 0.0, 1.0, 1.0], 0.5),
                    &c.draw_state,
                    c.transform,
                    g
                );
            }

            // Draw rect of the original grid.
            graphics::Rectangle::border([1.0, 0.0, 0.0, 1.0], 1.5)
                .draw([0.0, 0.0, width, height], &c.draw_state, c.transform, g);

            // Draw control points.
            let original = graphics::Ellipse::new([1.0, 0.0, 0.0, 0.5]);
            let current = graphics::Ellipse::new([0.0, 0.0, 0.0, 0.5]);
            for i in (0..grid.ps.len()) {
                use graphics::ellipse::circle;

                // Original positions.
                let x = grid.ps[i][0];
                let y = grid.ps[i][1];
                original.draw(circle(x, y, 3.0), &c.draw_state, c.transform, g);

                // Current positions.
                let x = grid.qs[i][0];
                let y = grid.qs[i][1];
                current.draw(circle(x, y, 3.0), &c.draw_state, c.transform, g);
            };

            });
        } // end render
    }
}

<commit_msg>Upgraded deform example to latest piston<commit_after>extern crate piston;
extern crate graphics;
extern crate opengl_graphics;
extern crate drag_controller;
extern crate sdl2_window;

use std::path::Path;
use std::rc::Rc;
use std::cell::RefCell;

use graphics::ImageSize;
use drag_controller::{
    DragController,
    Drag
};
use opengl_graphics::{ GlGraphics, OpenGL, Texture };
use sdl2_window::Sdl2Window as Window;
use piston::window::{ WindowSettings, Size };

fn main() {
    println!("Click in the red square and drag.");
    println!("Toggle grid with G.");
    println!("Reset grid with R.");

    let opengl = OpenGL::_3_2;
    let window = Window::new(
        opengl,
        WindowSettings::new(
            "piston-example-deform".to_string(),
            Size { width: 300, height: 300 }
        )
        .exit_on_esc(true)
        .samples(4)
    );

    let image = Path::new("./bin/assets/rust-logo.png");
    let image = Texture::from_path(&image).unwrap();

    let (width, height) = image.get_size();
    let width = width as f64;
    let height = height as f64;
    let mut grid = graphics::deform::DeformGrid::new(
        [0.0, 0.0, width, height],
        20, 20
    );

    let mut drag = DragController::new();
    let mut draw_grid = true;

    let mut gl = GlGraphics::new(opengl);
    let window = Rc::new(RefCell::new(window));
    for e in piston::events(window) {
        use piston::event::*;

        drag.event(&e, |action| {
            match action {
                Drag::Start(x, y) => {
                    match grid.hit([x, y]) {
                        None => {
                            // Did not hit deformed grid.
                            grid.add_control_point([x, y]);
                        }
                        Some(pos) => {
                            // Add point to deformed grid.
                            grid.add_control_point(pos);
                            let n = grid.ps.len();
                            grid.set_current(n - 1, [x, y]);
                        }
                    }
                    grid.update();
                    true
                }
                Drag::Move(x, y) => {
                    let n = grid.ps.len();
                    grid.set_current(n - 1, [x, y]);
                    grid.update();
                    true
                }
                Drag::End(_, _) => false,
                // Continue dragging when receiving focus.
                Drag::Interrupt => true,
            }
        });
        if let Some(button) = e.press_args() {
            use piston::input::Button::Keyboard;
            use piston::input::keyboard::Key;

            if button == Keyboard(Key::G) {
                draw_grid = !draw_grid;
                println!("Draw grid {}", draw_grid);
            } else if button == Keyboard(Key::R) {
                grid.reset_control_points();
                grid.reset_vertices_and_texture_coords();
                grid.update();
                println!("Reset grid");
            }
        }
        if let Some(args) = e.render_args() {
            gl.draw(
                [0, 0, args.width as i32, args.height as i32],
                |c, g| {

            graphics::clear(graphics::color::WHITE, g);

            // Draw deformed image.
            grid.draw_image(&image, &c.draw_state, c.transform, g);

            if draw_grid {
                // Draw grid.
                grid.draw_vertical_lines(
                    &graphics::Line::new([0.0, 1.0, 0.0, 1.0], 0.5),
                    &c.draw_state,
                    c.transform,
                    g
                );
                grid.draw_horizontal_lines(
                    &graphics::Line::new([0.0, 0.0, 1.0, 1.0], 0.5),
                    &c.draw_state,
                    c.transform,
                    g
                );
            }

            // Draw rect of the original grid.
            graphics::Rectangle::new_border([1.0, 0.0, 0.0, 1.0], 1.5)
                .draw([0.0, 0.0, width, height], &c.draw_state, c.transform, g);

            // Draw control points.
            let original = graphics::Ellipse::new([1.0, 0.0, 0.0, 0.5]);
            let current = graphics::Ellipse::new([0.0, 0.0, 0.0, 0.5]);
            for i in (0..grid.ps.len()) {
                use graphics::ellipse::circle;

                // Original positions.
                let x = grid.ps[i][0];
                let y = grid.ps[i][1];
                original.draw(circle(x, y, 3.0), &c.draw_state, c.transform, g);

                // Current positions.
                let x = grid.qs[i][0];
                let y = grid.qs[i][1];
                current.draw(circle(x, y, 3.0), &c.draw_state, c.transform, g);
            };

            });
        } // end render
    }
}

<|endoftext|>


<commit_before>// Copyright 2016 Google Inc. All Rights Reserved.
//
// Licensed under the MIT License, <LICENSE or http://opensource.org/licenses/MIT>.
// This file may not be copied, modified, or distributed except according to those terms.

#![feature(inclusive_range_syntax, conservative_impl_trait, plugin)]
#![plugin(tarpc_plugins)]

extern crate chrono;
extern crate env_logger;
extern crate futures;
#[macro_use]
extern crate log;
#[macro_use]
extern crate tarpc;
extern crate tokio_core;
extern crate futures_cpupool;

use futures::Future;
use futures_cpupool::{CpuFuture, CpuPool};
use std::thread;
use std::time::{Duration, Instant, SystemTime};
use tarpc::future::{Connect};
use tarpc::util::{FirstSocketAddr, Never};
use tokio_core::reactor;

service! {
    rpc read(size: u32) -> Vec<u8>;
}

#[derive(Clone)]
struct Server(CpuPool);

impl Server {
    fn new() -> Self {
        Server(CpuPool::new_num_cpus())
    }
}

impl FutureService for Server {
    type ReadFut = CpuFuture<Vec<u8>, Never>;

    fn read(&self, size: u32) -> Self::ReadFut {
        self.0
            .spawn(futures::lazy(move || {
                let mut vec: Vec<u8> = Vec::with_capacity(size as usize);
                for i in 0..size {
                    vec.push((i % 1 << 8) as u8);
                }
                futures::finished::<_, Never>(vec)
            }))
    }
}

fn run_once<'a>(clients: &'a [FutureClient], concurrency: u32, print: bool)
    -> Box<Future<Item=(), Error=()> + 'a>
{
    let start = Instant::now();
    let futs = clients.iter()
        .cycle()
        .take(concurrency as usize)
        .map(|client| {
            let start = SystemTime::now();
            let future = client.read(CHUNK_SIZE).map(move |_| start.elapsed().unwrap());
            future
        });
    let futs = futures::collect(futs);

    Box::new(futs.map(move |latencies| {
        let total_time = start.elapsed();

        let sum_latencies = latencies.iter().fold(Duration::new(0, 0), |sum, &dur| sum + dur);
        let mean = sum_latencies / latencies.len() as u32;
        let min_latency = *latencies.iter().min().unwrap();
        let max_latency = *latencies.iter().max().unwrap();

        if print {
            println!("{} requests => Mean={}µs, Min={}µs, Max={}µs, Total={}µs",
                     latencies.len(),
                     mean.microseconds(),
                     min_latency.microseconds(),
                     max_latency.microseconds(),
                     total_time.microseconds());
        }
    }).map_err(|e| panic!(e)))
}

trait Microseconds {
    fn microseconds(&self) -> i64;
}

impl Microseconds for Duration {
    fn microseconds(&self) -> i64 {
         chrono::Duration::from_std(*self)
             .unwrap()
             .num_microseconds()
             .unwrap()
    }
}

const CHUNK_SIZE: u32 = 1 << 10;
const MAX_CONCURRENCY: u32 = 100;

fn main() {
    let _ = env_logger::init();

    let server = Server::new().listen("localhost:0".first_socket_addr()).wait().unwrap();
    println!("Server listening on {}.", server.local_addr());

    let mut core = reactor::Core::new().unwrap();
    let handle = core.handle();
    let clients = (1...5)
        .map(|i| {
            println!("Client {} connecting...", i);
            FutureClient::connect_with(server.local_addr(), &handle).map_err(|e| panic!(e))
        });
    let clients = core.run(futures::collect(clients)).unwrap();
    println!("Starting...");
    let runs = ::std::iter::once(run_once(&clients, MAX_CONCURRENCY, false))
        .chain((1...MAX_CONCURRENCY).map(|concurrency| run_once(&clients, concurrency, true)));
    core.run(futures::collect(runs)).unwrap();
}
<commit_msg>Run clients in parallel (that's the whole point!).<commit_after>// Copyright 2016 Google Inc. All Rights Reserved.
//
// Licensed under the MIT License, <LICENSE or http://opensource.org/licenses/MIT>.
// This file may not be copied, modified, or distributed except according to those terms.

#![feature(inclusive_range_syntax, conservative_impl_trait, plugin)]
#![plugin(tarpc_plugins)]

extern crate chrono;
extern crate env_logger;
extern crate futures;
#[macro_use]
extern crate log;
#[macro_use]
extern crate tarpc;
extern crate tokio_core;
extern crate futures_cpupool;

use futures::Future;
use futures_cpupool::{CpuFuture, CpuPool};
use std::time::{Duration, Instant, SystemTime};
use tarpc::future::{Connect};
use tarpc::util::{FirstSocketAddr, Never, spawn_core};
use tokio_core::reactor;

service! {
    rpc read(size: u32) -> Vec<u8>;
}

#[derive(Clone)]
struct Server(CpuPool);

impl Server {
    fn new() -> Self {
        Server(CpuPool::new_num_cpus())
    }
}

impl FutureService for Server {
    type ReadFut = CpuFuture<Vec<u8>, Never>;

    fn read(&self, size: u32) -> Self::ReadFut {
        self.0
            .spawn(futures::lazy(move || {
                let mut vec: Vec<u8> = Vec::with_capacity(size as usize);
                for i in 0..size {
                    vec.push((i % 1 << 8) as u8);
                }
                futures::finished::<_, Never>(vec)
            }))
    }
}

fn run_once<'a>(clients: &'a [FutureClient], concurrency: u32, print: bool)
    -> Box<Future<Item=(), Error=()> + 'a>
{
    let start = Instant::now();
    let futs = clients.iter()
        .cycle()
        .take(concurrency as usize)
        .map(|client| {
            let start = SystemTime::now();
            let future = client.read(CHUNK_SIZE).map(move |_| start.elapsed().unwrap());
            future
        })
        // Need an intermediate collection to kick off each future,
        // because futures::collect will iterate sequentially.
        .collect::<Vec<_>>();
    let futs = futures::collect(futs);

    Box::new(futs.map(move |latencies| {
        let total_time = start.elapsed();

        let sum_latencies = latencies.iter().fold(Duration::new(0, 0), |sum, &dur| sum + dur);
        let mean = sum_latencies / latencies.len() as u32;
        let min_latency = *latencies.iter().min().unwrap();
        let max_latency = *latencies.iter().max().unwrap();

        if print {
            println!("{} requests => Mean={}µs, Min={}µs, Max={}µs, Total={}µs",
                     latencies.len(),
                     mean.microseconds(),
                     min_latency.microseconds(),
                     max_latency.microseconds(),
                     total_time.microseconds());
        }
    }).map_err(|e| panic!(e)))
}

trait Microseconds {
    fn microseconds(&self) -> i64;
}

impl Microseconds for Duration {
    fn microseconds(&self) -> i64 {
         chrono::Duration::from_std(*self)
             .unwrap()
             .num_microseconds()
             .unwrap()
    }
}

const CHUNK_SIZE: u32 = 1 << 10;
const MAX_CONCURRENCY: u32 = 100;

fn main() {
    let _ = env_logger::init();

    let server = Server::new().listen("localhost:0".first_socket_addr()).wait().unwrap();
    println!("Server listening on {}.", server.local_addr());

    // The driver of the main future.
    let mut core = reactor::Core::new().unwrap();

    let clients = (1...5)
        // Spin up a couple threads to drive the clients.
        .map(|i| (i, spawn_core()))
        .map(|(i, remote)| {
            println!("Client {} connecting...", i);
            FutureClient::connect_remotely(server.local_addr(), &remote)
                .map_err(|e| panic!(e))
        })
        // Need an intermediate collection to connect the clients in parallel,
        // because `futures::collect` iterates sequentially.
        .collect::<Vec<_>>();

    let clients = core.run(futures::collect(clients)).unwrap();
    println!("Starting...");
    let runs = ::std::iter::once(run_once(&clients, MAX_CONCURRENCY, false))
        .chain((1...MAX_CONCURRENCY).map(|concurrency| run_once(&clients, concurrency, true)));
    core.run(futures::collect(runs)).unwrap();
}
<|endoftext|>


<commit_before>//
// Copyright 2021 The Project Oak Authors
//
// 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 crate::proto::EncryptedData;
use anyhow::anyhow;
use ring::{
    aead::{self, BoundKey},
    agreement, rand,
};
use sha2::{digest::Digest, Sha256};

// `ring::aead` uses 96-bit (12-byte) nonces.
// https://briansmith.org/rustdoc/ring/aead/constant.NONCE_LEN.html
//
// TODO(#2087): Use a unique nonce for each message.
const NONCE: [u8; aead::NONCE_LEN] = [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0];
// Algorithm used for encrypting/decrypting messages.
// https://datatracker.ietf.org/doc/html/rfc5288
static AEAD_ALGORITHM: &aead::Algorithm = &aead::AES_256_GCM;
// Algorithm used for negotiating a session key.
// https://datatracker.ietf.org/doc/html/rfc7748
static KEY_AGREEMENT_ALGORITHM: &agreement::Algorithm = &agreement::X25519;

/// Nonce implementation used by [`AeadEncryptor`].
/// It returns a single nonce once and then only returns errors.
struct OneNonceSequence(Option<aead::Nonce>);

impl OneNonceSequence {
    fn new(nonce: [u8; aead::NONCE_LEN]) -> Self {
        let nonce = aead::Nonce::assume_unique_for_key(nonce);
        Self(Some(nonce))
    }
}

impl aead::NonceSequence for OneNonceSequence {
    fn advance(&mut self) -> Result<aead::Nonce, ring::error::Unspecified> {
        self.0.take().ok_or(ring::error::Unspecified)
    }
}

/// Implementation of Authenticated Encryption with Associated Data (AEAD).
/// https://datatracker.ietf.org/doc/html/rfc5116
pub struct AeadEncryptor {
    /// Key used for encrypting and decrypting data.
    ///
    /// AES-GCM algorithm uses the same key for both encryption and decryption.
    /// https://datatracker.ietf.org/doc/html/rfc5288
    key: Vec<u8>,
}

impl AeadEncryptor {
    pub fn new(key: Vec<u8>) -> Self {
        Self { key }
    }

    /// Encrypts `data` using [`AeadEncryptor::key`].
    pub fn encrypt(&mut self, data: &[u8]) -> anyhow::Result<EncryptedData> {
        // Bind `AeadEncryptor::key` to a `NONCE`.
        let unbound_sealing_key = aead::UnboundKey::new(AEAD_ALGORITHM, &self.key)
            .map_err(|error| anyhow!("Couldn't create sealing key: {:?}", error))?;
        let mut sealing_key =
            ring::aead::SealingKey::new(unbound_sealing_key, OneNonceSequence::new(NONCE));

        let mut encrypted_data = data.to_vec();
        sealing_key
            // Additional authenticated data is not required for the remotely attested channel,
            // since after session key is established client and server exchange messages with a
            // single encrypted field.
            // And the nonce is authenticated by the AEAD algorithm itself.
            // https://datatracker.ietf.org/doc/html/rfc5116#section-2.1
            .seal_in_place_append_tag(aead::Aad::empty(), &mut encrypted_data)
            .map_err(|error| anyhow!("Couldn't encrypt data: {:?}", error))?;

        Ok(EncryptedData {
            nonce: NONCE.to_vec(),
            data: encrypted_data,
        })
    }

    /// Decrypts and authenticates `data` using [`AeadEncryptor::key`].
    /// `data` must contain an encrypted message prefixed with a random nonce of [`aead::NONCE_LEN`]
    /// length.
    pub fn decrypt(&mut self, data: &EncryptedData) -> anyhow::Result<Vec<u8>> {
        // Extract nonce from `data`.
        let mut nonce: [u8; aead::NONCE_LEN] = Default::default();
        nonce.copy_from_slice(&data.nonce[0..NONCE.len()]);

        // Bind `AeadEncryptor::key` to the extracted `nonce`.
        let unbound_opening_key = aead::UnboundKey::new(AEAD_ALGORITHM, &self.key).unwrap();
        let mut opening_key =
            ring::aead::OpeningKey::new(unbound_opening_key, OneNonceSequence::new(nonce));

        let mut decrypted_data = data.data.to_vec();
        let decrypted_data = opening_key
            // Additional authenticated data is not required for the remotely attested channel,
            // since after session key is established client and server exchange messages with a
            // single encrypted field.
            // And the nonce is authenticated by the AEAD algorithm itself.
            // https://datatracker.ietf.org/doc/html/rfc5116#section-2.1
            .open_in_place(aead::Aad::empty(), &mut decrypted_data)
            .map_err(|error| anyhow!("Couldn't decrypt data: {:?}", error))?;
        Ok(decrypted_data.to_vec())
    }
}

/// Implementation of the X25519 Elliptic Curve Diffie-Hellman (ECDH) key negotiation.
/// https://datatracker.ietf.org/doc/html/rfc7748#section-6.1
pub struct KeyNegotiator {
    private_key: agreement::EphemeralPrivateKey,
}

impl KeyNegotiator {
    pub fn create() -> anyhow::Result<Self> {
        let rng = rand::SystemRandom::new();
        let private_key = agreement::EphemeralPrivateKey::generate(KEY_AGREEMENT_ALGORITHM, &rng)
            .map_err(|error| anyhow!("Couldn't generate private key: {:?}", error))?;
        Ok(Self { private_key })
    }

    pub fn public_key(&self) -> anyhow::Result<Vec<u8>> {
        let public_key = self
            .private_key
            .compute_public_key()
            .map_err(|error| anyhow!("Couldn't get public key: {:?}", error))?;
        Ok(public_key.as_ref().to_vec())
    }

    pub fn derive_session_key(self, peer_public_key: &[u8]) -> anyhow::Result<Vec<u8>> {
        let peer_public_key =
            agreement::UnparsedPublicKey::new(KEY_AGREEMENT_ALGORITHM, peer_public_key);
        let session_key = agreement::agree_ephemeral(
            self.private_key,
            &peer_public_key,
            ring::error::Unspecified,
            Self::key_derivation_function,
        )
        .map_err(|error| anyhow!("Couldn't derive session key: {:?}", error))?;
        Ok(session_key)
    }

    /// Derives a session key from `key_material`.
    fn key_derivation_function(key_material: &[u8]) -> Result<Vec<u8>, ring::error::Unspecified> {
        // TODO(#2100): Derive a session key not just from `key_material` but also from server's and
        // client's public keys.
        // https://datatracker.ietf.org/doc/html/rfc7748#section-6.1
        Ok(key_material.to_vec())
    }
}

/// Computes a SHA-256 digest of `input` and returns it in a form of raw bytes.
pub fn get_sha256(input: &[u8]) -> Vec<u8> {
    let mut hasher = Sha256::new();
    hasher.update(&input);
    hasher.finalize().to_vec()
}
<commit_msg>Use random nonce in Remote Attestation (#2182)<commit_after>//
// Copyright 2021 The Project Oak Authors
//
// 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 crate::proto::EncryptedData;
use anyhow::anyhow;
use ring::{
    aead::{self, BoundKey},
    agreement,
    rand::{SecureRandom, SystemRandom},
};
use sha2::{digest::Digest, Sha256};

// Algorithm used for encrypting/decrypting messages.
// https://datatracker.ietf.org/doc/html/rfc5288
static AEAD_ALGORITHM: &aead::Algorithm = &aead::AES_256_GCM;
// Algorithm used for negotiating a session key.
// https://datatracker.ietf.org/doc/html/rfc7748
static KEY_AGREEMENT_ALGORITHM: &agreement::Algorithm = &agreement::X25519;

/// Nonce implementation used by [`AeadEncryptor`].
/// It returns a single nonce once and then only returns errors.
struct OneNonceSequence(Option<aead::Nonce>);

impl OneNonceSequence {
    fn new(nonce: [u8; aead::NONCE_LEN]) -> Self {
        let nonce = aead::Nonce::assume_unique_for_key(nonce);
        Self(Some(nonce))
    }
}

impl aead::NonceSequence for OneNonceSequence {
    fn advance(&mut self) -> Result<aead::Nonce, ring::error::Unspecified> {
        self.0.take().ok_or(ring::error::Unspecified)
    }
}

/// Implementation of Authenticated Encryption with Associated Data (AEAD).
/// https://datatracker.ietf.org/doc/html/rfc5116
pub struct AeadEncryptor {
    /// Key used for encrypting and decrypting data.
    ///
    /// AES-GCM algorithm uses the same key for both encryption and decryption.
    /// https://datatracker.ietf.org/doc/html/rfc5288
    key: Vec<u8>,
}

impl AeadEncryptor {
    pub fn new(key: Vec<u8>) -> Self {
        Self { key }
    }

    /// Encrypts `data` using [`AeadEncryptor::key`].
    pub fn encrypt(&mut self, data: &[u8]) -> anyhow::Result<EncryptedData> {
        // Generate a random nonce.
        let nonce = Self::generate_nonce();

        // Bind `AeadEncryptor::key` to a `NONCE`.
        let unbound_sealing_key = aead::UnboundKey::new(AEAD_ALGORITHM, &self.key)
            .map_err(|error| anyhow!("Couldn't create sealing key: {:?}", error))?;
        let mut sealing_key =
            ring::aead::SealingKey::new(unbound_sealing_key, OneNonceSequence::new(nonce));

        let mut encrypted_data = data.to_vec();
        sealing_key
            // Additional authenticated data is not required for the remotely attested channel,
            // since after session key is established client and server exchange messages with a
            // single encrypted field.
            // And the nonce is authenticated by the AEAD algorithm itself.
            // https://datatracker.ietf.org/doc/html/rfc5116#section-2.1
            .seal_in_place_append_tag(aead::Aad::empty(), &mut encrypted_data)
            .map_err(|error| anyhow!("Couldn't encrypt data: {:?}", error))?;

        Ok(EncryptedData {
            nonce: nonce.to_vec(),
            data: encrypted_data,
        })
    }

    /// Decrypts and authenticates `data` using [`AeadEncryptor::key`].
    /// `data` must contain an encrypted message prefixed with a random nonce of [`aead::NONCE_LEN`]
    /// length.
    pub fn decrypt(&mut self, data: &EncryptedData) -> anyhow::Result<Vec<u8>> {
        // Extract nonce from `data`.
        let mut nonce: [u8; aead::NONCE_LEN] = Default::default();
        nonce.copy_from_slice(&data.nonce[0..aead::NONCE_LEN]);

        // Bind `AeadEncryptor::key` to the extracted `nonce`.
        let unbound_opening_key = aead::UnboundKey::new(AEAD_ALGORITHM, &self.key).unwrap();
        let mut opening_key =
            ring::aead::OpeningKey::new(unbound_opening_key, OneNonceSequence::new(nonce));

        let mut decrypted_data = data.data.to_vec();
        let decrypted_data = opening_key
            // Additional authenticated data is not required for the remotely attested channel,
            // since after session key is established client and server exchange messages with a
            // single encrypted field.
            // And the nonce is authenticated by the AEAD algorithm itself.
            // https://datatracker.ietf.org/doc/html/rfc5116#section-2.1
            .open_in_place(aead::Aad::empty(), &mut decrypted_data)
            .map_err(|error| anyhow!("Couldn't decrypt data: {:?}", error))?;
        Ok(decrypted_data.to_vec())
    }

    /// Generate a random nonce.
    /// `ring::aead` uses 96-bit (12-byte) nonces.
    /// https://briansmith.org/rustdoc/ring/aead/constant.NONCE_LEN.html
    fn generate_nonce() -> [u8; aead::NONCE_LEN] {
        let mut nonce: [u8; aead::NONCE_LEN] = Default::default();
        let rng = SystemRandom::new();
        rng.fill(&mut nonce).unwrap();
        nonce
    }
}

/// Implementation of the X25519 Elliptic Curve Diffie-Hellman (ECDH) key negotiation.
/// https://datatracker.ietf.org/doc/html/rfc7748#section-6.1
pub struct KeyNegotiator {
    private_key: agreement::EphemeralPrivateKey,
}

impl KeyNegotiator {
    pub fn create() -> anyhow::Result<Self> {
        let rng = ring::rand::SystemRandom::new();
        let private_key = agreement::EphemeralPrivateKey::generate(KEY_AGREEMENT_ALGORITHM, &rng)
            .map_err(|error| anyhow!("Couldn't generate private key: {:?}", error))?;
        Ok(Self { private_key })
    }

    pub fn public_key(&self) -> anyhow::Result<Vec<u8>> {
        let public_key = self
            .private_key
            .compute_public_key()
            .map_err(|error| anyhow!("Couldn't get public key: {:?}", error))?;
        Ok(public_key.as_ref().to_vec())
    }

    pub fn derive_session_key(self, peer_public_key: &[u8]) -> anyhow::Result<Vec<u8>> {
        let peer_public_key =
            agreement::UnparsedPublicKey::new(KEY_AGREEMENT_ALGORITHM, peer_public_key);
        let session_key = agreement::agree_ephemeral(
            self.private_key,
            &peer_public_key,
            ring::error::Unspecified,
            Self::key_derivation_function,
        )
        .map_err(|error| anyhow!("Couldn't derive session key: {:?}", error))?;
        Ok(session_key)
    }

    /// Derives a session key from `key_material`.
    fn key_derivation_function(key_material: &[u8]) -> Result<Vec<u8>, ring::error::Unspecified> {
        // TODO(#2100): Derive a session key not just from `key_material` but also from server's and
        // client's public keys.
        // https://datatracker.ietf.org/doc/html/rfc7748#section-6.1
        Ok(key_material.to_vec())
    }
}

/// Computes a SHA-256 digest of `input` and returns it in a form of raw bytes.
pub fn get_sha256(input: &[u8]) -> Vec<u8> {
    let mut hasher = Sha256::new();
    hasher.update(&input);
    hasher.finalize().to_vec()
}
<|endoftext|>


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

use dom::bindings::codegen::BindingDeclarations::HTMLDirectoryElementBinding;
use dom::bindings::codegen::InheritTypes::HTMLDirectoryElementDerived;
use dom::bindings::js::{JSRef, Temporary};
use dom::bindings::error::ErrorResult;
use dom::document::Document;
use dom::element::HTMLDirectoryElementTypeId;
use dom::eventtarget::{EventTarget, NodeTargetTypeId};
use dom::htmlelement::HTMLElement;
use dom::node::{Node, ElementNodeTypeId};
use servo_util::str::DOMString;

#[deriving(Encodable)]
pub struct HTMLDirectoryElement {
    pub htmlelement: HTMLElement
}

impl HTMLDirectoryElementDerived for EventTarget {
    fn is_htmldirectoryelement(&self) -> bool {
        self.type_id == NodeTargetTypeId(ElementNodeTypeId(HTMLDirectoryElementTypeId))
    }
}

impl HTMLDirectoryElement {
    pub fn new_inherited(localName: DOMString, document: &JSRef<Document>) -> HTMLDirectoryElement {
        HTMLDirectoryElement {
            htmlelement: HTMLElement::new_inherited(HTMLDirectoryElementTypeId, localName, document)
        }
    }

    pub fn new(localName: DOMString, document: &JSRef<Document>) -> Temporary<HTMLDirectoryElement> {
        let element = HTMLDirectoryElement::new_inherited(localName, document);
        Node::reflect_node(box element, document, HTMLDirectoryElementBinding::Wrap)
    }
}

pub trait HTMLDirectoryElementMethods {
    fn Compact(&self) -> bool;
    fn SetCompact(&mut self, _compact: bool) -> ErrorResult;
}

impl<'a> HTMLDirectoryElementMethods for JSRef<'a, HTMLDirectoryElement> {
    fn Compact(&self) -> bool {
        false
    }

    fn SetCompact(&mut self, _compact: bool) -> ErrorResult {
        Ok(())
    }
}
<commit_msg>Remove needless '&mut self' from HTMLDirectryElementMethods.<commit_after>/* This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at http://mozilla.org/MPL/2.0/. */

use dom::bindings::codegen::BindingDeclarations::HTMLDirectoryElementBinding;
use dom::bindings::codegen::InheritTypes::HTMLDirectoryElementDerived;
use dom::bindings::js::{JSRef, Temporary};
use dom::bindings::error::ErrorResult;
use dom::document::Document;
use dom::element::HTMLDirectoryElementTypeId;
use dom::eventtarget::{EventTarget, NodeTargetTypeId};
use dom::htmlelement::HTMLElement;
use dom::node::{Node, ElementNodeTypeId};
use servo_util::str::DOMString;

#[deriving(Encodable)]
pub struct HTMLDirectoryElement {
    pub htmlelement: HTMLElement
}

impl HTMLDirectoryElementDerived for EventTarget {
    fn is_htmldirectoryelement(&self) -> bool {
        self.type_id == NodeTargetTypeId(ElementNodeTypeId(HTMLDirectoryElementTypeId))
    }
}

impl HTMLDirectoryElement {
    pub fn new_inherited(localName: DOMString, document: &JSRef<Document>) -> HTMLDirectoryElement {
        HTMLDirectoryElement {
            htmlelement: HTMLElement::new_inherited(HTMLDirectoryElementTypeId, localName, document)
        }
    }

    pub fn new(localName: DOMString, document: &JSRef<Document>) -> Temporary<HTMLDirectoryElement> {
        let element = HTMLDirectoryElement::new_inherited(localName, document);
        Node::reflect_node(box element, document, HTMLDirectoryElementBinding::Wrap)
    }
}

pub trait HTMLDirectoryElementMethods {
    fn Compact(&self) -> bool;
    fn SetCompact(&self, _compact: bool) -> ErrorResult;
}

impl<'a> HTMLDirectoryElementMethods for JSRef<'a, HTMLDirectoryElement> {
    fn Compact(&self) -> bool {
        false
    }

    fn SetCompact(&self, _compact: bool) -> ErrorResult {
        Ok(())
    }
}
<|endoftext|>


<commit_before>use collections::BTreeMap;

use redox::{self, env, BMPFile};
use redox::event::{self, EventOption, MouseEvent};
use redox::fs::file::File;
use redox::io::{Read, Seek, SeekFrom};
use redox::orbital::Window;
use redox::time::{self, Duration};
use redox::vec::Vec;
use redox::string::{String, ToString};

pub struct FileType {
    description: String,
    icon: BMPFile,
}

impl FileType {
    pub fn new(desc: &str, icon: &str) -> FileType {
        FileType { description: desc.to_string(), icon: load_icon(icon) }
    }

}

pub struct FileManager {
    file_types: BTreeMap<String, FileType>,
    /*
    folder_icon: BMPFile,
    audio_icon: BMPFile,
    bin_icon: BMPFile,
    image_icon: BMPFile,
    source_icon: BMPFile,
    script_icon: BMPFile,
    text_icon: BMPFile,
    file_icon: BMPFile,
    */
    files: Vec<String>,
    file_sizes: Vec<String>,
    selected: isize,
    last_mouse_event: MouseEvent,
    click_time: Duration,
}

fn load_icon(path: &str) -> BMPFile {
    let mut vec: Vec<u8> = Vec::new();
    if let Some(mut file) = File::open(&("file:///ui/mimetypes/".to_string() + path + ".bmp")) {
        file.read_to_end(&mut vec);
    }
    BMPFile::from_data(&vec)
}

impl FileManager {
    pub fn new() -> Self {
        FileManager {
            file_types: {
                let mut file_types = BTreeMap::<String, FileType>::new();
                file_types.insert("/".to_string(),
                                  FileType::new("Folder", "inode-directory"));
                file_types.insert("wav".to_string(),
                                  FileType::new("WAV audio", "audio-x-wav"));
                file_types.insert("bin".to_string(),
                                  FileType::new("Executable", "application-x-executable"));
                file_types.insert("bmp".to_string(),
                                  FileType::new("Bitmap Image", "image-x-generic"));
                file_types.insert("rs".to_string(),
                                  FileType::new("Rust source code", "image-x-makefile"));
                file_types.insert("crate".to_string(),
                                  FileType::new("Rust crate", "image-x-makefile"));
                file_types.insert("asm".to_string(),
                                  FileType::new("Assembly source", "image-x-makefile"));
                file_types.insert("list".to_string(),
                                  FileType::new("Disassembly source", "image-x-makefile"));
                file_types.insert("c".to_string(),
                                  FileType::new("C source code", "image-x-makefile"));
                file_types.insert("cpp".to_string(),
                                  FileType::new("C++ source code", "image-x-makefile"));
                file_types.insert("sh".to_string(),
                                  FileType::new("Shell script", "image-x-script"));
                file_types.insert("lua".to_string(),
                                  FileType::new("Lua script", "image-x-script"));
                file_types.insert("txt".to_string(),
                                  FileType::new("plain text document", "text-x-generic"));
                file_types.insert("md".to_string(),
                                  FileType::new("Markdown", "text-x-generic"));
                file_types.insert(String::new(),
                                  FileType::new("Unknown file", "unknown"));
                file_types
            },
            files: Vec::new(),
            file_sizes: Vec::new(),
            selected: -1,
            last_mouse_event: MouseEvent {
                x: 0,
                y: 0,
                left_button: false,
                middle_button: false,
                right_button: false,
            },
            click_time: Duration::new(0, 0),
        }
    }

    fn load_icon_with(&self, file_name: &str, row: isize, window: &mut Window) {
        if let Some(pos) = file_name.find('.') {
            match self.file_types.get(&file_name[(pos + 1)..]) {
                Some(file_type) => {
                    window.image(0,
                                 32 * row,
                                 file_type.icon.width(),
                                 file_type.icon.height(),
                                 file_type.icon.as_slice());
                }
                None => {
                    window.image(0,
                                 32 * row,
                                 self.file_types[""].icon.width(),
                                 self.file_types[""].icon.height(),
                                 self.file_types[""].icon.as_slice());
                }
            }
        }
        else if file_name.ends_with("/") {
            window.image(0,
                         32 * row as isize,
                         self.file_types["/"].icon.width(),
                         self.file_types["/"].icon.height(),
                         self.file_types["/"].icon.as_slice());
        }
    }

    fn draw_content(&mut self, window: &mut Window) {
        window.set([255, 255, 255, 255]);

        let mut i = 0;
        let mut row = 0;
        let column1 = {
            let mut tmp = 0;
            for string in self.files.iter() {
                if tmp < string.len() {
                    tmp = string.len();
                }
            }
            tmp + 1
        };
        for (file_name, file_size) in self.files.iter().zip(self.file_sizes.iter()) {
            if i == self.selected {
                let width = window.width();
                window.rect(0, 32 * row as isize, width, 32, [224, 224, 224, 255]);
            }

            self.load_icon_with(&file_name, row as isize, window);

            let mut col = 0;
            for c in file_name.chars() {
                if c == '\n' {
                    col = 0;
                    row += 1;
                } else if c == '\t' {
                    col += 8 - col % 8;
                } else {
                    if col < window.width() / 8 && row < window.height() / 32 {
                        window.char(8 * col as isize + 40,
                                    32 * row as isize + 8,
                                    c,
                                    [0, 0, 0, 255]);
                        col += 1;
                    }
                }
                if col >= window.width() / 8 {
                    col = 0;
                    row += 1;
                }
            }

            col = column1;

            for c in file_size.chars() {
                if c == '\n' {
                    col = 0;
                    row += 1;
                } else if c == '\t' {
                    col += 8 - col % 8;
                } else {
                    if col < window.width() / 8 && row < window.height() / 32 {
                        window.char(8 * col as isize + 40,
                                    32 * row as isize + 8,
                                    c,
                                    [0, 0, 0, 255]);
                        col += 1;
                    }
                }
                if col >= window.width() / 8 {
                    col = 0;
                    row += 1;
                }
            }

            row += 1;
            i += 1;
        }

        window.sync();
    }

    fn main(&mut self, path: &str) {
        let mut width = 160;
        let mut height = 0;
        if let Some(mut file) = File::open(path) {
            let mut list = String::new();
            file.read_to_string(&mut list);

            for entry in list.split('\n') {
                self.files.push(entry.to_string());
                self.file_sizes.push(
                    match File::open(&(path.to_string() + entry)) {
                        Some(mut file) => {
                            // When the entry is a folder
                            if entry.ends_with('/') {
                                let mut string = String::new();
                                file.read_to_string(&mut string);

                                let count = string.split('\n').count();
                                if count == 1 {
                                    "1 entry".to_string()
                                } else {
                                    format!("{} entries", count)
                                }
                            } else {
                                match file.seek(SeekFrom::End(0)) {
                                    Some(size) => {
                                        if size >= 1_000_000_000 {
                                            format!("{:.1} GB", (size as f64)/1_000_000_000.0)
                                        } else if size >= 1_000_000 {
                                            format!("{:.1} MB", (size as f64)/1_000_000.0)
                                        } else if size >= 1_000 {
                                            format!("{:.1} KB", (size as f64)/1_000.0)
                                        } else {
                                            format!("{:.1} bytes", size)
                                        }
                                    }
                                    None => "Failed to seek".to_string()
                                }
                            }
                        },
                        None => "Failed to open".to_string(),
                    }
                );
                // Unwrapping the last file size will not panic since it has
                // been at least pushed once in the vector
                let current_width = (40 + (entry.len() + 1) * 8) +
                                    (8 + (self.file_sizes.last().unwrap().len() + 1) * 8);
                if width < current_width {
                    width = current_width;
                }
            }

            if height < self.files.len() * 32 {
                height = self.files.len() * 32;
            }
        }

        let mut window = Window::new((redox::rand() % 400 + 50) as isize,
                                     (redox::rand() % 300 + 50) as isize,
                                     width,
                                     height,
                                     &path).unwrap();

        self.draw_content(&mut window);

        while let Some(event) = window.poll() {
            match event.to_option() {
                EventOption::Key(key_event) => {
                    if key_event.pressed {
                        match key_event.scancode {
                            event::K_ESC => break,
                            event::K_HOME => self.selected = 0,
                            event::K_UP => if self.selected > 0 {
                                self.selected -= 1;
                            },
                            event::K_END => self.selected = self.files.len() as isize - 1,
                            event::K_DOWN => if self.selected < self.files.len() as isize - 1 {
                                self.selected += 1;
                            },
                            _ => match key_event.character {
                                '\0' => (),
                                '\n' => {
                                    if self.selected >= 0 &&
                                       self.selected < self.files.len() as isize {
                                        match self.files.get(self.selected as usize) {
                                            Some(file) => {
                                                File::exec(&(path.to_string() + &file));
                                            },
                                            None => (),
                                        }
                                    }
                                }
                                _ => {
                                    let mut i = 0;
                                    for file in self.files.iter() {
                                        if file.starts_with(key_event.character) {
                                            self.selected = i;
                                            break;
                                        }
                                        i += 1;
                                    }
                                }
                            },
                        }

                        self.draw_content(&mut window);
                    }
                }
                EventOption::Mouse(mouse_event) => {
                    let mut redraw = false;
                    let mut i = 0;
                    let mut row = 0;
                    for file in self.files.iter() {
                        let mut col = 0;
                        for c in file.chars() {
                            if mouse_event.y >= 32 * row as isize &&
                               mouse_event.y < 32 * row as isize + 32 {
                                self.selected = i;
                                redraw = true;
                            }

                            if c == '\n' {
                                col = 0;
                                row += 1;
                            } else if c == '\t' {
                                col += 8 - col % 8;
                            } else {
                                if col < window.width() / 8 && row < window.height() / 32 {
                                    col += 1;
                                }
                            }
                            if col >= window.width() / 8 {
                                col = 0;
                                row += 1;
                            }
                        }
                        row += 1;
                        i += 1;
                    }

                    if redraw {
                        self.draw_content(&mut window);
                    }

                    //Check for double click
                    if mouse_event.left_button {
                        let click_time = Duration::realtime();

                        if click_time - self.click_time < Duration::new(0, 500 * time::NANOS_PER_MILLI)
                            && self.last_mouse_event.x == mouse_event.x
                            && self.last_mouse_event.y == mouse_event.y {
                            if self.selected >= 0 && self.selected < self.files.len() as isize {
                                match self.files.get(self.selected as usize) {
                                    Some(file) => {
                                        File::exec(&(path.to_string() + &file));
                                    },
                                    None => (),
                                }
                            }
                            self.click_time = Duration::new(0, 0);
                        } else {
                            self.click_time = click_time;
                        }
                    }
                    self.last_mouse_event = mouse_event;
                }
                _ => (),
            }
        }
    }
}

pub fn main() {
    match env::args().get(1) {
        Some(arg) => FileManager::new().main(arg),
        None => FileManager::new().main("file:///"),
    }
}
<commit_msg>Add file description column in file manager<commit_after>use collections::BTreeMap;

use redox::{self, env, BMPFile};
use redox::event::{self, EventOption, MouseEvent};
use redox::fs::file::File;
use redox::io::{Read, Seek, SeekFrom};
use redox::orbital::Window;
use redox::time::{self, Duration};
use redox::vec::Vec;
use redox::string::{String, ToString};

pub struct FileType {
    description: String,
    icon: BMPFile,
}

impl FileType {
    pub fn new(desc: &str, icon: &str) -> FileType {
        FileType { description: desc.to_string(), icon: load_icon(icon) }
    }

}

pub struct FileManager {
    file_types: BTreeMap<String, FileType>,
    /*
    folder_icon: BMPFile,
    audio_icon: BMPFile,
    bin_icon: BMPFile,
    image_icon: BMPFile,
    source_icon: BMPFile,
    script_icon: BMPFile,
    text_icon: BMPFile,
    file_icon: BMPFile,
    */
    files: Vec<String>,
    file_sizes: Vec<String>,
    selected: isize,
    last_mouse_event: MouseEvent,
    click_time: Duration,
}

fn load_icon(path: &str) -> BMPFile {
    let mut vec: Vec<u8> = Vec::new();
    if let Some(mut file) = File::open(&("file:///ui/mimetypes/".to_string() + path + ".bmp")) {
        file.read_to_end(&mut vec);
    }
    BMPFile::from_data(&vec)
}

impl FileManager {
    pub fn new() -> Self {
        FileManager {
            file_types: {
                let mut file_types = BTreeMap::<String, FileType>::new();
                file_types.insert("/".to_string(),
                                  FileType::new("Folder", "inode-directory"));
                file_types.insert("wav".to_string(),
                                  FileType::new("WAV audio", "audio-x-wav"));
                file_types.insert("bin".to_string(),
                                  FileType::new("Executable", "application-x-executable"));
                file_types.insert("bmp".to_string(),
                                  FileType::new("Bitmap Image", "image-x-generic"));
                file_types.insert("rs".to_string(),
                                  FileType::new("Rust source code", "image-x-makefile"));
                file_types.insert("crate".to_string(),
                                  FileType::new("Rust crate", "image-x-makefile"));
                file_types.insert("asm".to_string(),
                                  FileType::new("Assembly source", "image-x-makefile"));
                file_types.insert("list".to_string(),
                                  FileType::new("Disassembly source", "image-x-makefile"));
                file_types.insert("c".to_string(),
                                  FileType::new("C source code", "image-x-makefile"));
                file_types.insert("cpp".to_string(),
                                  FileType::new("C++ source code", "image-x-makefile"));
                file_types.insert("sh".to_string(),
                                  FileType::new("Shell script", "image-x-script"));
                file_types.insert("lua".to_string(),
                                  FileType::new("Lua script", "image-x-script"));
                file_types.insert("txt".to_string(),
                                  FileType::new("plain text document", "text-x-generic"));
                file_types.insert("md".to_string(),
                                  FileType::new("Markdown", "text-x-generic"));
                file_types.insert(String::new(),
                                  FileType::new("Unknown file", "unknown"));
                file_types
            },
            files: Vec::new(),
            file_sizes: Vec::new(),
            selected: -1,
            last_mouse_event: MouseEvent {
                x: 0,
                y: 0,
                left_button: false,
                middle_button: false,
                right_button: false,
            },
            click_time: Duration::new(0, 0),
        }
    }

    fn load_icon_with(&self, file_name: &str, row: isize, window: &mut Window) {
        if let Some(pos) = file_name.find('.') {
            match self.file_types.get(&file_name[(pos + 1)..]) {
                Some(file_type) => {
                    window.image(0,
                                 32 * row,
                                 file_type.icon.width(),
                                 file_type.icon.height(),
                                 file_type.icon.as_slice());
                }
                None => {
                    window.image(0,
                                 32 * row,
                                 self.file_types[""].icon.width(),
                                 self.file_types[""].icon.height(),
                                 self.file_types[""].icon.as_slice());
                }
            }
        }
        else if file_name.ends_with("/") {
            window.image(0,
                         32 * row as isize,
                         self.file_types["/"].icon.width(),
                         self.file_types["/"].icon.height(),
                         self.file_types["/"].icon.as_slice());
        }
    }

    fn get_description(&self, file_name: &str) -> String {
        if let Some(pos) = file_name.find('.') {
            match self.file_types.get(&file_name[(pos + 1)..]) {
                Some(file_type) => file_type.description.clone(),
                None => self.file_types[""].description.clone(),
            }
        }
        else { self.file_types["/"].description.clone() }
    }

    fn draw_content(&mut self, window: &mut Window) {
        window.set([255, 255, 255, 255]);

        let mut i = 0;
        let mut row = 0;
        let (column1, column2) = {
            let mut tmp1 = 0;
            for string in self.files.iter() {
                if tmp1 < string.len() {
                    tmp1 = string.len();
                }
            }

            let mut tmp2 = 0;
            for file_size in self.file_sizes.iter() {
                if tmp2 < file_size.len() {
                    tmp2 = file_size.len();
                }
            }
            (tmp1 + 1, tmp1 + 1 + tmp2 + 1)
        };
        for (file_name, file_size) in self.files.iter().zip(self.file_sizes.iter()) {
            if i == self.selected {
                let width = window.width();
                window.rect(0, 32 * row as isize, width, 32, [224, 224, 224, 255]);
            }

            self.load_icon_with(&file_name, row as isize, window);

            let mut col = 0;
            for c in file_name.chars() {
                if c == '\n' {
                    col = 0;
                    row += 1;
                } else if c == '\t' {
                    col += 8 - col % 8;
                } else {
                    if col < window.width() / 8 && row < window.height() / 32 {
                        window.char(8 * col as isize + 40,
                                    32 * row as isize + 8,
                                    c,
                                    [0, 0, 0, 255]);
                        col += 1;
                    }
                }
                if col >= window.width() / 8 {
                    col = 0;
                    row += 1;
                }
            }

            col = column1;

            for c in file_size.chars() {
                if c == '\n' {
                    col = 0;
                    row += 1;
                } else if c == '\t' {
                    col += 8 - col % 8;
                } else {
                    if col < window.width() / 8 && row < window.height() / 32 {
                        window.char(8 * col as isize + 40,
                                    32 * row as isize + 8,
                                    c,
                                    [0, 0, 0, 255]);
                        col += 1;
                    }
                }
                if col >= window.width() / 8 {
                    col = 0;
                    row += 1;
                }
            }

            col = column2;

            for c in self.get_description(file_name).chars() {
                if c == '\n' {
                    col = 0;
                    row += 1;
                } else if c == '\t' {
                    col += 8 - col % 8;
                } else {
                    if col < window.width() / 8 && row < window.height() / 32 {
                        window.char(8 * col as isize + 40,
                                    32 * row as isize + 8,
                                    c,
                                    [0, 0, 0, 255]);
                        col += 1;
                    }
                }
                if col >= window.width() / 8 {
                    col = 0;
                    row += 1;
                }
            }

            row += 1;
            i += 1;
        }

        window.sync();
    }

    fn main(&mut self, path: &str) {
        let mut width = 160;
        let mut height = 0;
        if let Some(mut file) = File::open(path) {
            let mut list = String::new();
            file.read_to_string(&mut list);

            for entry in list.split('\n') {
                self.files.push(entry.to_string());
                self.file_sizes.push(
                    match File::open(&(path.to_string() + entry)) {
                        Some(mut file) => {
                            // When the entry is a folder
                            if entry.ends_with('/') {
                                let mut string = String::new();
                                file.read_to_string(&mut string);

                                let count = string.split('\n').count();
                                if count == 1 {
                                    "1 entry".to_string()
                                } else {
                                    format!("{} entries", count)
                                }
                            } else {
                                match file.seek(SeekFrom::End(0)) {
                                    Some(size) => {
                                        if size >= 1_000_000_000 {
                                            format!("{:.1} GB", (size as f64)/1_000_000_000.0)
                                        } else if size >= 1_000_000 {
                                            format!("{:.1} MB", (size as f64)/1_000_000.0)
                                        } else if size >= 1_000 {
                                            format!("{:.1} KB", (size as f64)/1_000.0)
                                        } else {
                                            format!("{:.1} bytes", size)
                                        }
                                    }
                                    None => "Failed to seek".to_string()
                                }
                            }
                        },
                        None => "Failed to open".to_string(),
                    }
                );
                // Unwrapping the last file size will not panic since it has
                // been at least pushed once in the vector
                let tmp_width = (40 + (entry.len() + 1) * 8) +
                                (8 + (self.file_sizes.last().unwrap().len() + 1) * 8) +
                                (8 + (self.get_description(entry).len() + 1) * 8);
                if width < tmp_width {
                    width = tmp_width;
                }
            }

            if height < self.files.len() * 32 {
                height = self.files.len() * 32;
            }
        }

        let mut window = Window::new((redox::rand() % 400 + 50) as isize,
                                     (redox::rand() % 300 + 50) as isize,
                                     width,
                                     height,
                                     &path).unwrap();

        self.draw_content(&mut window);

        while let Some(event) = window.poll() {
            match event.to_option() {
                EventOption::Key(key_event) => {
                    if key_event.pressed {
                        match key_event.scancode {
                            event::K_ESC => break,
                            event::K_HOME => self.selected = 0,
                            event::K_UP => if self.selected > 0 {
                                self.selected -= 1;
                            },
                            event::K_END => self.selected = self.files.len() as isize - 1,
                            event::K_DOWN => if self.selected < self.files.len() as isize - 1 {
                                self.selected += 1;
                            },
                            _ => match key_event.character {
                                '\0' => (),
                                '\n' => {
                                    if self.selected >= 0 &&
                                       self.selected < self.files.len() as isize {
                                        match self.files.get(self.selected as usize) {
                                            Some(file) => {
                                                File::exec(&(path.to_string() + &file));
                                            },
                                            None => (),
                                        }
                                    }
                                }
                                _ => {
                                    let mut i = 0;
                                    for file in self.files.iter() {
                                        if file.starts_with(key_event.character) {
                                            self.selected = i;
                                            break;
                                        }
                                        i += 1;
                                    }
                                }
                            },
                        }

                        self.draw_content(&mut window);
                    }
                }
                EventOption::Mouse(mouse_event) => {
                    let mut redraw = false;
                    let mut i = 0;
                    let mut row = 0;
                    for file in self.files.iter() {
                        let mut col = 0;
                        for c in file.chars() {
                            if mouse_event.y >= 32 * row as isize &&
                               mouse_event.y < 32 * row as isize + 32 {
                                self.selected = i;
                                redraw = true;
                            }

                            if c == '\n' {
                                col = 0;
                                row += 1;
                            } else if c == '\t' {
                                col += 8 - col % 8;
                            } else {
                                if col < window.width() / 8 && row < window.height() / 32 {
                                    col += 1;
                                }
                            }
                            if col >= window.width() / 8 {
                                col = 0;
                                row += 1;
                            }
                        }
                        row += 1;
                        i += 1;
                    }

                    if redraw {
                        self.draw_content(&mut window);
                    }

                    //Check for double click
                    if mouse_event.left_button {
                        let click_time = Duration::realtime();

                        if click_time - self.click_time < Duration::new(0, 500 * time::NANOS_PER_MILLI)
                            && self.last_mouse_event.x == mouse_event.x
                            && self.last_mouse_event.y == mouse_event.y {
                            if self.selected >= 0 && self.selected < self.files.len() as isize {
                                match self.files.get(self.selected as usize) {
                                    Some(file) => {
                                        File::exec(&(path.to_string() + &file));
                                    },
                                    None => (),
                                }
                            }
                            self.click_time = Duration::new(0, 0);
                        } else {
                            self.click_time = click_time;
                        }
                    }
                    self.last_mouse_event = mouse_event;
                }
                _ => (),
            }
        }
    }
}

pub fn main() {
    match env::args().get(1) {
        Some(arg) => FileManager::new().main(arg),
        None => FileManager::new().main("file:///"),
    }
}
<|endoftext|>


<commit_before>//! Connection pooling via r2d2.
//!
//! Note: This module requires enabling the `r2d2` feature

extern crate r2d2;

pub use self::r2d2::*;

/// A re-export of [`r2d2::Error`], which is only used by methods on [`r2d2::Pool`].
///
/// [`r2d2::Error`]: r2d2::Error
/// [`r2d2::Pool`]: r2d2::Pool
pub type PoolError = self::r2d2::Error;

use std::convert::Into;
use std::fmt;
use std::marker::PhantomData;

use crate::connection::{SimpleConnection, TransactionManager};
use crate::deserialize::FromSqlRow;
use crate::expression::QueryMetadata;
use crate::prelude::*;
use crate::query_builder::{AsQuery, QueryFragment, QueryId};
use crate::query_dsl::load_dsl::CompatibleType;

/// An r2d2 connection manager for use with Diesel.
///
/// See the [r2d2 documentation] for usage examples.
///
/// [r2d2 documentation]: r2d2
#[derive(Clone)]
pub struct ConnectionManager<T> {
    database_url: String,
    _marker: PhantomData<T>,
}

impl<T> fmt::Debug for ConnectionManager<T> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "ConnectionManager<{}>", std::any::type_name::<T>())
    }
}

unsafe impl<T: Send + 'static> Sync for ConnectionManager<T> {}

impl<T> ConnectionManager<T> {
    /// Returns a new connection manager,
    /// which establishes connections to the given database URL.
    pub fn new<S: Into<String>>(database_url: S) -> Self {
        ConnectionManager {
            database_url: database_url.into(),
            _marker: PhantomData,
        }
    }
}

/// The error used when managing connections with `r2d2`.
#[derive(Debug)]
pub enum Error {
    /// An error occurred establishing the connection
    ConnectionError(ConnectionError),

    /// An error occurred pinging the database
    QueryError(crate::result::Error),
}

impl fmt::Display for Error {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match *self {
            Error::ConnectionError(ref e) => e.fmt(f),
            Error::QueryError(ref e) => e.fmt(f),
        }
    }
}

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

/// A trait indicating a connection could be used inside a r2d2 pool
pub trait R2D2Connection: Connection {
    /// Check if a connection is still valid
    fn ping(&mut self) -> QueryResult<()>;
}

#[cfg(feature = "postgres")]
impl R2D2Connection for crate::pg::PgConnection {
    fn ping(&mut self) -> QueryResult<()> {
        self.execute("SELECT 1").map(|_| ())
    }
}

#[cfg(feature = "mysql")]
impl R2D2Connection for crate::mysql::MysqlConnection {
    fn ping(&mut self) -> QueryResult<()> {
        self.execute("SELECT 1").map(|_| ())
    }
}

#[cfg(feature = "sqlite")]
impl R2D2Connection for crate::sqlite::SqliteConnection {
    fn ping(&mut self) -> QueryResult<()> {
        self.execute("SELECT 1").map(|_| ())
    }
}

impl<T> ManageConnection for ConnectionManager<T>
where
    T: R2D2Connection + Send + 'static,
{
    type Connection = T;
    type Error = Error;

    fn connect(&self) -> Result<T, Error> {
        T::establish(&self.database_url).map_err(Error::ConnectionError)
    }

    fn is_valid(&self, conn: &mut T) -> Result<(), Error> {
        conn.ping().map_err(Error::QueryError)
    }

    fn has_broken(&self, _conn: &mut T) -> bool {
        false
    }
}

impl<M> SimpleConnection for PooledConnection<M>
where
    M: ManageConnection,
    M::Connection: R2D2Connection + Send + 'static,
{
    fn batch_execute(&mut self, query: &str) -> QueryResult<()> {
        (&mut **self).batch_execute(query)
    }
}

impl<M> Connection for PooledConnection<M>
where
    M: ManageConnection,
    M::Connection: Connection + R2D2Connection + Send + 'static,
{
    type Backend = <M::Connection as Connection>::Backend;
    type TransactionManager =
        PoolTransactionManager<<M::Connection as Connection>::TransactionManager>;

    fn establish(_: &str) -> ConnectionResult<Self> {
        Err(ConnectionError::BadConnection(String::from(
            "Cannot directly establish a pooled connection",
        )))
    }

    fn execute(&mut self, query: &str) -> QueryResult<usize> {
        (&mut **self).execute(query)
    }

    fn load<T, U, ST>(&mut self, source: T) -> QueryResult<Vec<U>>
    where
        T: AsQuery,
        T::Query: QueryFragment<Self::Backend> + QueryId,
        T::SqlType: CompatibleType<U, Self::Backend, SqlType = ST>,
        U: FromSqlRow<ST, Self::Backend>,
        Self::Backend: QueryMetadata<T::SqlType>,
    {
        (&mut **self).load(source)
    }

    fn execute_returning_count<T>(&mut self, source: &T) -> QueryResult<usize>
    where
        T: QueryFragment<Self::Backend> + QueryId,
    {
        (&mut **self).execute_returning_count(source)
    }

    fn transaction_state(
        &mut self,
    ) -> &mut <Self::TransactionManager as TransactionManager<Self>>::TransactionStateData {
        (&mut **self).transaction_state()
    }

    fn begin_test_transaction(&mut self) -> QueryResult<()> {
        (&mut **self).begin_test_transaction()
    }
}

#[doc(hidden)]
#[allow(missing_debug_implementations)]
pub struct PoolTransactionManager<T>(std::marker::PhantomData<T>);

impl<M, T> TransactionManager<PooledConnection<M>> for PoolTransactionManager<T>
where
    M: ManageConnection,
    M::Connection: Connection<TransactionManager = T> + R2D2Connection,
    T: TransactionManager<M::Connection>,
{
    type TransactionStateData = T::TransactionStateData;

    fn begin_transaction(conn: &mut PooledConnection<M>) -> QueryResult<()> {
        T::begin_transaction(&mut **conn)
    }

    fn rollback_transaction(conn: &mut PooledConnection<M>) -> QueryResult<()> {
        T::begin_transaction(&mut **conn)
    }

    fn commit_transaction(conn: &mut PooledConnection<M>) -> QueryResult<()> {
        T::commit_transaction(&mut **conn)
    }

    fn get_transaction_depth(conn: &mut PooledConnection<M>) -> u32 {
        T::get_transaction_depth(&mut **conn)
    }
}

impl<M> crate::migration::MigrationConnection for PooledConnection<M>
where
    M: ManageConnection,
    M::Connection: crate::migration::MigrationConnection,
    Self: Connection,
{
    fn setup(&mut self) -> QueryResult<usize> {
        (&mut **self).setup()
    }
}

impl<Changes, Output, M> crate::query_dsl::UpdateAndFetchResults<Changes, Output>
    for PooledConnection<M>
where
    M: ManageConnection,
    M::Connection: crate::query_dsl::UpdateAndFetchResults<Changes, Output>,
    Self: Connection,
{
    fn update_and_fetch(&mut self, changeset: Changes) -> QueryResult<Output> {
        (&mut **self).update_and_fetch(changeset)
    }
}

#[cfg(test)]
mod tests {
    use std::sync::mpsc;
    use std::sync::Arc;
    use std::thread;

    use crate::r2d2::*;
    use crate::test_helpers::*;

    #[test]
    fn establish_basic_connection() {
        let manager = ConnectionManager::<TestConnection>::new(database_url());
        let pool = Arc::new(Pool::builder().max_size(2).build(manager).unwrap());

        let (s1, r1) = mpsc::channel();
        let (s2, r2) = mpsc::channel();

        let pool1 = Arc::clone(&pool);
        let t1 = thread::spawn(move || {
            let conn = pool1.get().unwrap();
            s1.send(()).unwrap();
            r2.recv().unwrap();
            drop(conn);
        });

        let pool2 = Arc::clone(&pool);
        let t2 = thread::spawn(move || {
            let conn = pool2.get().unwrap();
            s2.send(()).unwrap();
            r1.recv().unwrap();
            drop(conn);
        });

        t1.join().unwrap();
        t2.join().unwrap();

        pool.get().unwrap();
    }

    #[test]
    fn is_valid() {
        let manager = ConnectionManager::<TestConnection>::new(database_url());
        let pool = Pool::builder()
            .max_size(1)
            .test_on_check_out(true)
            .build(manager)
            .unwrap();

        pool.get().unwrap();
    }

    #[test]
    fn pooled_connection_impls_connection() {
        use crate::select;
        use crate::sql_types::Text;

        let manager = ConnectionManager::<TestConnection>::new(database_url());
        let pool = Pool::builder()
            .max_size(1)
            .test_on_check_out(true)
            .build(manager)
            .unwrap();
        let mut conn = pool.get().unwrap();

        let query = select("foo".into_sql::<Text>());
        assert_eq!("foo", query.get_result::<String>(&mut conn).unwrap());
    }
}
<commit_msg>Call rollback_transaction instead of begin_transaction when rolling back<commit_after>//! Connection pooling via r2d2.
//!
//! Note: This module requires enabling the `r2d2` feature

extern crate r2d2;

pub use self::r2d2::*;

/// A re-export of [`r2d2::Error`], which is only used by methods on [`r2d2::Pool`].
///
/// [`r2d2::Error`]: r2d2::Error
/// [`r2d2::Pool`]: r2d2::Pool
pub type PoolError = self::r2d2::Error;

use std::convert::Into;
use std::fmt;
use std::marker::PhantomData;

use crate::connection::{SimpleConnection, TransactionManager};
use crate::deserialize::FromSqlRow;
use crate::expression::QueryMetadata;
use crate::prelude::*;
use crate::query_builder::{AsQuery, QueryFragment, QueryId};
use crate::query_dsl::load_dsl::CompatibleType;

/// An r2d2 connection manager for use with Diesel.
///
/// See the [r2d2 documentation] for usage examples.
///
/// [r2d2 documentation]: r2d2
#[derive(Clone)]
pub struct ConnectionManager<T> {
    database_url: String,
    _marker: PhantomData<T>,
}

impl<T> fmt::Debug for ConnectionManager<T> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "ConnectionManager<{}>", std::any::type_name::<T>())
    }
}

unsafe impl<T: Send + 'static> Sync for ConnectionManager<T> {}

impl<T> ConnectionManager<T> {
    /// Returns a new connection manager,
    /// which establishes connections to the given database URL.
    pub fn new<S: Into<String>>(database_url: S) -> Self {
        ConnectionManager {
            database_url: database_url.into(),
            _marker: PhantomData,
        }
    }
}

/// The error used when managing connections with `r2d2`.
#[derive(Debug)]
pub enum Error {
    /// An error occurred establishing the connection
    ConnectionError(ConnectionError),

    /// An error occurred pinging the database
    QueryError(crate::result::Error),
}

impl fmt::Display for Error {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match *self {
            Error::ConnectionError(ref e) => e.fmt(f),
            Error::QueryError(ref e) => e.fmt(f),
        }
    }
}

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

/// A trait indicating a connection could be used inside a r2d2 pool
pub trait R2D2Connection: Connection {
    /// Check if a connection is still valid
    fn ping(&mut self) -> QueryResult<()>;
}

#[cfg(feature = "postgres")]
impl R2D2Connection for crate::pg::PgConnection {
    fn ping(&mut self) -> QueryResult<()> {
        self.execute("SELECT 1").map(|_| ())
    }
}

#[cfg(feature = "mysql")]
impl R2D2Connection for crate::mysql::MysqlConnection {
    fn ping(&mut self) -> QueryResult<()> {
        self.execute("SELECT 1").map(|_| ())
    }
}

#[cfg(feature = "sqlite")]
impl R2D2Connection for crate::sqlite::SqliteConnection {
    fn ping(&mut self) -> QueryResult<()> {
        self.execute("SELECT 1").map(|_| ())
    }
}

impl<T> ManageConnection for ConnectionManager<T>
where
    T: R2D2Connection + Send + 'static,
{
    type Connection = T;
    type Error = Error;

    fn connect(&self) -> Result<T, Error> {
        T::establish(&self.database_url).map_err(Error::ConnectionError)
    }

    fn is_valid(&self, conn: &mut T) -> Result<(), Error> {
        conn.ping().map_err(Error::QueryError)
    }

    fn has_broken(&self, _conn: &mut T) -> bool {
        false
    }
}

impl<M> SimpleConnection for PooledConnection<M>
where
    M: ManageConnection,
    M::Connection: R2D2Connection + Send + 'static,
{
    fn batch_execute(&mut self, query: &str) -> QueryResult<()> {
        (&mut **self).batch_execute(query)
    }
}

impl<M> Connection for PooledConnection<M>
where
    M: ManageConnection,
    M::Connection: Connection + R2D2Connection + Send + 'static,
{
    type Backend = <M::Connection as Connection>::Backend;
    type TransactionManager =
        PoolTransactionManager<<M::Connection as Connection>::TransactionManager>;

    fn establish(_: &str) -> ConnectionResult<Self> {
        Err(ConnectionError::BadConnection(String::from(
            "Cannot directly establish a pooled connection",
        )))
    }

    fn execute(&mut self, query: &str) -> QueryResult<usize> {
        (&mut **self).execute(query)
    }

    fn load<T, U, ST>(&mut self, source: T) -> QueryResult<Vec<U>>
    where
        T: AsQuery,
        T::Query: QueryFragment<Self::Backend> + QueryId,
        T::SqlType: CompatibleType<U, Self::Backend, SqlType = ST>,
        U: FromSqlRow<ST, Self::Backend>,
        Self::Backend: QueryMetadata<T::SqlType>,
    {
        (&mut **self).load(source)
    }

    fn execute_returning_count<T>(&mut self, source: &T) -> QueryResult<usize>
    where
        T: QueryFragment<Self::Backend> + QueryId,
    {
        (&mut **self).execute_returning_count(source)
    }

    fn transaction_state(
        &mut self,
    ) -> &mut <Self::TransactionManager as TransactionManager<Self>>::TransactionStateData {
        (&mut **self).transaction_state()
    }

    fn begin_test_transaction(&mut self) -> QueryResult<()> {
        (&mut **self).begin_test_transaction()
    }
}

#[doc(hidden)]
#[allow(missing_debug_implementations)]
pub struct PoolTransactionManager<T>(std::marker::PhantomData<T>);

impl<M, T> TransactionManager<PooledConnection<M>> for PoolTransactionManager<T>
where
    M: ManageConnection,
    M::Connection: Connection<TransactionManager = T> + R2D2Connection,
    T: TransactionManager<M::Connection>,
{
    type TransactionStateData = T::TransactionStateData;

    fn begin_transaction(conn: &mut PooledConnection<M>) -> QueryResult<()> {
        T::begin_transaction(&mut **conn)
    }

    fn rollback_transaction(conn: &mut PooledConnection<M>) -> QueryResult<()> {
        T::rollback_transaction(&mut **conn)
    }

    fn commit_transaction(conn: &mut PooledConnection<M>) -> QueryResult<()> {
        T::commit_transaction(&mut **conn)
    }

    fn get_transaction_depth(conn: &mut PooledConnection<M>) -> u32 {
        T::get_transaction_depth(&mut **conn)
    }
}

impl<M> crate::migration::MigrationConnection for PooledConnection<M>
where
    M: ManageConnection,
    M::Connection: crate::migration::MigrationConnection,
    Self: Connection,
{
    fn setup(&mut self) -> QueryResult<usize> {
        (&mut **self).setup()
    }
}

impl<Changes, Output, M> crate::query_dsl::UpdateAndFetchResults<Changes, Output>
    for PooledConnection<M>
where
    M: ManageConnection,
    M::Connection: crate::query_dsl::UpdateAndFetchResults<Changes, Output>,
    Self: Connection,
{
    fn update_and_fetch(&mut self, changeset: Changes) -> QueryResult<Output> {
        (&mut **self).update_and_fetch(changeset)
    }
}

#[cfg(test)]
mod tests {
    use std::sync::mpsc;
    use std::sync::Arc;
    use std::thread;

    use crate::r2d2::*;
    use crate::test_helpers::*;

    #[test]
    fn establish_basic_connection() {
        let manager = ConnectionManager::<TestConnection>::new(database_url());
        let pool = Arc::new(Pool::builder().max_size(2).build(manager).unwrap());

        let (s1, r1) = mpsc::channel();
        let (s2, r2) = mpsc::channel();

        let pool1 = Arc::clone(&pool);
        let t1 = thread::spawn(move || {
            let conn = pool1.get().unwrap();
            s1.send(()).unwrap();
            r2.recv().unwrap();
            drop(conn);
        });

        let pool2 = Arc::clone(&pool);
        let t2 = thread::spawn(move || {
            let conn = pool2.get().unwrap();
            s2.send(()).unwrap();
            r1.recv().unwrap();
            drop(conn);
        });

        t1.join().unwrap();
        t2.join().unwrap();

        pool.get().unwrap();
    }

    #[test]
    fn is_valid() {
        let manager = ConnectionManager::<TestConnection>::new(database_url());
        let pool = Pool::builder()
            .max_size(1)
            .test_on_check_out(true)
            .build(manager)
            .unwrap();

        pool.get().unwrap();
    }

    #[test]
    fn pooled_connection_impls_connection() {
        use crate::select;
        use crate::sql_types::Text;

        let manager = ConnectionManager::<TestConnection>::new(database_url());
        let pool = Pool::builder()
            .max_size(1)
            .test_on_check_out(true)
            .build(manager)
            .unwrap();
        let mut conn = pool.get().unwrap();

        let query = select("foo".into_sql::<Text>());
        assert_eq!("foo", query.get_result::<String>(&mut conn).unwrap());
    }
}
<|endoftext|>


<commit_before>use orbclient::Color;

pub const BAR_COLOR: Color = Color::rgba(47, 52, 63, 224);
pub const BAR_HIGHLIGHT_COLOR: Color = Color::rgb(80, 86, 102);
pub const TEXT_COLOR: Color = Color::rgb(204, 210, 224);
pub const TEXT_HIGHLIGHT_COLOR: Color = Color::rgb(235, 241, 255);
<commit_msg>Adjust theme to match orbital<commit_after>use orbclient::Color;

pub const BAR_COLOR: Color = Color::rgba(47, 52, 63, 224);
pub const BAR_HIGHLIGHT_COLOR: Color = Color::rgba(80, 86, 102, 224);
pub const TEXT_COLOR: Color = Color::rgb(204, 210, 224);
pub const TEXT_HIGHLIGHT_COLOR: Color = Color::rgb(235, 241, 255);
<|endoftext|>