<commit_before>//
// Tests for the pwx binary
//

use std::process::Command;

// TODO: Find a better way to get pwx binary path
const PWXBIN: &'static str = "target/debug/pwx";

macro_rules! pwxrun {
    ($($arg:expr),*) => {
        Command::new(PWXBIN)
            .env("PWX_PASSWORD", "test")
            .env("PWX_DATABASE", "tests/test.psafe3")
            $(.arg($arg))*
            .output().unwrap_or_else(|e| {panic!("Failed to execute pwx({}) {}", PWXBIN, e)})
    };
}

/* Show all entries in the DB */
#[test]
fn list() {
    let output = pwxrun!("list");
    let sout = String::from_utf8_lossy(&output.stdout);
    println!("{}", sout);
    assert_eq!(sout.trim().split('\n').count(), 2);
}

/* Filter by title */
#[test]
fn list_title() {
    let output = pwxrun!("list", "-T", "face");
    let sout = String::from_utf8_lossy(&output.stdout);
    println!("{}", sout);
    assert_eq!(sout.trim().split('\n').count(), 1);
}

/* Filter by User */
#[test]
fn list_user() {
    let output = pwxrun!("list", "-U", "testuser");
    let sout = String::from_utf8_lossy(&output.stdout);
    println!("{}", sout);
    assert_eq!(sout.trim().split('\n').count(), 1);
}

/* Filter by Group */
#[test]
fn list_group() {
    let output = pwxrun!("list", "-G", "social");
    let sout = String::from_utf8_lossy(&output.stdout);
    println!("{}", sout);
    assert_eq!(sout.trim().split('\n').count(), 1);
    assert!(sout.contains("facebook"));
}

/* Filter by URL */
#[test]
fn list_url() {
    let output = pwxrun!("list", "-R", "facebook.com");
    let sout = String::from_utf8_lossy(&output.stdout);
    println!("{}", sout);
    assert_eq!(sout.trim().split('\n').count(), 1);
}

/* list <filter> searches all text fields */
#[test]
fn list_filter() {
    let output = pwxrun!("list", "some");
    let sout = String::from_utf8_lossy(&output.stdout);
    println!("{}", sout);
    assert_eq!(sout.trim().split('\n').count(), 2);

    // Filters
    let output = pwxrun!("list", "-U", "some", "some");
    let sout = String::from_utf8_lossy(&output.stdout);
    println!("{}", sout);
    assert_eq!(sout.trim().split('\n').count(), 1);
}

#[test]
fn get() {
    // URL
    let output = pwxrun!("get", "43fe1d0e-b65f-4e48-9abf-a1c5a1beeee8", "url");
    assert!(output.status.success());
    let sout = String::from_utf8_lossy(&output.stdout);
    println!("{}", sout);
    assert_eq!(sout.trim(), "https://facebook.com");

    // Group
    let output = pwxrun!("get", "43fe1d0e-b65f-4e48-9abf-a1c5a1beeee8", "group");
    assert!(output.status.success());
    let sout = String::from_utf8_lossy(&output.stdout);
    println!("{}", sout);
    assert_eq!(sout.trim(), "social");

    // Group
    let output = pwxrun!("get", "43fe1d0e-b65f-4e48-9abf-a1c5a1beeee8", "notes");
    assert!(output.status.success());
    let sout = String::from_utf8_lossy(&output.stdout);
    println!("{}", sout);
    assert_eq!(sout.trim(), "Some notes");

    // Command is not set so the command fails
    let output = pwxrun!("get", "43fe1d0e-b65f-4e48-9abf-a1c5a1beeee8", "command");
    assert!(!output.status.success());
}

<commit_msg>Fix path detection for pwx functional tests<commit_after>//
// Tests for the pwx binary
//

use std::process::Command;
use std::path::PathBuf;
use std::env::current_exe;

macro_rules! pwxrun {
    ($($arg:expr),*) => {{
        let binpath = current_exe().unwrap()
                        .parent().expect("executable path")
                        .to_path_buf().join("pwx");
        Command::new(&binpath)
            .env("PWX_PASSWORD", "test")
            .env("PWX_DATABASE", "tests/test.psafe3")
            $(.arg($arg))*
            .output().unwrap_or_else(|e| {panic!("Failed to execute pwx({}) {}", binpath.to_string_lossy(), e)})
    }};
}

/* Show all entries in the DB */
#[test]
fn list() {
    let output = pwxrun!("list");
    let sout = String::from_utf8_lossy(&output.stdout);
    println!("{}", sout);
    assert_eq!(sout.trim().split('\n').count(), 2);
}

/* Filter by title */
#[test]
fn list_title() {
    let output = pwxrun!("list", "-T", "face");
    let sout = String::from_utf8_lossy(&output.stdout);
    println!("{}", sout);
    assert_eq!(sout.trim().split('\n').count(), 1);
}

/* Filter by User */
#[test]
fn list_user() {
    let output = pwxrun!("list", "-U", "testuser");
    let sout = String::from_utf8_lossy(&output.stdout);
    println!("{}", sout);
    assert_eq!(sout.trim().split('\n').count(), 1);
}

/* Filter by Group */
#[test]
fn list_group() {
    let output = pwxrun!("list", "-G", "social");
    let sout = String::from_utf8_lossy(&output.stdout);
    println!("{}", sout);
    assert_eq!(sout.trim().split('\n').count(), 1);
    assert!(sout.contains("facebook"));
}

/* Filter by URL */
#[test]
fn list_url() {
    let output = pwxrun!("list", "-R", "facebook.com");
    let sout = String::from_utf8_lossy(&output.stdout);
    println!("{}", sout);
    assert_eq!(sout.trim().split('\n').count(), 1);
}

/* list <filter> searches all text fields */
#[test]
fn list_filter() {
    let output = pwxrun!("list", "some");
    let sout = String::from_utf8_lossy(&output.stdout);
    println!("{}", sout);
    assert_eq!(sout.trim().split('\n').count(), 2);

    // Filters
    let output = pwxrun!("list", "-U", "some", "some");
    let sout = String::from_utf8_lossy(&output.stdout);
    println!("{}", sout);
    assert_eq!(sout.trim().split('\n').count(), 1);
}

#[test]
fn get() {
    // URL
    let output = pwxrun!("get", "43fe1d0e-b65f-4e48-9abf-a1c5a1beeee8", "url");
    assert!(output.status.success());
    let sout = String::from_utf8_lossy(&output.stdout);
    println!("{}", sout);
    assert_eq!(sout.trim(), "https://facebook.com");

    // Group
    let output = pwxrun!("get", "43fe1d0e-b65f-4e48-9abf-a1c5a1beeee8", "group");
    assert!(output.status.success());
    let sout = String::from_utf8_lossy(&output.stdout);
    println!("{}", sout);
    assert_eq!(sout.trim(), "social");

    // Group
    let output = pwxrun!("get", "43fe1d0e-b65f-4e48-9abf-a1c5a1beeee8", "notes");
    assert!(output.status.success());
    let sout = String::from_utf8_lossy(&output.stdout);
    println!("{}", sout);
    assert_eq!(sout.trim(), "Some notes");

    // Command is not set so the command fails
    let output = pwxrun!("get", "43fe1d0e-b65f-4e48-9abf-a1c5a1beeee8", "command");
    assert!(!output.status.success());
}

<|endoftext|>


<commit_before>use bytes::{Buf, Bytes};
use h2::{Reason, RecvStream, SendStream};
use http::header::{HeaderName, CONNECTION, TE, TRAILER, TRANSFER_ENCODING, UPGRADE};
use http::HeaderMap;
use pin_project_lite::pin_project;
use std::error::Error as StdError;
use std::io::{self, Cursor, IoSlice};
use std::mem;
use std::task::Context;
use tokio::io::{AsyncRead, AsyncWrite, ReadBuf};

use crate::body::HttpBody;
use crate::common::{task, Future, Pin, Poll};
use crate::proto::h2::ping::Recorder;

pub(crate) mod ping;

cfg_client! {
    pub(crate) mod client;
    pub(crate) use self::client::ClientTask;
}

cfg_server! {
    pub(crate) mod server;
    pub(crate) use self::server::Server;
}

/// Default initial stream window size defined in HTTP2 spec.
pub(crate) const SPEC_WINDOW_SIZE: u32 = 65_535;

fn strip_connection_headers(headers: &mut HeaderMap, is_request: bool) {
    // List of connection headers from:
    // https://developer.mozilla.org/en-US/docs/Web/HTTP/Headers/Connection
    //
    // TE headers are allowed in HTTP/2 requests as long as the value is "trailers", so they're
    // tested separately.
    let connection_headers = [
        HeaderName::from_lowercase(b"keep-alive").unwrap(),
        HeaderName::from_lowercase(b"proxy-connection").unwrap(),
        TRAILER,
        TRANSFER_ENCODING,
        UPGRADE,
    ];

    for header in connection_headers.iter() {
        if headers.remove(header).is_some() {
            warn!("Connection header illegal in HTTP/2: {}", header.as_str());
        }
    }

    if is_request {
        if headers
            .get(TE)
            .map(|te_header| te_header != "trailers")
            .unwrap_or(false)
        {
            warn!("TE headers not set to \"trailers\" are illegal in HTTP/2 requests");
            headers.remove(TE);
        }
    } else if headers.remove(TE).is_some() {
        warn!("TE headers illegal in HTTP/2 responses");
    }

    if let Some(header) = headers.remove(CONNECTION) {
        warn!(
            "Connection header illegal in HTTP/2: {}",
            CONNECTION.as_str()
        );
        let header_contents = header.to_str().unwrap();

        // A `Connection` header may have a comma-separated list of names of other headers that
        // are meant for only this specific connection.
        //
        // Iterate these names and remove them as headers. Connection-specific headers are
        // forbidden in HTTP2, as that information has been moved into frame types of the h2
        // protocol.
        for name in header_contents.split(',') {
            let name = name.trim();
            headers.remove(name);
        }
    }
}

// body adapters used by both Client and Server

pin_project! {
    struct PipeToSendStream<S>
    where
        S: HttpBody,
    {
        body_tx: SendStream<SendBuf<S::Data>>,
        data_done: bool,
        #[pin]
        stream: S,
    }
}

impl<S> PipeToSendStream<S>
where
    S: HttpBody,
{
    fn new(stream: S, tx: SendStream<SendBuf<S::Data>>) -> PipeToSendStream<S> {
        PipeToSendStream {
            body_tx: tx,
            data_done: false,
            stream,
        }
    }
}

