<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 std::{error::Error, fmt, io, str, sync};

pub type StratisResult<T> = Result<T, StratisError>;

#[derive(Debug, Clone)]
pub enum ErrorEnum {
    Error,

    AlreadyExists,
    Busy,
    Invalid,
    NotFound,
}

#[derive(Debug)]
pub enum StratisError {
    Error(String),
    Engine(ErrorEnum, String),
    Io(io::Error),
    Nix(nix::Error),
    Uuid(uuid::Error),
    Utf8(str::Utf8Error),
    Serde(serde_json::error::Error),
    Decode(data_encoding::DecodeError),
    DM(devicemapper::DmError),
    Crypt(libcryptsetup_rs::LibcryptErr),
    Recv(sync::mpsc::RecvError),
    Null(std::ffi::NulError),
    Join(tokio::task::JoinError),

    #[cfg(feature = "dbus_enabled")]
    Dbus(dbus::Error),
    Udev(libudev::Error),
}

impl fmt::Display for StratisError {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match *self {
            StratisError::Error(ref s) => write!(f, "Error: {}", s),
            StratisError::Engine(_, ref msg) => write!(f, "Engine error: {}", msg),
            StratisError::Io(ref err) => write!(f, "IO error: {}", err),
            StratisError::Nix(ref err) => write!(f, "Nix error: {}", err),
            StratisError::Uuid(ref err) => write!(f, "Uuid error: {}", err),
            StratisError::Utf8(ref err) => write!(f, "Utf8 error: {}", err),
            StratisError::Serde(ref err) => write!(f, "Serde error: {}", err),
            StratisError::Decode(ref err) => write!(f, "Data encoding error: {}", err),
            StratisError::DM(ref err) => write!(f, "DM error: {}", err),
            StratisError::Crypt(ref err) => write!(f, "Cryptsetup error: {}", err),
            StratisError::Recv(ref err) => write!(f, "Synchronization channel error: {}", err),
            StratisError::Null(ref err) => write!(f, "C string conversion error: {}", err),
            StratisError::Join(ref err) => write!(f, "Failed to join thread: {}", err),

            #[cfg(feature = "dbus_enabled")]
            StratisError::Dbus(ref err) => {
                write!(f, "Dbus error: {}", err.message().unwrap_or("Unknown"))
            }
            StratisError::Udev(ref err) => write!(f, "Udev error: {}", err),
        }
    }
}

impl Error for StratisError {
    fn cause(&self) -> Option<&dyn Error> {
        match *self {
            StratisError::Error(_) | StratisError::Engine(_, _) => None,
            StratisError::Io(ref err) => Some(err),
            StratisError::Nix(ref err) => Some(err),
            StratisError::Uuid(ref err) => Some(err),
            StratisError::Utf8(ref err) => Some(err),
            StratisError::Serde(ref err) => Some(err),
            StratisError::Decode(ref err) => Some(err),
            StratisError::DM(ref err) => Some(err),
            StratisError::Crypt(ref err) => Some(err),
            StratisError::Recv(ref err) => Some(err),
            StratisError::Null(ref err) => Some(err),
            StratisError::Join(ref err) => Some(err),

            #[cfg(feature = "dbus_enabled")]
            StratisError::Dbus(ref err) => Some(err),
            StratisError::Udev(ref err) => Some(err),
        }
    }
}

impl From<tokio::task::JoinError> for StratisError {
    fn from(err: tokio::task::JoinError) -> StratisError {
        StratisError::Join(err)
    }
}

impl From<std::ffi::NulError> for StratisError {
    fn from(err: std::ffi::NulError) -> StratisError {
        StratisError::Null(err)
    }
}

impl From<io::Error> for StratisError {
    fn from(err: io::Error) -> StratisError {
        StratisError::Io(err)
    }
}

impl From<nix::Error> for StratisError {
    fn from(err: nix::Error) -> StratisError {
        StratisError::Nix(err)
    }
}

impl From<uuid::Error> for StratisError {
    fn from(err: uuid::Error) -> StratisError {
        StratisError::Uuid(err)
    }
}

impl From<str::Utf8Error> for StratisError {
    fn from(err: str::Utf8Error) -> StratisError {
        StratisError::Utf8(err)
    }
}

impl From<serde_json::error::Error> for StratisError {
    fn from(err: serde_json::error::Error) -> StratisError {
        StratisError::Serde(err)
    }
}

impl From<data_encoding::DecodeError> for StratisError {
    fn from(err: data_encoding::DecodeError) -> StratisError {
        StratisError::Decode(err)
    }
}

impl From<devicemapper::DmError> for StratisError {
    fn from(err: devicemapper::DmError) -> StratisError {
        StratisError::DM(err)
    }
}

impl From<libcryptsetup_rs::LibcryptErr> for StratisError {
    fn from(err: libcryptsetup_rs::LibcryptErr) -> StratisError {
        StratisError::Crypt(err)
    }
}

#[cfg(feature = "dbus_enabled")]
impl From<dbus::Error> for StratisError {
    fn from(err: dbus::Error) -> StratisError {
        StratisError::Dbus(err)
    }
}

impl From<libudev::Error> for StratisError {
    fn from(err: libudev::Error) -> StratisError {
        StratisError::Udev(err)
    }
}

impl<T> From<sync::PoisonError<T>> for StratisError {
    fn from(err: sync::PoisonError<T>) -> StratisError {
        StratisError::Error(err.to_string())
    }
}

impl From<sync::mpsc::RecvError> for StratisError {
    fn from(err: sync::mpsc::RecvError) -> StratisError {
        StratisError::Recv(err)
    }
}
<commit_msg>Drop implementation of From for PoisonError<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 std::{error::Error, fmt, io, str, sync};

pub type StratisResult<T> = Result<T, StratisError>;

#[derive(Debug, Clone)]
pub enum ErrorEnum {
    Error,

    AlreadyExists,
    Busy,
    Invalid,
    NotFound,
}

#[derive(Debug)]
pub enum StratisError {
    Error(String),
    Engine(ErrorEnum, String),
    Io(io::Error),
    Nix(nix::Error),
    Uuid(uuid::Error),
    Utf8(str::Utf8Error),
    Serde(serde_json::error::Error),
    Decode(data_encoding::DecodeError),
    DM(devicemapper::DmError),
    Crypt(libcryptsetup_rs::LibcryptErr),
    Recv(sync::mpsc::RecvError),
    Null(std::ffi::NulError),
    Join(tokio::task::JoinError),

    #[cfg(feature = "dbus_enabled")]
    Dbus(dbus::Error),
    Udev(libudev::Error),
}

impl fmt::Display for StratisError {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match *self {
            StratisError::Error(ref s) => write!(f, "Error: {}", s),
            StratisError::Engine(_, ref msg) => write!(f, "Engine error: {}", msg),
            StratisError::Io(ref err) => write!(f, "IO error: {}", err),
            StratisError::Nix(ref err) => write!(f, "Nix error: {}", err),
            StratisError::Uuid(ref err) => write!(f, "Uuid error: {}", err),
            StratisError::Utf8(ref err) => write!(f, "Utf8 error: {}", err),
            StratisError::Serde(ref err) => write!(f, "Serde error: {}", err),
            StratisError::Decode(ref err) => write!(f, "Data encoding error: {}", err),
            StratisError::DM(ref err) => write!(f, "DM error: {}", err),
            StratisError::Crypt(ref err) => write!(f, "Cryptsetup error: {}", err),
            StratisError::Recv(ref err) => write!(f, "Synchronization channel error: {}", err),
            StratisError::Null(ref err) => write!(f, "C string conversion error: {}", err),
            StratisError::Join(ref err) => write!(f, "Failed to join thread: {}", err),

            #[cfg(feature = "dbus_enabled")]
            StratisError::Dbus(ref err) => {
                write!(f, "Dbus error: {}", err.message().unwrap_or("Unknown"))
            }
            StratisError::Udev(ref err) => write!(f, "Udev error: {}", err),
        }
    }
}

impl Error for StratisError {
    fn cause(&self) -> Option<&dyn Error> {
        match *self {
            StratisError::Error(_) | StratisError::Engine(_, _) => None,
            StratisError::Io(ref err) => Some(err),
            StratisError::Nix(ref err) => Some(err),
            StratisError::Uuid(ref err) => Some(err),
            StratisError::Utf8(ref err) => Some(err),
            StratisError::Serde(ref err) => Some(err),
            StratisError::Decode(ref err) => Some(err),
            StratisError::DM(ref err) => Some(err),
            StratisError::Crypt(ref err) => Some(err),
            StratisError::Recv(ref err) => Some(err),
            StratisError::Null(ref err) => Some(err),
            StratisError::Join(ref err) => Some(err),

            #[cfg(feature = "dbus_enabled")]
            StratisError::Dbus(ref err) => Some(err),
            StratisError::Udev(ref err) => Some(err),
        }
    }
}

impl From<tokio::task::JoinError> for StratisError {
    fn from(err: tokio::task::JoinError) -> StratisError {
        StratisError::Join(err)
    }
}

impl From<std::ffi::NulError> for StratisError {
    fn from(err: std::ffi::NulError) -> StratisError {
        StratisError::Null(err)
    }
}

impl From<io::Error> for StratisError {
    fn from(err: io::Error) -> StratisError {
        StratisError::Io(err)
    }
}

impl From<nix::Error> for StratisError {
    fn from(err: nix::Error) -> StratisError {
        StratisError::Nix(err)
    }
}

impl From<uuid::Error> for StratisError {
    fn from(err: uuid::Error) -> StratisError {
        StratisError::Uuid(err)
    }
}

impl From<str::Utf8Error> for StratisError {
    fn from(err: str::Utf8Error) -> StratisError {
        StratisError::Utf8(err)
    }
}

impl From<serde_json::error::Error> for StratisError {
    fn from(err: serde_json::error::Error) -> StratisError {
        StratisError::Serde(err)
    }
}

impl From<data_encoding::DecodeError> for StratisError {
    fn from(err: data_encoding::DecodeError) -> StratisError {
        StratisError::Decode(err)
    }
}

impl From<devicemapper::DmError> for StratisError {
    fn from(err: devicemapper::DmError) -> StratisError {
        StratisError::DM(err)
    }
}

impl From<libcryptsetup_rs::LibcryptErr> for StratisError {
    fn from(err: libcryptsetup_rs::LibcryptErr) -> StratisError {
        StratisError::Crypt(err)
    }
}

#[cfg(feature = "dbus_enabled")]
impl From<dbus::Error> for StratisError {
    fn from(err: dbus::Error) -> StratisError {
        StratisError::Dbus(err)
    }
}

impl From<libudev::Error> for StratisError {
    fn from(err: libudev::Error) -> StratisError {
        StratisError::Udev(err)
    }
}

impl From<sync::mpsc::RecvError> for StratisError {
    fn from(err: sync::mpsc::RecvError) -> StratisError {
        StratisError::Recv(err)
    }
}
<|endoftext|>


<commit_before>// Copyright 2015 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

use std::ops::{Add, Sub};

type Test = Add +
            //~^ ERROR the type parameter `RHS` must be explicitly specified in an object type because its default value `Self` references the type `Self`
            Sub;
            //~^ ERROR only the builtin traits can be used as closure or object bounds

fn main() {
}
<commit_msg>Add `ignore-tidy-linelength` and fix formatting<commit_after>// Copyright 2015 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

// ignore-tidy-linelength

use std::ops::{Add, Sub};

type Test = Add +
            //~^ ERROR the type parameter `RHS` must be explicitly specified in an object type because its default value `Self` references the type `Self`
            Sub;
            //~^ ERROR only the builtin traits can be used as closure or object bounds

fn main() { }
<|endoftext|>


