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

mod bridge_converter;
mod byvalue_checker;
mod cpp_postprocessor;
mod preprocessor_parse_callbacks;

#[cfg(test)]
mod integration_tests;

use proc_macro2::Span;
use proc_macro2::TokenStream as TokenStream2;
use std::{fmt::Display, path::PathBuf};

use cxx_gen::GeneratedCode;
use indoc::indoc;
use quote::ToTokens;
use syn::parse::{Parse, ParseStream, Result as ParseResult};
use syn::{parse_quote, Ident, ItemMod, Macro, TypePath};

use cpp_postprocessor::EncounteredType;
use log::{debug, info, warn};
use osstrtools::OsStrTools;
use preprocessor_parse_callbacks::{PreprocessorDefinitions, PreprocessorParseCallbacks};
use std::rc::Rc;
use std::sync::Mutex;

/// Turn off and remove this constant when
/// https://github.com/rust-lang/rust/pull/75857 is fixed.
/// This should allow removal of the entire cpp_postprocessor module.
const TEMPORARY_HACK_TO_AVOID_REDEFINITIONS: bool = true;

const BINDGEN_BLOCKLIST: &[&str] = &["std.*", ".*mbstate_t.*"];

/// Any time we store a type name, we should use this.
/// At the moment it's just a string, but one day it will need to become
/// sufficiently intelligent to handle namespaces.
#[derive(Debug, PartialEq, PartialOrd, Eq, Hash, Clone)]
pub struct TypeName(String);

impl TypeName {
    fn from_ident(id: &Ident) -> Self {
        TypeName(id.to_string())
    }

    fn from_type_path(p: &TypePath) -> Self {
        // TODO better handle generics, multi-segment paths, etc.
        TypeName::from_ident(&p.path.segments.last().unwrap().ident)
    }

    fn new(id: &str) -> Self {
        TypeName(id.to_string())
    }

    fn to_ident(&self) -> Ident {
        Ident::new(&self.0, Span::call_site())
    }
}

impl Display for TypeName {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.write_str(&self.0)
    }
}

#[derive(Debug)]
pub enum Error {
    Io(std::io::Error),
    Bindgen(()),
    CxxGen(cxx_gen::Error),
    Parsing(syn::Error),
    NoAutoCxxInc,
    CouldNotCanoncalizeIncludeDir(PathBuf),
    Conversion(bridge_converter::ConvertError),
}

pub type Result<T, E = Error> = std::result::Result<T, E>;

pub enum CppInclusion {
    Define(String),
    Header(String),
}

/// Core of the autocxx engine.
/// TODO - consider whether this 'engine' crate should actually be a
/// directory of source symlinked from all the other sub-crates, so that
/// we avoid exposing an external interface from this code.
pub struct IncludeCpp {
    inclusions: Vec<CppInclusion>,
    allowlist: Vec<String>, // not TypeName as it may be functions or whatever.
    pod_types: Vec<TypeName>,
    preconfigured_inc_dirs: Option<std::ffi::OsString>,
    parse_only: bool,
    preprocessor_definitions: Rc<Mutex<PreprocessorDefinitions>>,
}

impl Parse for IncludeCpp {
    fn parse(input: ParseStream) -> ParseResult<Self> {
        Self::new_from_parse_stream(input)
    }
}

fn dump_generated_code(gen: GeneratedCode) -> Result<GeneratedCode> {
    info!(
        "CXX:\n{}",
        String::from_utf8(gen.implementation.clone()).unwrap()
    );
    info!(
        "header:\n{}",
        String::from_utf8(gen.header.clone()).unwrap()
    );
    Ok(gen)
}

/// Prelude of C++ for squirting into bindgen. This configures
/// bindgen to output simpler types to replace some STL types
/// that bindgen just can't cope with. Although we then replace
/// those types with cxx types (e.g. UniquePtr), this intermediate
/// step is still necessary because bindgen can't otherwise
/// give us the templated types (e.g. when faced with the STL
/// unique_ptr, bindgen would normally give us std_unique_ptr
/// as opposed to std_unique_ptr<T>.)
static PRELUDE: &str = indoc! {"
    /**
    *

    */
    template<typename T> class UniquePtr {
        T* ptr;
    };

    /**
    *

    */
    class CxxString {
        char* str_data;
    };
    \n"};

impl IncludeCpp {
    fn new_from_parse_stream(input: ParseStream) -> syn::Result<Self> {
        // Takes as inputs:
        // 1. List of headers to include
        // 2. List of #defines to include
        // 3. Allowlist

        let mut inclusions = Vec::new();
        let mut allowlist = Vec::new();
        let mut pod_types = Vec::new();
        let mut parse_only = false;

        while !input.is_empty() {
            let ident: syn::Ident = input.parse()?;
            if ident == "Header" {
                let args;
                syn::parenthesized!(args in input);
                let hdr: syn::LitStr = args.parse()?;
                inclusions.push(CppInclusion::Header(hdr.value()));
            } else if ident == "Allow" || ident == "AllowPOD" {
                let args;
                syn::parenthesized!(args in input);
                let allow: syn::LitStr = args.parse()?;
                allowlist.push(allow.value());
                if ident == "AllowPOD" {
                    pod_types.push(TypeName::new(&allow.value()));
                }
            } else if ident == "ParseOnly" {
                parse_only = true;
            } else {
                return Err(syn::Error::new(
                    ident.span(),
                    "expected Header, Allow or AllowPOD",
                ));
            }
            if input.is_empty() {
                break;
            }
            input.parse::<syn::Token![,]>()?;
        }

        Ok(IncludeCpp {
            inclusions,
            allowlist,
            pod_types,
            preconfigured_inc_dirs: None,
            parse_only,
            preprocessor_definitions: Rc::new(Mutex::new(PreprocessorDefinitions::new())),
        })
    }

    pub fn new_from_syn(mac: Macro) -> Result<Self> {
        mac.parse_body::<IncludeCpp>().map_err(Error::Parsing)
    }

    pub fn set_include_dirs<P: AsRef<std::ffi::OsStr>>(&mut self, include_dirs: P) {
        self.preconfigured_inc_dirs = Some(include_dirs.as_ref().into());
    }

    fn build_header(&self) -> String {
        let mut s = PRELUDE.to_string();
        for incl in &self.inclusions {
            let text = match incl {
                CppInclusion::Define(symbol) => format!("#define {}\n", symbol),
                CppInclusion::Header(path) => format!("#include \"{}\"\n", path),
            };
            s.push_str(&text);
        }
        s
    }

    fn determine_incdirs(&self) -> Result<Vec<PathBuf>> {
        let inc_dirs = match &self.preconfigured_inc_dirs {
            Some(d) => d.clone(),
            None => std::env::var_os("AUTOCXX_INC").ok_or(Error::NoAutoCxxInc)?,
        };
        // TODO consider if we can or should look up the include path automatically
        // instead of requiring callers always to set AUTOCXX_INC.
        let multi_path_separator = if std::path::MAIN_SEPARATOR == '/' {
            b':'
        } else {
            b';'
        }; // there's probably a crate for this
        let splitter = [multi_path_separator];
        let inc_dirs = inc_dirs.split(&splitter[0..1]);
        let mut inc_dir_paths = Vec::new();
        for inc_dir in inc_dirs {
            let p: PathBuf = inc_dir.into();
            let p = p
                .canonicalize()
                .map_err(|_| Error::CouldNotCanoncalizeIncludeDir(p))?;
            inc_dir_paths.push(p);
        }
        Ok(inc_dir_paths)
    }

    fn make_bindgen_builder(&self) -> Result<bindgen::Builder> {
        let inc_dirs = self.determine_incdirs()?;

        let full_header = self.build_header();
        debug!("Full header: {}", full_header);
        debug!("Inc dir: {:?}", inc_dirs);

        // TODO support different C++ versions
        let mut builder = bindgen::builder()
            .clang_args(&["-x", "c++", "-std=c++14"])
            .derive_copy(false)
            .derive_debug(false)
            .parse_callbacks(Box::new(PreprocessorParseCallbacks::new(
                self.preprocessor_definitions.clone(),
            )))
            .default_enum_style(bindgen::EnumVariation::Rust {
                non_exhaustive: false,
            })
            .layout_tests(false); // TODO revisit later
        for item in BINDGEN_BLOCKLIST.iter() {
            builder = builder.blacklist_item(*item);
        }

        for inc_dir in inc_dirs {
            // TODO work with OsStrs here to avoid the .display()
            builder = builder.clang_arg(format!("-I{}", inc_dir.display()));
        }
        builder = builder.header_contents("example.hpp", &full_header);
        // 3. Passes allowlist and other options to the bindgen::Builder equivalent
        //    to --output-style=cxx --allowlist=<as passed in>
        for a in &self.allowlist {
            // TODO - allowlist type/functions/separately
            builder = builder
                .whitelist_type(a)
                .whitelist_function(a)
                .whitelist_var(a);
        }
        Ok(builder)
    }

    pub fn generate_rs(&self) -> Result<TokenStream2> {
        let (itemmod, _) = self.do_generation()?;
        Ok(itemmod.to_token_stream())
    }

    fn get_preprocessor_defs_mod(&self) -> Option<ItemMod> {
        let another_ref = self.preprocessor_definitions.clone();
        let m = another_ref.try_lock().unwrap().to_mod();
        m
    }

    fn do_generation(&self) -> Result<(Option<ItemMod>, Vec<EncounteredType>)> {
        // If we are in parse only mode, do nothing. This is used for
        // doc tests to ensure the parsing is valid, but we can't expect
        // valid C++ header files or linkers to allow a complete build.
        if self.parse_only {
            return Ok((None, Vec::new()));
        }
        // 4. (also respects environment variables to pick up more headers,
        //     include paths and #defines)
        // Then:
        // 1. Builds an overall C++ header with all those #defines and #includes
        // 2. Passes it to bindgen::Builder::header
        let bindings = self
            .make_bindgen_builder()?
            .generate()
            .map_err(Error::Bindgen)?;
        // This bindings object is actually a TokenStream internally and we're wasting
        // effort converting to and from string. We could enhance the bindgen API
        // in future.
        let bindings = bindings.to_string();
        // Manually add the mod ffi {} so that we can ask syn to parse
        // into a single construct.
        let bindings = format!("mod ffi {{ {} }}", bindings);
        info!("Bindings: {}", bindings);
        let bindings = syn::parse_str::<ItemMod>(&bindings).map_err(Error::Parsing)?;

        let mut include_list = Vec::new();
        for incl in &self.inclusions {
            match incl {
                CppInclusion::Header(ref hdr) => {
                    include_list.push(hdr.clone());
                }
                CppInclusion::Define(_) => warn!("Currently no way to define! within cxx"),
            }
        }

        let mut converter = bridge_converter::BridgeConverter::new(
            include_list,
            self.pod_types.clone(), // TODO take self by value to avoid clone.
            TEMPORARY_HACK_TO_AVOID_REDEFINITIONS,
        );
        let results = converter.convert(bindings).map_err(Error::Conversion)?;
        let mut items = results.items;
        let types_encountered = results.types_encountered;
        if let Some(itemmod) = self.get_preprocessor_defs_mod() {
            items.push(syn::Item::Mod(itemmod));
        }
        let new_bindings: ItemMod = parse_quote! {
            mod ffi {
                #(#items)*
            }
        };
        info!(
            "New bindings: {}",
            new_bindings.to_token_stream().to_string()
        );
        Ok((Some(new_bindings), types_encountered))
    }

    pub fn generate_h_and_cxx(self) -> Result<GeneratedCode> {
        let (rs, types_encountered) = self.do_generation()?;
        let rs = match rs {
            Some(itemmod) => itemmod.into_token_stream(),
            None => TokenStream2::new(),
        };
        let opt = cxx_gen::Opt::default();
        cxx_gen::generate_header_and_cc(rs, &opt)
            .map_err(Error::CxxGen)
            .and_then(dump_generated_code)
            .and_then(|gen| {
                if TEMPORARY_HACK_TO_AVOID_REDEFINITIONS {
                    Ok(GeneratedCode {
                        header: cpp_postprocessor::disable_types(gen.header, &types_encountered),
                        implementation: cpp_postprocessor::disable_types(
                            gen.implementation,
                            &types_encountered,
                        ),
                    })
                } else {
                    Ok(gen)
                }
            })
            .and_then(dump_generated_code)
    }

    pub fn include_dirs(&self) -> Result<Vec<PathBuf>> {
        self.determine_incdirs()
    }
}
<commit_msg>Adding test.<commit_after>// Copyright 2020 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//    https://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

mod bridge_converter;
mod byvalue_checker;
mod cpp_postprocessor;
mod preprocessor_parse_callbacks;

#[cfg(test)]
mod integration_tests;

use proc_macro2::Span;
use proc_macro2::TokenStream as TokenStream2;
use std::{fmt::Display, path::PathBuf};

use cxx_gen::GeneratedCode;
use indoc::indoc;
use quote::ToTokens;
use syn::parse::{Parse, ParseStream, Result as ParseResult};
use syn::{parse_quote, Ident, ItemMod, Macro, TypePath};

use cpp_postprocessor::EncounteredType;
use log::{debug, info, warn};
use osstrtools::OsStrTools;
use preprocessor_parse_callbacks::{PreprocessorDefinitions, PreprocessorParseCallbacks};
use std::rc::Rc;
use std::sync::Mutex;

/// Turn off and remove this constant when
/// https://github.com/rust-lang/rust/pull/75857 is fixed.
/// This should allow removal of the entire cpp_postprocessor module.
const TEMPORARY_HACK_TO_AVOID_REDEFINITIONS: bool = true;

const BINDGEN_BLOCKLIST: &[&str] = &["std.*", ".*mbstate_t.*"];

/// Any time we store a type name, we should use this.
/// At the moment it's just a string, but one day it will need to become
/// sufficiently intelligent to handle namespaces.
#[derive(Debug, PartialEq, PartialOrd, Eq, Hash, Clone)]
pub struct TypeName(String);

impl TypeName {
    fn from_ident(id: &Ident) -> Self {
        TypeName(id.to_string())
    }

    fn from_type_path(p: &TypePath) -> Self {
        // TODO better handle generics, multi-segment paths, etc.
        TypeName::from_ident(&p.path.segments.last().unwrap().ident)
    }

    fn new(id: &str) -> Self {
        TypeName(id.to_string())
    }

    fn to_ident(&self) -> Ident {
        Ident::new(&self.0, Span::call_site())
    }
}

impl Display for TypeName {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.write_str(&self.0)
    }
}