impl<S> Future for PipeToSendStream<S>
where
    S: HttpBody,
    S::Error: Into<Box<dyn StdError + Send + Sync>>,
{
    type Output = crate::Result<()>;

    fn poll(self: Pin<&mut Self>, cx: &mut task::Context<'_>) -> Poll<Self::Output> {
        let mut me = self.project();
        loop {
            if !*me.data_done {
                // we don't have the next chunk of data yet, so just reserve 1 byte to make
                // sure there's some capacity available. h2 will handle the capacity management
                // for the actual body chunk.
                me.body_tx.reserve_capacity(1);

                if me.body_tx.capacity() == 0 {
                    loop {
                        match ready!(me.body_tx.poll_capacity(cx)) {
                            Some(Ok(0)) => {}
                            Some(Ok(_)) => break,
                            Some(Err(e)) => {
                                return Poll::Ready(Err(crate::Error::new_body_write(e)))
                            }
                            None => {
                                // None means the stream is no longer in a
                                // streaming state, we either finished it
                                // somehow, or the remote reset us.
                                return Poll::Ready(Err(crate::Error::new_body_write(
                                    "send stream capacity unexpectedly closed",
                                )));
                            }
                        }
                    }
                } else if let Poll::Ready(reason) = me
                    .body_tx
                    .poll_reset(cx)
                    .map_err(crate::Error::new_body_write)?
                {
                    debug!("stream received RST_STREAM: {:?}", reason);
                    return Poll::Ready(Err(crate::Error::new_body_write(::h2::Error::from(
                        reason,
                    ))));
                }

                match ready!(me.stream.as_mut().poll_data(cx)) {
                    Some(Ok(chunk)) => {
                        let is_eos = me.stream.is_end_stream();
                        trace!(
                            "send body chunk: {} bytes, eos={}",
                            chunk.remaining(),
                            is_eos,
                        );

                        let buf = SendBuf::Buf(chunk);
                        me.body_tx
                            .send_data(buf, is_eos)
                            .map_err(crate::Error::new_body_write)?;

                        if is_eos {
                            return Poll::Ready(Ok(()));
                        }
                    }
                    Some(Err(e)) => return Poll::Ready(Err(me.body_tx.on_user_err(e))),
                    None => {
                        me.body_tx.reserve_capacity(0);
                        let is_eos = me.stream.is_end_stream();
                        if is_eos {
                            return Poll::Ready(me.body_tx.send_eos_frame());
                        } else {
                            *me.data_done = true;
                            // loop again to poll_trailers
                        }
                    }
                }
            } else {
                if let Poll::Ready(reason) = me
                    .body_tx
                    .poll_reset(cx)
                    .map_err(crate::Error::new_body_write)?
                {
                    debug!("stream received RST_STREAM: {:?}", reason);
                    return Poll::Ready(Err(crate::Error::new_body_write(::h2::Error::from(
                        reason,
                    ))));
                }

                match ready!(me.stream.poll_trailers(cx)) {
                    Ok(Some(trailers)) => {
                        me.body_tx
                            .send_trailers(trailers)
                            .map_err(crate::Error::new_body_write)?;
                        return Poll::Ready(Ok(()));
                    }
                    Ok(None) => {
                        // There were no trailers, so send an empty DATA frame...
                        return Poll::Ready(me.body_tx.send_eos_frame());
                    }
                    Err(e) => return Poll::Ready(Err(me.body_tx.on_user_err(e))),
                }
            }
        }
    }
}

trait SendStreamExt {
    fn on_user_err<E>(&mut self, err: E) -> crate::Error
    where
        E: Into<Box<dyn std::error::Error + Send + Sync>>;
    fn send_eos_frame(&mut self) -> crate::Result<()>;
}

impl<B: Buf> SendStreamExt for SendStream<SendBuf<B>> {
    fn on_user_err<E>(&mut self, err: E) -> crate::Error
    where
        E: Into<Box<dyn std::error::Error + Send + Sync>>,
    {
        let err = crate::Error::new_user_body(err);
        debug!("send body user stream error: {}", err);
        self.send_reset(err.h2_reason());
        err
    }

    fn send_eos_frame(&mut self) -> crate::Result<()> {
        trace!("send body eos");
        self.send_data(SendBuf::None, true)
            .map_err(crate::Error::new_body_write)
    }
}

#[repr(usize)]
enum SendBuf<B> {
    Buf(B),
    Cursor(Cursor<Box<[u8]>>),
    None,
}

impl<B: Buf> Buf for SendBuf<B> {
    #[inline]
    fn remaining(&self) -> usize {
        match *self {
            Self::Buf(ref b) => b.remaining(),
            Self::Cursor(ref c) => Buf::remaining(c),
            Self::None => 0,
        }
    }

    #[inline]
    fn chunk(&self) -> &[u8] {
        match *self {
            Self::Buf(ref b) => b.chunk(),
            Self::Cursor(ref c) => c.chunk(),
            Self::None => &[],
        }
    }

    #[inline]
    fn advance(&mut self, cnt: usize) {
        match *self {
            Self::Buf(ref mut b) => b.advance(cnt),
            Self::Cursor(ref mut c) => c.advance(cnt),
            Self::None => {}
        }
    }

    fn chunks_vectored<'a>(&'a self, dst: &mut [IoSlice<'a>]) -> usize {
        match *self {
            Self::Buf(ref b) => b.chunks_vectored(dst),
            Self::Cursor(ref c) => c.chunks_vectored(dst),
            Self::None => 0,
        }
    }
}

struct H2Upgraded<B>
where
    B: Buf,
{
    ping: Recorder,
    send_stream: UpgradedSendStream<B>,
    recv_stream: RecvStream,
    buf: Bytes,
}

impl<B> AsyncRead for H2Upgraded<B>
where
    B: Buf,
{
    fn poll_read(
        mut self: Pin<&mut Self>,
        cx: &mut Context<'_>,
        read_buf: &mut ReadBuf<'_>,
    ) -> Poll<Result<(), io::Error>> {
        if self.buf.is_empty() {
            self.buf = loop {
                match ready!(self.recv_stream.poll_data(cx)) {
                    None => return Poll::Ready(Ok(())),
                    Some(Ok(buf)) if buf.is_empty() && !self.recv_stream.is_end_stream() => {
                        continue
                    }
                    Some(Ok(buf)) => {
                        self.ping.record_data(buf.len());
                        break buf;
                    }
                    Some(Err(e)) => {
                        return Poll::Ready(match e.reason() {
                            Some(Reason::NO_ERROR) | Some(Reason::CANCEL) => Ok(()),
                            Some(Reason::STREAM_CLOSED) => Err(io::ErrorKind::BrokenPipe.into()),
                            _ => Err(h2_to_io_error(e)),
                        })
                    }
                }
            };
        }
        let cnt = std::cmp::min(self.buf.len(), read_buf.remaining());
        read_buf.put_slice(&self.buf[..cnt]);
        self.buf.advance(cnt);
        let _ = self.recv_stream.flow_control().release_capacity(cnt);
        Poll::Ready(Ok(()))
    }
}