<commit_before>pub extern crate berg_compiler;
pub extern crate num;
pub use crate::compiler_test::berg_compiler::test::expect;
pub use crate::compiler_test::berg_compiler::ErrorCode::*;
pub use crate::compiler_test::berg_compiler::*;
pub use crate::compiler_test::num::BigRational;
pub use std::str::FromStr;
<commit_msg>Use edition 2018 idioms<commit_after>pub use berg_compiler;
pub use num;
pub use crate::compiler_test::berg_compiler::test::expect;
pub use crate::compiler_test::berg_compiler::ErrorCode::*;
pub use crate::compiler_test::berg_compiler::*;
pub use crate::compiler_test::num::BigRational;
pub use std::str::FromStr;
<|endoftext|>


<commit_before>// `cargo test --benches` and `#[feature(test)]` work only in nightly
#![cfg(rustc_nightly)]
#![cfg(feature = "bytes")]
#![feature(test)]

extern crate bytes;
extern crate test;

extern crate protobuf;

extern crate perftest_bytes;

use std::fmt::Write;

use bytes::Bytes;

use perftest_bytes::messages;
use protobuf::Message;

fn make_string_of_len(len: usize) -> String {
    let mut s = String::new();
    for i in 0..len {
        write!(s, "{}", i % 10).expect("unreachable");
    }
    s
}

fn make_repeated(len: usize) -> Vec<u8> {
    let mut m = messages::TestMessage::new();

    for i in 0..100 {
        m.sr.push(make_string_of_len(i % len));
        m.br.push(make_string_of_len((i + len / 2) % len).into_bytes());
    }

    m.write_to_bytes().expect("write")
}

#[bench]
fn parse_repeated_small_regular(b: &mut test::Bencher) {
    let bs = make_repeated(30);
    b.iter(|| messages::TestMessage::parse_from_bytes(&bs).expect("parse"))
}

#[bench]
fn parse_repeated_small_bytes(b: &mut test::Bencher) {
    let bs = Bytes::from(make_repeated(30));
    b.iter(|| messages::TestMessageWithBytes::parse_from_carllerche_bytes(&bs).expect("parse"))
}

#[bench]
fn parse_repeated_medium_regular(b: &mut test::Bencher) {
    let bs = make_repeated(300);
    b.iter(|| protobuf::parse_from_bytes::<messages::TestMessage>(&bs).expect("parse"))
}

#[bench]
fn parse_repeated_medium_bytes(b: &mut test::Bencher) {
    let bs = Bytes::from(make_repeated(300));
    b.iter(|| messages::TestMessageWithBytes::parse_from_carllerche_bytes(&bs).expect("parse"))
}

#[bench]
fn parse_repeated_large_regular(b: &mut test::Bencher) {
    let bs = make_repeated(3000);
    b.iter(|| messages::TestMessage::parse_from_bytes(&bs).expect("parse"))
}

#[bench]
fn parse_repeated_large_bytes(b: &mut test::Bencher) {
    let bs = Bytes::from(make_repeated(3000));
    b.iter(|| messages::TestMessageWithBytes::parse_from_carllerche_bytes(&bs).expect("parse"))
}

#[bench]
fn parse_repeated_huge_regular(b: &mut test::Bencher) {
    let bs = make_repeated(30000);
    b.iter(|| messages::TestMessage::parse_from_bytes(&bs).expect("parse"))
}

#[bench]
fn parse_repeated_huge_bytes(b: &mut test::Bencher) {
    let bs = Bytes::from(make_repeated(30000));
    b.iter(|| messages::TestMessageWithBytes::parse_from_carllerche_bytes(&bs).expect("parse"))
}
<commit_msg>Fix bench compilation<commit_after>// `cargo test --benches` and `#[feature(test)]` work only in nightly
#![cfg(rustc_nightly)]
#![cfg(feature = "bytes")]
#![feature(test)]

extern crate bytes;
extern crate test;

extern crate protobuf;

extern crate perftest_bytes;

use std::fmt::Write;

use bytes::Bytes;

use perftest_bytes::messages;
use protobuf::Message;

fn make_string_of_len(len: usize) -> String {
    let mut s = String::new();
    for i in 0..len {
        write!(s, "{}", i % 10).expect("unreachable");
    }
    s
}

fn make_repeated(len: usize) -> Vec<u8> {
    let mut m = messages::TestMessage::new();

    for i in 0..100 {
        m.sr.push(make_string_of_len(i % len));
        m.br.push(make_string_of_len((i + len / 2) % len).into_bytes());
    }

    m.write_to_bytes().expect("write")
}

#[bench]
fn parse_repeated_small_regular(b: &mut test::Bencher) {
    let bs = make_repeated(30);
    b.iter(|| messages::TestMessage::parse_from_bytes(&bs).expect("parse"))
}

#[bench]
fn parse_repeated_small_bytes(b: &mut test::Bencher) {
    let bs = Bytes::from(make_repeated(30));
    b.iter(|| messages::TestMessageWithBytes::parse_from_carllerche_bytes(&bs).expect("parse"))
}

#[bench]
fn parse_repeated_medium_regular(b: &mut test::Bencher) {
    let bs = make_repeated(300);
    b.iter(|| messages::TestMessage::parse_from_bytes(&bs).expect("parse"))
}

#[bench]
fn parse_repeated_medium_bytes(b: &mut test::Bencher) {
    let bs = Bytes::from(make_repeated(300));
    b.iter(|| messages::TestMessageWithBytes::parse_from_carllerche_bytes(&bs).expect("parse"))
}

#[bench]
fn parse_repeated_large_regular(b: &mut test::Bencher) {
    let bs = make_repeated(3000);
    b.iter(|| messages::TestMessage::parse_from_bytes(&bs).expect("parse"))
}

#[bench]
fn parse_repeated_large_bytes(b: &mut test::Bencher) {
    let bs = Bytes::from(make_repeated(3000));
    b.iter(|| messages::TestMessageWithBytes::parse_from_carllerche_bytes(&bs).expect("parse"))
}

#[bench]
fn parse_repeated_huge_regular(b: &mut test::Bencher) {
    let bs = make_repeated(30000);
    b.iter(|| messages::TestMessage::parse_from_bytes(&bs).expect("parse"))
}

#[bench]
fn parse_repeated_huge_bytes(b: &mut test::Bencher) {
    let bs = Bytes::from(make_repeated(30000));
    b.iter(|| messages::TestMessageWithBytes::parse_from_carllerche_bytes(&bs).expect("parse"))
}
<|endoftext|>


<commit_before>//! TCP connection stream for local server with remote (proxy server or remote target)

use std::{
    fmt::{self, Display, Formatter},
    io::{self, Error},
    net::SocketAddr,
    pin::Pin,
    task::{Context as TaskContext, Poll},
    time::Duration,
};

use bytes::BytesMut;
use log::{debug, error, trace};
use pin_project::pin_project;
use tokio::io::{AsyncRead, AsyncWrite, AsyncWriteExt, ReadHalf, WriteHalf};

use crate::{
    config::{ConfigType, ServerAddr, ServerConfig},
    context::{Context, SharedContext},
    relay::{socks5::Address, sys::tcp_stream_connect, utils::try_timeout},
};

use super::{connection::Connection, CryptoStream, STcpStream};

/// Stream wrapper for both direct connections and proxied connections
#[allow(clippy::large_enum_variant)]
#[pin_project]
pub enum ProxyStream {
    Direct {
        #[pin]
        stream: STcpStream,
        context: SharedContext,
    },
    Proxied {
        #[pin]
        stream: CryptoStream<STcpStream>,
        context: SharedContext,
    },
}

#[derive(Debug)]
pub struct ProxyStreamError {
    inner: Error,
    bypassed: bool,
}

impl ProxyStreamError {
    fn new(inner: Error, bypassed: bool) -> ProxyStreamError {
        ProxyStreamError { inner, bypassed }
    }

    /// Check if it is proxied
    pub fn is_proxied(&self) -> bool {
        self.bypassed
    }

    /// Into internal `std::io::Error`
    pub fn into_inner(self) -> Error {
        self.inner
    }
}

impl From<ProxyStreamError> for Error {
    fn from(err: ProxyStreamError) -> Error {
        err.inner
    }
}

impl Display for ProxyStreamError {
    fn fmt(&self, f: &mut Formatter) -> fmt::Result {
        Display::fmt(&self.inner, f)
    }
}

impl ProxyStream {
    /// Connect to remote by ACL rules
    pub async fn connect(
        context: SharedContext,
        svr_cfg: &ServerConfig,
        addr: &Address,
    ) -> Result<ProxyStream, ProxyStreamError> {
        if context.check_target_bypassed(addr).await {
            ProxyStream::connect_direct_wrapped(context, addr).await
        } else {
            ProxyStream::connect_proxied_wrapped(context, svr_cfg, addr).await
        }
    }

    /// Connect to remote directly (without proxy)
    ///
    /// This is used for hosts that matches ACL bypassed rules
    pub async fn connect_direct(context: SharedContext, addr: &Address) -> io::Result<ProxyStream> {
        debug!("connect to {} directly (bypassed)", addr);

        // NOTE: Direct connection's timeout is controlled by the global key
        let timeout = context.config().timeout;

        let stream = match *addr {
            Address::SocketAddress(ref saddr) => try_timeout(tcp_stream_connect(&saddr, &context), timeout).await?,
            Address::DomainNameAddress(ref domain, port) => {
                lookup_then!(context, domain, port, |saddr| {
                    try_timeout(tcp_stream_connect(&saddr, &context), timeout).await
                })?
                .1
            }
        };

        Ok(ProxyStream::Direct {
            stream: Connection::new(stream, timeout),
            context,
        })
    }

    async fn connect_direct_wrapped(context: SharedContext, addr: &Address) -> Result<ProxyStream, ProxyStreamError> {
        match ProxyStream::connect_direct(context, addr).await {
            Ok(s) => Ok(s),
            Err(err) => Err(ProxyStreamError::new(err, true)),
        }
    }

    /// Connect to remote via proxy server
    ///
    /// This is used for hosts that matches ACL proxied rules
    pub async fn connect_proxied(
        context: SharedContext,
        svr_cfg: &ServerConfig,
        addr: &Address,
    ) -> io::Result<ProxyStream> {
        debug!(
            "connect to {} via {} ({}) (proxied)",
            addr,
            svr_cfg.addr(),
            svr_cfg.external_addr()
        );

        let server_stream = connect_proxy_server(&context, svr_cfg).await?;
        let proxy_stream = proxy_server_handshake(context.clone(), server_stream, svr_cfg, addr).await?;

        Ok(ProxyStream::Proxied {
            stream: proxy_stream,
            context,
        })
    }

    async fn connect_proxied_wrapped(
        context: SharedContext,
        svr_cfg: &ServerConfig,
        addr: &Address,
    ) -> Result<ProxyStream, ProxyStreamError> {
        match ProxyStream::connect_proxied(context, svr_cfg, addr).await {
            Ok(s) => Ok(s),
            Err(err) => Err(ProxyStreamError::new(err, false)),
        }
    }

    /// Split into reader and writer
    pub fn split(self) -> (ReadHalf<ProxyStream>, WriteHalf<ProxyStream>) {
        use tokio::io::split;
        split(self)
    }

    /// Returns the local socket address of this stream socket
    pub fn local_addr(&self) -> io::Result<SocketAddr> {
        match *self {
            ProxyStream::Direct { ref stream, .. } => stream.get_ref().local_addr(),
            ProxyStream::Proxied { ref stream, .. } => stream.get_ref().get_ref().local_addr(),
        }
    }

    /// Check if the underlying connection is proxied
    pub fn is_proxied(&self) -> bool {
        match *self {
            ProxyStream::Proxied { .. } => true,
            _ => false,
        }
    }