#[derive(Debug)]
pub enum Error {
    Io(std::io::Error),
    Bindgen(()),
    CxxGen(cxx_gen::Error),
    Parsing(syn::Error),
    NoAutoCxxInc,
    CouldNotCanoncalizeIncludeDir(PathBuf),
    Conversion(bridge_converter::ConvertError),
}

pub type Result<T, E = Error> = std::result::Result<T, E>;

pub enum CppInclusion {
    Define(String),
    Header(String),
}

/// Core of the autocxx engine.
/// TODO - consider whether this 'engine' crate should actually be a
/// directory of source symlinked from all the other sub-crates, so that
/// we avoid exposing an external interface from this code.
pub struct IncludeCpp {
    inclusions: Vec<CppInclusion>,
    allowlist: Vec<String>, // not TypeName as it may be functions or whatever.
    pod_types: Vec<TypeName>,
    preconfigured_inc_dirs: Option<std::ffi::OsString>,
    parse_only: bool,
    preprocessor_definitions: Rc<Mutex<PreprocessorDefinitions>>,
}

impl Parse for IncludeCpp {
    fn parse(input: ParseStream) -> ParseResult<Self> {
        Self::new_from_parse_stream(input)
    }
}

fn dump_generated_code(gen: GeneratedCode) -> Result<GeneratedCode> {
    info!(
        "CXX:\n{}",
        String::from_utf8(gen.implementation.clone()).unwrap()
    );
    info!(
        "header:\n{}",
        String::from_utf8(gen.header.clone()).unwrap()
    );
    Ok(gen)
}