impl<B> AsyncWrite for H2Upgraded<B>
where
    B: Buf,
{
    fn poll_write(
        mut self: Pin<&mut Self>,
        cx: &mut Context<'_>,
        buf: &[u8],
    ) -> Poll<Result<usize, io::Error>> {
        if buf.is_empty() {
            return Poll::Ready(Ok(0));
        }
        self.send_stream.reserve_capacity(buf.len());

        // We ignore all errors returned by `poll_capacity` and `write`, as we
        // will get the correct from `poll_reset` anyway.
        let cnt = match ready!(self.send_stream.poll_capacity(cx)) {
            None => Some(0),
            Some(Ok(cnt)) => self
                .send_stream
                .write(&buf[..cnt], false)
                .ok()
                .map(|()| cnt),
            Some(Err(_)) => None,
        };

        if let Some(cnt) = cnt {
            return Poll::Ready(Ok(cnt));
        }

        Poll::Ready(Err(h2_to_io_error(
            match ready!(self.send_stream.poll_reset(cx)) {
                Ok(Reason::NO_ERROR) | Ok(Reason::CANCEL) | Ok(Reason::STREAM_CLOSED) => {
                    return Poll::Ready(Err(io::ErrorKind::BrokenPipe.into()))
                }
                Ok(reason) => reason.into(),
                Err(e) => e,
            },
        )))
    }

    fn poll_flush(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<Result<(), io::Error>> {
        Poll::Ready(Ok(()))
    }

    fn poll_shutdown(
        mut self: Pin<&mut Self>,
        _cx: &mut Context<'_>,
    ) -> Poll<Result<(), io::Error>> {
        Poll::Ready(self.send_stream.write(&[], true))
    }
}

fn h2_to_io_error(e: h2::Error) -> io::Error {
    if e.is_io() {
        e.into_io().unwrap()
    } else {
        io::Error::new(io::ErrorKind::Other, e)
    }
}

struct UpgradedSendStream<B>(SendStream<SendBuf<Neutered<B>>>);

impl<B> UpgradedSendStream<B>
where
    B: Buf,
{
    unsafe fn new(inner: SendStream<SendBuf<B>>) -> Self {
        assert_eq!(mem::size_of::<B>(), mem::size_of::<Neutered<B>>());
        Self(mem::transmute(inner))
    }

    fn reserve_capacity(&mut self, cnt: usize) {
        unsafe { self.as_inner_unchecked().reserve_capacity(cnt) }
    }

    fn poll_capacity(&mut self, cx: &mut Context<'_>) -> Poll<Option<Result<usize, h2::Error>>> {
        unsafe { self.as_inner_unchecked().poll_capacity(cx) }
    }

    fn poll_reset(&mut self, cx: &mut Context<'_>) -> Poll<Result<h2::Reason, h2::Error>> {
        unsafe { self.as_inner_unchecked().poll_reset(cx) }
    }

    fn write(&mut self, buf: &[u8], end_of_stream: bool) -> Result<(), io::Error> {
        let send_buf = SendBuf::Cursor(Cursor::new(buf.into()));
        unsafe {
            self.as_inner_unchecked()
                .send_data(send_buf, end_of_stream)
                .map_err(h2_to_io_error)
        }
    }

    unsafe fn as_inner_unchecked(&mut self) -> &mut SendStream<SendBuf<B>> {
        &mut *(&mut self.0 as *mut _ as *mut _)
    }
}

#[repr(transparent)]
struct Neutered<B> {
    _inner: B,
    impossible: Impossible,
}

enum Impossible {}

unsafe impl<B> Send for Neutered<B> {}

impl<B> Buf for Neutered<B> {
    fn remaining(&self) -> usize {
        match self.impossible {}
    }

    fn chunk(&self) -> &[u8] {
        match self.impossible {}
    }

    fn advance(&mut self, _cnt: usize) {
        match self.impossible {}
    }
}
<commit_msg>fix(http2): improve errors emitted by HTTP2 `Upgraded` stream shutdown (#2622)<commit_after>use bytes::{Buf, Bytes};
use h2::{Reason, RecvStream, SendStream};
use http::header::{HeaderName, CONNECTION, TE, TRAILER, TRANSFER_ENCODING, UPGRADE};
use http::HeaderMap;
use pin_project_lite::pin_project;
use std::error::Error as StdError;
use std::io::{self, Cursor, IoSlice};
use std::mem;
use std::task::Context;
use tokio::io::{AsyncRead, AsyncWrite, ReadBuf};

use crate::body::HttpBody;
use crate::common::{task, Future, Pin, Poll};
use crate::proto::h2::ping::Recorder;

pub(crate) mod ping;

cfg_client! {
    pub(crate) mod client;
    pub(crate) use self::client::ClientTask;
}

cfg_server! {
    pub(crate) mod server;
    pub(crate) use self::server::Server;
}

/// Default initial stream window size defined in HTTP2 spec.
pub(crate) const SPEC_WINDOW_SIZE: u32 = 65_535;

fn strip_connection_headers(headers: &mut HeaderMap, is_request: bool) {
    // List of connection headers from:
    // https://developer.mozilla.org/en-US/docs/Web/HTTP/Headers/Connection
    //
    // TE headers are allowed in HTTP/2 requests as long as the value is "trailers", so they're
    // tested separately.
    let connection_headers = [
        HeaderName::from_lowercase(b"keep-alive").unwrap(),
        HeaderName::from_lowercase(b"proxy-connection").unwrap(),
        TRAILER,
        TRANSFER_ENCODING,
        UPGRADE,
    ];

    for header in connection_headers.iter() {
        if headers.remove(header).is_some() {
            warn!("Connection header illegal in HTTP/2: {}", header.as_str());
        }
    }

    if is_request {
        if headers
            .get(TE)
            .map(|te_header| te_header != "trailers")
            .unwrap_or(false)
        {
            warn!("TE headers not set to \"trailers\" are illegal in HTTP/2 requests");
            headers.remove(TE);
        }
    } else if headers.remove(TE).is_some() {
        warn!("TE headers illegal in HTTP/2 responses");
    }

    if let Some(header) = headers.remove(CONNECTION) {
        warn!(
            "Connection header illegal in HTTP/2: {}",
            CONNECTION.as_str()
        );
        let header_contents = header.to_str().unwrap();

        // A `Connection` header may have a comma-separated list of names of other headers that
        // are meant for only this specific connection.
        //
        // Iterate these names and remove them as headers. Connection-specific headers are
        // forbidden in HTTP2, as that information has been moved into frame types of the h2
        // protocol.
        for name in header_contents.split(',') {
            let name = name.trim();
            headers.remove(name);
        }
    }
}

// body adapters used by both Client and Server

pin_project! {
    struct PipeToSendStream<S>
    where
        S: HttpBody,
    {
        body_tx: SendStream<SendBuf<S::Data>>,
        data_done: bool,
        #[pin]
        stream: S,
    }
}

impl<S> PipeToSendStream<S>
where
    S: HttpBody,
{
    fn new(stream: S, tx: SendStream<SendBuf<S::Data>>) -> PipeToSendStream<S> {
        PipeToSendStream {
            body_tx: tx,
            data_done: false,
            stream,
        }
    }
}

impl<S> Future for PipeToSendStream<S>
where
    S: HttpBody,
    S::Error: Into<Box<dyn StdError + Send + Sync>>,
{
    type Output = crate::Result<()>;

    fn poll(self: Pin<&mut Self>, cx: &mut task::Context<'_>) -> Poll<Self::Output> {
        let mut me = self.project();
        loop {
            if !*me.data_done {
                // we don't have the next chunk of data yet, so just reserve 1 byte to make
                // sure there's some capacity available. h2 will handle the capacity management
                // for the actual body chunk.
                me.body_tx.reserve_capacity(1);

                if me.body_tx.capacity() == 0 {
                    loop {
                        match ready!(me.body_tx.poll_capacity(cx)) {
                            Some(Ok(0)) => {}
                            Some(Ok(_)) => break,
                            Some(Err(e)) => {
                                return Poll::Ready(Err(crate::Error::new_body_write(e)))
                            }
                            None => {
                                // None means the stream is no longer in a
                                // streaming state, we either finished it
                                // somehow, or the remote reset us.
                                return Poll::Ready(Err(crate::Error::new_body_write(
                                    "send stream capacity unexpectedly closed",
                                )));
                            }
                        }
                    }
                } else if let Poll::Ready(reason) = me
                    .body_tx
                    .poll_reset(cx)
                    .map_err(crate::Error::new_body_write)?
                {
                    debug!("stream received RST_STREAM: {:?}", reason);
                    return Poll::Ready(Err(crate::Error::new_body_write(::h2::Error::from(
                        reason,
                    ))));
                }

                match ready!(me.stream.as_mut().poll_data(cx)) {
                    Some(Ok(chunk)) => {
                        let is_eos = me.stream.is_end_stream();
                        trace!(
                            "send body chunk: {} bytes, eos={}",
                            chunk.remaining(),
                            is_eos,
                        );

                        let buf = SendBuf::Buf(chunk);
                        me.body_tx
                            .send_data(buf, is_eos)
                            .map_err(crate::Error::new_body_write)?;

                        if is_eos {
                            return Poll::Ready(Ok(()));
                        }
                    }
                    Some(Err(e)) => return Poll::Ready(Err(me.body_tx.on_user_err(e))),
                    None => {
                        me.body_tx.reserve_capacity(0);
                        let is_eos = me.stream.is_end_stream();
                        if is_eos {
                            return Poll::Ready(me.body_tx.send_eos_frame());
                        } else {
                            *me.data_done = true;
                            // loop again to poll_trailers
                        }
                    }
                }
            } else {
                if let Poll::Ready(reason) = me
                    .body_tx
                    .poll_reset(cx)
                    .map_err(crate::Error::new_body_write)?
                {
                    debug!("stream received RST_STREAM: {:?}", reason);
                    return Poll::Ready(Err(crate::Error::new_body_write(::h2::Error::from(
                        reason,
                    ))));
                }

                match ready!(me.stream.poll_trailers(cx)) {
                    Ok(Some(trailers)) => {
                        me.body_tx
                            .send_trailers(trailers)
                            .map_err(crate::Error::new_body_write)?;
                        return Poll::Ready(Ok(()));
                    }
                    Ok(None) => {
                        // There were no trailers, so send an empty DATA frame...
                        return Poll::Ready(me.body_tx.send_eos_frame());
                    }
                    Err(e) => return Poll::Ready(Err(me.body_tx.on_user_err(e))),
                }
            }
        }
    }
}

trait SendStreamExt {
    fn on_user_err<E>(&mut self, err: E) -> crate::Error
    where
        E: Into<Box<dyn std::error::Error + Send + Sync>>;
    fn send_eos_frame(&mut self) -> crate::Result<()>;
}

impl<B: Buf> SendStreamExt for SendStream<SendBuf<B>> {
    fn on_user_err<E>(&mut self, err: E) -> crate::Error
    where
        E: Into<Box<dyn std::error::Error + Send + Sync>>,
    {
        let err = crate::Error::new_user_body(err);
        debug!("send body user stream error: {}", err);
        self.send_reset(err.h2_reason());
        err
    }

    fn send_eos_frame(&mut self) -> crate::Result<()> {
        trace!("send body eos");
        self.send_data(SendBuf::None, true)
            .map_err(crate::Error::new_body_write)
    }
}

#[repr(usize)]
enum SendBuf<B> {
    Buf(B),
    Cursor(Cursor<Box<[u8]>>),
    None,
}

impl<B: Buf> Buf for SendBuf<B> {
    #[inline]
    fn remaining(&self) -> usize {
        match *self {
            Self::Buf(ref b) => b.remaining(),
            Self::Cursor(ref c) => Buf::remaining(c),
            Self::None => 0,
        }
    }

    #[inline]
    fn chunk(&self) -> &[u8] {
        match *self {
            Self::Buf(ref b) => b.chunk(),
            Self::Cursor(ref c) => c.chunk(),
            Self::None => &[],
        }
    }

    #[inline]
    fn advance(&mut self, cnt: usize) {
        match *self {
            Self::Buf(ref mut b) => b.advance(cnt),
            Self::Cursor(ref mut c) => c.advance(cnt),
            Self::None => {}
        }
    }

    fn chunks_vectored<'a>(&'a self, dst: &mut [IoSlice<'a>]) -> usize {
        match *self {
            Self::Buf(ref b) => b.chunks_vectored(dst),
            Self::Cursor(ref c) => c.chunks_vectored(dst),
            Self::None => 0,
        }
    }
}

struct H2Upgraded<B>
where
    B: Buf,
{
    ping: Recorder,
    send_stream: UpgradedSendStream<B>,
    recv_stream: RecvStream,
    buf: Bytes,
}

impl<B> AsyncRead for H2Upgraded<B>
where
    B: Buf,
{
    fn poll_read(
        mut self: Pin<&mut Self>,
        cx: &mut Context<'_>,
        read_buf: &mut ReadBuf<'_>,
    ) -> Poll<Result<(), io::Error>> {
        if self.buf.is_empty() {
            self.buf = loop {
                match ready!(self.recv_stream.poll_data(cx)) {
                    None => return Poll::Ready(Ok(())),
                    Some(Ok(buf)) if buf.is_empty() && !self.recv_stream.is_end_stream() => {
                        continue
                    }
                    Some(Ok(buf)) => {
                        self.ping.record_data(buf.len());
                        break buf;
                    }
                    Some(Err(e)) => {
                        return Poll::Ready(match e.reason() {
                            Some(Reason::NO_ERROR) | Some(Reason::CANCEL) => Ok(()),
                            Some(Reason::STREAM_CLOSED) => Err(io::ErrorKind::BrokenPipe.into()),
                            _ => Err(h2_to_io_error(e)),
                        })
                    }
                }
            };
        }
        let cnt = std::cmp::min(self.buf.len(), read_buf.remaining());
        read_buf.put_slice(&self.buf[..cnt]);
        self.buf.advance(cnt);
        let _ = self.recv_stream.flow_control().release_capacity(cnt);
        Poll::Ready(Ok(()))
    }
}

impl<B> AsyncWrite for H2Upgraded<B>
where
    B: Buf,
{
    fn poll_write(
        mut self: Pin<&mut Self>,
        cx: &mut Context<'_>,
        buf: &[u8],
    ) -> Poll<Result<usize, io::Error>> {
        if buf.is_empty() {
            return Poll::Ready(Ok(0));
        }
        self.send_stream.reserve_capacity(buf.len());

        // We ignore all errors returned by `poll_capacity` and `write`, as we
        // will get the correct from `poll_reset` anyway.
        let cnt = match ready!(self.send_stream.poll_capacity(cx)) {
            None => Some(0),
            Some(Ok(cnt)) => self
                .send_stream
                .write(&buf[..cnt], false)
                .ok()
                .map(|()| cnt),
            Some(Err(_)) => None,
        };

        if let Some(cnt) = cnt {
            return Poll::Ready(Ok(cnt));
        }

        Poll::Ready(Err(h2_to_io_error(
            match ready!(self.send_stream.poll_reset(cx)) {
                Ok(Reason::NO_ERROR) | Ok(Reason::CANCEL) | Ok(Reason::STREAM_CLOSED) => {
                    return Poll::Ready(Err(io::ErrorKind::BrokenPipe.into()))
                }
                Ok(reason) => reason.into(),
                Err(e) => e,
            },
        )))
    }

    fn poll_flush(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<Result<(), io::Error>> {
        Poll::Ready(Ok(()))
    }

    fn poll_shutdown(
        mut self: Pin<&mut Self>,
        cx: &mut Context<'_>,
    ) -> Poll<Result<(), io::Error>> {
        if self.send_stream.write(&[], true).is_ok() {
            return Poll::Ready(Ok(()))
        }

        Poll::Ready(Err(h2_to_io_error(
            match ready!(self.send_stream.poll_reset(cx)) {
                Ok(Reason::NO_ERROR) => {
                    return Poll::Ready(Ok(()))
                }
                Ok(Reason::CANCEL) | Ok(Reason::STREAM_CLOSED) => {
                    return Poll::Ready(Err(io::ErrorKind::BrokenPipe.into()))
                }
                Ok(reason) => reason.into(),
                Err(e) => e,
            },
        )))
    }
}

fn h2_to_io_error(e: h2::Error) -> io::Error {
    if e.is_io() {
        e.into_io().unwrap()
    } else {
        io::Error::new(io::ErrorKind::Other, e)
    }
}

struct UpgradedSendStream<B>(SendStream<SendBuf<Neutered<B>>>);

impl<B> UpgradedSendStream<B>
where
    B: Buf,
{
    unsafe fn new(inner: SendStream<SendBuf<B>>) -> Self {
        assert_eq!(mem::size_of::<B>(), mem::size_of::<Neutered<B>>());
        Self(mem::transmute(inner))
    }

    fn reserve_capacity(&mut self, cnt: usize) {
        unsafe { self.as_inner_unchecked().reserve_capacity(cnt) }
    }

    fn poll_capacity(&mut self, cx: &mut Context<'_>) -> Poll<Option<Result<usize, h2::Error>>> {
        unsafe { self.as_inner_unchecked().poll_capacity(cx) }
    }

    fn poll_reset(&mut self, cx: &mut Context<'_>) -> Poll<Result<h2::Reason, h2::Error>> {
        unsafe { self.as_inner_unchecked().poll_reset(cx) }
    }

    fn write(&mut self, buf: &[u8], end_of_stream: bool) -> Result<(), io::Error> {
        let send_buf = SendBuf::Cursor(Cursor::new(buf.into()));
        unsafe {
            self.as_inner_unchecked()
                .send_data(send_buf, end_of_stream)
                .map_err(h2_to_io_error)
        }
    }

    unsafe fn as_inner_unchecked(&mut self) -> &mut SendStream<SendBuf<B>> {
        &mut *(&mut self.0 as *mut _ as *mut _)
    }
}

#[repr(transparent)]
struct Neutered<B> {
    _inner: B,
    impossible: Impossible,
}

enum Impossible {}

unsafe impl<B> Send for Neutered<B> {}

impl<B> Buf for Neutered<B> {
    fn remaining(&self) -> usize {
        match self.impossible {}
    }

    fn chunk(&self) -> &[u8] {
        match self.impossible {}
    }

    fn advance(&mut self, _cnt: usize) {
        match self.impossible {}
    }
}
<|endoftext|>


<commit_before>// Bloom
//
// HTTP REST API caching middleware
// Copyright: 2017, Valerian Saliou <valerian@valeriansaliou.name>
// License: Mozilla Public License v2.0 (MPL v2.0)

use hyper::{Error, Client, Method, Uri, Headers, Body, Request};
use hyper::client::HttpConnector;
use hyper::server::Response;
use tokio_core::reactor::Core;

use APP_CONF;

const MAX_SHARDS: u8 = 1;

lazy_static! {
    static ref SHARD_URI: Uri = format!("http://{}:{}", APP_CONF.proxy.inet.ip(),
        APP_CONF.proxy.inet.port()).parse().unwrap();
}

pub struct ProxyTunnelBuilder;

pub struct ProxyTunnel {
    core: Core,
    client: Client<HttpConnector>,
    shards: [Option<&'static Uri>; MAX_SHARDS as usize],
}

impl ProxyTunnelBuilder {
    pub fn new() -> ProxyTunnel {
        // TODO: keep a pool of connections active? (re-use existing connectors)
        let core = Core::new().unwrap();
        let handle = core.handle();
        let client = Client::configure()
            .connector(HttpConnector::new(APP_CONF.proxy.tunnel_threads, &handle))
            .build(&handle);

        // We support only 1 shard for now.
        ProxyTunnel {
            core: core,
            client: client,
            shards: [Some(&*SHARD_URI)],
        }
    }
}

impl ProxyTunnel {
    pub fn run(
        &mut self,
        method: &Method,
        uri: &Uri,
        headers: &Headers,
        body: Body,
        shard: u8,
    ) -> Result<Response, Error> {
        if shard < MAX_SHARDS {
            // Route to target shard
            match self.shards[shard as usize] {
                Some(ref shard_uri) => {
                    let tunnel_uri = format!("{}{}", shard_uri, uri.path()).parse().unwrap();

                    let mut tunnel_req = Request::new(method.to_owned(), tunnel_uri);

                    // Forward headers
                    {
                        let tunnel_headers = tunnel_req.headers_mut();

                        tunnel_headers.clone_from(headers);
                    }

                    // Forward body
                    // TODO: blocking, wtf?
                    // tunnel_req.set_body(Body::from(body));

                    // TODO: debug
                    // tunnel_req.set_body(Body::from(
                    //     "{\"type\":\"online\",\"time\":{\"for\":60}}"));

                    self.core.run(self.client.request(tunnel_req))
                }
                None => Err(Error::Header),
            }
        } else {
            // Shard out of bounds
            Err(Error::Header)
        }
    }
}
<commit_msg>Debugs<commit_after>// Bloom
//
// HTTP REST API caching middleware
// Copyright: 2017, Valerian Saliou <valerian@valeriansaliou.name>
// License: Mozilla Public License v2.0 (MPL v2.0)

use hyper::{Error, Client, Method, Uri, Headers, Body, Request};
use hyper::client::HttpConnector;
use hyper::server::Response;
use tokio_core::reactor::Core;

use APP_CONF;

const MAX_SHARDS: u8 = 1;

lazy_static! {
    static ref SHARD_URI: Uri = format!("http://{}:{}", APP_CONF.proxy.inet.ip(),
        APP_CONF.proxy.inet.port()).parse().unwrap();
}

pub struct ProxyTunnelBuilder;

pub struct ProxyTunnel {
    core: Core,
    client: Client<HttpConnector>,
    shards: [Option<&'static Uri>; MAX_SHARDS as usize],
}

impl ProxyTunnelBuilder {
    pub fn new() -> ProxyTunnel {
        // TODO: keep a pool of connections active? (re-use existing connectors)
        let core = Core::new().unwrap();
        let handle = core.handle();
        let client = Client::configure()
            .connector(HttpConnector::new(APP_CONF.proxy.tunnel_threads, &handle))
            .build(&handle);

        // We support only 1 shard for now.
        ProxyTunnel {
            core: core,
            client: client,
            shards: [Some(&*SHARD_URI)],
        }
    }
}

impl ProxyTunnel {
    pub fn run(
        &mut self,
        method: &Method,
        uri: &Uri,
        headers: &Headers,
        body: Body,
        shard: u8,
    ) -> Result<Response, Error> {
        if shard < MAX_SHARDS {
            // Route to target shard
            match self.shards[shard as usize] {
                Some(ref shard_uri) => {
                    let tunnel_uri = format!("{}{}", shard_uri, uri.path()).parse().unwrap();

                    let mut tunnel_req = Request::new(method.to_owned(), tunnel_uri);

                    // Forward headers
                    {
                        let tunnel_headers = tunnel_req.headers_mut();

                        tunnel_headers.clone_from(headers);
                    }

                    // TODO: ignore the body if not POST, PATCH, etc (no need to fwd it)

                    // Forward body
                    // TODO: blocking if non-empty, eg. if PATCH, why?
                    tunnel_req.set_body(body);

                    // TODO: debug (this works when the content-length is ==)
                    // tunnel_req.set_body(Body::from(
                    //     "{\"type\":\"online\",\"time\":{\"for\":60}}"));

                    self.core.run(self.client.request(tunnel_req))
                }
                None => Err(Error::Header),
            }
        } else {
            // Shard out of bounds
            Err(Error::Header)
        }
    }
}
<|endoftext|>


<commit_before>use ::color::Color;
use ::powerline::*;
use ::part::*;
use std::process::Command;
use std::str;
use regex::Regex;

pub struct GitInfo {
    untracked: u32,
    conflicted: u32,
    non_staged: u32,
    pub ahead: u32,
    pub behind: u32,
    staged: u32,
}

impl GitInfo{
    pub fn new() -> GitInfo{
        GitInfo{ untracked: 0, conflicted: 0, non_staged: 0, staged: 0, ahead: 0, behind: 0}
    }

    fn is_dirty(&self) -> bool {
        (self.untracked + self.conflicted + self.staged + self.non_staged) > 0
    }

    fn add_file(&mut self, begin: &str) -> Result<(), Error> {
        match begin{
            "??" => self.untracked += 1,
            "DD" => self.conflicted += 1,
            "AU" => self.conflicted += 1,
            "UD" => self.conflicted += 1,
            "UA" => self.conflicted += 1,
            "UU" => self.conflicted += 1,
            "DU" => self.conflicted += 1,
            "AA" => self.conflicted += 1,
            _ => {
                let mut chars = begin.chars();
                let a = chars.next().ok_or(Error::from_str("Invalid file status"))?;
                let b = chars.next().ok_or(Error::from_str("Invalid file status"))?;
                if b != ' '{ self.non_staged += 1; }
                if a != ' '{ self.staged += 1; }
            },
        };
        Ok(())
    }

}

fn get_detached_branch_name() -> Result<String, Error> {
    let child = Command::new("git").args(&["describe", "--tags", "--always"]).output().map_err(|e| Error::wrap(e, "Failed to run git"))?;
    Ok(if child.status.success() {
        String::from("Big Bang")
    } else {
        let branch = str::from_utf8(&child.stdout)?.split("\n").next().ok_or(Error::from_str("Empty git output"))?;
        format!("\u{2693}{}", branch)
    })
}

fn  quantity(val: u32) -> String{
    if val  > 1 {
        format!("{}",val)
    } else {
        String::from("")
    }
}

impl Part for GitInfo {
    fn get_segments(mut self) -> Result<Vec<Segment>, Error> {
        let output = Command::new("git").args(&["status", "--porcelain", "-b"]).output().map_err(|e| Error::wrap(e, "Failed to run git"))?;
        let data = str::from_utf8(&output.stdout)?;
        if data == "" { return Ok(Vec::new()); }
        let mut lines: Vec<&str> = data.split("\n").collect();
        lines.pop();

        let mut iter = lines.into_iter();
        let branch_line = iter.next().ok_or(Error::from_str("Empty git output"))?;
        let re = Regex::new(r"^## (?P<local>[^\.]+)?")?;

        let branch = {
            if let Some(caps) = re.captures(branch_line) {
                caps["local"].to_owned()
            } else {
                get_detached_branch_name()?
            }
        };
        for x in iter {
            let file = x.get(..2).ok_or(Error::from_str("Invalid file status line"))?;
            self.add_file(file)?;
        }

        let (branch_fg, branch_bg) = if self.is_dirty() {
            (Color::REPO_DIRTY_FG, Color::REPO_DIRTY_BG)
        } else {
            (Color::REPO_CLEAN_FG, Color::REPO_CLEAN_BG)
        };
        let mut segments = Vec::new();
        segments.push(Segment::simple(&format!(" {} ", branch), branch_fg, branch_bg));
        {
            let mut add_elem = |count, symbol, fg, bg| {
                if count > 0 {
                    let text = format!(" {}{} ", quantity(count), symbol);
                    segments.push(Segment::simple(&text, fg, bg));
                }
            };
            add_elem(self.ahead, '\u{2B06}', Color::GIT_AHEAD_FG, Color::GIT_AHEAD_BG);
            add_elem(self.behind, '\u{2B07}', Color::GIT_BEHIND_FG, Color::GIT_BEHIND_BG);
            add_elem(self.staged, '\u{2714}', Color::GIT_STAGED_FG, Color::GIT_STAGED_BG);
            add_elem(self.non_staged, '\u{270E}', Color::GIT_NOTSTAGED_FG, Color::GIT_NOTSTAGED_BG);
            add_elem(self.untracked, '\u{2753}', Color::GIT_UNTRACKED_FG, Color::GIT_UNTRACKED_BG);
            add_elem(self.conflicted, '\u{273C}', Color::GIT_CONFLICTED_FG, Color::GIT_CONFLICTED_BG);
        }
        Ok(segments)
    }
}<commit_msg>added git ahead/behind support<commit_after>use ::color::Color;
use ::powerline::*;
use ::part::*;
use std::process::Command;
use std::str;
use regex::Regex;

pub struct GitInfo {
    untracked: u32,
    conflicted: u32,
    non_staged: u32,
    pub ahead: u32,
    pub behind: u32,
    staged: u32,
}

impl GitInfo{
    pub fn new() -> GitInfo{
        GitInfo{ untracked: 0, conflicted: 0, non_staged: 0, staged: 0, ahead: 0, behind: 0}
    }

    fn is_dirty(&self) -> bool {
        (self.untracked + self.conflicted + self.staged + self.non_staged) > 0
    }

    fn add_file(&mut self, begin: &str) -> Result<(), Error> {
        match begin{
            "??" => self.untracked += 1,
            "DD" => self.conflicted += 1,
            "AU" => self.conflicted += 1,
            "UD" => self.conflicted += 1,
            "UA" => self.conflicted += 1,
            "UU" => self.conflicted += 1,
            "DU" => self.conflicted += 1,
            "AA" => self.conflicted += 1,
            _ => {
                let mut chars = begin.chars();
                let a = chars.next().ok_or(Error::from_str("Invalid file status"))?;
                let b = chars.next().ok_or(Error::from_str("Invalid file status"))?;
                if b != ' '{ self.non_staged += 1; }
                if a != ' '{ self.staged += 1; }
            },
        };
        Ok(())
    }

}

fn get_detached_branch_name() -> Result<String, Error> {
    let child = Command::new("git").args(&["describe", "--tags", "--always"]).output().map_err(|e| Error::wrap(e, "Failed to run git"))?;
    Ok(if child.status.success() {
        String::from("Big Bang")
    } else {
        let branch = str::from_utf8(&child.stdout)?.split("\n").next().ok_or(Error::from_str("Empty git output"))?;
        format!("\u{2693}{}", branch)
    })
}

fn  quantity(val: u32) -> String{
    if val  > 1 {
        format!("{}",val)
    } else {
        String::from("")
    }
}

impl Part for GitInfo {
    fn get_segments(mut self) -> Result<Vec<Segment>, Error> {
        let output = Command::new("git").args(&["status", "--porcelain", "-b"]).output().map_err(|e| Error::wrap(e, "Failed to run git"))?;
        let data = str::from_utf8(&output.stdout)?;
        if data == "" { return Ok(Vec::new()); }
        let mut lines: Vec<&str> = data.split("\n").collect();
        lines.pop();

        let mut iter = lines.into_iter();
        let branch_line = iter.next().ok_or(Error::from_str("Empty git output"))?;
        let re = Regex::new(r"^## (?P<local>[^\.]+)?")?;
        let aheadbehind_re = Regex::new(r"\[(ahead (\d+))?(, )?(behind (\d+))?\]")?;

        if let Some(ab_caps) = aheadbehind_re.captures(branch_line) {
            if let Some(ahead) = ab_caps.get(2) {
                self.ahead += ahead.as_str().parse()?;
            }
            if let Some(behind) = ab_caps.get(5) {
                self.behind += behind.as_str().parse()?;
            }
        }
        let branch = {
            if let Some(caps) = re.captures(branch_line) {
                caps["local"].to_owned()
            } else {
                get_detached_branch_name()?
            }
        };
        for x in iter {
            let file = x.get(..2).ok_or(Error::from_str("Invalid file status line"))?;
            self.add_file(file)?;
        }

        let (branch_fg, branch_bg) = if self.is_dirty() {
            (Color::REPO_DIRTY_FG, Color::REPO_DIRTY_BG)
        } else {
            (Color::REPO_CLEAN_FG, Color::REPO_CLEAN_BG)
        };
        let mut segments = Vec::new();
        segments.push(Segment::simple(&format!(" {} ", branch), branch_fg, branch_bg));
        {
            let mut add_elem = |count, symbol, fg, bg| {
                if count > 0 {
                    let text = format!(" {}{} ", quantity(count), symbol);
                    segments.push(Segment::simple(&text, fg, bg));
                }
            };
            add_elem(self.ahead, '\u{2B06}', Color::GIT_AHEAD_FG, Color::GIT_AHEAD_BG);
            add_elem(self.behind, '\u{2B07}', Color::GIT_BEHIND_FG, Color::GIT_BEHIND_BG);
            add_elem(self.staged, '\u{2714}', Color::GIT_STAGED_FG, Color::GIT_STAGED_BG);
            add_elem(self.non_staged, '\u{270E}', Color::GIT_NOTSTAGED_FG, Color::GIT_NOTSTAGED_BG);
            add_elem(self.untracked, '\u{2753}', Color::GIT_UNTRACKED_FG, Color::GIT_UNTRACKED_BG);
            add_elem(self.conflicted, '\u{273C}', Color::GIT_CONFLICTED_FG, Color::GIT_CONFLICTED_BG);
        }
        Ok(segments)
    }
}<|endoftext|>


<commit_before>//! Contains result and error type for FBX reader.

extern crate byteorder;

use std::io;
use std::string;
use std::fmt;

/// A specialized `std::result::Result` type for FBX parsing.
pub type Result<T> = ::std::result::Result<T, Error>;

/// An FBX parsing error.
#[derive(Debug)]
pub struct Error {
    /// Last position of successfully read data when an error detected.
    pos: u64,
    /// Error type.
    kind: ErrorKind,
}

impl Error {
    /// Constructs `Error` with position and objects which can be converted to `ErrorKind`.
    pub fn new<K: Into<ErrorKind>>(pos: u64, kind: K) -> Self {
        Error {
            pos: pos,
            kind: kind.into(),
        }
    }
}

impl fmt::Display for Error {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match self.kind {
            ErrorKind::FromUtf8Error(ref err) => write!(f, "UTF-8 conversion error at pos={}: {}", self.pos, err),
            ErrorKind::InvalidMagic => write!(f, "Invalid magic header at pos={}: Non-FBX or corrupted data?", self.pos),
            ErrorKind::Io(ref err) => write!(f, "I/O error at pos={}: {}", self.pos, err),
            ErrorKind::DataError(ref err) => write!(f, "Invalid data at pos={}: {}", self.pos, err),
            ErrorKind::UnexpectedValue(ref err) => write!(f, "Got an unexpected value at pos={}: {}", self.pos, err),
            ErrorKind::UnexpectedEof => write!(f, "Unexpected EOF at pos={}", self.pos),
            ErrorKind::Unimplemented(ref err) => write!(f, "Unimplemented feature: {}", err),
        }
    }
}

impl ::std::error::Error for Error {
    fn description(&self) -> &str {
        match self.kind {
            ErrorKind::FromUtf8Error(ref err) => err.description(),
            ErrorKind::InvalidMagic => "Got an invalid magic header",
            ErrorKind::Io(ref err) => err.description(),
            ErrorKind::DataError(_) => "Got an invalid data",
            ErrorKind::UnexpectedValue(_) => "Invalid value in FBX data",
            ErrorKind::UnexpectedEof => "Unexpected EOF",
            ErrorKind::Unimplemented(_) => "Attempt to use unimplemented feature",
        }
    }
}

/// Error type.
#[derive(Debug)]
pub enum ErrorKind {
    /// Conversion from array of u8 to String failed.
    FromUtf8Error(string::FromUtf8Error),
    /// Invalid magic binary detected.
    InvalidMagic,
    /// I/O operation error.
    Io(io::Error),
    /// Corrupted or inconsistent FBX data detected.
    DataError(String),
    /// Got an unexpected value, and cannot continue parsing.
    ///
    /// This is specialization of `DataError`.
    UnexpectedValue(String),
    /// Reached unexpected EOF.
    UnexpectedEof,
    /// Attempted to use unimplemented feature.
    Unimplemented(String),
}

impl From<string::FromUtf8Error> for ErrorKind {
    fn from(err: string::FromUtf8Error) -> ErrorKind {
        ErrorKind::FromUtf8Error(err)
    }
}

impl From<io::Error> for ErrorKind {
    fn from(err: io::Error) -> ErrorKind {
        // TODO: `io::Error::UnexpectedEof` should be converted to `ErrorKind::UnexpectedEof`.
        ErrorKind::Io(err)
    }
}

impl From<byteorder::Error> for ErrorKind {
    fn from(err: byteorder::Error) -> ErrorKind {
        match err {
            byteorder::Error::UnexpectedEOF => ErrorKind::UnexpectedEof,
            byteorder::Error::Io(err) => ErrorKind::Io(err),
        }
    }
}
<commit_msg>Implement `Clone` trait for reader `Error` related types<commit_after>//! Contains result and error type for FBX reader.

extern crate byteorder;

use std::io;
use std::string;
use std::str;
use std::fmt;

/// A specialized `std::result::Result` type for FBX parsing.
pub type Result<T> = ::std::result::Result<T, Error>;

/// An FBX parsing error.
#[derive(Debug, Clone)]
pub struct Error {
    /// Last position of successfully read data when an error detected.
    pos: u64,
    /// Error type.
    kind: ErrorKind,
}

impl Error {
    /// Constructs `Error` with position and objects which can be converted to `ErrorKind`.
    pub fn new<K: Into<ErrorKind>>(pos: u64, kind: K) -> Self {
        Error {
            pos: pos,
            kind: kind.into(),
        }
    }
}

impl fmt::Display for Error {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match self.kind {
            ErrorKind::Utf8Error(ref err) => write!(f, "UTF-8 conversion error at pos={}: {}", self.pos, err),
            ErrorKind::InvalidMagic => write!(f, "Invalid magic header at pos={}: Non-FBX or corrupted data?", self.pos),
            ErrorKind::Io(ref err) => write!(f, "I/O error at pos={}: {}", self.pos, err),
            ErrorKind::DataError(ref err) => write!(f, "Invalid data at pos={}: {}", self.pos, err),
            ErrorKind::UnexpectedValue(ref err) => write!(f, "Got an unexpected value at pos={}: {}", self.pos, err),
            ErrorKind::UnexpectedEof => write!(f, "Unexpected EOF at pos={}", self.pos),
            ErrorKind::Unimplemented(ref err) => write!(f, "Unimplemented feature: {}", err),
        }
    }
}

impl ::std::error::Error for Error {
    fn description(&self) -> &str {
        match self.kind {
            ErrorKind::Utf8Error(ref err) => err.description(),
            ErrorKind::InvalidMagic => "Got an invalid magic header",
            ErrorKind::Io(ref err) => err.description(),
            ErrorKind::DataError(_) => "Got an invalid data",
            ErrorKind::UnexpectedValue(_) => "Invalid value in FBX data",
            ErrorKind::UnexpectedEof => "Unexpected EOF",
            ErrorKind::Unimplemented(_) => "Attempt to use unimplemented feature",
        }
    }
}

/// Error type.
#[derive(Debug)]
pub enum ErrorKind {
    /// Conversion from array of u8 to String failed.
    Utf8Error(str::Utf8Error),
    /// Invalid magic binary detected.
    InvalidMagic,
    /// I/O operation error.
    Io(io::Error),
    /// Corrupted or inconsistent FBX data detected.
    DataError(String),
    /// Got an unexpected value, and cannot continue parsing.
    ///
    /// This is specialization of `DataError`.
    UnexpectedValue(String),
    /// Reached unexpected EOF.
    UnexpectedEof,
    /// Attempted to use unimplemented feature.
    Unimplemented(String),
}

impl Clone for ErrorKind {
    fn clone(&self) -> Self {
        use self::ErrorKind::*;
        use std::error::Error;
        match *self {
            Utf8Error(ref e) => Utf8Error(e.clone()),
            InvalidMagic => InvalidMagic,
            // `io::Error` (and an error wrapped by `io::Error`) cannot be cloned.
            Io(ref e) => Io(io::Error::new(e.kind(), e.description())),
            DataError(ref e) => DataError(e.clone()),
            UnexpectedValue(ref e) => UnexpectedValue(e.clone()),
            UnexpectedEof => UnexpectedEof,
            Unimplemented(ref e) => Unimplemented(e.clone()),
        }
    }
}

impl From<string::FromUtf8Error> for ErrorKind {
    fn from(err: string::FromUtf8Error) -> ErrorKind {
        ErrorKind::Utf8Error(err.utf8_error())
    }
}

impl From<io::Error> for ErrorKind {
    fn from(err: io::Error) -> ErrorKind {
        // TODO: `io::Error::UnexpectedEof` should be converted to `ErrorKind::UnexpectedEof`.
        ErrorKind::Io(err)
    }
}

impl From<byteorder::Error> for ErrorKind {
    fn from(err: byteorder::Error) -> ErrorKind {
        match err {
            byteorder::Error::UnexpectedEOF => ErrorKind::UnexpectedEof,
            byteorder::Error::Io(err) => ErrorKind::Io(err),
        }
    }
}
<|endoftext|>


<commit_before>use bytes::{BufMut, BytesMut};
use postgres_protocol::message::backend;
use postgres_protocol::message::frontend;
use socket2::{Domain, SockAddr, Socket, Type};
use std::io::{self, BufWriter, Read, Write};
use std::net::{SocketAddr, ToSocketAddrs};
#[cfg(unix)]
use std::os::unix::io::{AsRawFd, RawFd};
#[cfg(windows)]
use std::os::windows::io::{AsRawSocket, RawSocket};
use std::result;
use std::time::Duration;

use error;
use params::{ConnectParams, Host};
use tls::TlsStream;
use {Result, TlsMode};

const INITIAL_CAPACITY: usize = 8 * 1024;

pub struct MessageStream {
    stream: BufWriter<Box<TlsStream>>,
    in_buf: BytesMut,
    out_buf: Vec<u8>,
}

impl MessageStream {
    pub fn new(stream: Box<TlsStream>) -> MessageStream {
        MessageStream {
            stream: BufWriter::new(stream),
            in_buf: BytesMut::with_capacity(INITIAL_CAPACITY),
            out_buf: vec![],
        }
    }

    pub fn get_ref(&self) -> &Box<TlsStream> {
        self.stream.get_ref()
    }

    pub fn write_message<F, E>(&mut self, f: F) -> result::Result<(), E>
    where
        F: FnOnce(&mut Vec<u8>) -> result::Result<(), E>,
        E: From<io::Error>,
    {
        self.out_buf.clear();
        f(&mut self.out_buf)?;
        self.stream.write_all(&self.out_buf).map_err(From::from)
    }

    pub fn read_message(&mut self) -> io::Result<backend::Message> {
        loop {
            match backend::Message::parse(&mut self.in_buf) {
                Ok(Some(message)) => return Ok(message),
                Ok(None) => self.read_in()?,
                Err(e) => return Err(e),
            }
        }
    }

    fn read_in(&mut self) -> io::Result<()> {
        self.in_buf.reserve(1);
        match self.stream
            .get_mut()
            .read(unsafe { self.in_buf.bytes_mut() })
        {
            Ok(0) => Err(io::Error::new(
                io::ErrorKind::UnexpectedEof,
                "unexpected EOF",
            )),
            Ok(n) => {
                unsafe { self.in_buf.advance_mut(n) };
                Ok(())
            }
            Err(e) => Err(e),
        }
    }

    pub fn read_message_timeout(
        &mut self,
        timeout: Duration,
    ) -> io::Result<Option<backend::Message>> {
        if self.in_buf.is_empty() {
            self.set_read_timeout(Some(timeout))?;
            let r = self.read_in();
            self.set_read_timeout(None)?;

            match r {
                Ok(()) => {}
                Err(ref e)
                    if e.kind() == io::ErrorKind::WouldBlock
                        || e.kind() == io::ErrorKind::TimedOut =>
                {
                    return Ok(None)
                }
                Err(e) => return Err(e),
            }
        }

        self.read_message().map(Some)
    }

    pub fn read_message_nonblocking(&mut self) -> io::Result<Option<backend::Message>> {
        if self.in_buf.is_empty() {
            self.set_nonblocking(true)?;
            let r = self.read_in();
            self.set_nonblocking(false)?;

            match r {
                Ok(()) => {}
                Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => return Ok(None),
                Err(e) => return Err(e),
            }
        }

        self.read_message().map(Some)
    }

    pub fn flush(&mut self) -> io::Result<()> {
        self.stream.flush()
    }

    fn set_read_timeout(&self, timeout: Option<Duration>) -> io::Result<()> {
        self.stream.get_ref().get_ref().0.set_read_timeout(timeout)
    }

    fn set_nonblocking(&self, nonblock: bool) -> io::Result<()> {
        self.stream.get_ref().get_ref().0.set_nonblocking(nonblock)
    }
}

/// A connection to the Postgres server.
///
/// It implements `Read`, `Write` and `TlsStream`, as well as `AsRawFd` on
/// Unix platforms and `AsRawSocket` on Windows platforms.
#[derive(Debug)]
pub struct Stream(Socket);

impl Read for Stream {
    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
        self.0.read(buf)
    }
}

impl Write for Stream {
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
        self.0.write(buf)
    }

    fn flush(&mut self) -> io::Result<()> {
        self.0.flush()
    }
}

impl TlsStream for Stream {
    fn get_ref(&self) -> &Stream {
        self
    }

    fn get_mut(&mut self) -> &mut Stream {
        self
    }
}

#[cfg(unix)]
impl AsRawFd for Stream {
    fn as_raw_fd(&self) -> RawFd {
        self.0.as_raw_fd()
    }
}

#[cfg(windows)]
impl AsRawSocket for Stream {
    fn as_raw_socket(&self) -> RawSocket {
        self.0.as_raw_socket()
    }
}

fn open_socket(params: &ConnectParams) -> Result<Socket> {
    let port = params.port();
    match *params.host() {
        Host::Tcp(ref host) => {
            let mut error = None;
            for addr in (&**host, port).to_socket_addrs()? {
                let domain = match addr {
                    SocketAddr::V4(_) => Domain::ipv4(),
                    SocketAddr::V6(_) => Domain::ipv6(),
                };
                let socket = Socket::new(domain, Type::stream(), None)?;
                socket.set_keepalive(params.keepalive())?;
                let addr = SockAddr::from(addr);
                let r = match params.connect_timeout() {
                    Some(timeout) => socket.connect_timeout(&addr, timeout),
                    None => socket.connect(&addr),
                };
                match r {
                    Ok(()) => return Ok(socket),
                    Err(e) => error = Some(e),
                }
            }

            Err(error
                .unwrap_or_else(|| {
                    io::Error::new(
                        io::ErrorKind::InvalidInput,
                        "could not resolve any addresses",
                    )
                })
                .into())
        }
        #[cfg(unix)]
        Host::Unix(ref path) => {
            let path = path.join(&format!(".s.PGSQL.{}", port));
            let socket = Socket::new(Domain::unix(), Type::stream(), None)?;
            let addr = SockAddr::unix(path)?;
            socket.connect(&addr)?;
            Ok(socket)
        }
        #[cfg(not(unix))]
        Host::Unix(..) => Err(io::Error::new(
            io::ErrorKind::InvalidInput,
            "unix sockets are not supported on this system",
        ).into()),
    }
}

pub fn initialize_stream(params: &ConnectParams, tls: TlsMode) -> Result<Box<TlsStream>> {
    let mut socket = Stream(open_socket(params)?);

    let (tls_required, handshaker) = match tls {
        TlsMode::None => return Ok(Box::new(socket)),
        TlsMode::Prefer(handshaker) => (false, handshaker),
        TlsMode::Require(handshaker) => (true, handshaker),
    };

    let mut buf = vec![];
    frontend::ssl_request(&mut buf);
    socket.write_all(&buf)?;
    socket.flush()?;

    let mut b = [0; 1];
    socket.read_exact(&mut b)?;
    if b[0] == b'N' {
        if tls_required {
            return Err(error::tls("the server does not support TLS".into()));
        } else {
            return Ok(Box::new(socket));
        }
    }

    let host = match *params.host() {
        Host::Tcp(ref host) => host,
        // Postgres doesn't support TLS over unix sockets
        Host::Unix(_) => return Err(::bad_response().into()),
    };

    handshaker.tls_handshake(host, socket).map_err(error::tls)
}
<commit_msg>Don't explicity disable keepalive<commit_after>use bytes::{BufMut, BytesMut};
use postgres_protocol::message::backend;
use postgres_protocol::message::frontend;
use socket2::{Domain, SockAddr, Socket, Type};
use std::io::{self, BufWriter, Read, Write};
use std::net::{SocketAddr, ToSocketAddrs};
#[cfg(unix)]
use std::os::unix::io::{AsRawFd, RawFd};
#[cfg(windows)]
use std::os::windows::io::{AsRawSocket, RawSocket};
use std::result;
use std::time::Duration;

use error;
use params::{ConnectParams, Host};
use tls::TlsStream;
use {Result, TlsMode};

const INITIAL_CAPACITY: usize = 8 * 1024;

pub struct MessageStream {
    stream: BufWriter<Box<TlsStream>>,
    in_buf: BytesMut,
    out_buf: Vec<u8>,
}

impl MessageStream {
    pub fn new(stream: Box<TlsStream>) -> MessageStream {
        MessageStream {
            stream: BufWriter::new(stream),
            in_buf: BytesMut::with_capacity(INITIAL_CAPACITY),
            out_buf: vec![],
        }
    }

    pub fn get_ref(&self) -> &Box<TlsStream> {
        self.stream.get_ref()
    }

    pub fn write_message<F, E>(&mut self, f: F) -> result::Result<(), E>
    where
        F: FnOnce(&mut Vec<u8>) -> result::Result<(), E>,
        E: From<io::Error>,
    {
        self.out_buf.clear();
        f(&mut self.out_buf)?;
        self.stream.write_all(&self.out_buf).map_err(From::from)
    }

    pub fn read_message(&mut self) -> io::Result<backend::Message> {
        loop {
            match backend::Message::parse(&mut self.in_buf) {
                Ok(Some(message)) => return Ok(message),
                Ok(None) => self.read_in()?,
                Err(e) => return Err(e),
            }
        }
    }

    fn read_in(&mut self) -> io::Result<()> {
        self.in_buf.reserve(1);
        match self.stream
            .get_mut()
            .read(unsafe { self.in_buf.bytes_mut() })
        {
            Ok(0) => Err(io::Error::new(
                io::ErrorKind::UnexpectedEof,
                "unexpected EOF",
            )),
            Ok(n) => {
                unsafe { self.in_buf.advance_mut(n) };
                Ok(())
            }
            Err(e) => Err(e),
        }
    }

    pub fn read_message_timeout(
        &mut self,
        timeout: Duration,
    ) -> io::Result<Option<backend::Message>> {
        if self.in_buf.is_empty() {
            self.set_read_timeout(Some(timeout))?;
            let r = self.read_in();
            self.set_read_timeout(None)?;

            match r {
                Ok(()) => {}
                Err(ref e)
                    if e.kind() == io::ErrorKind::WouldBlock
                        || e.kind() == io::ErrorKind::TimedOut =>
                {
                    return Ok(None)
                }
                Err(e) => return Err(e),
            }
        }

        self.read_message().map(Some)
    }

    pub fn read_message_nonblocking(&mut self) -> io::Result<Option<backend::Message>> {
        if self.in_buf.is_empty() {
            self.set_nonblocking(true)?;
            let r = self.read_in();
            self.set_nonblocking(false)?;

            match r {
                Ok(()) => {}
                Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => return Ok(None),
                Err(e) => return Err(e),
            }
        }

        self.read_message().map(Some)
    }

    pub fn flush(&mut self) -> io::Result<()> {
        self.stream.flush()
    }

    fn set_read_timeout(&self, timeout: Option<Duration>) -> io::Result<()> {
        self.stream.get_ref().get_ref().0.set_read_timeout(timeout)
    }

    fn set_nonblocking(&self, nonblock: bool) -> io::Result<()> {
        self.stream.get_ref().get_ref().0.set_nonblocking(nonblock)
    }
}

/// A connection to the Postgres server.
///
/// It implements `Read`, `Write` and `TlsStream`, as well as `AsRawFd` on
/// Unix platforms and `AsRawSocket` on Windows platforms.
#[derive(Debug)]
pub struct Stream(Socket);

impl Read for Stream {
    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
        self.0.read(buf)
    }
}