    /// Get reference to context
    pub fn context(&self) -> &Context {
        match *self {
            ProxyStream::Direct { ref context, .. } => &context,
            ProxyStream::Proxied { ref context, .. } => &context,
        }
    }
}

macro_rules! forward_call {
    ($self:expr, $method:ident $(, $param:expr)*) => {
        match $self.as_mut().project() {
            __ProxyStreamProjection::Direct { stream, .. } => stream.$method($($param),*),
            __ProxyStreamProjection::Proxied { stream, .. } => stream.$method($($param),*),
        }
    };
}

impl AsyncRead for ProxyStream {
    fn poll_read(mut self: Pin<&mut Self>, cx: &mut TaskContext<'_>, buf: &mut [u8]) -> Poll<io::Result<usize>> {
        let p = forward_call!(self, poll_read, cx, buf);

        // Flow statistic for Android client
        #[cfg(feature = "local-flow-stat")]
        {
            if self.is_proxied() {
                if let Poll::Ready(Ok(n)) = p {
                    self.context().local_flow_statistic().tcp().incr_tx(n as u64);
                }
            }
        }

        p
    }
}

impl AsyncWrite for ProxyStream {
    fn poll_write(mut self: Pin<&mut Self>, cx: &mut TaskContext<'_>, buf: &[u8]) -> Poll<io::Result<usize>> {
        let p = forward_call!(self, poll_write, cx, buf);

        // Flow statistic for Android client
        #[cfg(feature = "local-flow-stat")]
        {
            if self.is_proxied() {
                if let Poll::Ready(Ok(n)) = p {
                    self.context().local_flow_statistic().tcp().incr_rx(n as u64);
                }
            }
        }

        p
    }

    fn poll_flush(mut self: Pin<&mut Self>, cx: &mut TaskContext<'_>) -> Poll<io::Result<()>> {
        forward_call!(self, poll_flush, cx)
    }

    fn poll_shutdown(mut self: Pin<&mut Self>, cx: &mut TaskContext<'_>) -> Poll<io::Result<()>> {
        forward_call!(self, poll_shutdown, cx)
    }
}

async fn connect_proxy_server_internal(
    context: &Context,
    orig_svr_addr: &ServerAddr,
    svr_addr: &ServerAddr,
    timeout: Option<Duration>,
) -> io::Result<STcpStream> {
    match svr_addr {
        ServerAddr::SocketAddr(ref addr) => {
            let stream = try_timeout(tcp_stream_connect(&addr, &context), timeout).await?;
            trace!("connected proxy {} ({})", orig_svr_addr, addr);
            Ok(STcpStream::new(stream, timeout))
        }
        ServerAddr::DomainName(ref domain, port) => {
            let result = lookup_then!(context, domain.as_str(), *port, |addr| {
                match try_timeout(tcp_stream_connect(&addr, &context), timeout).await {
                    Ok(s) => Ok(STcpStream::new(s, timeout)),
                    Err(e) => {
                        debug!(
                            "failed to connect proxy {} ({}:{} ({})) try another (err: {})",
                            orig_svr_addr, domain, port, addr, e
                        );
                        Err(e)
                    }
                }
            });

            match result {
                Ok((addr, s)) => {
                    trace!("connected proxy {} ({}:{} ({}))", orig_svr_addr, domain, port, addr);
                    Ok(s)
                }
                Err(err) => {
                    error!(
                        "failed to connect proxy {} ({}:{}), {}",
                        orig_svr_addr, domain, port, err
                    );
                    Err(err)
                }
            }
        }
    }
}

/// Connect to proxy server with `ServerConfig`
async fn connect_proxy_server(context: &Context, svr_cfg: &ServerConfig) -> io::Result<STcpStream> {
    let timeout = svr_cfg.timeout().or(context.config().timeout);

    let svr_addr = match context.config().config_type {
        ConfigType::Server => svr_cfg.addr(),
        ConfigType::Socks5Local | ConfigType::TunnelLocal | ConfigType::DnsLocal => svr_cfg.external_addr(),
        #[cfg(feature = "local-http")]
        ConfigType::HttpLocal => svr_cfg.external_addr(),
        #[cfg(feature = "local-redir")]
        ConfigType::RedirLocal => svr_cfg.external_addr(),
        ConfigType::Manager => unreachable!("ConfigType::Manager shouldn't need to connect to proxy server"),
    };

    // Retry if connect failed
    //
    // FIXME: This won't work if server is actually down.
    //        Probably we should retry with another server.
    //
    // Also works if plugin is starting
    const RETRY_TIMES: i32 = 3;

    let orig_svr_addr = svr_cfg.addr();
    trace!(
        "connecting to proxy {} ({}), timeout: {:?}",
        orig_svr_addr,
        svr_addr,
        timeout
    );

    let mut last_err = None;
    for retry_time in 0..RETRY_TIMES {
        match connect_proxy_server_internal(context, orig_svr_addr, svr_addr, timeout).await {
            Ok(mut s) => {
                // IMPOSSIBLE, won't fail, but just a guard
                if let Err(err) = s.set_nodelay(context.config().no_delay) {
                    error!("failed to set TCP_NODELAY on remote socket, error: {:?}", err);
                }

                return Ok(s);
            }
            Err(err) => {
                // Connection failure, retry
                debug!(
                    "failed to connect {}, retried {} times (last err: {})",
                    svr_addr, retry_time, err
                );
                last_err = Some(err);

                // Yield and let the others' run
                //
                // It may take some time for scheduler to resume this coroutine.
                tokio::task::yield_now().await;
            }
        }
    }

    let last_err = last_err.unwrap();
    error!(
        "failed to connect {}, retried {} times, last_err: {}",
        svr_addr, RETRY_TIMES, last_err
    );
    Err(last_err)
}