/// Prelude of C++ for squirting into bindgen. This configures
/// bindgen to output simpler types to replace some STL types
/// that bindgen just can't cope with. Although we then replace
/// those types with cxx types (e.g. UniquePtr), this intermediate
/// step is still necessary because bindgen can't otherwise
/// give us the templated types (e.g. when faced with the STL
/// unique_ptr, bindgen would normally give us std_unique_ptr
/// as opposed to std_unique_ptr<T>.)
static PRELUDE: &str = indoc! {"
    /**
    *

    */
    template<typename T> class UniquePtr {
        T* ptr;
    };

    /**
    *

    */
    class CxxString {
        char* str_data;
    };
    \n"};

impl IncludeCpp {
    fn new_from_parse_stream(input: ParseStream) -> syn::Result<Self> {
        // Takes as inputs:
        // 1. List of headers to include
        // 2. List of #defines to include
        // 3. Allowlist

        let mut inclusions = Vec::new();
        let mut allowlist = Vec::new();
        let mut pod_types = Vec::new();
        let mut parse_only = false;

        while !input.is_empty() {
            let ident: syn::Ident = input.parse()?;
            if ident == "Header" {
                let args;
                syn::parenthesized!(args in input);
                let hdr: syn::LitStr = args.parse()?;
                inclusions.push(CppInclusion::Header(hdr.value()));
            } else if ident == "Allow" || ident == "AllowPOD" {
                let args;
                syn::parenthesized!(args in input);
                let allow: syn::LitStr = args.parse()?;
                allowlist.push(allow.value());
                if ident == "AllowPOD" {
                    pod_types.push(TypeName::new(&allow.value()));
                }
            } else if ident == "ParseOnly" {
                parse_only = true;
            } else {
                return Err(syn::Error::new(
                    ident.span(),
                    "expected Header, Allow or AllowPOD",
                ));
            }
            if input.is_empty() {
                break;
            }
            input.parse::<syn::Token![,]>()?;
        }

        Ok(IncludeCpp {
            inclusions,
            allowlist,
            pod_types,
            preconfigured_inc_dirs: None,
            parse_only,
            preprocessor_definitions: Rc::new(Mutex::new(PreprocessorDefinitions::new())),
        })
    }

    pub fn new_from_syn(mac: Macro) -> Result<Self> {
        mac.parse_body::<IncludeCpp>().map_err(Error::Parsing)
    }

    pub fn set_include_dirs<P: AsRef<std::ffi::OsStr>>(&mut self, include_dirs: P) {
        self.preconfigured_inc_dirs = Some(include_dirs.as_ref().into());
    }

    fn build_header(&self) -> String {
        let mut s = PRELUDE.to_string();
        for incl in &self.inclusions {
            let text = match incl {
                CppInclusion::Define(symbol) => format!("#define {}\n", symbol),
                CppInclusion::Header(path) => format!("#include \"{}\"\n", path),
            };
            s.push_str(&text);
        }
        s
    }

    fn determine_incdirs(&self) -> Result<Vec<PathBuf>> {
        let inc_dirs = match &self.preconfigured_inc_dirs {
            Some(d) => d.clone(),
            None => std::env::var_os("AUTOCXX_INC").ok_or(Error::NoAutoCxxInc)?,
        };
        // TODO consider if we can or should look up the include path automatically
        // instead of requiring callers always to set AUTOCXX_INC.
        let multi_path_separator = if std::path::MAIN_SEPARATOR == '/' {
            b':'
        } else {
            b';'
        }; // there's probably a crate for this
        let splitter = [multi_path_separator];
        let inc_dirs = inc_dirs.split(&splitter[0..1]);
        let mut inc_dir_paths = Vec::new();
        for inc_dir in inc_dirs {
            let p: PathBuf = inc_dir.into();
            let p = p
                .canonicalize()
                .map_err(|_| Error::CouldNotCanoncalizeIncludeDir(p))?;
            inc_dir_paths.push(p);
        }
        Ok(inc_dir_paths)
    }

    fn make_bindgen_builder(&self) -> Result<bindgen::Builder> {
        let inc_dirs = self.determine_incdirs()?;

        let full_header = self.build_header();
        debug!("Full header: {}", full_header);
        debug!("Inc dir: {:?}", inc_dirs);

        // TODO support different C++ versions
        let mut builder = bindgen::builder()
            .clang_args(&["-x", "c++", "-std=c++14"])
            .derive_copy(false)
            .derive_debug(false)
            .parse_callbacks(Box::new(PreprocessorParseCallbacks::new(
                self.preprocessor_definitions.clone(),
            )))
            .default_enum_style(bindgen::EnumVariation::Rust {
                non_exhaustive: false,
            })
            .layout_tests(false); // TODO revisit later
        for item in BINDGEN_BLOCKLIST.iter() {
            builder = builder.blacklist_item(*item);
        }

        for inc_dir in inc_dirs {
            // TODO work with OsStrs here to avoid the .display()
            builder = builder.clang_arg(format!("-I{}", inc_dir.display()));
        }
        builder = builder.header_contents("example.hpp", &full_header);
        // 3. Passes allowlist and other options to the bindgen::Builder equivalent
        //    to --output-style=cxx --allowlist=<as passed in>
        for a in &self.allowlist {
            // TODO - allowlist type/functions/separately
            builder = builder
                .whitelist_type(a)
                .whitelist_function(a)
                .whitelist_var(a);
        }
        Ok(builder)
    }

    pub fn generate_rs(&self) -> Result<TokenStream2> {
        let (itemmod, _) = self.do_generation()?;
        Ok(itemmod.to_token_stream())
    }

    fn get_preprocessor_defs_mod(&self) -> Option<ItemMod> {
        let another_ref = self.preprocessor_definitions.clone();
        let m = another_ref.try_lock().unwrap().to_mod();
        m
    }

    fn do_generation(&self) -> Result<(Option<ItemMod>, Vec<EncounteredType>)> {
        // If we are in parse only mode, do nothing. This is used for
        // doc tests to ensure the parsing is valid, but we can't expect
        // valid C++ header files or linkers to allow a complete build.
        if self.parse_only {
            return Ok((None, Vec::new()));
        }
        // 4. (also respects environment variables to pick up more headers,
        //     include paths and #defines)
        // Then:
        // 1. Builds an overall C++ header with all those #defines and #includes
        // 2. Passes it to bindgen::Builder::header
        let bindings = self
            .make_bindgen_builder()?
            .generate()
            .map_err(Error::Bindgen)?;
        // This bindings object is actually a TokenStream internally and we're wasting
        // effort converting to and from string. We could enhance the bindgen API
        // in future.
        let bindings = bindings.to_string();
        // Manually add the mod ffi {} so that we can ask syn to parse
        // into a single construct.
        let bindings = format!("mod ffi {{ {} }}", bindings);
        info!("Bindings: {}", bindings);
        let bindings = syn::parse_str::<ItemMod>(&bindings).map_err(Error::Parsing)?;

        let mut include_list = Vec::new();
        for incl in &self.inclusions {
            match incl {
                CppInclusion::Header(ref hdr) => {
                    include_list.push(hdr.clone());
                }
                CppInclusion::Define(_) => warn!("Currently no way to define! within cxx"),
            }
        }

        let mut converter = bridge_converter::BridgeConverter::new(
            include_list,
            self.pod_types.clone(), // TODO take self by value to avoid clone.
            TEMPORARY_HACK_TO_AVOID_REDEFINITIONS,
        );
        let results = converter.convert(bindings).map_err(Error::Conversion)?;
        let mut items = results.items;
        let types_encountered = results.types_encountered;
        if let Some(itemmod) = self.get_preprocessor_defs_mod() {
            items.push(syn::Item::Mod(itemmod));
        }
        let new_bindings: ItemMod = parse_quote! {
            mod ffi {
                #(#items)*
            }
        };
        info!(
            "New bindings: {}",
            new_bindings.to_token_stream().to_string()
        );
        Ok((Some(new_bindings), types_encountered))
    }

    pub fn generate_h_and_cxx(self) -> Result<GeneratedCode> {
        let (rs, types_encountered) = self.do_generation()?;
        let rs = match rs {
            Some(itemmod) => itemmod.into_token_stream(),
            None => TokenStream2::new(),
        };
        let opt = cxx_gen::Opt::default();
        cxx_gen::generate_header_and_cc(rs, &opt)
            .map_err(Error::CxxGen)
            .and_then(dump_generated_code)
            .and_then(|gen| {
                if TEMPORARY_HACK_TO_AVOID_REDEFINITIONS {
                    Ok(GeneratedCode {
                        header: cpp_postprocessor::disable_types(gen.header, &types_encountered),
                        implementation: cpp_postprocessor::disable_types(
                            gen.implementation,
                            &types_encountered,
                        ),
                    })
                } else {
                    Ok(gen)
                }
            })
            .and_then(dump_generated_code)
    }

    pub fn include_dirs(&self) -> Result<Vec<PathBuf>> {
        self.determine_incdirs()
    }
}

#[cfg(test)]
mod tests {
    use crate::TypeName;

    #[test]
    fn test_typename() {
        let s = proc_macro2::Span::call_site();
        let id = syn::Ident::new("Bob", s);
        let tn = TypeName::from_ident(&id);
        assert_eq!(tn.to_ident(), id);
    }
}<|endoftext|>


<commit_before>extern crate elektra_sys;

use crate::{KeySetError, ReadableKey, StringKey, WriteableKey};
use bitflags::bitflags;
use std::convert::TryInto;
use std::ffi::CString;

#[derive(Debug)]
pub struct KeySet {
    ptr: std::ptr::NonNull<elektra_sys::KeySet>,
}

pub type Cursor = elektra_sys::cursor_t;

bitflags! {
    // Results in default value with bits: 0
    #[derive(Default)]
    pub struct LookupOption: elektra_sys::option_t {
        const KDB_O_NONE = elektra_sys::KDB_O_NONE as elektra_sys::option_t;
        const KDB_O_DEL = elektra_sys::KDB_O_DEL as elektra_sys::option_t;
        const KDB_O_POP = elektra_sys::KDB_O_POP as elektra_sys::option_t;
        const KDB_O_NODIR = elektra_sys::KDB_O_NODIR as elektra_sys::option_t;
        const KDB_O_DIRONLY = elektra_sys::KDB_O_DIRONLY as elektra_sys::option_t;
        const KDB_O_NOREMOVE = elektra_sys::KDB_O_NOREMOVE as elektra_sys::option_t;
        const KDB_O_REMOVEONLY = elektra_sys::KDB_O_REMOVEONLY as elektra_sys::option_t;
        const KDB_O_INACTIVE = elektra_sys::KDB_O_INACTIVE as elektra_sys::option_t;
        const KDB_O_SYNC = elektra_sys::KDB_O_SYNC as elektra_sys::option_t;
        const KDB_O_SORT = elektra_sys::KDB_O_SORT as elektra_sys::option_t;
        const KDB_O_NORECURSIVE = elektra_sys::KDB_O_NORECURSIVE as elektra_sys::option_t;
        const KDB_O_NOCASE = elektra_sys::KDB_O_NOCASE as elektra_sys::option_t;
        const KDB_O_WITHOWNER = elektra_sys::KDB_O_WITHOWNER as elektra_sys::option_t;
        const KDB_O_NOALL = elektra_sys::KDB_O_NOALL as elektra_sys::option_t;
    }
}

impl Drop for KeySet {
    fn drop(&mut self) {
        println!("Drop {:?}", self);
        unsafe {
            elektra_sys::ksDel(self.as_ptr());
        }
    }
}

impl Iterator for KeySet {
    type Item = StringKey;
    fn next(&mut self) -> Option<Self::Item> {
        let key_ptr = unsafe { elektra_sys::ksNext(self.as_ptr()) };
        if (key_ptr as *const elektra_sys::Key) == std::ptr::null() {
            None
        } else {
            Some(StringKey::from_ptr(key_ptr as *mut elektra_sys::Key))
        }
    }
}

impl KeySet {
    fn as_ptr(&mut self) -> *mut elektra_sys::KeySet {
        self.ptr.as_ptr()
    }

    fn as_ref(&self) -> &elektra_sys::KeySet {
        unsafe { self.ptr.as_ref() }
    }

    /// Create a new empty KeySet
    pub fn new() -> KeySet {
        let ks_ptr = unsafe { elektra_sys::ksNew(0, elektra_sys::KEY_END) };
        KeySet::from_ptr(ks_ptr)
    }

    /// Create a new KeySet that allocates enough space for the
    /// given capacity of keys
    pub fn with_capacity(capacity: usize) -> Self {
        let ks_ptr = unsafe { elektra_sys::ksNew(capacity, elektra_sys::KEY_END) };
        Self::from_ptr(ks_ptr)
    }

    fn from_ptr(keyset_ptr: *mut elektra_sys::KeySet) -> KeySet {
        KeySet {
            ptr: std::ptr::NonNull::new(keyset_ptr).unwrap(),
        }
    }

    /// Append a key to the keyset
    /// May return an InsertionFailure error, if the key was already added
    pub fn append_key<T: WriteableKey>(&mut self, mut key: T) -> Result<(), KeySetError> {
        let ret_val = unsafe { elektra_sys::ksAppendKey(self.as_ptr(), key.as_ptr()) };
        if ret_val == -1 {
            Err(KeySetError::InsertionFailure)
        } else {
            Ok(())
        }
    }

    /// Append a vector of keys to the KeySet.
    /// It will call append_key on every key. If any key returns an error it will be returned,
    /// but only after iterating over all keys.
    pub fn append_all<T: WriteableKey>(&mut self, keys: Vec<T>) -> Result<(), KeySetError> {
        let mut err = None;
        for key in keys {
            if let Err(ks_err) = self.append_key(key) {
                err = Some(ks_err);
            }
        }
        if let Some(error) = err {
            Err(error)
        } else {
            Ok(())
        }
    }

    /// Return a duplicate of a keyset.
    pub fn duplicate(&self) -> Self {
        let ks_ptr = unsafe { elektra_sys::ksDup(self.as_ref()) };
        Self::from_ptr(ks_ptr)
    }

    /// Replace the content of a keyset with another one.
    /// Copies the contents of source into self
    pub fn copy(&mut self, source: &Self) {
        unsafe { elektra_sys::ksCopy(self.as_ptr(), source.as_ref()) };
    }

    /// Return the number of keys that ks contains.
    pub fn get_size(&self) -> usize {
        unsafe { elektra_sys::ksGetSize(self.as_ref()).try_into().unwrap() }
    }

    /// Rewinds the KeySet internal cursor.
    pub fn rewind(&mut self) {
        unsafe {
            elektra_sys::ksRewind(self.as_ptr());
        }
    }

    /// Return the first key in the KeySet
    /// or None if the KeySet is empty.
    pub fn head(&self) -> Option<StringKey> {
        let key_ptr = unsafe { elektra_sys::ksHead(self.as_ref()) };
        if key_ptr as *const elektra_sys::Key == std::ptr::null() {
            None
        } else {
            Some(StringKey::from_ptr(key_ptr))
        }
    }

    /// Return the key pointed at by the internal cursor
    /// or None if the end is reached or after [`rewind`].
    ///
    /// [`rewind`]: #method.rewind
    pub fn current(&self) -> Option<StringKey> {
        let key_ptr = unsafe { elektra_sys::ksCurrent(self.as_ref()) };
        if key_ptr as *const elektra_sys::Key == std::ptr::null() {
            None
        } else {
            Some(StringKey::from_ptr(key_ptr))
        }
    }

    /// Return the last key in the KeySet
    /// or None if the KeySet is empty.
    pub fn tail(&self) -> Option<StringKey> {
        let key_ptr = unsafe { elektra_sys::ksTail(self.as_ref()) };
        if key_ptr as *const elektra_sys::Key == std::ptr::null() {
            None
        } else {
            Some(StringKey::from_ptr(key_ptr))
        }
    }

    /// Get the KeySet internal cursor.
    pub fn get_cursor(&self) -> Cursor {
        unsafe { elektra_sys::ksGetCursor(self.as_ref()) }
    }

    /// Return key at given cursor position or None if the cursor
    /// has a negative position or a position that does not lie within the keyset.
    pub fn at_cursor(&mut self, cursor: Cursor) -> Option<StringKey> {
        let key_ptr = unsafe { elektra_sys::ksAtCursor(self.as_ptr(), cursor) };
        if key_ptr as *const elektra_sys::Key == std::ptr::null() {
            None
        } else {
            Some(StringKey::from_ptr(key_ptr))
        }
    }

    /// Set the KeySet cursor.
    pub fn set_cursor(&mut self, cursor: Cursor) {
        unsafe {
            elektra_sys::ksSetCursor(self.as_ptr(), cursor);
        }
    }

    pub fn lookup(&mut self, mut key: StringKey, options: LookupOption) -> Option<StringKey> {
        println!("Got flags {:?}", options.bits() as elektra_sys::option_t);
        let key_ptr = unsafe {
            elektra_sys::ksLookup(
                self.as_ptr(),
                key.as_ptr(),
                options.bits() as elektra_sys::option_t,
            )
        };
        if key_ptr as *const elektra_sys::Key == std::ptr::null() {
            None
        } else {
            Some(StringKey::from_ptr(key_ptr))
        }
    }

    pub fn lookup_by_name(&mut self, name: &str, options: LookupOption) -> Option<StringKey> {
        println!("Got flags {:?}", options.bits() as elektra_sys::option_t);
        let cname = CString::new(name).unwrap();
        let key_ptr = unsafe {
            elektra_sys::ksLookupByName(
                self.as_ptr(),
                cname.as_ptr(),
                options.bits() as elektra_sys::option_t,
            )
        };
        if key_ptr as *const elektra_sys::Key == std::ptr::null() {
            None
        } else {
            Some(StringKey::from_ptr(key_ptr))
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::KeyBuilder;

    #[test]
    fn can_build_simple_keyset() {
        let mut ks = KeySet::new();
        ks.append_key(
            KeyBuilder::<StringKey>::new("user/sw/org/app/bool")
                .value("true")
                .build(),
        )
        .unwrap();
        // Make sure that the returned key from head can be dropped
        // without deleting the actual key
        {
            let head = ks.head().unwrap();
            assert_eq!(head.get_value(), "true");
        }
        {
            let head = ks.head().unwrap();
            assert_eq!(head.get_value(), "true");
        }
    }

    #[test]
    fn can_iterate_simple_keyset() {
        let names = ["user/test/key1", "user/test/key2", "user/test/key3"];
        let values = ["value1", "value2", "value3"];

        let mut ks = KeySet::with_capacity(3);

        ks.append_all(vec![
            KeyBuilder::<StringKey>::new(names[0])
                .value(values[0])
                .build(),
            KeyBuilder::<StringKey>::new(names[1])
                .value(values[1])
                .build(),
            KeyBuilder::<StringKey>::new(names[2])
                .value(values[2])
                .build(),
        ])
        .unwrap();
        ks.rewind();

        let mut did_iterate = false;
        for (i, key) in ks.enumerate() {
            did_iterate = true;
            assert_eq!(key.get_value(), values[i]);
            assert_eq!(key.get_name(), names[i]);
        }
        assert!(did_iterate);
    }

    fn setup_keyset() -> KeySet {
        let names = ["system/test/key", "user/test/key"];
        let values = ["value1", "value3"];

        let mut ks = KeySet::with_capacity(3);

        ks.append_all(vec![
            KeyBuilder::<StringKey>::new(names[0])
                .value(values[0])
                .build(),
            KeyBuilder::<StringKey>::new(names[1])
                .value(values[1])
                .build(),
            // KeyBuilder::<StringKey>::new(names[2])
            //     .value(values[2])
            //     .build(),
        ])
        .unwrap();
        ks
    }

    // #[test]
    // fn can_lookup_key() {
    //     let ret_val;
    //     {
    //         let mut ks = setup_keyset();
    //         let lookup_key = StringKey::new("/test/key").unwrap();
    //         println!("Lookup_key is {:?}", lookup_key);
    //         ret_val = ks.lookup(lookup_key, LookupOption::KDB_O_NONE);
    //         println!(
    //             "retval {:?}, size {}, name {}",
    //             ret_val,
    //             ks.get_size(),
    //             ks.tail().unwrap().get_name()
    //         );
    //         assert_eq!(ks.get_size(), 2);
    //         assert_eq!(ks.tail().unwrap().get_name(), "user/test/key");
    //     }
    //     assert_eq!(ret_val.unwrap().get_name(), "user/test/key");
    // }

    // #[test]
    // fn can_lookup_key_by_name() {
    //     let key;
    //     {
    //         let mut ks = setup_keyset();
    //         key = ks.lookup_by_name("/test/key", LookupOption::KDB_O_POP);
    //         println!(
    //             "retval {:?}, size {}, name {}",
    //             key,
    //             ks.get_size(),
    //             ks.tail().unwrap().get_name()
    //         );
    //         assert_eq!(ks.get_size(), 1);
    //         assert_eq!(ks.head().unwrap().get_name(), "system/test/key");
    //     } // <-- ks is dropped here
    //       // key is still alive, because of KDB_O_POP
    //     assert_eq!(key.unwrap().get_name(), "user/test/key");
    // }

}
<commit_msg>Rust: Further implementation details for lookup<commit_after>extern crate elektra_sys;

use crate::{KeySetError, ReadableKey, StringKey, WriteableKey};
use bitflags::bitflags;
use std::convert::TryInto;
use std::ffi::CString;

#[derive(Debug)]
pub struct KeySet {
    ptr: std::ptr::NonNull<elektra_sys::KeySet>,
}

pub type Cursor = elektra_sys::cursor_t;

bitflags! {
    // Results in default value with bits: 0
    #[derive(Default)]
    pub struct LookupOption: elektra_sys::option_t {
        const KDB_O_NONE = elektra_sys::KDB_O_NONE as elektra_sys::option_t;
        const KDB_O_DEL = elektra_sys::KDB_O_DEL as elektra_sys::option_t;
        const KDB_O_POP = elektra_sys::KDB_O_POP as elektra_sys::option_t;
        const KDB_O_NODIR = elektra_sys::KDB_O_NODIR as elektra_sys::option_t;
        const KDB_O_DIRONLY = elektra_sys::KDB_O_DIRONLY as elektra_sys::option_t;
        const KDB_O_NOREMOVE = elektra_sys::KDB_O_NOREMOVE as elektra_sys::option_t;
        const KDB_O_REMOVEONLY = elektra_sys::KDB_O_REMOVEONLY as elektra_sys::option_t;
        const KDB_O_INACTIVE = elektra_sys::KDB_O_INACTIVE as elektra_sys::option_t;
        const KDB_O_SYNC = elektra_sys::KDB_O_SYNC as elektra_sys::option_t;
        const KDB_O_SORT = elektra_sys::KDB_O_SORT as elektra_sys::option_t;
        const KDB_O_NORECURSIVE = elektra_sys::KDB_O_NORECURSIVE as elektra_sys::option_t;
        const KDB_O_NOCASE = elektra_sys::KDB_O_NOCASE as elektra_sys::option_t;
        const KDB_O_WITHOWNER = elektra_sys::KDB_O_WITHOWNER as elektra_sys::option_t;
        const KDB_O_NOALL = elektra_sys::KDB_O_NOALL as elektra_sys::option_t;
    }
}

impl Drop for KeySet {
    fn drop(&mut self) {
        println!("Drop {:?}", self);
        unsafe {
            elektra_sys::ksDel(self.as_ptr());
        }
    }
}

impl Iterator for KeySet {
    type Item = StringKey;
    fn next(&mut self) -> Option<Self::Item> {
        let key_ptr = unsafe { elektra_sys::ksNext(self.as_ptr()) };
        if (key_ptr as *const elektra_sys::Key) == std::ptr::null() {
            None
        } else {
            Some(StringKey::from_ptr(key_ptr as *mut elektra_sys::Key))
        }
    }
}

impl KeySet {
    fn as_ptr(&mut self) -> *mut elektra_sys::KeySet {
        self.ptr.as_ptr()
    }

    fn as_ref(&self) -> &elektra_sys::KeySet {
        unsafe { self.ptr.as_ref() }
    }

    /// Create a new empty KeySet
    pub fn new() -> KeySet {
        let ks_ptr = unsafe { elektra_sys::ksNew(0, elektra_sys::KEY_END) };
        KeySet::from_ptr(ks_ptr)
    }

    /// Create a new KeySet that allocates enough space for the
    /// given capacity of keys
    pub fn with_capacity(capacity: usize) -> Self {
        let ks_ptr = unsafe { elektra_sys::ksNew(capacity, elektra_sys::KEY_END) };
        Self::from_ptr(ks_ptr)
    }

    fn from_ptr(keyset_ptr: *mut elektra_sys::KeySet) -> KeySet {
        KeySet {
            ptr: std::ptr::NonNull::new(keyset_ptr).unwrap(),
        }
    }

    /// Append a key to the keyset
    /// May return an InsertionFailure error, if the key was already added
    pub fn append_key<T: WriteableKey>(&mut self, mut key: T) -> Result<(), KeySetError> {
        let ret_val = unsafe { elektra_sys::ksAppendKey(self.as_ptr(), key.as_ptr()) };
        if ret_val == -1 {
            Err(KeySetError::InsertionFailure)
        } else {
            Ok(())
        }
    }

    /// Append a vector of keys to the KeySet.
    /// It will call append_key on every key. If any key returns an error it will be returned,
    /// but only after iterating over all keys.
    pub fn append_all<T: WriteableKey>(&mut self, keys: Vec<T>) -> Result<(), KeySetError> {
        let mut err = None;
        for key in keys {
            if let Err(ks_err) = self.append_key(key) {
                err = Some(ks_err);
            }
        }
        if let Some(error) = err {
            Err(error)
        } else {
            Ok(())
        }
    }

    /// Return a duplicate of a keyset.
    pub fn duplicate(&self) -> Self {
        let ks_ptr = unsafe { elektra_sys::ksDup(self.as_ref()) };
        Self::from_ptr(ks_ptr)
    }

    /// Replace the content of a keyset with another one.
    /// Copies the contents of source into self
    pub fn copy(&mut self, source: &Self) {
        unsafe { elektra_sys::ksCopy(self.as_ptr(), source.as_ref()) };
    }

    /// Return the number of keys that ks contains.
    pub fn get_size(&self) -> usize {
        unsafe { elektra_sys::ksGetSize(self.as_ref()).try_into().unwrap() }
    }

    /// Rewinds the KeySet internal cursor.
    pub fn rewind(&mut self) {
        unsafe {
            elektra_sys::ksRewind(self.as_ptr());
        }
    }

    /// Return the first key in the KeySet
    /// or None if the KeySet is empty.
    pub fn head(&self) -> Option<StringKey> {
        let key_ptr = unsafe { elektra_sys::ksHead(self.as_ref()) };
        if key_ptr as *const elektra_sys::Key == std::ptr::null() {
            None
        } else {
            Some(StringKey::from_ptr(key_ptr))
        }
    }

    /// Return the key pointed at by the internal cursor
    /// or None if the end is reached or after [`rewind`].
    ///
    /// [`rewind`]: #method.rewind
    pub fn current(&self) -> Option<StringKey> {
        let key_ptr = unsafe { elektra_sys::ksCurrent(self.as_ref()) };
        if key_ptr as *const elektra_sys::Key == std::ptr::null() {
            None
        } else {
            Some(StringKey::from_ptr(key_ptr))
        }
    }

    /// Return the last key in the KeySet
    /// or None if the KeySet is empty.
    pub fn tail(&self) -> Option<StringKey> {
        let key_ptr = unsafe { elektra_sys::ksTail(self.as_ref()) };
        if key_ptr as *const elektra_sys::Key == std::ptr::null() {
            None
        } else {
            Some(StringKey::from_ptr(key_ptr))
        }
    }

    /// Get the KeySet internal cursor.
    pub fn get_cursor(&self) -> Cursor {
        unsafe { elektra_sys::ksGetCursor(self.as_ref()) }
    }

    /// Return key at given cursor position or None if the cursor
    /// has a negative position or a position that does not lie within the keyset.
    pub fn at_cursor(&mut self, cursor: Cursor) -> Option<StringKey> {
        let key_ptr = unsafe { elektra_sys::ksAtCursor(self.as_ptr(), cursor) };
        if key_ptr as *const elektra_sys::Key == std::ptr::null() {
            None
        } else {
            Some(StringKey::from_ptr(key_ptr))
        }
    }

    /// Set the KeySet cursor.
    pub fn set_cursor(&mut self, cursor: Cursor) {
        unsafe {
            elektra_sys::ksSetCursor(self.as_ptr(), cursor);
        }
    }

    pub fn lookup(&mut self, mut key: StringKey, options: LookupOption) -> Option<StringKey> {
        println!("Got flags {:?}", options.bits() as elektra_sys::option_t);
        let key_ptr = unsafe {
            elektra_sys::ksLookup(
                self.as_ptr(),
                key.as_ptr(),
                options.bits() as elektra_sys::option_t,
            )
        };
        if options.contains(LookupOption::KDB_O_DEL) {
            println!("contains del");
            std::mem::forget(key);
        }

        println!("keyptr is {:?}", key_ptr);
        if key_ptr as *const elektra_sys::Key == std::ptr::null() {
            println!("is null");
            None
        } else {
            let dup_ptr = unsafe { elektra_sys::keyDup(key_ptr as *const elektra_sys::Key) };
            Some(StringKey::from_ptr(dup_ptr))
        }
    }

    pub fn lookup_by_name(&mut self, name: &str, options: LookupOption) -> Option<StringKey> {
        println!("Got flags {:?}", options.bits() as elektra_sys::option_t);
        let cname = CString::new(name).unwrap();
        let key_ptr = unsafe {
            elektra_sys::ksLookupByName(
                self.as_ptr(),
                cname.as_ptr(),
                options.bits() as elektra_sys::option_t,
            )
        };
        println!("keyptr is {:?}", key_ptr);

        if key_ptr as *const elektra_sys::Key == std::ptr::null() {
            println!("is null");
            None
        } else {
            if options.contains(LookupOption::KDB_O_DEL) {
                println!("contains del");
                None
            } else {
                let dup_ptr = unsafe { elektra_sys::keyDup(key_ptr as *const elektra_sys::Key) };
                Some(StringKey::from_ptr(dup_ptr))
            }
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::KeyBuilder;

    #[test]
    fn can_build_simple_keyset() {
        let mut ks = KeySet::new();
        ks.append_key(
            KeyBuilder::<StringKey>::new("user/sw/org/app/bool")
                .value("true")
                .build(),
        )
        .unwrap();
        // Make sure that the returned key from head can be dropped
        // without deleting the actual key
        {
            let head = ks.head().unwrap();
            assert_eq!(head.get_value(), "true");
        }
        {
            let head = ks.head().unwrap();
            assert_eq!(head.get_value(), "true");
        }
    }

    #[test]
    fn can_iterate_simple_keyset() {
        let names = ["user/test/key1", "user/test/key2", "user/test/key3"];
        let values = ["value1", "value2", "value3"];

        let mut ks = KeySet::with_capacity(3);

        ks.append_all(vec![
            KeyBuilder::<StringKey>::new(names[0])
                .value(values[0])
                .build(),
            KeyBuilder::<StringKey>::new(names[1])
                .value(values[1])
                .build(),
            KeyBuilder::<StringKey>::new(names[2])
                .value(values[2])
                .build(),
        ])
        .unwrap();
        ks.rewind();

        let mut did_iterate = false;
        for (i, key) in ks.enumerate() {
            did_iterate = true;
            assert_eq!(key.get_value(), values[i]);
            assert_eq!(key.get_name(), names[i]);
        }
        assert!(did_iterate);
    }

    fn setup_keyset() -> KeySet {
        let names = ["system/test/key", "user/test/key"];
        let values = ["value1", "value3"];

        let mut ks = KeySet::with_capacity(3);

        ks.append_all(vec![
            KeyBuilder::<StringKey>::new(names[0])
                .value(values[0])
                .build(),
            KeyBuilder::<StringKey>::new(names[1])
                .value(values[1])
                .build(),
            // KeyBuilder::<StringKey>::new(names[2])
            //     .value(values[2])
            //     .build(),
        ])
        .unwrap();
        ks
    }

    #[test]
    fn can_lookup_key_with_none_option() {
        let ret_val;
        {
            let mut ks = setup_keyset();
            let lookup_key = StringKey::new("/test/key").unwrap();
            println!("Lookup_key is {:?}", lookup_key);
            ret_val = ks.lookup(lookup_key, LookupOption::KDB_O_NONE);
            println!(
                "retval {:?}, size {}, name {}",
                ret_val,
                ks.get_size(),
                ks.tail().unwrap().get_name()
            );
            assert_eq!(ks.get_size(), 2);
            assert_eq!(ks.tail().unwrap().get_name(), "user/test/key");
        }
        assert_eq!(ret_val.unwrap().get_name(), "user/test/key");
    }

    #[test]
    fn can_lookup_key_with_del_option() {
        let key;
        {
            let mut ks = setup_keyset();
            let lookup_key = StringKey::new("/test/key").unwrap();
            key = ks.lookup(lookup_key, LookupOption::KDB_O_DEL);
            println!(
                "retval {:?}, size {}, name {}",
                key,
                ks.get_size(),
                ks.tail().unwrap().get_name()
            );
            assert_eq!(ks.get_size(), 2);
            assert_eq!(ks.head().unwrap().get_name(), "system/test/key");
        } // <-- ks is dropped here
          // Check that key is still valid
        assert_eq!(key.unwrap().get_name(), "user/test/key");
    }

}
<|endoftext|>


<commit_before>// CITA
// Copyright 2016-2019 Cryptape Technologies LLC.

// This program is free software: you can redistribute it
// and/or modify it under the terms of the GNU General Public
// License as published by the Free Software Foundation,
// either version 3 of the License, or (at your option) any
// later version.

// This program is distributed in the hope that it will be
// useful, but WITHOUT ANY WARRANTY; without even the implied
// warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR
// PURPOSE. See the GNU General Public License for more details.

// You should have received a copy of the GNU General Public License
// along with this program.  If not, see <http://www.gnu.org/licenses/>.

use cita_database::{Config, Database, RocksDB, NUM_COLUMNS};
use cita_directories::DataPath;
use cita_types::H256;
use libproto::blockchain::SignedTransaction;
use libproto::{TryFrom, TryInto};
use std::sync::Arc;

/// Wal means write ahead log
/// used to persist transaction pools message
#[derive(Clone)]
pub struct TxWal {
    db: Arc<Database>,
}

impl TxWal {
    pub fn new(path: &str) -> Self {
        let nosql_path = DataPath::root_node_path() + path;
        // TODO: Can remove NUM_COLUMNS(useless)
        let config = Config::with_category_num(NUM_COLUMNS);
        let db = RocksDB::open(&nosql_path, &config).unwrap();
        TxWal { db: Arc::new(db) }
    }

    pub fn regenerate(&mut self, path: &str) {
        let nosql_path = DataPath::root_node_path() + path;
        let _ = self.db.restore(&nosql_path);
    }

    pub fn write(&self, tx: &SignedTransaction) {
        // TODO Fix the block_binary. tx_binary?
        let block_binary: Vec<u8> = tx.try_into().unwrap();
        self.db
            .insert(None, tx.get_tx_hash().to_vec(), block_binary)
            .expect("insert tx");
    }

    pub fn write_batch(&self, txs: &[SignedTransaction]) {
        let mut values: Vec<Vec<u8>> = Vec::new();
        let mut keys: Vec<Vec<u8>> = Vec::new();
        for tx in txs {
            let block_binary: Vec<u8> = tx.try_into().unwrap();
            values.push(block_binary);
            keys.push(tx.get_tx_hash().to_vec());
        }
        self.db
            .insert_batch(None, values, keys)
            .expect("insert batch txs");
    }

    pub fn delete_with_hash(&mut self, tx_hash: &H256) {
        self.db.remove(None, tx_hash).expect("delete with hash");
    }

    pub fn delete_with_hashes(&mut self, tx_hashes: &[H256]) {
        let mut keys: Vec<Vec<u8>> = Vec::new();
        for tx_hash in tx_hashes {
            keys.push(tx_hash.to_vec());
        }
        self.db
            .remove_batch(None, &keys)
            .expect("delete with hashes");
    }

    pub fn read_all(&self) -> Vec<SignedTransaction> {
        // TODO fix the unwrap
        let items = self.db.iterator(None).unwrap();

        items
            .map(|item| SignedTransaction::try_from(item.1.as_ref()).unwrap())
            .collect()
    }

    // FIXME Implement it.
    pub fn get(&self, _tx_hash: &[u8]) -> Option<SignedTransaction> {
        unimplemented!()
        // let result = self.db.get(None, tx_hash).unwrap();
        // result.map(|item| SignedTransaction::try_from(item).unwrap())
    }
}

// FIXME
#[cfg(test)]
mod tests {
    extern crate tempdir;
    use self::tempdir::TempDir;
    use super::*;
    use crypto::{CreateKey, KeyPair};
    use libproto::blockchain::Transaction;

    fn tx_wal() -> TxWal {
        let tempdir = TempDir::new("").unwrap().into_path();
        let config = DatabaseConfig::with_columns(None);
        let db = Database::open(&config, &tempdir.to_str().unwrap()).unwrap();
        TxWal { db: Arc::new(db) }
    }

    #[test]
    fn test_write_delete() {
        let mut wal = tx_wal();
        let keypair = KeyPair::gen_keypair();
        let privkey = keypair.privkey();
        let mut raw_tx = Transaction::new();
        raw_tx.quota = 1000;
        let tx = raw_tx.sign(*privkey);
        wal.write(&tx);

        let tx1 = wal.get(tx.get_tx_hash());

        assert_eq!(Some(tx.clone()), tx1);

        wal.delete_with_hash(&H256::from(tx.get_tx_hash()));
        let tx2 = wal.get(tx.get_tx_hash());

        assert_eq!(None, tx2);
    }

    #[test]
    fn test_batch() {
        let mut wal = tx_wal();

        let keypair = KeyPair::gen_keypair();
        let privkey = keypair.privkey();
        let mut raw_tx = Transaction::new();
        raw_tx.quota = 1000;
        let tx1 = raw_tx.sign(*privkey);
        let mut raw_tx2 = Transaction::new();
        raw_tx2.quota = 1000;
        let tx2 = raw_tx2.sign(*privkey);
        wal.write_batch(&vec![tx1.clone(), tx2.clone()]);

        let tx11 = wal.get(tx1.get_tx_hash());
        assert_eq!(Some(tx1.clone()), tx11);

        let tx11 = wal.get(tx2.get_tx_hash());
        assert_eq!(Some(tx2.clone()), tx11);

        wal.delete_with_hashes(&vec![
            H256::from(tx1.get_tx_hash()),
            H256::from(tx2.get_tx_hash()),
        ]);
        let tx12 = wal.get(tx1.get_tx_hash());
        let tx22 = wal.get(tx2.get_tx_hash());

        assert_eq!(None, tx12);
        assert_eq!(None, tx22);
    }
}
<commit_msg>Implement get function of cita-auth. [skip ci]<commit_after>// CITA
// Copyright 2016-2019 Cryptape Technologies LLC.

// This program is free software: you can redistribute it
// and/or modify it under the terms of the GNU General Public
// License as published by the Free Software Foundation,
// either version 3 of the License, or (at your option) any
// later version.

// This program is distributed in the hope that it will be
// useful, but WITHOUT ANY WARRANTY; without even the implied
// warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR
// PURPOSE. See the GNU General Public License for more details.

// You should have received a copy of the GNU General Public License
// along with this program.  If not, see <http://www.gnu.org/licenses/>.

use cita_database::{Config, Database, RocksDB, NUM_COLUMNS};
use cita_directories::DataPath;
use cita_types::H256;
use libproto::blockchain::SignedTransaction;
use libproto::{TryFrom, TryInto};
use std::sync::Arc;

/// Wal means write ahead log
/// used to persist transaction pools message
#[derive(Clone)]
pub struct TxWal {
    db: Arc<Database>,
}

impl TxWal {
    pub fn new(path: &str) -> Self {
        let nosql_path = DataPath::root_node_path() + path;
        // TODO: Can remove NUM_COLUMNS(useless)
        let config = Config::with_category_num(NUM_COLUMNS);
        let db = RocksDB::open(&nosql_path, &config).unwrap();
        TxWal { db: Arc::new(db) }
    }

    pub fn regenerate(&mut self, path: &str) {
        let nosql_path = DataPath::root_node_path() + path;
        let _ = self.db.restore(&nosql_path);
    }

    pub fn write(&self, tx: &SignedTransaction) {
        // TODO Fix the block_binary. tx_binary?
        let block_binary: Vec<u8> = tx.try_into().unwrap();
        self.db
            .insert(None, tx.get_tx_hash().to_vec(), block_binary)
            .expect("insert tx");
    }

    pub fn write_batch(&self, txs: &[SignedTransaction]) {
        let mut values: Vec<Vec<u8>> = Vec::new();
        let mut keys: Vec<Vec<u8>> = Vec::new();
        for tx in txs {
            let block_binary: Vec<u8> = tx.try_into().unwrap();
            values.push(block_binary);
            keys.push(tx.get_tx_hash().to_vec());
        }
        self.db
            .insert_batch(None, values, keys)
            .expect("insert batch txs");
    }

    pub fn delete_with_hash(&mut self, tx_hash: &H256) {
        self.db.remove(None, tx_hash).expect("delete with hash");
    }

    pub fn delete_with_hashes(&mut self, tx_hashes: &[H256]) {
        let mut keys: Vec<Vec<u8>> = Vec::new();
        for tx_hash in tx_hashes {
            keys.push(tx_hash.to_vec());
        }
        self.db
            .remove_batch(None, &keys)
            .expect("delete with hashes");
    }

    pub fn read_all(&self) -> Vec<SignedTransaction> {
        // TODO fix the unwrap
        let items = self.db.iterator(None).unwrap();

        items
            .map(|item| SignedTransaction::try_from(item.1.as_ref()).unwrap())
            .collect()
    }

    pub fn get(&self, _tx_hash: &[u8]) -> Option<SignedTransaction> {
        // TODO fix the unwrap
        let result = self.db.get(None, tx_hash).unwrap();
        result.map(|item| SignedTransaction::try_from(&item).unwrap())
    }
}

// FIXME
#[cfg(test)]
mod tests {
    extern crate tempdir;
    use self::tempdir::TempDir;
    use super::*;
    use crypto::{CreateKey, KeyPair};
    use libproto::blockchain::Transaction;

    fn tx_wal() -> TxWal {
        let tempdir = TempDir::new("").unwrap().into_path();
        let config = DatabaseConfig::with_columns(None);
        let db = Database::open(&config, &tempdir.to_str().unwrap()).unwrap();
        TxWal { db: Arc::new(db) }
    }

    #[test]
    fn test_write_delete() {
        let mut wal = tx_wal();
        let keypair = KeyPair::gen_keypair();
        let privkey = keypair.privkey();
        let mut raw_tx = Transaction::new();
        raw_tx.quota = 1000;
        let tx = raw_tx.sign(*privkey);
        wal.write(&tx);

        let tx1 = wal.get(tx.get_tx_hash());

        assert_eq!(Some(tx.clone()), tx1);

        wal.delete_with_hash(&H256::from(tx.get_tx_hash()));
        let tx2 = wal.get(tx.get_tx_hash());

        assert_eq!(None, tx2);
    }

    #[test]
    fn test_batch() {
        let mut wal = tx_wal();

        let keypair = KeyPair::gen_keypair();
        let privkey = keypair.privkey();
        let mut raw_tx = Transaction::new();
        raw_tx.quota = 1000;
        let tx1 = raw_tx.sign(*privkey);
        let mut raw_tx2 = Transaction::new();
        raw_tx2.quota = 1000;
        let tx2 = raw_tx2.sign(*privkey);
        wal.write_batch(&vec![tx1.clone(), tx2.clone()]);

        let tx11 = wal.get(tx1.get_tx_hash());
        assert_eq!(Some(tx1.clone()), tx11);

        let tx11 = wal.get(tx2.get_tx_hash());
        assert_eq!(Some(tx2.clone()), tx11);

        wal.delete_with_hashes(&vec![
            H256::from(tx1.get_tx_hash()),
            H256::from(tx2.get_tx_hash()),
        ]);
        let tx12 = wal.get(tx1.get_tx_hash());
        let tx22 = wal.get(tx2.get_tx_hash());

        assert_eq!(None, tx12);
        assert_eq!(None, tx22);
    }
}
<|endoftext|>


<commit_before>// Copyright 2018 The Chromium OS Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

#[allow(unused_imports)]
#[macro_use]
extern crate msg_on_socket_derive;
extern crate data_model;
#[macro_use]
extern crate sys_util;

mod msg_on_socket;

use std::io::Result;
use std::marker::PhantomData;
use std::os::unix::io::{AsRawFd, RawFd};
use std::os::unix::net::UnixDatagram;
use sys_util::{Error as SysError, ScmSocket, UnlinkUnixDatagram};

pub use msg_on_socket::*;
pub use msg_on_socket_derive::*;

/// Create a pair of socket. Request is send in one direction while response is in the other
/// direction.
pub fn pair<Request: MsgOnSocket, Response: MsgOnSocket>(
) -> Option<(MsgSocket<Request, Response>, MsgSocket<Response, Request>)> {
    let (sock1, sock2) = match UnixDatagram::pair() {
        Ok((sock1, sock2)) => (sock1, sock2),
        _ => {
            return None;
        }
    };
    let requester = MsgSocket {
        sock: sock1,
        _i: PhantomData,
        _o: PhantomData,
    };
    let responder = MsgSocket {
        sock: sock2,
        _i: PhantomData,
        _o: PhantomData,
    };
    Some((requester, responder))
}

/// Bidirection sock that support both send and recv.
pub struct MsgSocket<I: MsgOnSocket, O: MsgOnSocket> {
    sock: UnixDatagram,
    _i: PhantomData<I>,
    _o: PhantomData<O>,
}

impl<I: MsgOnSocket, O: MsgOnSocket> MsgSocket<I, O> {
    // Create a new MsgSocket.
    pub fn new(s: UnixDatagram) -> MsgSocket<I, O> {
        MsgSocket {
            sock: s,
            _i: PhantomData,
            _o: PhantomData,
        }
    }
}

/// Bidirection sock that support both send and recv.
pub struct UnlinkMsgSocket<I: MsgOnSocket, O: MsgOnSocket> {
    sock: UnlinkUnixDatagram,
    _i: PhantomData<I>,
    _o: PhantomData<O>,
}

impl<I: MsgOnSocket, O: MsgOnSocket> UnlinkMsgSocket<I, O> {
    // Create a new MsgSocket.
    pub fn new(s: UnlinkUnixDatagram) -> UnlinkMsgSocket<I, O> {
        UnlinkMsgSocket {
            sock: s,
            _i: PhantomData,
            _o: PhantomData,
        }
    }
}

/// One direction socket that only supports sending.
pub struct Sender<M: MsgOnSocket> {
    sock: UnixDatagram,
    _m: PhantomData<M>,
}

impl<M: MsgOnSocket> Sender<M> {
    /// Create a new sender sock.
    pub fn new(s: UnixDatagram) -> Sender<M> {
        Sender {
            sock: s,
            _m: PhantomData,
        }
    }
}

/// One direction socket that only supports receiving.
pub struct Receiver<M: MsgOnSocket> {
    sock: UnixDatagram,
    _m: PhantomData<M>,
}

impl<M: MsgOnSocket> Receiver<M> {
    /// Create a new receiver sock.
    pub fn new(s: UnixDatagram) -> Receiver<M> {
        Receiver {
            sock: s,
            _m: PhantomData,
        }
    }
}

impl<I: MsgOnSocket, O: MsgOnSocket> AsRef<UnixDatagram> for MsgSocket<I, O> {
    fn as_ref(&self) -> &UnixDatagram {
        &self.sock
    }
}

impl<I: MsgOnSocket, O: MsgOnSocket> AsRawFd for MsgSocket<I, O> {
    fn as_raw_fd(&self) -> RawFd {
        self.sock.as_raw_fd()
    }
}

impl<I: MsgOnSocket, O: MsgOnSocket> AsRef<UnixDatagram> for UnlinkMsgSocket<I, O> {
    fn as_ref(&self) -> &UnixDatagram {
        self.sock.as_ref()
    }
}

impl<I: MsgOnSocket, O: MsgOnSocket> AsRawFd for UnlinkMsgSocket<I, O> {
    fn as_raw_fd(&self) -> RawFd {
        self.as_ref().as_raw_fd()
    }
}

impl<M: MsgOnSocket> AsRef<UnixDatagram> for Sender<M> {
    fn as_ref(&self) -> &UnixDatagram {
        &self.sock
    }
}

impl<M: MsgOnSocket> AsRawFd for Sender<M> {
    fn as_raw_fd(&self) -> RawFd {
        self.sock.as_raw_fd()
    }
}

impl<M: MsgOnSocket> AsRef<UnixDatagram> for Receiver<M> {
    fn as_ref(&self) -> &UnixDatagram {
        &self.sock
    }
}

impl<M: MsgOnSocket> AsRawFd for Receiver<M> {
    fn as_raw_fd(&self) -> RawFd {
        self.sock.as_raw_fd()
    }
}

/// Types that could send a message.
pub trait MsgSender<M: MsgOnSocket>: AsRef<UnixDatagram> {
    fn send(&self, msg: &M) -> MsgResult<()> {
        let msg_size = M::msg_size();
        let fd_size = M::max_fd_count();
        let mut msg_buffer: Vec<u8> = vec![0; msg_size];
        let mut fd_buffer: Vec<RawFd> = vec![0; fd_size];

        let fd_size = msg.write_to_buffer(&mut msg_buffer, &mut fd_buffer)?;
        let sock: &UnixDatagram = self.as_ref();
        if fd_size == 0 {
            handle_eintr!(sock.send(&msg_buffer))
                .map_err(|e| MsgError::Send(SysError::new(e.raw_os_error().unwrap_or(0))))?;
        } else {
            sock.send_with_fds(&msg_buffer[..], &fd_buffer[0..fd_size])
                .map_err(|e| MsgError::Send(e))?;
        }
        Ok(())
    }
}

/// Types that could receive a message.
pub trait MsgReceiver<M: MsgOnSocket>: AsRef<UnixDatagram> {
    fn recv(&self) -> MsgResult<M> {
        let msg_size = M::msg_size();
        let fd_size = M::max_fd_count();
        let mut msg_buffer: Vec<u8> = vec![0; msg_size];
        let mut fd_buffer: Vec<RawFd> = vec![0; fd_size];

        let sock: &UnixDatagram = self.as_ref();

        let (recv_msg_size, recv_fd_size) = {
            if fd_size == 0 {
                let size = sock
                    .recv(&mut msg_buffer)
                    .map_err(|e| MsgError::Recv(SysError::new(e.raw_os_error().unwrap_or(0))))?;
                (size, 0)
            } else {
                sock.recv_with_fds(&mut msg_buffer, &mut fd_buffer)
                    .map_err(|e| MsgError::Recv(e))?
            }
        };
        if msg_size != recv_msg_size {
            return Err(MsgError::BadRecvSize(msg_size));
        }
        // Safe because fd buffer is read from socket.
        let (v, read_fd_size) = unsafe {
            M::read_from_buffer(&msg_buffer[0..recv_msg_size], &fd_buffer[0..recv_fd_size])?
        };
        if recv_fd_size != read_fd_size {
            return Err(MsgError::NotExpectFd);
        }
        Ok(v)
    }
}

impl<I: MsgOnSocket, O: MsgOnSocket> MsgSender<I> for MsgSocket<I, O> {}
impl<I: MsgOnSocket, O: MsgOnSocket> MsgReceiver<O> for MsgSocket<I, O> {}

impl<I: MsgOnSocket, O: MsgOnSocket> MsgSender<I> for UnlinkMsgSocket<I, O> {}
impl<I: MsgOnSocket, O: MsgOnSocket> MsgReceiver<O> for UnlinkMsgSocket<I, O> {}

impl<M: MsgOnSocket> MsgSender<M> for Sender<M> {}
impl<M: MsgOnSocket> MsgReceiver<M> for Receiver<M> {}
<commit_msg>msg_socket: return io:Error from pair() instead of Option<commit_after>// Copyright 2018 The Chromium OS Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

#[allow(unused_imports)]
#[macro_use]
extern crate msg_on_socket_derive;
extern crate data_model;
#[macro_use]
extern crate sys_util;

mod msg_on_socket;

use std::io::Result;
use std::marker::PhantomData;
use std::os::unix::io::{AsRawFd, RawFd};
use std::os::unix::net::UnixDatagram;
use sys_util::{Error as SysError, ScmSocket, UnlinkUnixDatagram};

pub use msg_on_socket::*;
pub use msg_on_socket_derive::*;

/// Create a pair of socket. Request is send in one direction while response is in the other
/// direction.
pub fn pair<Request: MsgOnSocket, Response: MsgOnSocket>(
) -> Result<(MsgSocket<Request, Response>, MsgSocket<Response, Request>)> {
    let (sock1, sock2) = UnixDatagram::pair()?;
    let requester = MsgSocket {
        sock: sock1,
        _i: PhantomData,
        _o: PhantomData,
    };
    let responder = MsgSocket {
        sock: sock2,
        _i: PhantomData,
        _o: PhantomData,
    };
    Ok((requester, responder))
}

/// Bidirection sock that support both send and recv.
pub struct MsgSocket<I: MsgOnSocket, O: MsgOnSocket> {
    sock: UnixDatagram,
    _i: PhantomData<I>,
    _o: PhantomData<O>,
}

impl<I: MsgOnSocket, O: MsgOnSocket> MsgSocket<I, O> {
    // Create a new MsgSocket.
    pub fn new(s: UnixDatagram) -> MsgSocket<I, O> {
        MsgSocket {
            sock: s,
            _i: PhantomData,
            _o: PhantomData,
        }
    }
}

/// Bidirection sock that support both send and recv.
pub struct UnlinkMsgSocket<I: MsgOnSocket, O: MsgOnSocket> {
    sock: UnlinkUnixDatagram,
    _i: PhantomData<I>,
    _o: PhantomData<O>,
}

impl<I: MsgOnSocket, O: MsgOnSocket> UnlinkMsgSocket<I, O> {
    // Create a new MsgSocket.
    pub fn new(s: UnlinkUnixDatagram) -> UnlinkMsgSocket<I, O> {
        UnlinkMsgSocket {
            sock: s,
            _i: PhantomData,
            _o: PhantomData,
        }
    }
}

/// One direction socket that only supports sending.
pub struct Sender<M: MsgOnSocket> {
    sock: UnixDatagram,
    _m: PhantomData<M>,
}

impl<M: MsgOnSocket> Sender<M> {
    /// Create a new sender sock.
    pub fn new(s: UnixDatagram) -> Sender<M> {
        Sender {
            sock: s,
            _m: PhantomData,
        }
    }
}

/// One direction socket that only supports receiving.
pub struct Receiver<M: MsgOnSocket> {
    sock: UnixDatagram,
    _m: PhantomData<M>,
}

impl<M: MsgOnSocket> Receiver<M> {
    /// Create a new receiver sock.
    pub fn new(s: UnixDatagram) -> Receiver<M> {
        Receiver {
            sock: s,
            _m: PhantomData,
        }
    }
}

impl<I: MsgOnSocket, O: MsgOnSocket> AsRef<UnixDatagram> for MsgSocket<I, O> {
    fn as_ref(&self) -> &UnixDatagram {
        &self.sock
    }
}

impl<I: MsgOnSocket, O: MsgOnSocket> AsRawFd for MsgSocket<I, O> {
    fn as_raw_fd(&self) -> RawFd {
        self.sock.as_raw_fd()
    }
}

impl<I: MsgOnSocket, O: MsgOnSocket> AsRef<UnixDatagram> for UnlinkMsgSocket<I, O> {
    fn as_ref(&self) -> &UnixDatagram {
        self.sock.as_ref()
    }
}

impl<I: MsgOnSocket, O: MsgOnSocket> AsRawFd for UnlinkMsgSocket<I, O> {
    fn as_raw_fd(&self) -> RawFd {
        self.as_ref().as_raw_fd()
    }
}

impl<M: MsgOnSocket> AsRef<UnixDatagram> for Sender<M> {
    fn as_ref(&self) -> &UnixDatagram {
        &self.sock
    }
}

impl<M: MsgOnSocket> AsRawFd for Sender<M> {
    fn as_raw_fd(&self) -> RawFd {
        self.sock.as_raw_fd()
    }
}

impl<M: MsgOnSocket> AsRef<UnixDatagram> for Receiver<M> {
    fn as_ref(&self) -> &UnixDatagram {
        &self.sock
    }
}

impl<M: MsgOnSocket> AsRawFd for Receiver<M> {
    fn as_raw_fd(&self) -> RawFd {
        self.sock.as_raw_fd()
    }
}

/// Types that could send a message.
pub trait MsgSender<M: MsgOnSocket>: AsRef<UnixDatagram> {
    fn send(&self, msg: &M) -> MsgResult<()> {
        let msg_size = M::msg_size();
        let fd_size = M::max_fd_count();
        let mut msg_buffer: Vec<u8> = vec![0; msg_size];
        let mut fd_buffer: Vec<RawFd> = vec![0; fd_size];

        let fd_size = msg.write_to_buffer(&mut msg_buffer, &mut fd_buffer)?;
        let sock: &UnixDatagram = self.as_ref();
        if fd_size == 0 {
            handle_eintr!(sock.send(&msg_buffer))
                .map_err(|e| MsgError::Send(SysError::new(e.raw_os_error().unwrap_or(0))))?;
        } else {
            sock.send_with_fds(&msg_buffer[..], &fd_buffer[0..fd_size])
                .map_err(|e| MsgError::Send(e))?;
        }
        Ok(())
    }
}

/// Types that could receive a message.
pub trait MsgReceiver<M: MsgOnSocket>: AsRef<UnixDatagram> {
    fn recv(&self) -> MsgResult<M> {
        let msg_size = M::msg_size();
        let fd_size = M::max_fd_count();
        let mut msg_buffer: Vec<u8> = vec![0; msg_size];
        let mut fd_buffer: Vec<RawFd> = vec![0; fd_size];

        let sock: &UnixDatagram = self.as_ref();

        let (recv_msg_size, recv_fd_size) = {
            if fd_size == 0 {
                let size = sock
                    .recv(&mut msg_buffer)
                    .map_err(|e| MsgError::Recv(SysError::new(e.raw_os_error().unwrap_or(0))))?;
                (size, 0)
            } else {
                sock.recv_with_fds(&mut msg_buffer, &mut fd_buffer)
                    .map_err(|e| MsgError::Recv(e))?
            }
        };
        if msg_size != recv_msg_size {
            return Err(MsgError::BadRecvSize(msg_size));
        }
        // Safe because fd buffer is read from socket.
        let (v, read_fd_size) = unsafe {
            M::read_from_buffer(&msg_buffer[0..recv_msg_size], &fd_buffer[0..recv_fd_size])?
        };
        if recv_fd_size != read_fd_size {
            return Err(MsgError::NotExpectFd);
        }
        Ok(v)
    }
}

impl<I: MsgOnSocket, O: MsgOnSocket> MsgSender<I> for MsgSocket<I, O> {}
impl<I: MsgOnSocket, O: MsgOnSocket> MsgReceiver<O> for MsgSocket<I, O> {}

impl<I: MsgOnSocket, O: MsgOnSocket> MsgSender<I> for UnlinkMsgSocket<I, O> {}
impl<I: MsgOnSocket, O: MsgOnSocket> MsgReceiver<O> for UnlinkMsgSocket<I, O> {}

impl<M: MsgOnSocket> MsgSender<M> for Sender<M> {}
impl<M: MsgOnSocket> MsgReceiver<M> for Receiver<M> {}
<|endoftext|>


<commit_before>use std::rand::{task_rng, Rng};

fn rnd_permu<T>(list: ~[T]) -> ~[T] {
    task_rng().shuffle(list)
}

fn main() {
    let list = ~['a', 'b', 'c', 'd', 'e', 'f'];
    println!("{:?}", rnd_permu(list));
}

<commit_msg>Cleanup Problem 25<commit_after>// The author of this work hereby waives all claim of copyright (economic and
// moral) in this work and immediately places it in the public domain; it may
// be used, distorted or destroyed in any manner whatsoever without further
// attribution or notice to the creator.

//! Problem 25: Vectors: rnd_permu
//!
//! Generate a random permutation of the elements of a vector.
//!
//! Your function could have this signature:
//! `fn rnd_permu<T>(vec: ~[T]) -> Vec<T>`

extern crate rand;
use rand::{task_rng, Rng};

fn rnd_permu<T>(vec: ~[T]) -> ~[T] {
    let mut vec = vec;
    task_rng().shuffle(vec);
    vec
}
<|endoftext|>


<commit_before>extern crate cc;
extern crate metadeps;

use std::env;
use std::ffi::OsString;
use std::fs;
use std::path::PathBuf;
use std::process::Command;

macro_rules! t {
    ($e:expr) => (match $e {
        Ok(t) => t,
        Err(e) => panic!("{} return the error {}", stringify!($e), e),
    })
}

fn main() {
    // Use the installed libcares if it is available.
    if metadeps::probe().is_ok() {
        return;
    }

    // Rerun if the c-ares source code has changed.
    println!("cargo:rerun-if-changed=c-ares");

    // MSVC builds are different.
    let target = env::var("TARGET").unwrap();
    if target.contains("msvc") {
        build_msvc(&target);
        return;
    }

    let dst = PathBuf::from(env::var_os("OUT_DIR").unwrap());
    let build = dst.join("build");
    let _ = fs::create_dir(&build);

    // Prepare.
    let src = env::current_dir().unwrap();
    run(
        Command::new("sh")
            .current_dir(&src.join("c-ares"))
            .arg("buildconf"),
    );

    // Configure.
    let cfg = cc::Build::new();
    let compiler = cfg.get_compiler();
    let mut cflags = OsString::new();
    for arg in compiler.args() {
        cflags.push(arg);
        cflags.push(" ");
    }

    let mut cmd = Command::new("sh");
    cmd.env("CFLAGS", &cflags)
        .env("CC", compiler.path())
        .current_dir(&build)
        .arg(format!("{}", src.join("c-ares/configure").display()))
        .arg("--enable-static=yes")
        .arg("--enable-shared=no")
        .arg("--enable-optimize")
        .arg(format!("--prefix={}", dst.display()));

    let host = env::var("HOST").unwrap();
    if target != host &&
       (!target.contains("windows") || !host.contains("windows")) {
        // NOTE GNU terminology
        // BUILD = machine where we are (cross) compiling curl
        // HOST = machine where the compiled curl will be used
        // TARGET = only relevant when compiling compilers
        if target.contains("windows") {
            // curl's configure can't parse `-windows-` triples when used
            // as `--host`s. In those cases we use this combination of
            // `host` and `target` that appears to do the right thing.
            cmd.arg(format!("--host={}", host));
            cmd.arg(format!("--target={}", target));
        } else {
            cmd.arg(format!("--build={}", host));
            cmd.arg(format!("--host={}", target));
        }
    }
    run(&mut cmd);

    // Compile.
    run(
        Command::new(make())
            .arg(format!("-j{}", env::var("NUM_JOBS").unwrap()))
            .current_dir(&build),
    );

    // Link to compiled library.
    println!("cargo:rustc-link-search={}/.libs", build.display());
    println!("cargo:rustc-link-lib=static=cares");
}

fn run(cmd: &mut Command) {
    println!("running: {:?}", cmd);
    assert!(t!(cmd.status()).success());
}

fn make() -> &'static str {
    if cfg!(target_os = "freebsd") {
        "gmake"
    } else {
        "make"
    }
}

fn nmake(target: &str) -> Command {
    // cargo messes with the environment in a way that nmake does not like -
    // see https://github.com/rust-lang/cargo/issues/4156.  Explicitly remove
    // the unwanted variables.
    let mut cmd = cc::windows_registry::find(target, "nmake.exe").unwrap();
    cmd.env_remove("MAKEFLAGS").env_remove("MFLAGS");
    cmd
}

fn build_msvc(target: &str) {
    // Prepare.
    let src = env::current_dir().unwrap();
    let c_ares_dir = &src.join("c-ares");
    run(
        Command::new("cmd")
            .current_dir(c_ares_dir)
            .arg("/c")
            .arg("buildconf.bat"),
    );

    // Compile.
    let mut cmd = nmake(target);
    cmd.current_dir(c_ares_dir);
    cmd.args(&["/f", "Makefile.msvc", "CFG=lib-release", "c-ares"]);
    run(&mut cmd);

    // Install library.
    let dst = PathBuf::from(env::var_os("OUT_DIR").unwrap());
    let build = dst.join("build");
    let mut cmd = nmake(target);
    cmd.current_dir(c_ares_dir);
    cmd.args(&["/f", "Makefile.msvc", "/a", "CFG=lib-release", "install"]);
    cmd.env("INSTALL_DIR", format!("{}", build.display()));
    run(&mut cmd);

    // Link to compiled library.
    println!("cargo:rustc-link-search={}/lib", build.display());
    println!("cargo:rustc-link-lib=static=libcares");
}
<commit_msg>make curl references deliberate, not accidental<commit_after>extern crate cc;
extern crate metadeps;

use std::env;
use std::ffi::OsString;
use std::fs;
use std::path::PathBuf;
use std::process::Command;

macro_rules! t {
    ($e:expr) => (match $e {
        Ok(t) => t,
        Err(e) => panic!("{} return the error {}", stringify!($e), e),
    })
}

fn main() {
    // Use the installed libcares if it is available.
    if metadeps::probe().is_ok() {
        return;
    }

    // Rerun if the c-ares source code has changed.
    println!("cargo:rerun-if-changed=c-ares");

    // MSVC builds are different.
    let target = env::var("TARGET").unwrap();
    if target.contains("msvc") {
        build_msvc(&target);
        return;
    }

    let dst = PathBuf::from(env::var_os("OUT_DIR").unwrap());
    let build = dst.join("build");
    let _ = fs::create_dir(&build);

    // Prepare.
    let src = env::current_dir().unwrap();
    run(
        Command::new("sh")
            .current_dir(&src.join("c-ares"))
            .arg("buildconf"),
    );

    // Configure.
    let cfg = cc::Build::new();
    let compiler = cfg.get_compiler();
    let mut cflags = OsString::new();
    for arg in compiler.args() {
        cflags.push(arg);
        cflags.push(" ");
    }

    let mut cmd = Command::new("sh");
    cmd.env("CFLAGS", &cflags)
        .env("CC", compiler.path())
        .current_dir(&build)
        .arg(format!("{}", src.join("c-ares/configure").display()))
        .arg("--enable-static=yes")
        .arg("--enable-shared=no")
        .arg("--enable-optimize")
        .arg(format!("--prefix={}", dst.display()));

    // This code fragment copied from curl-rust... c-ares and curl come from
    // the same developer so are usually pretty similar, and this seems to
    // work.
    //
    // NOTE GNU terminology
    // BUILD = machine where we are (cross) compiling
    // HOST = machine where the compiled binary will be used
    // TARGET = only relevant when compiling compilers
    let host = env::var("HOST").unwrap();
    if target != host &&
       (!target.contains("windows") || !host.contains("windows")) {
        if target.contains("windows") {
            cmd.arg(format!("--host={}", host));
            cmd.arg(format!("--target={}", target));
        } else {
            cmd.arg(format!("--build={}", host));
            cmd.arg(format!("--host={}", target));
        }
    }
    run(&mut cmd);

    // Compile.
    run(
        Command::new(make())
            .arg(format!("-j{}", env::var("NUM_JOBS").unwrap()))
            .current_dir(&build),
    );

    // Link to compiled library.
    println!("cargo:rustc-link-search={}/.libs", build.display());
    println!("cargo:rustc-link-lib=static=cares");
}

fn run(cmd: &mut Command) {
    println!("running: {:?}", cmd);
    assert!(t!(cmd.status()).success());
}

fn make() -> &'static str {
    if cfg!(target_os = "freebsd") {
        "gmake"
    } else {
        "make"
    }
}

fn nmake(target: &str) -> Command {
    // cargo messes with the environment in a way that nmake does not like -
    // see https://github.com/rust-lang/cargo/issues/4156.  Explicitly remove
    // the unwanted variables.
    let mut cmd = cc::windows_registry::find(target, "nmake.exe").unwrap();
    cmd.env_remove("MAKEFLAGS").env_remove("MFLAGS");
    cmd
}

fn build_msvc(target: &str) {
    // Prepare.
    let src = env::current_dir().unwrap();
    let c_ares_dir = &src.join("c-ares");
    run(
        Command::new("cmd")
            .current_dir(c_ares_dir)
            .arg("/c")
            .arg("buildconf.bat"),
    );

    // Compile.
    let mut cmd = nmake(target);
    cmd.current_dir(c_ares_dir);
    cmd.args(&["/f", "Makefile.msvc", "CFG=lib-release", "c-ares"]);
    run(&mut cmd);

    // Install library.
    let dst = PathBuf::from(env::var_os("OUT_DIR").unwrap());
    let build = dst.join("build");
    let mut cmd = nmake(target);
    cmd.current_dir(c_ares_dir);
    cmd.args(&["/f", "Makefile.msvc", "/a", "CFG=lib-release", "install"]);
    cmd.env("INSTALL_DIR", format!("{}", build.display()));
    run(&mut cmd);

    // Link to compiled library.
    println!("cargo:rustc-link-search={}/lib", build.display());
    println!("cargo:rustc-link-lib=static=libcares");
}
<|endoftext|>


<commit_before>use nth_prime as np;

#[test]
fn test_first_prime() {
    assert_eq!(np::nth(0), 2);
}

#[test]
#[ignore]
fn test_second_prime() {
    assert_eq!(np::nth(1), 3);
}

#[test]
#[ignore]
fn test_sixth_prime() {
    assert_eq!(np::nth(5), 13);
}

#[test]
#[ignore]
fn test_big_prime() {
    assert_eq!(np::nth(10000), 104743);
}
<commit_msg>Wrong tests<commit_after>use nth_prime as np;

#[test]
fn test_first_prime() {
    assert_eq!(np::nth(1), 2);
}

#[test]
#[ignore]
fn test_second_prime() {
    assert_eq!(np::nth(2), 3);
}

#[test]
#[ignore]
fn test_sixth_prime() {
    assert_eq!(np::nth(6), 13);
}

#[test]
#[ignore]
fn test_big_prime() {
    assert_eq!(np::nth(10001), 104743);
}
<|endoftext|>


<commit_before>//! A proxy that forwards data to another server and forwards that server's
//! responses back to clients.
//!
//! Because the Tokio runtime uses a thread pool, each TCP connection is
//! processed concurrently with all other TCP connections across multiple
//! threads.
//!
//! You can showcase this by running this in one terminal:
//!
//!     cargo run --example proxy
//!
//! This in another terminal
//!
//!     cargo run --example echo
//!
//! And finally this in another terminal
//!
//!     cargo run --example connect 127.0.0.1:8081
//!
//! This final terminal will connect to our proxy, which will in turn connect to
//! the echo server, and you'll be able to see data flowing between them.

#![warn(rust_2018_idioms)]

use tokio::io;
use tokio::net::{TcpListener, TcpStream};

use futures::future::try_join;
use futures::FutureExt;
use std::env;
use std::error::Error;

#[tokio::main]
async fn main() -> Result<(), Box<dyn Error>> {
    let listen_addr = env::args()
        .nth(1)
        .unwrap_or_else(|| "127.0.0.1:8081".to_string());
    let server_addr = env::args()
        .nth(2)
        .unwrap_or_else(|| "127.0.0.1:8080".to_string());

    println!("Listening on: {}", listen_addr);
    println!("Proxying to: {}", server_addr);

    let mut listener = TcpListener::bind(listen_addr).await?;

    while let Ok((inbound, _)) = listener.accept().await {
        let transfer = transfer(inbound, server_addr.clone()).map(|r| {
            if let Err(e) = r {
                println!("Failed to transfer; error={}", e);
            }
        });

        tokio::spawn(transfer);
    }

    Ok(())
}

async fn transfer(mut inbound: TcpStream, proxy_addr: String) -> Result<(), Box<dyn Error>> {
    let mut outbound = TcpStream::connect(proxy_addr).await?;

    let (mut ri, mut wi) = inbound.split();
    let (mut ro, mut wo) = outbound.split();

    let client_to_server = io::copy(&mut ri, &mut wo);
    let server_to_client = io::copy(&mut ro, &mut wi);

    try_join(client_to_server, server_to_client).await?;

    Ok(())
}
<commit_msg>example: close pending connection on proxy exemple  (#2590)<commit_after>//! A proxy that forwards data to another server and forwards that server's
//! responses back to clients.
//!
//! Because the Tokio runtime uses a thread pool, each TCP connection is
//! processed concurrently with all other TCP connections across multiple
//! threads.
//!
//! You can showcase this by running this in one terminal:
//!
//!     cargo run --example proxy
//!
//! This in another terminal
//!
//!     cargo run --example echo
//!
//! And finally this in another terminal
//!
//!     cargo run --example connect 127.0.0.1:8081
//!
//! This final terminal will connect to our proxy, which will in turn connect to
//! the echo server, and you'll be able to see data flowing between them.

#![warn(rust_2018_idioms)]

use tokio::io;
use tokio::io::AsyncWriteExt;
use tokio::net::{TcpListener, TcpStream};

use futures::future::try_join;
use futures::FutureExt;
use std::env;
use std::error::Error;

#[tokio::main]
async fn main() -> Result<(), Box<dyn Error>> {
    let listen_addr = env::args()
        .nth(1)
        .unwrap_or_else(|| "127.0.0.1:8081".to_string());
    let server_addr = env::args()
        .nth(2)
        .unwrap_or_else(|| "127.0.0.1:8080".to_string());

    println!("Listening on: {}", listen_addr);
    println!("Proxying to: {}", server_addr);

    let mut listener = TcpListener::bind(listen_addr).await?;

    while let Ok((inbound, _)) = listener.accept().await {
        let transfer = transfer(inbound, server_addr.clone()).map(|r| {
            if let Err(e) = r {
                println!("Failed to transfer; error={}", e);
            }
        });

        tokio::spawn(transfer);
    }

    Ok(())
}

async fn transfer(mut inbound: TcpStream, proxy_addr: String) -> Result<(), Box<dyn Error>> {
    let mut outbound = TcpStream::connect(proxy_addr).await?;

    let (mut ri, mut wi) = inbound.split();
    let (mut ro, mut wo) = outbound.split();

    let client_to_server = async {
        io::copy(&mut ri, &mut wo).await?;
        wo.shutdown().await
    };

    let server_to_client = async {
        io::copy(&mut ro, &mut wi).await?;
        wi.shutdown().await
    };

    try_join(client_to_server, server_to_client).await?;

    Ok(())
}
<|endoftext|>


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

#[macro_use]
extern crate gj;
extern crate gjio;

use std::net::ToSocketAddrs;
use gjio::{AsyncRead, AsyncWrite};
use gj::{EventLoop, Promise};

fn forward<R,W,B>(mut src: R, mut dst: W, buf: B) -> Promise<(R, W), ::std::io::Error>
    where R: AsyncRead + 'static, W: AsyncWrite + 'static, B: AsMut<[u8]> + AsRef<[u8]> + 'static
{
    src.try_read(buf, 1).then(move |(buf, n)| {
        if n == 0 { // EOF
            Promise::ok((src, dst))
        } else {
            dst.write(gjio::BufferPrefix::new(buf, n)).then(move |prefix| {
                forward(src, dst, prefix.buf)
            })
        }
    })
}

fn accept_loop(receiver: gjio::SocketListener,
               outbound_addr: gjio::SocketAddress,
               timer: gjio::Timer,
               mut task_set: gj::TaskSet<(), ::std::io::Error>)
               -> Promise<(), ::std::io::Error>
{
    receiver.accept().then(move |src_stream| {
        println!("handling connection");

        timer.timeout_after(::std::time::Duration::from_secs(3), outbound_addr.connect())
           .then_else(move |r| match r {
               Ok(dst_stream) =>  {
                   let task1 = forward(src_stream.clone(), dst_stream.clone(), vec![0; 1024])
                       .map(|(src, mut dst)| dst.shutdown(::std::net::Shutdown::Write));
                   let task2 =
                       forward(dst_stream.clone(), src_stream.clone(), vec![0; 1024]).map(|_| Ok(()))
                   task_set.add(task1);
                   task_set.add(task2);
                   accept_loop(receiver, outbound_addr, timer, task_set)
               }
               Err(e) => {
                   println!("failed to connect: {}", e);
                   Promise::err(e.into())
               }
           })
    })
}

pub struct Reporter;

impl gj::TaskReaper<(), ::std::io::Error> for Reporter {
    fn task_failed(&mut self, error: ::std::io::Error) {
        println!("Task failed: {}", error);
    }
}

pub fn main() {
    let args: Vec<String> = ::std::env::args().collect();
    if args.len() != 3 {
        println!("usage: {} <LISTEN_ADDR> <CONNECT_ADDR>", args[0]);
        return;
    }

    EventLoop::top_level(|wait_scope| -> Result<(), ::std::io::Error> {
        let mut event_port = try!(gjio::EventPort::new());
        let network = event_port.get_network();
        let addr = try!(args[1].to_socket_addrs()).next().expect("could not parse address");
        let mut address = network.get_tcp_address(addr);
        let listener = try!(address.listen());

        let outbound_addr = try!(args[2].to_socket_addrs()).next().expect("could not parse address");
        let outbound_address = network.get_tcp_address(outbound_addr);
        accept_loop(listener, outbound_address, event_port.get_timer(),
                    gj::TaskSet::new(Box::new(Reporter))).wait(wait_scope, &mut event_port)
    }).expect("top level");
}
<commit_msg>add missing semicolon<commit_after>// Copyright (c) 2013-2016 Sandstorm Development Group, Inc. and contributors
// Licensed under the MIT License:
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.

#[macro_use]
extern crate gj;
extern crate gjio;

use std::net::ToSocketAddrs;
use gjio::{AsyncRead, AsyncWrite};
use gj::{EventLoop, Promise};

fn forward<R,W,B>(mut src: R, mut dst: W, buf: B) -> Promise<(R, W), ::std::io::Error>
    where R: AsyncRead + 'static, W: AsyncWrite + 'static, B: AsMut<[u8]> + AsRef<[u8]> + 'static
{
    src.try_read(buf, 1).then(move |(buf, n)| {
        if n == 0 { // EOF
            Promise::ok((src, dst))
        } else {
            dst.write(gjio::BufferPrefix::new(buf, n)).then(move |prefix| {
                forward(src, dst, prefix.buf)
            })
        }
    })
}

fn accept_loop(receiver: gjio::SocketListener,
               outbound_addr: gjio::SocketAddress,
               timer: gjio::Timer,
               mut task_set: gj::TaskSet<(), ::std::io::Error>)
               -> Promise<(), ::std::io::Error>
{
    receiver.accept().then(move |src_stream| {
        println!("handling connection");

        timer.timeout_after(::std::time::Duration::from_secs(3), outbound_addr.connect())
           .then_else(move |r| match r {
               Ok(dst_stream) =>  {
                   let task1 = forward(src_stream.clone(), dst_stream.clone(), vec![0; 1024])
                       .map(|(_src, mut dst)| dst.shutdown(::std::net::Shutdown::Write));
                   let task2 =
                       forward(dst_stream.clone(), src_stream.clone(), vec![0; 1024]).map(|_| Ok(()));
                   task_set.add(task1);
                   task_set.add(task2);
                   accept_loop(receiver, outbound_addr, timer, task_set)
               }
               Err(e) => {
                   println!("failed to connect: {}", e);
                   Promise::err(e.into())
               }
           })
    })
}

pub struct Reporter;

impl gj::TaskReaper<(), ::std::io::Error> for Reporter {
    fn task_failed(&mut self, error: ::std::io::Error) {
        println!("Task failed: {}", error);
    }
}

pub fn main() {
    let args: Vec<String> = ::std::env::args().collect();
    if args.len() != 3 {
        println!("usage: {} <LISTEN_ADDR> <CONNECT_ADDR>", args[0]);
        return;
    }

    EventLoop::top_level(|wait_scope| -> Result<(), ::std::io::Error> {
        let mut event_port = try!(gjio::EventPort::new());
        let network = event_port.get_network();
        let addr = try!(args[1].to_socket_addrs()).next().expect("could not parse address");
        let mut address = network.get_tcp_address(addr);
        let listener = try!(address.listen());

        let outbound_addr = try!(args[2].to_socket_addrs()).next().expect("could not parse address");
        let outbound_address = network.get_tcp_address(outbound_addr);
        accept_loop(listener, outbound_address, event_port.get_timer(),
                    gj::TaskSet::new(Box::new(Reporter))).wait(wait_scope, &mut event_port)
    }).expect("top level");
}
<|endoftext|>