impl Write for Stream {
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
        self.0.write(buf)
    }

    fn flush(&mut self) -> io::Result<()> {
        self.0.flush()
    }
}

impl TlsStream for Stream {
    fn get_ref(&self) -> &Stream {
        self
    }

    fn get_mut(&mut self) -> &mut Stream {
        self
    }
}

#[cfg(unix)]
impl AsRawFd for Stream {
    fn as_raw_fd(&self) -> RawFd {
        self.0.as_raw_fd()
    }
}

#[cfg(windows)]
impl AsRawSocket for Stream {
    fn as_raw_socket(&self) -> RawSocket {
        self.0.as_raw_socket()
    }
}

fn open_socket(params: &ConnectParams) -> Result<Socket> {
    let port = params.port();
    match *params.host() {
        Host::Tcp(ref host) => {
            let mut error = None;
            for addr in (&**host, port).to_socket_addrs()? {
                let domain = match addr {
                    SocketAddr::V4(_) => Domain::ipv4(),
                    SocketAddr::V6(_) => Domain::ipv6(),
                };
                let socket = Socket::new(domain, Type::stream(), None)?;
                if let Some(keepalive) = params.keepalive() {
                    socket.set_keepalive(Some(keepalive))?;
                }
                let addr = SockAddr::from(addr);
                let r = match params.connect_timeout() {
                    Some(timeout) => socket.connect_timeout(&addr, timeout),
                    None => socket.connect(&addr),
                };
                match r {
                    Ok(()) => return Ok(socket),
                    Err(e) => error = Some(e),
                }
            }

            Err(error
                .unwrap_or_else(|| {
                    io::Error::new(
                        io::ErrorKind::InvalidInput,
                        "could not resolve any addresses",
                    )
                })
                .into())
        }
        #[cfg(unix)]
        Host::Unix(ref path) => {
            let path = path.join(&format!(".s.PGSQL.{}", port));
            let socket = Socket::new(Domain::unix(), Type::stream(), None)?;
            let addr = SockAddr::unix(path)?;
            socket.connect(&addr)?;
            Ok(socket)
        }
        #[cfg(not(unix))]
        Host::Unix(..) => Err(io::Error::new(
            io::ErrorKind::InvalidInput,
            "unix sockets are not supported on this system",
        ).into()),
    }
}