/// Handshake logic for ShadowSocks Client
async fn proxy_server_handshake(
    context: SharedContext,
    remote_stream: STcpStream,
    svr_cfg: &ServerConfig,
    relay_addr: &Address,
) -> io::Result<CryptoStream<STcpStream>> {
    let mut stream = CryptoStream::new(context, remote_stream, svr_cfg);

    trace!("got encrypt stream and going to send addr: {:?}", relay_addr);

    // Send relay address to remote
    //
    // NOTE: `Address` handshake packets are very small in most cases,
    // so it will be sent with the IV/Nonce data (implemented inside `CryptoStream`).
    //
    // For lower latency, first packet should be sent back quickly,
    // so TCP_NODELAY should be kept enabled until the first data packet is received.
    let mut addr_buf = BytesMut::with_capacity(relay_addr.serialized_len());
    relay_addr.write_to_buf(&mut addr_buf);
    stream.write_all(&addr_buf).await?;

    // Here we should keep the TCP_NODELAY set until the first packet is received.
    // https://github.com/shadowsocks/shadowsocks-libev/pull/746
    //
    // Reset TCP_NODELAY after the first packet is received and sent back.

    Ok(stream)
}
<commit_msg>[#232] Send data along with header (handshake)<commit_after>//! TCP connection stream for local server with remote (proxy server or remote target)

use std::{
    fmt::{self, Display, Formatter},
    io::{self, Error},
    net::SocketAddr,
    pin::Pin,
    task::{self, Poll},
    time::Duration,
};

use bytes::{Buf, BytesMut};
use futures::ready;
use log::{debug, error, trace};
use pin_project::{pin_project, project};
use tokio::io::{AsyncRead, AsyncWrite, ReadHalf, WriteHalf};

use crate::{
    config::{ConfigType, ServerAddr, ServerConfig},
    context::{Context, SharedContext},
    relay::{socks5::Address, sys::tcp_stream_connect, utils::try_timeout},
};

use super::{connection::Connection, CryptoStream, STcpStream};

enum ProxiedConnectState {
    Connected(Address),
    Handshaking { buf: BytesMut, data_len: usize },
    Established,
}

#[pin_project]
struct ProxiedConnection {
    #[pin]
    stream: CryptoStream<STcpStream>,
    state: ProxiedConnectState,
}

impl ProxiedConnection {
    fn connected(stream: CryptoStream<STcpStream>, addr: Address) -> ProxiedConnection {
        ProxiedConnection {
            stream,
            state: ProxiedConnectState::Connected(addr),
        }
    }

    fn local_addr(&self) -> io::Result<SocketAddr> {
        self.stream.get_ref().get_ref().local_addr()
    }
}

impl AsyncRead for ProxiedConnection {
    #[project]
    fn poll_read(self: Pin<&mut Self>, cx: &mut task::Context<'_>, buf: &mut [u8]) -> Poll<io::Result<usize>> {
        self.project().stream.poll_read(cx, buf)
    }
}

impl AsyncWrite for ProxiedConnection {
    fn poll_write(mut self: Pin<&mut Self>, cx: &mut task::Context<'_>, data: &[u8]) -> Poll<io::Result<usize>> {
        loop {
            let this = self.as_mut().project();

            match this.state {
                ProxiedConnectState::Connected(ref addr) => {
                    assert_ne!(data.len(), 0);

                    // Send relay address to remote
                    //
                    // NOTE: `Address` handshake packets are very small in most cases,
                    // so
                    // 1. it will be sent with the IV/Nonce data (implemented inside `CryptoStream`).
                    // 2. concatenating target's Address with the first data buffer (#232)
                    //
                    // For lower latency, first packet should be sent back quickly,
                    // so TCP_NODELAY should be kept enabled until the first data packet is received.
                    let addr_len = addr.serialized_len();
                    let mut buf = BytesMut::with_capacity(addr_len + data.len());
                    addr.write_to_buf(&mut buf);
                    buf.extend_from_slice(data);

                    trace!("sending handshake address {} with data {} bytes", addr, data.len());

                    // Fast path
                    //
                    // For CryptoStream (Stream and AEAD), poll_write will return Ready(..) until all data have been sent out
                    match this.stream.poll_write(cx, &buf) {
                        Poll::Ready(Err(err)) => return Poll::Ready(Err(err)),
                        Poll::Ready(Ok(n)) => {
                            buf.advance(n);

                            let remaining = buf.remaining();
                            if remaining < data.len() {
                                // Ok, written some data with Address
                                let written_len = data.len() - remaining;

                                trace!(
                                    "sent handshake address {} with {} bytes of data, data len {} bytes",
                                    addr,
                                    written_len,
                                    data.len(),
                                );

                                self.state = ProxiedConnectState::Established;
                                return Poll::Ready(Ok(written_len));
                            }

                            // FALLTHROUGH
                            // Handshaking branch will try to poll_write again
                            self.state = ProxiedConnectState::Handshaking {
                                buf,
                                data_len: data.len(),
                            };
                        }
                        Poll::Pending => {
                            // poll_write is not ready, let Handshaking branch try again later
                            self.state = ProxiedConnectState::Handshaking {
                                buf,
                                data_len: data.len(),
                            };

                            return Poll::Pending;
                        }
                    }
                }
                ProxiedConnectState::Handshaking { ref mut buf, data_len } => {
                    let data_len = *data_len;

                    // Try to write at least addr_len size
                    let n = ready!(this.stream.poll_write(cx, buf))?;
                    buf.advance(n);

                    let remaining = buf.remaining();
                    if remaining < data_len {
                        // Ok, written some data with Address
                        let written_len = data_len - remaining;

                        trace!(
                            "sent handshake address with {} bytes of data, data len {} bytes",
                            written_len,
                            data_len
                        );

                        self.state = ProxiedConnectState::Established;
                        return Poll::Ready(Ok(written_len));
                    }
                }
                ProxiedConnectState::Established => {
                    break;
                }
            }
        }

        self.project().stream.poll_write(cx, data)
    }

    fn poll_flush(self: Pin<&mut Self>, cx: &mut task::Context<'_>) -> Poll<io::Result<()>> {
        self.project().stream.poll_flush(cx)
    }

    fn poll_shutdown(self: Pin<&mut Self>, cx: &mut task::Context<'_>) -> Poll<io::Result<()>> {
        self.project().stream.poll_shutdown(cx)
    }
}

#[pin_project]
enum ProxyConnection {
    Direct(#[pin] STcpStream),
    Proxied(#[pin] ProxiedConnection),
}

impl ProxyConnection {
    /// Check if the underlying connection is proxied
    fn is_proxied(&self) -> bool {
        match *self {
            ProxyConnection::Proxied { .. } => true,
            _ => false,
        }
    }

    fn local_addr(&self) -> io::Result<SocketAddr> {
        match *self {
            ProxyConnection::Direct(ref stream) => stream.get_ref().local_addr(),
            ProxyConnection::Proxied(ref stream) => stream.local_addr(),
        }
    }
}

macro_rules! forward_call {
    ($self:expr, $method:ident $(, $param:expr)*) => {
        // #[project]
        match $self.as_mut().project() {
            // ProxyConnection::Direct(stream) => stream.$method($($param),*),
            __ProxyConnectionProjection::Direct(stream) => stream.$method($($param),*),
            // ProxyConnection::Proxied(stream) => stream.$method($($param),*),
            __ProxyConnectionProjection::Proxied(stream) => stream.$method($($param),*),
        }
    };
}

impl AsyncRead for ProxyConnection {
    #[project]
    fn poll_read(mut self: Pin<&mut Self>, cx: &mut task::Context<'_>, buf: &mut [u8]) -> Poll<io::Result<usize>> {
        forward_call!(self, poll_read, cx, buf)
    }
}

impl AsyncWrite for ProxyConnection {
    #[project]
    fn poll_write(mut self: Pin<&mut Self>, cx: &mut task::Context<'_>, buf: &[u8]) -> Poll<io::Result<usize>> {
        // let p = forward_call!(self, poll_write, cx, buf);
        forward_call!(self, poll_write, cx, buf)
    }

    #[project]
    fn poll_flush(mut self: Pin<&mut Self>, cx: &mut task::Context<'_>) -> Poll<io::Result<()>> {
        forward_call!(self, poll_flush, cx)
    }

    #[project]
    fn poll_shutdown(mut self: Pin<&mut Self>, cx: &mut task::Context<'_>) -> Poll<io::Result<()>> {
        forward_call!(self, poll_shutdown, cx)
    }
}

#[derive(Debug)]
pub struct ProxyStreamError {
    inner: Error,
    bypassed: bool,
}

impl ProxyStreamError {
    fn new(inner: Error, bypassed: bool) -> ProxyStreamError {
        ProxyStreamError { inner, bypassed }
    }

    /// Check if it is proxied
    pub fn is_proxied(&self) -> bool {
        self.bypassed
    }

    /// Into internal `std::io::Error`
    pub fn into_inner(self) -> Error {
        self.inner
    }
}

impl From<ProxyStreamError> for Error {
    fn from(err: ProxyStreamError) -> Error {
        err.inner
    }
}

impl Display for ProxyStreamError {
    fn fmt(&self, f: &mut Formatter) -> fmt::Result {
        Display::fmt(&self.inner, f)
    }
}

/// Stream wrapper for both direct connections and proxied connections
#[pin_project]
pub struct ProxyStream {
    #[pin]
    connection: ProxyConnection,
    context: SharedContext,
}

impl ProxyStream {
    /// Connect to remote by ACL rules
    pub async fn connect(
        context: SharedContext,
        svr_cfg: &ServerConfig,
        addr: &Address,
    ) -> Result<ProxyStream, ProxyStreamError> {
        if context.check_target_bypassed(addr).await {
            ProxyStream::connect_direct_wrapped(context, addr).await
        } else {
            ProxyStream::connect_proxied_wrapped(context, svr_cfg, addr).await
        }
    }

    /// Connect to remote directly (without proxy)
    ///
    /// This is used for hosts that matches ACL bypassed rules
    pub async fn connect_direct(context: SharedContext, addr: &Address) -> io::Result<ProxyStream> {
        debug!("connect to {} directly (bypassed)", addr);

        // NOTE: Direct connection's timeout is controlled by the global key
        let timeout = context.config().timeout;

        let stream = match *addr {
            Address::SocketAddress(ref saddr) => try_timeout(tcp_stream_connect(&saddr, &context), timeout).await?,
            Address::DomainNameAddress(ref domain, port) => {
                lookup_then!(context, domain, port, |saddr| {
                    try_timeout(tcp_stream_connect(&saddr, &context), timeout).await
                })?
                .1
            }
        };

        Ok(ProxyStream {
            context,
            connection: ProxyConnection::Direct(Connection::new(stream, timeout)),
        })
    }

    async fn connect_direct_wrapped(context: SharedContext, addr: &Address) -> Result<ProxyStream, ProxyStreamError> {
        match ProxyStream::connect_direct(context, addr).await {
            Ok(s) => Ok(s),
            Err(err) => Err(ProxyStreamError::new(err, true)),
        }
    }

    /// Connect to remote via proxy server
    ///
    /// This is used for hosts that matches ACL proxied rules
    pub async fn connect_proxied(
        context: SharedContext,
        svr_cfg: &ServerConfig,
        addr: &Address,
    ) -> io::Result<ProxyStream> {
        debug!(
            "connect to {} via {} ({}) (proxied)",
            addr,
            svr_cfg.addr(),
            svr_cfg.external_addr()
        );

        let server_stream = connect_proxy_server(&context, svr_cfg).await?;
        let proxy_stream = CryptoStream::new(context.clone(), server_stream, svr_cfg);

        Ok(ProxyStream {
            context,
            connection: ProxyConnection::Proxied(ProxiedConnection::connected(proxy_stream, addr.clone())),
        })
    }

    async fn connect_proxied_wrapped(
        context: SharedContext,
        svr_cfg: &ServerConfig,
        addr: &Address,
    ) -> Result<ProxyStream, ProxyStreamError> {
        match ProxyStream::connect_proxied(context, svr_cfg, addr).await {
            Ok(s) => Ok(s),
            Err(err) => Err(ProxyStreamError::new(err, false)),
        }
    }

    /// Split into reader and writer
    pub fn split(self) -> (ReadHalf<ProxyStream>, WriteHalf<ProxyStream>) {
        use tokio::io::split;
        split(self)
    }

    /// Returns the local socket address of this stream socket
    pub fn local_addr(&self) -> io::Result<SocketAddr> {
        self.connection.local_addr()
    }

    /// Check if the underlying connection is proxied
    pub fn is_proxied(&self) -> bool {
        self.connection.is_proxied()
    }

    /// Get reference to context
    pub fn context(&self) -> &Context {
        &self.context
    }
}

impl AsyncRead for ProxyStream {
    fn poll_read(mut self: Pin<&mut Self>, cx: &mut task::Context<'_>, buf: &mut [u8]) -> Poll<io::Result<usize>> {
        let p = self.as_mut().project().connection.poll_read(cx, buf);

        // Flow statistic for Android client
        #[cfg(feature = "local-flow-stat")]
        {
            if self.is_proxied() {
                if let Poll::Ready(Ok(n)) = p {
                    self.context().local_flow_statistic().tcp().incr_tx(n as u64);
                }
            }
        }

        p
    }
}

impl AsyncWrite for ProxyStream {
    fn poll_write(mut self: Pin<&mut Self>, cx: &mut task::Context<'_>, buf: &[u8]) -> Poll<io::Result<usize>> {
        let p = self.as_mut().project().connection.poll_write(cx, buf);

        // Flow statistic for Android client
        #[cfg(feature = "local-flow-stat")]
        {
            if self.is_proxied() {
                if let Poll::Ready(Ok(n)) = p {
                    self.context().local_flow_statistic().tcp().incr_rx(n as u64);
                }
            }
        }

        p
    }

    fn poll_flush(self: Pin<&mut Self>, cx: &mut task::Context<'_>) -> Poll<io::Result<()>> {
        self.project().connection.poll_flush(cx)
    }

    fn poll_shutdown(self: Pin<&mut Self>, cx: &mut task::Context<'_>) -> Poll<io::Result<()>> {
        self.project().connection.poll_shutdown(cx)
    }
}

async fn connect_proxy_server_internal(
    context: &Context,
    orig_svr_addr: &ServerAddr,
    svr_addr: &ServerAddr,
    timeout: Option<Duration>,
) -> io::Result<STcpStream> {
    match svr_addr {
        ServerAddr::SocketAddr(ref addr) => {
            let stream = try_timeout(tcp_stream_connect(&addr, &context), timeout).await?;
            trace!("connected proxy {} ({})", orig_svr_addr, addr);
            Ok(STcpStream::new(stream, timeout))
        }
        ServerAddr::DomainName(ref domain, port) => {
            let result = lookup_then!(context, domain.as_str(), *port, |addr| {
                match try_timeout(tcp_stream_connect(&addr, &context), timeout).await {
                    Ok(s) => Ok(STcpStream::new(s, timeout)),
                    Err(e) => {
                        debug!(
                            "failed to connect proxy {} ({}:{} ({})) try another (err: {})",
                            orig_svr_addr, domain, port, addr, e
                        );
                        Err(e)
                    }
                }
            });

            match result {
                Ok((addr, s)) => {
                    trace!("connected proxy {} ({}:{} ({}))", orig_svr_addr, domain, port, addr);
                    Ok(s)
                }
                Err(err) => {
                    error!(
                        "failed to connect proxy {} ({}:{}), {}",
                        orig_svr_addr, domain, port, err
                    );
                    Err(err)
                }
            }
        }
    }
}

/// Connect to proxy server with `ServerConfig`
async fn connect_proxy_server(context: &Context, svr_cfg: &ServerConfig) -> io::Result<STcpStream> {
    let timeout = svr_cfg.timeout().or(context.config().timeout);

    let svr_addr = match context.config().config_type {
        ConfigType::Server => svr_cfg.addr(),
        ConfigType::Socks5Local | ConfigType::TunnelLocal | ConfigType::DnsLocal => svr_cfg.external_addr(),
        #[cfg(feature = "local-http")]
        ConfigType::HttpLocal => svr_cfg.external_addr(),
        #[cfg(feature = "local-redir")]
        ConfigType::RedirLocal => svr_cfg.external_addr(),
        ConfigType::Manager => unreachable!("ConfigType::Manager shouldn't need to connect to proxy server"),
    };

    // Retry if connect failed
    //
    // FIXME: This won't work if server is actually down.
    //        Probably we should retry with another server.
    //
    // Also works if plugin is starting
    const RETRY_TIMES: i32 = 3;

    let orig_svr_addr = svr_cfg.addr();
    trace!(
        "connecting to proxy {} ({}), timeout: {:?}",
        orig_svr_addr,
        svr_addr,
        timeout
    );

    let mut last_err = None;
    for retry_time in 0..RETRY_TIMES {
        match connect_proxy_server_internal(context, orig_svr_addr, svr_addr, timeout).await {
            Ok(mut s) => {
                // IMPOSSIBLE, won't fail, but just a guard
                if let Err(err) = s.set_nodelay(context.config().no_delay) {
                    error!("failed to set TCP_NODELAY on remote socket, error: {:?}", err);
                }

                return Ok(s);
            }
            Err(err) => {
                // Connection failure, retry
                debug!(
                    "failed to connect {}, retried {} times (last err: {})",
                    svr_addr, retry_time, err
                );
                last_err = Some(err);

                // Yield and let the others' run
                //
                // It may take some time for scheduler to resume this coroutine.
                tokio::task::yield_now().await;
            }
        }
    }

    let last_err = last_err.unwrap();
    error!(
        "failed to connect {}, retried {} times, last_err: {}",
        svr_addr, RETRY_TIMES, last_err
    );
    Err(last_err)
}
<|endoftext|>


<commit_before>use openvr_sys;

use tracking::*;
use error::*;
use subsystems::*;

pub struct IVRTrackedCamera(pub *const ());

#[derive(Debug, Copy, Clone)]
pub struct CameraFrameSize {
    pub width: u32,
    pub height: u32,
    pub buffer: u32
}

#[derive(Debug, Copy, Clone)]
pub enum CameraFrameType {
    Distorted,
    Undistorted,
    MaximumUndistorted,
    MaximumFrameTypes
}

#[derive(Debug, Copy, Clone)]
pub struct CameraIntriniscs {
    pub focal_length: [f32; 2],
    pub center: [f32; 2]
}

#[derive(Debug, Copy, Clone)]
pub struct CameraFrameHeader {
    pub width: u32,
    pub height: u32,
    pub bytes_per_pixel: u32,
    pub frame_sequence: u32,
    pub pose: TrackedDevicePose
}

#[derive(Debug)]
pub struct CameraFrame {
    pub framebuffer: Vec<u8>,
    pub header: CameraFrameHeader
}

pub struct CameraStream {
    pub handle: u64,
    pub owner: TrackedDevicePose
}

impl Drop for CameraStream {
    fn drop(&mut self) {
        unsafe {
            let cam = *{ tracked_camera().unwrap().0 as *mut openvr_sys::VR_IVRTrackedCamera_FnTable };

            let error = Error::from_raw(
                cam.ReleaseVideoStreamingService.unwrap()(self.handle));

            if error.is_err() {
                println!("Failed to drop camera stream! Possible memory leak! {}", error.message());
            }
        }
    }
}

impl CameraStream {
    /// reads current camera frame
    pub fn read(&self, ctype: CameraFrameType) -> Result<CameraFrame, Error<openvr_sys::EVRTrackedCameraError>> {
        use std::mem;
        use std;

        unsafe {
            // get subsystems
            let cam = *{ tracked_camera().unwrap().0 as *mut openvr_sys::VR_IVRTrackedCamera_FnTable };
            let size = tracked_camera().unwrap().frame_size(&self.owner, ctype).unwrap();

            // create raw buffer where openvr can store it's data into
            let mut buffer = Vec::<u8>::with_capacity(size.buffer as usize);
            let mut raw_buffer = buffer.as_mut_ptr();
            mem::forget(buffer);

            // create header
            let mut header = openvr_sys::CameraVideoStreamFrameHeader_t::default();

            let error = Error::from_raw(
                cam.GetVideoStreamFrameBuffer.unwrap()(
                    self.handle,
                    ctype.to_raw(),
                    raw_buffer as *mut std::os::raw::c_void,
                    size.buffer,
                    &mut header,
                    mem::size_of::<openvr_sys::CameraVideoStreamFrameHeader_t>() as u32
                ));

            if error.is_ok() {
                // bring framebuffer back into rusts controll
                let buffer = Vec::from_raw_parts(raw_buffer, size.buffer as usize, size.buffer as usize);

                return Ok(CameraFrame {
                    framebuffer: buffer,
                    header: CameraFrameHeader {
                        width: header.nWidth,
                        height: header.nHeight,
                        bytes_per_pixel: header.nBytesPerPixel,
                        frame_sequence: header.nFrameSequence,
                        pose: TrackedDevicePose::from_raw(self.owner.index, header.standingTrackedDevicePose)
                    }
                });
            } else {
                return Err(error);
            }
        }
    }
}

impl CameraFrameType {
    pub fn to_raw(&self) -> openvr_sys::EVRTrackedCameraFrameType {
        use openvr_sys::EVRTrackedCameraFrameType::*;

        match self {
            &CameraFrameType::Distorted => EVRTrackedCameraFrameType_VRTrackedCameraFrameType_Distorted,
            &CameraFrameType::Undistorted => EVRTrackedCameraFrameType_VRTrackedCameraFrameType_Undistorted,
            &CameraFrameType::MaximumUndistorted => EVRTrackedCameraFrameType_VRTrackedCameraFrameType_MaximumUndistorted ,
            &CameraFrameType::MaximumFrameTypes => EVRTrackedCameraFrameType_MAX_CAMERA_FRAME_TYPES
        }
    }
}

impl IVRTrackedCamera {
    pub unsafe fn from_raw(ptr: *const ()) -> Self {
        IVRTrackedCamera(ptr as *mut ())
    }

    /// checks whether the current system has a camera
    pub fn has_camera(&self, device: &TrackedDevicePose) -> Result<bool, Error<openvr_sys::EVRTrackedCameraError>> {
        unsafe {
            let cam = * { self.0 as *mut openvr_sys::VR_IVRTrackedCamera_FnTable };
            let mut has_cam = 0i32;

            let error = Error::from_raw(
                cam.HasCamera.unwrap()(device.index as u32, &mut has_cam as *mut i32));

            if error.is_ok() {
                return Ok(has_cam > 0i32);
            } else {
                return Err(error);
            }
        }
    }

    /// gets frame buffer information of camera
    pub fn frame_size(&self, device: &TrackedDevicePose, ctype: CameraFrameType)
        -> Result<CameraFrameSize, Error<openvr_sys::EVRTrackedCameraError>>
    {
        unsafe {
            let mut result = CameraFrameSize {
                width: 0,
                height: 0,
                buffer: 0,
            };

            let cam = *{ self.0 as *mut openvr_sys::VR_IVRTrackedCamera_FnTable };

            let error = Error::from_raw(
                cam.GetCameraFrameSize.unwrap()(device.index as u32,
                                       ctype.to_raw(),
                                       &mut result.width,
                                       &mut result.height,
                                       &mut result.buffer));

            if error.is_ok() {
                return Ok(result);
            } else {
                return Err(error);
            }
        }
    }

    // gets camera intrinsic
    pub fn intrinisics(&self, device: &TrackedDevicePose, ctype: CameraFrameType)
        -> Result<CameraIntriniscs, Error<openvr_sys::EVRTrackedCameraError>>
    {
        unsafe {
            let mut focal = openvr_sys::HmdVector2_t { v: [0.0, 0.0] };
            let mut center = openvr_sys::HmdVector2_t { v: [0.0, 0.0] };

            let cam = *{ self.0 as *mut openvr_sys::VR_IVRTrackedCamera_FnTable };

            let error = Error::from_raw(
                cam.GetCameraIntrinisics.unwrap()(device.index as u32,
                                                ctype.to_raw(),
                                                &mut focal,
                                                &mut center));

            if error.is_ok() {
                return Ok(CameraIntriniscs {
                    focal_length: focal.v,
                    center: center.v
                });
            } else {
                return Err(error);
            }
        }
    }

    /// aquires a stream to the given camera device
    pub fn stream(&self, device: &TrackedDevicePose) -> Result<CameraStream, Error<openvr_sys::EVRTrackedCameraError>> {
        unsafe {
            let cam = *{ self.0 as *mut openvr_sys::VR_IVRTrackedCamera_FnTable };
            let mut handle = 0u64;

            let error = Error::from_raw(
                cam.AcquireVideoStreamingService.unwrap()(device.index as u32, &mut handle));

            if error.is_ok() {
                return Ok(CameraStream {
                            handle: handle,
                            owner: *device
                        });
            } else {
                return Err(error);
            }
        }
    }
}
<commit_msg>fix warning "variable does not need to be mutable"<commit_after>use openvr_sys;

use tracking::*;
use error::*;
use subsystems::*;

pub struct IVRTrackedCamera(pub *const ());

#[derive(Debug, Copy, Clone)]
pub struct CameraFrameSize {
    pub width: u32,
    pub height: u32,
    pub buffer: u32
}

#[derive(Debug, Copy, Clone)]
pub enum CameraFrameType {
    Distorted,
    Undistorted,
    MaximumUndistorted,
    MaximumFrameTypes
}

#[derive(Debug, Copy, Clone)]
pub struct CameraIntriniscs {
    pub focal_length: [f32; 2],
    pub center: [f32; 2]
}

#[derive(Debug, Copy, Clone)]
pub struct CameraFrameHeader {
    pub width: u32,
    pub height: u32,
    pub bytes_per_pixel: u32,
    pub frame_sequence: u32,
    pub pose: TrackedDevicePose
}

#[derive(Debug)]
pub struct CameraFrame {
    pub framebuffer: Vec<u8>,
    pub header: CameraFrameHeader
}

pub struct CameraStream {
    pub handle: u64,
    pub owner: TrackedDevicePose
}

impl Drop for CameraStream {
    fn drop(&mut self) {
        unsafe {
            let cam = *{ tracked_camera().unwrap().0 as *mut openvr_sys::VR_IVRTrackedCamera_FnTable };

            let error = Error::from_raw(
                cam.ReleaseVideoStreamingService.unwrap()(self.handle));

            if error.is_err() {
                println!("Failed to drop camera stream! Possible memory leak! {}", error.message());
            }
        }
    }
}

impl CameraStream {
    /// reads current camera frame
    pub fn read(&self, ctype: CameraFrameType) -> Result<CameraFrame, Error<openvr_sys::EVRTrackedCameraError>> {
        use std::mem;
        use std;

        unsafe {
            // get subsystems
            let cam = *{ tracked_camera().unwrap().0 as *mut openvr_sys::VR_IVRTrackedCamera_FnTable };
            let size = tracked_camera().unwrap().frame_size(&self.owner, ctype).unwrap();

            // create raw buffer where openvr can store it's data into
            let mut buffer = Vec::<u8>::with_capacity(size.buffer as usize);
            let raw_buffer = buffer.as_mut_ptr();
            mem::forget(buffer);

            // create header
            let mut header = openvr_sys::CameraVideoStreamFrameHeader_t::default();

            let error = Error::from_raw(
                cam.GetVideoStreamFrameBuffer.unwrap()(
                    self.handle,
                    ctype.to_raw(),
                    raw_buffer as *mut std::os::raw::c_void,
                    size.buffer,
                    &mut header,
                    mem::size_of::<openvr_sys::CameraVideoStreamFrameHeader_t>() as u32
                ));

            if error.is_ok() {
                // bring framebuffer back into rusts controll
                let buffer = Vec::from_raw_parts(raw_buffer, size.buffer as usize, size.buffer as usize);

                return Ok(CameraFrame {
                    framebuffer: buffer,
                    header: CameraFrameHeader {
                        width: header.nWidth,
                        height: header.nHeight,
                        bytes_per_pixel: header.nBytesPerPixel,
                        frame_sequence: header.nFrameSequence,
                        pose: TrackedDevicePose::from_raw(self.owner.index, header.standingTrackedDevicePose)
                    }
                });
            } else {
                return Err(error);
            }
        }
    }
}

impl CameraFrameType {
    pub fn to_raw(&self) -> openvr_sys::EVRTrackedCameraFrameType {
        use openvr_sys::EVRTrackedCameraFrameType::*;

        match self {
            &CameraFrameType::Distorted => EVRTrackedCameraFrameType_VRTrackedCameraFrameType_Distorted,
            &CameraFrameType::Undistorted => EVRTrackedCameraFrameType_VRTrackedCameraFrameType_Undistorted,
            &CameraFrameType::MaximumUndistorted => EVRTrackedCameraFrameType_VRTrackedCameraFrameType_MaximumUndistorted ,
            &CameraFrameType::MaximumFrameTypes => EVRTrackedCameraFrameType_MAX_CAMERA_FRAME_TYPES
        }
    }
}

impl IVRTrackedCamera {
    pub unsafe fn from_raw(ptr: *const ()) -> Self {
        IVRTrackedCamera(ptr as *mut ())
    }

    /// checks whether the current system has a camera
    pub fn has_camera(&self, device: &TrackedDevicePose) -> Result<bool, Error<openvr_sys::EVRTrackedCameraError>> {
        unsafe {
            let cam = * { self.0 as *mut openvr_sys::VR_IVRTrackedCamera_FnTable };
            let mut has_cam = 0i32;

            let error = Error::from_raw(
                cam.HasCamera.unwrap()(device.index as u32, &mut has_cam as *mut i32));

            if error.is_ok() {
                return Ok(has_cam > 0i32);
            } else {
                return Err(error);
            }
        }
    }

    /// gets frame buffer information of camera
    pub fn frame_size(&self, device: &TrackedDevicePose, ctype: CameraFrameType)
        -> Result<CameraFrameSize, Error<openvr_sys::EVRTrackedCameraError>>
    {
        unsafe {
            let mut result = CameraFrameSize {
                width: 0,
                height: 0,
                buffer: 0,
            };

            let cam = *{ self.0 as *mut openvr_sys::VR_IVRTrackedCamera_FnTable };

            let error = Error::from_raw(
                cam.GetCameraFrameSize.unwrap()(device.index as u32,
                                       ctype.to_raw(),
                                       &mut result.width,
                                       &mut result.height,
                                       &mut result.buffer));

            if error.is_ok() {
                return Ok(result);
            } else {
                return Err(error);
            }
        }
    }

    // gets camera intrinsic
    pub fn intrinisics(&self, device: &TrackedDevicePose, ctype: CameraFrameType)
        -> Result<CameraIntriniscs, Error<openvr_sys::EVRTrackedCameraError>>
    {
        unsafe {
            let mut focal = openvr_sys::HmdVector2_t { v: [0.0, 0.0] };
            let mut center = openvr_sys::HmdVector2_t { v: [0.0, 0.0] };

            let cam = *{ self.0 as *mut openvr_sys::VR_IVRTrackedCamera_FnTable };

            let error = Error::from_raw(
                cam.GetCameraIntrinisics.unwrap()(device.index as u32,
                                                ctype.to_raw(),
                                                &mut focal,
                                                &mut center));

            if error.is_ok() {
                return Ok(CameraIntriniscs {
                    focal_length: focal.v,
                    center: center.v
                });
            } else {
                return Err(error);
            }
        }
    }

    /// aquires a stream to the given camera device
    pub fn stream(&self, device: &TrackedDevicePose) -> Result<CameraStream, Error<openvr_sys::EVRTrackedCameraError>> {
        unsafe {
            let cam = *{ self.0 as *mut openvr_sys::VR_IVRTrackedCamera_FnTable };
            let mut handle = 0u64;

            let error = Error::from_raw(
                cam.AcquireVideoStreamingService.unwrap()(device.index as u32, &mut handle));

            if error.is_ok() {
                return Ok(CameraStream {
                            handle: handle,
                            owner: *device
                        });
            } else {
                return Err(error);
            }
        }
    }
}
<|endoftext|>


<commit_before>use std::clone::Clone;
use std::cmp::Ordering::Less;
use std::collections::VecDeque;
use std::error::Error;
use std::fs::{File, OpenOptions};
use std::io::SeekFrom::Start;
use std::io::{BufRead, BufReader, BufWriter, Seek, Write};
use std::marker::PhantomData;
use std::path::PathBuf;

use serde::de::DeserializeOwned;
use serde::Serialize;
use serde_json;
use tempdir::TempDir;

use super::{GlobalSettings, Line};

/// Trait for types that can be used by
/// [ExternalSorter](struct.ExternalSorter.html). Must be sortable, cloneable,
/// serializeable, and able to report on it's size
pub trait ExternallySortable: Clone + Serialize + DeserializeOwned {
    /// Get the size, in bytes, of this object (used to constrain the buffer
    /// used in the external sort).
    fn get_size(&self) -> u64;
}

/// Iterator that provides sorted `T`s
pub struct ExtSortedIterator<Line> {
    buffers: Vec<VecDeque<Line>>,
    chunk_offsets: Vec<u64>,
    max_per_chunk: u64,
    chunks: u64,
    tmp_dir: TempDir,
    settings: GlobalSettings,
    failed: bool,
}

impl Iterator for ExtSortedIterator<Line>
where
    Line: ExternallySortable,
{
    type Item = Result<Line, Box<dyn Error>>;

    /// # Errors
    ///
    /// This method can fail due to issues reading intermediate sorted chunks
    /// from disk, or due to serde deserialization issues
    fn next(&mut self) -> Option<Self::Item> {
        if self.failed {
            return None;
        }
        // fill up any empty buffers
        let mut empty = true;
        for chunk_num in 0..self.chunks {
            if self.buffers[chunk_num as usize].is_empty() {
                let mut f = match File::open(self.tmp_dir.path().join(chunk_num.to_string())) {
                    Ok(f) => f,
                    Err(e) => {
                        self.failed = true;
                        return Some(Err(Box::new(e)));
                    }
                };
                match f.seek(Start(self.chunk_offsets[chunk_num as usize])) {
                    Ok(_) => (),
                    Err(e) => {
                        self.failed = true;
                        return Some(Err(Box::new(e)));
                    }
                }
                let bytes_read =
                    match fill_buff(&mut self.buffers[chunk_num as usize], f, self.max_per_chunk) {
                        Ok(bytes_read) => bytes_read,
                        Err(e) => {
                            self.failed = true;
                            return Some(Err(e));
                        }
                    };
                self.chunk_offsets[chunk_num as usize] += bytes_read;
                if !self.buffers[chunk_num as usize].is_empty() {
                    empty = false;
                }
            } else {
                empty = false;
            }
        }
        if empty {
            return None;
        }

        // find the next record to write
        // check is_empty() before unwrap()ing
        let mut idx = 0;
        for chunk_num in 0..self.chunks as usize {
            if !self.buffers[chunk_num].is_empty() {
                if self.buffers[idx].is_empty()
                    || (super::compare_by)(
                        self.buffers[chunk_num].front().unwrap(),
                        self.buffers[idx].front().unwrap(),
                        &self.settings,
                    ) == Less
                {
                    idx = chunk_num;
                }
            }
        }

        // unwrap due to checks above
        let r = self.buffers[idx].pop_front().unwrap();
        Some(Ok(r))
    }
}

/// Perform an external sort on an unsorted stream of incoming data
pub struct ExternalSorter<Line>
where
    Line: ExternallySortable,
{
    tmp_dir: Option<PathBuf>,
    buffer_bytes: u64,
    phantom: PhantomData<Line>,
    settings: GlobalSettings,
}

impl ExternalSorter<Line>
where
    Line: ExternallySortable,
{
    /// Create a new `ExternalSorter` with a specified memory buffer and
    /// temporary directory
    pub fn new(
        buffer_bytes: u64,
        tmp_dir: Option<PathBuf>,
        settings: GlobalSettings,
    ) -> ExternalSorter<Line> {
        ExternalSorter {
            buffer_bytes,
            tmp_dir,
            phantom: PhantomData,
            settings,
        }
    }

    /// Sort (based on `compare`) the `T`s provided by `unsorted` and return an
    /// iterator
    ///
    /// # Errors
    ///
    /// This method can fail due to issues writing intermediate sorted chunks
    /// to disk, or due to serde serialization issues
    pub fn sort_by<I>(
        &self,
        unsorted: I,
        settings: GlobalSettings,
    ) -> Result<ExtSortedIterator<Line>, Box<dyn Error>>
    where
        I: Iterator<Item = Line>,
    {
        let tmp_dir = match self.tmp_dir {
            Some(ref p) => TempDir::new_in(p, "uutils_sort")?,
            None => TempDir::new("uutils_sort")?,
        };
        // creating the thing we need to return first due to the face that we need to
        // borrow tmp_dir and move it out
        let mut iter = ExtSortedIterator {
            buffers: Vec::new(),
            chunk_offsets: Vec::new(),
            max_per_chunk: 0,
            chunks: 0,
            tmp_dir,
            settings,
            failed: false,
        };

        {
            let mut total_read = 0;
            let mut chunk = Vec::new();
            // Initial buffer is specified by user
            let mut adjusted_buffer_size = self.buffer_bytes;

            // make the initial chunks on disk
            for seq in unsorted {
                let seq_size = seq.get_size();
                total_read += seq_size;
                // Grow buffer size for a struct/Line larger than buffer
                adjusted_buffer_size =
                    if adjusted_buffer_size < seq_size {
                        seq_size
                    } else {
                        adjusted_buffer_size
                    };
                chunk.push(seq);

                if total_read >= adjusted_buffer_size {
                    super::sort_by(&mut chunk, &self.settings);
                    self.write_chunk(
                        &iter.tmp_dir.path().join(iter.chunks.to_string()),
                        &mut chunk,
                    )?;
                    chunk.clear();
                    total_read = 0;
                    iter.chunks += 1;
                }
            }
            // write the last chunk
            if chunk.len() > 0 {
                super::sort_by(&mut chunk, &self.settings);
                self.write_chunk(
                    &iter.tmp_dir.path().join(iter.chunks.to_string()),
                    &mut chunk,
                )?;
                iter.chunks += 1;
            }

            // initialize buffers for each chunk
            //
            // Having a right sized buffer for each chunk for smallish values seems silly to me?
            //
            // We will have to have the entire iter in memory sometime right?
            // Set minimum to the size of the writer buffer, ~8K
            //
            const MINIMUM_READBACK_BUFFER: u64 = 8200;
            let right_sized_buffer = adjusted_buffer_size
                   .checked_div(iter.chunks)
                   .unwrap_or(adjusted_buffer_size);
            iter.max_per_chunk =
                if right_sized_buffer > MINIMUM_READBACK_BUFFER {
                    right_sized_buffer
                } else {
                    MINIMUM_READBACK_BUFFER
                };
            iter.buffers = vec![VecDeque::new(); iter.chunks as usize];
            iter.chunk_offsets = vec![0 as u64; iter.chunks as usize];
            for chunk_num in 0..iter.chunks {
                let offset = fill_buff(
                    &mut iter.buffers[chunk_num as usize],
                    File::open(iter.tmp_dir.path().join(chunk_num.to_string()))?,
                    iter.max_per_chunk,
                )?;
                iter.chunk_offsets[chunk_num as usize] = offset;
            }
        }

        Ok(iter)
    }

    fn write_chunk(&self, file: &PathBuf, chunk: &mut Vec<Line>) -> Result<(), Box<dyn Error>> {
        let new_file = OpenOptions::new().create(true).append(true).open(file)?;
        let mut buf_write = Box::new(BufWriter::new(new_file)) as Box<dyn Write>;
        for s in chunk {
            let mut serialized = serde_json::to_string(&s).expect("JSON write error: ");
            serialized.push_str("\n");
            buf_write.write(serialized.as_bytes())?;
        }
        buf_write.flush()?;

        Ok(())
    }
}

fn fill_buff<Line>(
    vec: &mut VecDeque<Line>,
    file: File,
    max_bytes: u64,
) -> Result<u64, Box<dyn Error>>
where
    Line: ExternallySortable,
{
    let mut total_read = 0;
    let mut bytes_read = 0;
    for line in BufReader::new(file).lines() {
        let line_s = line?;
        bytes_read += line_s.len() + 1;
        // This is where the bad stuff happens usually
        let deserialized: Line = serde_json::from_str(&line_s).expect("JSON read error: ");
        total_read += deserialized.get_size();
        vec.push_back(deserialized);
        if total_read > max_bytes {
            break;
        }
    }

    Ok(bytes_read as u64)
}
<commit_msg>Set a dynamic minimum buffer size<commit_after>use std::clone::Clone;
use std::cmp::Ordering::Less;
use std::collections::VecDeque;
use std::error::Error;
use std::fs::{File, OpenOptions};
use std::io::SeekFrom::Start;
use std::io::{BufRead, BufReader, BufWriter, Seek, Write};
use std::marker::PhantomData;
use std::path::PathBuf;

use serde::de::DeserializeOwned;
use serde::Serialize;
use serde_json;
use tempdir::TempDir;

use super::{GlobalSettings, Line};

/// Trait for types that can be used by
/// [ExternalSorter](struct.ExternalSorter.html). Must be sortable, cloneable,
/// serializeable, and able to report on it's size
pub trait ExternallySortable: Clone + Serialize + DeserializeOwned {
    /// Get the size, in bytes, of this object (used to constrain the buffer
    /// used in the external sort).
    fn get_size(&self) -> u64;
}

/// Iterator that provides sorted `T`s
pub struct ExtSortedIterator<Line> {
    buffers: Vec<VecDeque<Line>>,
    chunk_offsets: Vec<u64>,
    max_per_chunk: u64,
    chunks: u64,
    tmp_dir: TempDir,
    settings: GlobalSettings,
    failed: bool,
}

impl Iterator for ExtSortedIterator<Line>
where
    Line: ExternallySortable,
{
    type Item = Result<Line, Box<dyn Error>>;

    /// # Errors
    ///
    /// This method can fail due to issues reading intermediate sorted chunks
    /// from disk, or due to serde deserialization issues
    fn next(&mut self) -> Option<Self::Item> {
        if self.failed {
            return None;
        }
        // fill up any empty buffers
        let mut empty = true;
        for chunk_num in 0..self.chunks {
            if self.buffers[chunk_num as usize].is_empty() {
                let mut f = match File::open(self.tmp_dir.path().join(chunk_num.to_string())) {
                    Ok(f) => f,
                    Err(e) => {
                        self.failed = true;
                        return Some(Err(Box::new(e)));
                    }
                };
                match f.seek(Start(self.chunk_offsets[chunk_num as usize])) {
                    Ok(_) => (),
                    Err(e) => {
                        self.failed = true;
                        return Some(Err(Box::new(e)));
                    }
                }
                let bytes_read =
                    match fill_buff(&mut self.buffers[chunk_num as usize], f, self.max_per_chunk) {
                        Ok(bytes_read) => bytes_read,
                        Err(e) => {
                            self.failed = true;
                            return Some(Err(e));
                        }
                    };
                self.chunk_offsets[chunk_num as usize] += bytes_read;
                if !self.buffers[chunk_num as usize].is_empty() {
                    empty = false;
                }
            } else {
                empty = false;
            }
        }
        if empty {
            return None;
        }

        // find the next record to write
        // check is_empty() before unwrap()ing
        let mut idx = 0;
        for chunk_num in 0..self.chunks as usize {
            if !self.buffers[chunk_num].is_empty() {
                if self.buffers[idx].is_empty()
                    || (super::compare_by)(
                        self.buffers[chunk_num].front().unwrap(),
                        self.buffers[idx].front().unwrap(),
                        &self.settings,
                    ) == Less
                {
                    idx = chunk_num;
                }
            }
        }

        // unwrap due to checks above
        let r = self.buffers[idx].pop_front().unwrap();
        Some(Ok(r))
    }
}

/// Perform an external sort on an unsorted stream of incoming data
pub struct ExternalSorter<Line>
where
    Line: ExternallySortable,
{
    tmp_dir: Option<PathBuf>,
    buffer_bytes: u64,
    phantom: PhantomData<Line>,
    settings: GlobalSettings,
}

impl ExternalSorter<Line>
where
    Line: ExternallySortable,
{
    /// Create a new `ExternalSorter` with a specified memory buffer and
    /// temporary directory
    pub fn new(
        buffer_bytes: u64,
        tmp_dir: Option<PathBuf>,
        settings: GlobalSettings,
    ) -> ExternalSorter<Line> {
        ExternalSorter {
            buffer_bytes,
            tmp_dir,
            phantom: PhantomData,
            settings,
        }
    }

    /// Sort (based on `compare`) the `T`s provided by `unsorted` and return an
    /// iterator
    ///
    /// # Errors
    ///
    /// This method can fail due to issues writing intermediate sorted chunks
    /// to disk, or due to serde serialization issues
    pub fn sort_by<I>(
        &self,
        unsorted: I,
        settings: GlobalSettings,
    ) -> Result<ExtSortedIterator<Line>, Box<dyn Error>>
    where
        I: Iterator<Item = Line>,
    {
        let tmp_dir = match self.tmp_dir {
            Some(ref p) => TempDir::new_in(p, "uutils_sort")?,
            None => TempDir::new("uutils_sort")?,
        };
        // creating the thing we need to return first due to the face that we need to
        // borrow tmp_dir and move it out
        let mut iter = ExtSortedIterator {
            buffers: Vec::new(),
            chunk_offsets: Vec::new(),
            max_per_chunk: 0,
            chunks: 0,
            tmp_dir,
            settings,
            failed: false,
        };

        {
            let mut total_read = 0;
            let mut chunk = Vec::new();
            // Initial buffer is specified by user
            let mut adjusted_buffer_size = self.buffer_bytes;
            let (iter_size, _) = unsorted.size_hint();

            // make the initial chunks on disk
            for seq in unsorted {
                let seq_size = seq.get_size();
                total_read += seq_size;

                // GNU minimum is 16 * (sizeof struct + 2), but GNU uses about
                // 1/10 the memory that we do.  And GNU even says in the code it may
                // not work on small buffer sizes.
                //
                // The following seems to work pretty well, and has about the same max
                // RSS as lower minimum values.
                //
                let minimum_buffer_size: u64 = iter_size as u64 * seq_size / 8;

                adjusted_buffer_size =
                    // Grow buffer size for a struct/Line larger than buffer
                    if adjusted_buffer_size < seq_size {
                        seq_size
                    } else if adjusted_buffer_size < minimum_buffer_size {
                        minimum_buffer_size
                    } else {
                        adjusted_buffer_size
                    };
                chunk.push(seq);

                if total_read >= adjusted_buffer_size {
                    super::sort_by(&mut chunk, &self.settings);
                    self.write_chunk(
                        &iter.tmp_dir.path().join(iter.chunks.to_string()),
                        &mut chunk,
                    )?;
                    chunk.clear();
                    total_read = 0;
                    iter.chunks += 1;
                }
            }
            // write the last chunk
            if chunk.len() > 0 {
                super::sort_by(&mut chunk, &self.settings);
                self.write_chunk(
                    &iter.tmp_dir.path().join(iter.chunks.to_string()),
                    &mut chunk,
                )?;
                iter.chunks += 1;
            }

            // initialize buffers for each chunk
            //
            // Having a right sized buffer for each chunk for smallish values seems silly to me?
            //
            // We will have to have the entire iter in memory sometime right?
            // Set minimum to the size of the writer buffer, ~8K
            //
            const MINIMUM_READBACK_BUFFER: u64 = 8200;
            let right_sized_buffer = adjusted_buffer_size
                   .checked_div(iter.chunks)
                   .unwrap_or(adjusted_buffer_size);
            iter.max_per_chunk =
                if right_sized_buffer > MINIMUM_READBACK_BUFFER {
                    right_sized_buffer
                } else {
                    MINIMUM_READBACK_BUFFER
                };
            iter.buffers = vec![VecDeque::new(); iter.chunks as usize];
            iter.chunk_offsets = vec![0 as u64; iter.chunks as usize];
            for chunk_num in 0..iter.chunks {
                let offset = fill_buff(
                    &mut iter.buffers[chunk_num as usize],
                    File::open(iter.tmp_dir.path().join(chunk_num.to_string()))?,
                    iter.max_per_chunk,
                )?;
                iter.chunk_offsets[chunk_num as usize] = offset;
            }
        }

        Ok(iter)
    }

    fn write_chunk(&self, file: &PathBuf, chunk: &mut Vec<Line>) -> Result<(), Box<dyn Error>> {
        let new_file = OpenOptions::new().create(true).append(true).open(file)?;
        let mut buf_write = Box::new(BufWriter::new(new_file)) as Box<dyn Write>;
        for s in chunk {
            let mut serialized = serde_json::to_string(&s).expect("JSON write error: ");
            serialized.push_str("\n");
            buf_write.write(serialized.as_bytes())?;
        }
        buf_write.flush()?;

        Ok(())
    }
}

fn fill_buff<Line>(
    vec: &mut VecDeque<Line>,
    file: File,
    max_bytes: u64,
) -> Result<u64, Box<dyn Error>>
where
    Line: ExternallySortable,
{
    let mut total_read = 0;
    let mut bytes_read = 0;
    for line in BufReader::new(file).lines() {
        let line_s = line?;
        bytes_read += line_s.len() + 1;
        // This is where the bad stuff happens usually
        let deserialized: Line = serde_json::from_str(&line_s).expect("JSON read error: ");
        total_read += deserialized.get_size();
        vec.push_back(deserialized);
        if total_read > max_bytes {
            break;
        }
    }

    Ok(bytes_read as u64)
}
<|endoftext|>


<commit_before>#[cfg(unix)]
pub use self::unix::*;
#[cfg(target_os = "linux")]
pub use self::linux::*;
#[cfg(windows)]
pub use self::windows::*;

// Add any operating systems we support here
#[cfg(not(any(target_os = "linux")))]
pub use self::default::*;

#[cfg(unix)]
mod unix;
#[cfg(target_os = "linux")]
mod linux;
#[cfg(windows)]
mod windows;

// Add any operating systems we support here
#[cfg(not(any(target_os = "linux")))]
mod default;
<commit_msg>uucore: allow the Linux zero-copy setup on Android<commit_after>#[cfg(unix)]
pub use self::unix::*;
#[cfg(any(target_os = "linux", target_os = "android"))]
pub use self::linux::*;
#[cfg(windows)]
pub use self::windows::*;

// Add any operating systems we support here
#[cfg(not(any(target_os = "linux", target_os = "android")))]
pub use self::default::*;

#[cfg(unix)]
mod unix;
#[cfg(any(target_os = "linux", target_os = "android"))]
mod linux;
#[cfg(windows)]
mod windows;

// Add any operating systems we support here
#[cfg(not(any(target_os = "linux", target_os = "android")))]
mod default;
<|endoftext|>


<commit_before>#![deny(
    non_camel_case_types,
    non_snake_case,
    path_statements,
    trivial_numeric_casts,
    unstable_features,
    unused_allocation,
    unused_import_braces,
    unused_imports,
    unused_mut,
    unused_qualifications,
    while_true,
)]

#[macro_use] extern crate log;
#[macro_use] extern crate itertools;
#[cfg(unix)] extern crate xdg_basedir;
extern crate tempfile;

extern crate clap;
extern crate toml;

extern crate libimagstore;
extern crate libimagstorestdhook;
extern crate libimagutil;

mod configuration;
mod logger;

pub mod edit;
pub mod error;
pub mod runtime;

<commit_msg>Enable lints<commit_after>#![deny(
    dead_code,
    non_camel_case_types,
    non_snake_case,
    path_statements,
    trivial_numeric_casts,
    unstable_features,
    unused_allocation,
    unused_import_braces,
    unused_imports,
    unused_must_use,
    unused_mut,
    unused_qualifications,
    while_true,
)]

#[macro_use] extern crate log;
#[macro_use] extern crate itertools;
#[cfg(unix)] extern crate xdg_basedir;
extern crate tempfile;

extern crate clap;
extern crate toml;

extern crate libimagstore;
extern crate libimagstorestdhook;
extern crate libimagutil;

mod configuration;
mod logger;

pub mod edit;
pub mod error;
pub mod runtime;

<|endoftext|>


<commit_before>

use visit::IntoNeighbors;
use visit::{VisitMap, Visitable};

/// Strictly monotonically increasing event time for a depth first search.
#[derive(Copy, Clone, Debug, PartialEq, PartialOrd, Eq, Ord, Default)]
pub struct Time(pub usize);

/// A depth first search (DFS) visitor event.
#[derive(Copy, Clone, Debug)]
pub enum DfsEvent<N> {
    Discover(N, Time),
    /// An edge of the tree formed by the traversal.
    TreeEdge(N, N),
    /// An edge to an already visited node.
    BackEdge(N, N),
    /// A cross or forward edge.
    ///
    /// For an edge *(u, v)*, if the discover time of *v* is greater than *u*,
    /// then it is a forward edge, else a cross edge.
    CrossForwardEdge(N, N),
    Finish(N, Time),
}

/// Return if the expression is a break value.
macro_rules! try_control {
    ($e:expr) => {
        match $e {
            x => if x.should_break() {
                return x;
            }
        }
    }
}