pub fn initialize_stream(params: &ConnectParams, tls: TlsMode) -> Result<Box<TlsStream>> {
    let mut socket = Stream(open_socket(params)?);

    let (tls_required, handshaker) = match tls {
        TlsMode::None => return Ok(Box::new(socket)),
        TlsMode::Prefer(handshaker) => (false, handshaker),
        TlsMode::Require(handshaker) => (true, handshaker),
    };

    let mut buf = vec![];
    frontend::ssl_request(&mut buf);
    socket.write_all(&buf)?;
    socket.flush()?;

    let mut b = [0; 1];
    socket.read_exact(&mut b)?;
    if b[0] == b'N' {
        if tls_required {
            return Err(error::tls("the server does not support TLS".into()));
        } else {
            return Ok(Box::new(socket));
        }
    }

    let host = match *params.host() {
        Host::Tcp(ref host) => host,
        // Postgres doesn't support TLS over unix sockets
        Host::Unix(_) => return Err(::bad_response().into()),
    };

    handshaker.tls_handshake(host, socket).map_err(error::tls)
}
<|endoftext|>


<commit_before>macro_rules! run_count_test {
    ( $db:expr, $coll:expr, $filter:expr, $opt:expr, $outcome:expr ) => {{
        let n = $coll.count($filter, $opt).unwrap();
        assert!($outcome.result.int_eq(n));
        check_coll!($db, $coll, $outcome.collection);
    }};
}