/// Control flow for callbacks.
///
/// `Break` can carry a value.
#[derive(Copy, Clone, Debug)]
pub enum Control<B> {
    Continue,
    Break(B),
}

impl<B> Control<B> {
    pub fn breaking() -> Control<()> { Control::Break(()) }
    /// Get the value in `Control::Break(_)`, if present.
    pub fn break_value(self) -> Option<B> {
        match self {
            Control::Continue => None,
            Control::Break(b) => Some(b),
        }
    }
}

/// Control flow for callbacks.
///
/// The empty return value `()` is equivalent to continue.
pub trait ControlFlow {
    fn continuing() -> Self;
    fn should_break(&self) -> bool;
}

impl ControlFlow for () {
    fn continuing() { }
    #[inline]
    fn should_break(&self) -> bool { false }
}

impl<B> ControlFlow for Control<B> {
    fn continuing() -> Self { Control::Continue }
    fn should_break(&self) -> bool {
        if let Control::Break(_) = *self { true } else { false }
    }
}

impl<E> ControlFlow for Result<(), E> {
    fn continuing() -> Self { Ok(()) }
    fn should_break(&self) -> bool {
        if let Err(_) = *self { true } else { false }
    }
}

/// The default is `Continue`.
impl<B> Default for Control<B> {
    fn default() -> Self { Control::Continue }
}

/// A recursive depth first search.
///
/// Starting points are the nodes in the iterator `starts` (specify just one
/// start vertex *x* by using `Some(x)`).
///
/// The traversal emits discovery and finish events for each reachable vertex,
/// and edge classification of each reachable edge. `visitor` is called for each
/// event, see [`DfsEvent`][de] for possible values.
///
/// If the return value of the visitor is simply `()`, the visit runs until it
/// is finished. If the return value is a `Control<B>`, it can be used to
/// break the visit early, and the last control value is returned by the
/// function.
///
/// [de]: enum.DfsEvent.html
///
/// # Example
///
/// Find a path from vertex 0 to 5, and exit the visit as soon as we reach
/// the goal vertex.
///
/// ```
/// use petgraph::prelude::*;
/// use petgraph::graph::node_index as n;
/// use petgraph::visit::depth_first_search;
/// use petgraph::visit::{DfsEvent, Control};
///
/// let gr: Graph<(), ()> = Graph::from_edges(&[
///     (0, 1), (0, 2), (0, 3),
///     (1, 3),
///     (2, 3), (2, 4),
///     (4, 0), (4, 5),
/// ]);
///
/// // record each predecessor, mapping node → node
/// let mut predecessor = vec![NodeIndex::end(); gr.node_count()];
/// let start = n(0);
/// let goal = n(5);
/// depth_first_search(&gr, Some(start), |event| {
///     if let DfsEvent::TreeEdge(u, v) = event {
///         predecessor[v.index()] = u;
///         if v == goal {
///             return Control::Break(v);
///         }
///     }
///     Control::Continue
/// });
///
/// let mut next = goal;
/// let mut path = vec![next];
/// while next != start {
///     let pred = predecessor[next.index()];
///     path.push(pred);
///     next = pred;
/// }
/// path.reverse();
/// assert_eq!(&path, &[n(0), n(2), n(4), n(5)]);
/// ```
pub fn depth_first_search<G, I, F, C>(graph: G, starts: I, mut visitor: F) -> C
    where G: IntoNeighbors + Visitable,
          I: IntoIterator<Item=G::NodeId>,
          F: FnMut(DfsEvent<G::NodeId>) -> C,
          C: ControlFlow,
{
    let time = &mut Time(0);
    let discovered = &mut graph.visit_map();
    let finished = &mut graph.visit_map();

    for start in starts {
        try_control!(dfs_visitor(graph, start, &mut visitor, discovered, finished, time));
    }
    C::continuing()
}