macro_rules! run_delete_test {
    ( $db:expr, $coll:expr, $filter:expr, $outcome:expr, $many:expr ) => {{
        let count = if $many {
                        $coll.delete_many($filter, None)
                    } else {
                        $coll.delete_one($filter, None)
                    };

        let expected = count.unwrap().deleted_count;

        let actual = match $outcome.result {
            Bson::Document(ref doc) => doc.get("deletedCount").unwrap(),
            _ => panic!("`delete` test result should be a document")
        };

        assert!(actual.int_eq(expected as i64));
        check_coll!($db, $coll, $outcome.collection);
    }};
}

macro_rules! run_distinct_test {
    ( $db:expr, $coll:expr, $field_name:expr, $filter:expr, $outcome:expr ) => {{
        let actual = $coll.distinct(&$field_name, $filter, None).unwrap();

        let expected = match $outcome.result {
            Bson::Array(ref arr) => arr.clone(),
            _ => panic!("Invalid `result` of find test")
        };

        assert_eq!(actual.len(), expected.len());

        for i in 0..actual.len() {
            assert!(eq::bson_eq(&actual[i], &expected[i]));
        }

        check_coll!($db, $coll, $outcome.collection);
    }};
}

macro_rules! run_find_test {
    ( $db:expr, $coll:expr, $filter:expr, $opt:expr, $outcome:expr ) => {{
        let mut cursor = $coll.find($filter, $opt).unwrap();

        let array = match $outcome.result {
            Bson::Array(ref arr) => arr.clone(),
            _ => panic!("Invalid `result` of find test")
        };

        for bson in array {
            assert!(eq::bson_eq(&bson, &Bson::Document(cursor.next().unwrap().unwrap())));
        }

        assert!(!cursor.has_next().ok().expect("Failed to execute 'has_next()' on cursor."));
        check_coll!($db, $coll, $outcome.collection);
    }};
}