fn dfs_visitor<G, F, C>(graph: G, u: G::NodeId, visitor: &mut F,
                     discovered: &mut G::Map, finished: &mut G::Map,
                     time: &mut Time) -> C
    where G: IntoNeighbors + Visitable,
          F: FnMut(DfsEvent<G::NodeId>) -> C,
          C: ControlFlow,
{
    if !discovered.visit(u) {
        return C::continuing();
    }
    try_control!(visitor(DfsEvent::Discover(u, time_post_inc(time))));
    for v in graph.neighbors(u) {
        if !discovered.is_visited(&v) {
            try_control!(visitor(DfsEvent::TreeEdge(u, v)));
            try_control!(dfs_visitor(graph, v, visitor, discovered, finished, time));
        } else if !finished.is_visited(&v) {
            try_control!(visitor(DfsEvent::BackEdge(u, v)));
        } else {
            try_control!(visitor(DfsEvent::CrossForwardEdge(u, v)));
        }
    }
    let first_finish = finished.visit(u);
    debug_assert!(first_finish);
    try_control!(visitor(DfsEvent::Finish(u, time_post_inc(time))));
    C::continuing()
}

fn time_post_inc(x: &mut Time) -> Time {
    let v = *x;
    x.0 += 1;
    v
}
<commit_msg>FEAT: Add Hash impl for Time<commit_after>

use visit::IntoNeighbors;
use visit::{VisitMap, Visitable};

/// Strictly monotonically increasing event time for a depth first search.
#[derive(Copy, Clone, Debug, PartialEq, PartialOrd, Eq, Ord, Default, Hash)]
pub struct Time(pub usize);

/// A depth first search (DFS) visitor event.
#[derive(Copy, Clone, Debug)]
pub enum DfsEvent<N> {
    Discover(N, Time),
    /// An edge of the tree formed by the traversal.
    TreeEdge(N, N),
    /// An edge to an already visited node.
    BackEdge(N, N),
    /// A cross or forward edge.
    ///
    /// For an edge *(u, v)*, if the discover time of *v* is greater than *u*,
    /// then it is a forward edge, else a cross edge.
    CrossForwardEdge(N, N),
    Finish(N, Time),
}

/// Return if the expression is a break value.
macro_rules! try_control {
    ($e:expr) => {
        match $e {
            x => if x.should_break() {
                return x;
            }
        }
    }
}

/// Control flow for callbacks.
///
/// `Break` can carry a value.
#[derive(Copy, Clone, Debug)]
pub enum Control<B> {
    Continue,
    Break(B),
}

impl<B> Control<B> {
    pub fn breaking() -> Control<()> { Control::Break(()) }
    /// Get the value in `Control::Break(_)`, if present.
    pub fn break_value(self) -> Option<B> {
        match self {
            Control::Continue => None,
            Control::Break(b) => Some(b),
        }
    }
}

/// Control flow for callbacks.
///
/// The empty return value `()` is equivalent to continue.
pub trait ControlFlow {
    fn continuing() -> Self;
    fn should_break(&self) -> bool;
}

impl ControlFlow for () {
    fn continuing() { }
    #[inline]
    fn should_break(&self) -> bool { false }
}

impl<B> ControlFlow for Control<B> {
    fn continuing() -> Self { Control::Continue }
    fn should_break(&self) -> bool {
        if let Control::Break(_) = *self { true } else { false }
    }
}

impl<E> ControlFlow for Result<(), E> {
    fn continuing() -> Self { Ok(()) }
    fn should_break(&self) -> bool {
        if let Err(_) = *self { true } else { false }
    }
}

/// The default is `Continue`.
impl<B> Default for Control<B> {
    fn default() -> Self { Control::Continue }
}

/// A recursive depth first search.
///
/// Starting points are the nodes in the iterator `starts` (specify just one
/// start vertex *x* by using `Some(x)`).
///
/// The traversal emits discovery and finish events for each reachable vertex,
/// and edge classification of each reachable edge. `visitor` is called for each
/// event, see [`DfsEvent`][de] for possible values.
///
/// If the return value of the visitor is simply `()`, the visit runs until it
/// is finished. If the return value is a `Control<B>`, it can be used to
/// break the visit early, and the last control value is returned by the
/// function.
///
/// [de]: enum.DfsEvent.html
///
/// # Example
///
/// Find a path from vertex 0 to 5, and exit the visit as soon as we reach
/// the goal vertex.
///
/// ```
/// use petgraph::prelude::*;
/// use petgraph::graph::node_index as n;
/// use petgraph::visit::depth_first_search;
/// use petgraph::visit::{DfsEvent, Control};
///
/// let gr: Graph<(), ()> = Graph::from_edges(&[
///     (0, 1), (0, 2), (0, 3),
///     (1, 3),
///     (2, 3), (2, 4),
///     (4, 0), (4, 5),
/// ]);
///
/// // record each predecessor, mapping node → node
/// let mut predecessor = vec![NodeIndex::end(); gr.node_count()];
/// let start = n(0);
/// let goal = n(5);
/// depth_first_search(&gr, Some(start), |event| {
///     if let DfsEvent::TreeEdge(u, v) = event {
///         predecessor[v.index()] = u;
///         if v == goal {
///             return Control::Break(v);
///         }
///     }
///     Control::Continue
/// });
///
/// let mut next = goal;
/// let mut path = vec![next];
/// while next != start {
///     let pred = predecessor[next.index()];
///     path.push(pred);
///     next = pred;
/// }
/// path.reverse();
/// assert_eq!(&path, &[n(0), n(2), n(4), n(5)]);
/// ```
pub fn depth_first_search<G, I, F, C>(graph: G, starts: I, mut visitor: F) -> C
    where G: IntoNeighbors + Visitable,
          I: IntoIterator<Item=G::NodeId>,
          F: FnMut(DfsEvent<G::NodeId>) -> C,
          C: ControlFlow,
{
    let time = &mut Time(0);
    let discovered = &mut graph.visit_map();
    let finished = &mut graph.visit_map();

    for start in starts {
        try_control!(dfs_visitor(graph, start, &mut visitor, discovered, finished, time));
    }
    C::continuing()
}

fn dfs_visitor<G, F, C>(graph: G, u: G::NodeId, visitor: &mut F,
                     discovered: &mut G::Map, finished: &mut G::Map,
                     time: &mut Time) -> C
    where G: IntoNeighbors + Visitable,
          F: FnMut(DfsEvent<G::NodeId>) -> C,
          C: ControlFlow,
{
    if !discovered.visit(u) {
        return C::continuing();
    }
    try_control!(visitor(DfsEvent::Discover(u, time_post_inc(time))));
    for v in graph.neighbors(u) {
        if !discovered.is_visited(&v) {
            try_control!(visitor(DfsEvent::TreeEdge(u, v)));
            try_control!(dfs_visitor(graph, v, visitor, discovered, finished, time));
        } else if !finished.is_visited(&v) {
            try_control!(visitor(DfsEvent::BackEdge(u, v)));
        } else {
            try_control!(visitor(DfsEvent::CrossForwardEdge(u, v)));
        }
    }
    let first_finish = finished.visit(u);
    debug_assert!(first_finish);
    try_control!(visitor(DfsEvent::Finish(u, time_post_inc(time))));
    C::continuing()
}

fn time_post_inc(x: &mut Time) -> Time {
    let v = *x;
    x.0 += 1;
    v
}
<|endoftext|>