macro_rules! run_insert_many_test {
    ( $db:expr, $coll:expr, $docs:expr, $outcome:expr ) => {{
        let inserted = $coll.insert_many($docs, true, None).unwrap().inserted_ids.unwrap();
        let ids_bson = match $outcome.result {
            Bson::Document(ref doc) => doc.get("insertedIds").unwrap(),
            _ => panic!("`insert_many` test result should be a document")
        };

        let ids = match ids_bson {
            &Bson::Array(ref arr) => arr.into_iter(),
            _ => panic!("`insertedIds` test result should be an array")
        };

        let mut actual_ids = inserted.values();

        for expected_id in ids {
            assert!(eq::bson_eq(&expected_id, actual_ids.next().unwrap()));
        }

        check_coll!($db, $coll, $outcome.collection);
    }};
}

macro_rules! run_insert_one_test {
    ( $db:expr, $coll:expr, $doc:expr, $outcome:expr ) => {{
        let inserted = $coll.insert_one($doc, None).unwrap().inserted_id.unwrap();
        let id = match $outcome.result {
            Bson::Document(ref doc) => doc.get("insertedId").unwrap(),
            _ => panic!("`insert_one` test result should be a document")
        };

        assert!(eq::bson_eq(&id, &inserted));
        check_coll!($db, $coll, $outcome.collection);
    }};
}

macro_rules! run_replace_one_test {
    ( $db:expr, $coll:expr, $filter:expr, $replacement:expr, $upsert:expr,
        $outcome:expr ) => {{
            let actual = $coll.replace_one($filter, $replacement, $upsert,
                                           None).unwrap();

            let (matched, modified, upserted) = match $outcome.result {
                Bson::Document(ref doc) => (
                    doc.get("matchedCount").unwrap(),
                    doc.get("modifiedCount").unwrap(),
                    doc.get("upsertedId"),
                    ),
                _ => panic!("`delete` test result should be a document")
            };

            assert!(matched.int_eq(actual.matched_count as i64));
            assert!(modified.int_eq(actual.modified_count as i64));

            let id = match actual.upserted_id {
                Some(Bson::Document(ref doc)) => doc.get("_id"),
                _ => None
            };

            match (upserted, id) {
                (None, None) => (),
                (Some(ref bson1), Some(ref bson2)) =>
                assert!(eq::bson_eq(&bson1, &bson2)),
                _ => panic!("Wrong `upsertedId` returned")
            };

            check_coll!($db, $coll, $outcome.collection);
    }};
}

macro_rules! run_update_test {
    ( $db:expr, $coll:expr, $filter:expr, $update:expr, $upsert:expr,
      $many:expr, $outcome:expr ) => {{
          let result = if $many {
                           $coll.update_many($filter, $update, $upsert, None)
                       } else {
                           $coll.update_one($filter, $update, $upsert, None)
                       };

          let actual = result.unwrap();

          let (matched, modified, upserted) = match $outcome.result {
              Bson::Document(ref doc) => (
                  doc.get("matchedCount").unwrap(),
                  doc.get("modifiedCount").unwrap(),
                  doc.get("upsertedId"),
              ),
              _ => panic!("`update` test result should be a document")
          };

          assert!(matched.int_eq(actual.matched_count as i64));
          assert!(modified.int_eq(actual.modified_count as i64));

          let id = match actual.upserted_id {
	          Some(Bson::Document(ref doc)) => doc.get("_id"),
              _ => None
          };

          match (upserted, id) {
              (None, None) => (),
              (Some(ref bson1), Some(ref bson2)) =>
                  assert!(eq::bson_eq(&bson1, &bson2)),
              _ => panic!("Wrong `upsertedId` returned")
          };

          check_coll!($db, $coll, $outcome.collection);
    }};
}

#[macro_export]
macro_rules! run_suite {
    ( $file:expr, $coll:expr ) => {{
        let json = Json::from_file($file).unwrap();
        let suite = json.get_suite().unwrap();
        let client =  MongoClient::new("localhost", 27017).unwrap();
        let db = client.db("test");
        let coll = db.collection($coll);

        for test in suite.tests {
            coll.drop().unwrap();
            coll.insert_many(suite.data.clone(), true, None).unwrap();

            match test.operation {
                Arguments::Count { filter, options } =>
                    run_count_test!(db, coll, filter, Some(options), test.outcome),
                Arguments::Delete { filter, many } =>
                    run_delete_test!(db, coll, filter, test.outcome, many),
                Arguments::Distinct { field_name, filter } =>
                    run_distinct_test!(db, coll, field_name, filter, test.outcome),
                Arguments::Find { filter, options } =>
                    run_find_test!(db, coll, filter, Some(options), test.outcome),
                Arguments::InsertMany { documents } =>
                    run_insert_many_test!(db, coll, documents, test.outcome),
                Arguments::InsertOne { document } =>
                    run_insert_one_test!(db, coll, document, test.outcome),
                Arguments::ReplaceOne { filter, replacement, upsert } =>
                    run_replace_one_test!(db, coll, filter, replacement, upsert,
                                          test.outcome),
                Arguments::Update { filter, update, upsert, many } =>
                    run_update_test!(db, coll, filter, update, upsert, many,
                                     test.outcome),
            };
        }
    }};
}

#[macro_export]
macro_rules! check_coll {
    ( $db:expr, $coll:expr, $coll_opt:expr ) => {{
        let outcome_coll = match $coll_opt {
            Some(ref coll) => coll.clone(),
            None => continue
        };

        let coll = match outcome_coll.name {
            Some(ref str) => $db.collection(&str),
            None => $db.collection(&$coll.name())
        };

        let mut cursor = coll.find(None, None).unwrap();

        for doc in outcome_coll.data.iter() {
            assert!(eq::bson_eq(&Bson::Document(doc.clone()),
                                &Bson::Document(cursor.next().unwrap().unwrap())));
        }

        assert!(!cursor.has_next().ok().expect("Failed to execute 'has_next()' on cursor."));
    }};
}
<commit_msg>fixed mistaken error message<commit_after>macro_rules! run_count_test {
    ( $db:expr, $coll:expr, $filter:expr, $opt:expr, $outcome:expr ) => {{
        let n = $coll.count($filter, $opt).unwrap();
        assert!($outcome.result.int_eq(n));
        check_coll!($db, $coll, $outcome.collection);
    }};
}

macro_rules! run_delete_test {
    ( $db:expr, $coll:expr, $filter:expr, $outcome:expr, $many:expr ) => {{
        let count = if $many {
                        $coll.delete_many($filter, None)
                    } else {
                        $coll.delete_one($filter, None)
                    };

        let expected = count.unwrap().deleted_count;

        let actual = match $outcome.result {
            Bson::Document(ref doc) => doc.get("deletedCount").unwrap(),
            _ => panic!("`delete` test result should be a document")
        };

        assert!(actual.int_eq(expected as i64));
        check_coll!($db, $coll, $outcome.collection);
    }};
}

macro_rules! run_distinct_test {
    ( $db:expr, $coll:expr, $field_name:expr, $filter:expr, $outcome:expr ) => {{
        let actual = $coll.distinct(&$field_name, $filter, None).unwrap();

        let expected = match $outcome.result {
            Bson::Array(ref arr) => arr.clone(),
            _ => panic!("Invalid `result` of distinct test")
        };

        assert_eq!(actual.len(), expected.len());

        for i in 0..actual.len() {
            assert!(eq::bson_eq(&actual[i], &expected[i]));
        }

        check_coll!($db, $coll, $outcome.collection);
    }};
}

macro_rules! run_find_test {
    ( $db:expr, $coll:expr, $filter:expr, $opt:expr, $outcome:expr ) => {{
        let mut cursor = $coll.find($filter, $opt).unwrap();

        let array = match $outcome.result {
            Bson::Array(ref arr) => arr.clone(),
            _ => panic!("Invalid `result` of find test")
        };

        for bson in array {
            assert!(eq::bson_eq(&bson, &Bson::Document(cursor.next().unwrap().unwrap())));
        }

        assert!(!cursor.has_next().ok().expect("Failed to execute 'has_next()' on cursor."));
        check_coll!($db, $coll, $outcome.collection);
    }};
}

macro_rules! run_insert_many_test {
    ( $db:expr, $coll:expr, $docs:expr, $outcome:expr ) => {{
        let inserted = $coll.insert_many($docs, true, None).unwrap().inserted_ids.unwrap();
        let ids_bson = match $outcome.result {
            Bson::Document(ref doc) => doc.get("insertedIds").unwrap(),
            _ => panic!("`insert_many` test result should be a document")
        };

        let ids = match ids_bson {
            &Bson::Array(ref arr) => arr.into_iter(),
            _ => panic!("`insertedIds` test result should be an array")
        };

        let mut actual_ids = inserted.values();

        for expected_id in ids {
            assert!(eq::bson_eq(&expected_id, actual_ids.next().unwrap()));
        }

        check_coll!($db, $coll, $outcome.collection);
    }};
}

macro_rules! run_insert_one_test {
    ( $db:expr, $coll:expr, $doc:expr, $outcome:expr ) => {{
        let inserted = $coll.insert_one($doc, None).unwrap().inserted_id.unwrap();
        let id = match $outcome.result {
            Bson::Document(ref doc) => doc.get("insertedId").unwrap(),
            _ => panic!("`insert_one` test result should be a document")
        };

        assert!(eq::bson_eq(&id, &inserted));
        check_coll!($db, $coll, $outcome.collection);
    }};
}

macro_rules! run_replace_one_test {
    ( $db:expr, $coll:expr, $filter:expr, $replacement:expr, $upsert:expr,
        $outcome:expr ) => {{
            let actual = $coll.replace_one($filter, $replacement, $upsert,
                                           None).unwrap();

            let (matched, modified, upserted) = match $outcome.result {
                Bson::Document(ref doc) => (
                    doc.get("matchedCount").unwrap(),
                    doc.get("modifiedCount").unwrap(),
                    doc.get("upsertedId"),
                    ),
                _ => panic!("`delete` test result should be a document")
            };

            assert!(matched.int_eq(actual.matched_count as i64));
            assert!(modified.int_eq(actual.modified_count as i64));

            let id = match actual.upserted_id {
                Some(Bson::Document(ref doc)) => doc.get("_id"),
                _ => None
            };

            match (upserted, id) {
                (None, None) => (),
                (Some(ref bson1), Some(ref bson2)) =>
                assert!(eq::bson_eq(&bson1, &bson2)),
                _ => panic!("Wrong `upsertedId` returned")
            };

            check_coll!($db, $coll, $outcome.collection);
    }};
}

macro_rules! run_update_test {
    ( $db:expr, $coll:expr, $filter:expr, $update:expr, $upsert:expr,
      $many:expr, $outcome:expr ) => {{
          let result = if $many {
                           $coll.update_many($filter, $update, $upsert, None)
                       } else {
                           $coll.update_one($filter, $update, $upsert, None)
                       };

          let actual = result.unwrap();

          let (matched, modified, upserted) = match $outcome.result {
              Bson::Document(ref doc) => (
                  doc.get("matchedCount").unwrap(),
                  doc.get("modifiedCount").unwrap(),
                  doc.get("upsertedId"),
              ),
              _ => panic!("`update` test result should be a document")
          };

          assert!(matched.int_eq(actual.matched_count as i64));
          assert!(modified.int_eq(actual.modified_count as i64));

          let id = match actual.upserted_id {
	          Some(Bson::Document(ref doc)) => doc.get("_id"),
              _ => None
          };

          match (upserted, id) {
              (None, None) => (),
              (Some(ref bson1), Some(ref bson2)) =>
                  assert!(eq::bson_eq(&bson1, &bson2)),
              _ => panic!("Wrong `upsertedId` returned")
          };

          check_coll!($db, $coll, $outcome.collection);
    }};
}

#[macro_export]
macro_rules! run_suite {
    ( $file:expr, $coll:expr ) => {{
        let json = Json::from_file($file).unwrap();
        let suite = json.get_suite().unwrap();
        let client =  MongoClient::new("localhost", 27017).unwrap();
        let db = client.db("test");
        let coll = db.collection($coll);

        for test in suite.tests {
            coll.drop().unwrap();
            coll.insert_many(suite.data.clone(), true, None).unwrap();

            match test.operation {
                Arguments::Count { filter, options } =>
                    run_count_test!(db, coll, filter, Some(options), test.outcome),
                Arguments::Delete { filter, many } =>
                    run_delete_test!(db, coll, filter, test.outcome, many),
                Arguments::Distinct { field_name, filter } =>
                    run_distinct_test!(db, coll, field_name, filter, test.outcome),
                Arguments::Find { filter, options } =>
                    run_find_test!(db, coll, filter, Some(options), test.outcome),
                Arguments::InsertMany { documents } =>
                    run_insert_many_test!(db, coll, documents, test.outcome),
                Arguments::InsertOne { document } =>
                    run_insert_one_test!(db, coll, document, test.outcome),
                Arguments::ReplaceOne { filter, replacement, upsert } =>
                    run_replace_one_test!(db, coll, filter, replacement, upsert,
                                          test.outcome),
                Arguments::Update { filter, update, upsert, many } =>
                    run_update_test!(db, coll, filter, update, upsert, many,
                                     test.outcome),
            };
        }
    }};
}

#[macro_export]
macro_rules! check_coll {
    ( $db:expr, $coll:expr, $coll_opt:expr ) => {{
        let outcome_coll = match $coll_opt {
            Some(ref coll) => coll.clone(),
            None => continue
        };

        let coll = match outcome_coll.name {
            Some(ref str) => $db.collection(&str),
            None => $db.collection(&$coll.name())
        };

        let mut cursor = coll.find(None, None).unwrap();

        for doc in outcome_coll.data.iter() {
            assert!(eq::bson_eq(&Bson::Document(doc.clone()),
                                &Bson::Document(cursor.next().unwrap().unwrap())));
        }

        assert!(!cursor.has_next().ok().expect("Failed to execute 'has_next()' on cursor."));
    }};
}
<|endoftext|>