<commit_before>// Copyright (c) Facebook, Inc. and its affiliates
// SPDX-License-Identifier: MIT OR Apache-2.0

use crate::{
    analyzer,
    indent::{IndentConfig, IndentedWriter},
    CodeGeneratorConfig,
};
use serde_reflection::{ContainerFormat, Format, Named, Registry, VariantFormat};
use std::borrow::Cow;
use std::collections::{BTreeMap, HashSet};
use std::io::{Result, Write};
use std::path::PathBuf;

/// Main configuration object for code-generation in Rust.
pub struct CodeGenerator<'a> {
    /// Language-independent configuration.
    config: &'a CodeGeneratorConfig,
    /// Which derive macros should be added (independently from serialization).
    derive_macros: Vec<String>,
    /// Additional block of text added before each new container definition.
    custom_derive_block: Option<String>,
    /// Whether definitions and fields should be marked as `pub`.
    track_visibility: bool,
}

/// Shared state for the code generation of a Rust source file.
struct RustEmitter<'a, T> {
    /// Writer.
    out: IndentedWriter<T>,
    /// Generator.
    generator: &'a CodeGenerator<'a>,
    /// Track which definitions have a known size. (Used to add `Box` types.)
    known_sizes: Cow<'a, HashSet<&'a str>>,
    /// Current namespace (e.g. vec!["my_package", "my_module", "MyClass"])
    current_namespace: Vec<String>,
}

impl<'a> CodeGenerator<'a> {
    /// Create a Rust code generator for the given config.
    pub fn new(config: &'a CodeGeneratorConfig) -> Self {
        Self {
            config,
            derive_macros: vec!["Clone", "Debug", "PartialEq", "PartialOrd"]
                .into_iter()
                .map(String::from)
                .collect(),
            custom_derive_block: None,
            track_visibility: true,
        }
    }

    /// Which derive macros should be added (independently from serialization).
    pub fn with_derive_macros(mut self, derive_macros: Vec<String>) -> Self {
        self.derive_macros = derive_macros;
        self
    }

    /// Additional block of text added after `derive_macros` (if any), before each new
    /// container definition.
    pub fn with_custom_derive_block(mut self, custom_derive_block: Option<String>) -> Self {
        self.custom_derive_block = custom_derive_block;
        self
    }

    /// Whether definitions and fields should be marked as `pub`.
    pub fn with_track_visibility(mut self, track_visibility: bool) -> Self {
        self.track_visibility = track_visibility;
        self
    }

    /// Write container definitions in Rust.
    pub fn output(
        &self,
        out: &mut dyn Write,
        registry: &Registry,
    ) -> std::result::Result<(), Box<dyn std::error::Error>> {
        let external_names = self
            .config
            .external_definitions
            .values()
            .cloned()
            .flatten()
            .collect();
        let dependencies =
            analyzer::get_dependency_map_with_external_dependencies(registry, &external_names)?;
        let entries = analyzer::best_effort_topological_sort(&dependencies);

        let known_sizes = external_names
            .iter()
            .map(<String as std::ops::Deref>::deref)
            .collect::<HashSet<_>>();

        let current_namespace = self
            .config
            .module_name
            .split('.')
            .map(String::from)
            .collect();
        let mut emitter = RustEmitter {
            out: IndentedWriter::new(out, IndentConfig::Space(4)),
            generator: self,
            known_sizes: Cow::Owned(known_sizes),
            current_namespace,
        };

        emitter.output_preamble()?;
        for name in entries {
            let format = &registry[name];
            emitter.output_container(name, format)?;
            emitter.known_sizes.to_mut().insert(name);
        }
        Ok(())
    }

    /// For each container, generate a Rust definition.
    pub fn quote_container_definitions(
        &self,
        registry: &Registry,
    ) -> std::result::Result<BTreeMap<String, String>, Box<dyn std::error::Error>> {
        let dependencies = analyzer::get_dependency_map(registry)?;
        let entries = analyzer::best_effort_topological_sort(&dependencies);

        let mut result = BTreeMap::new();
        let mut known_sizes = HashSet::new();
        let current_namespace = self
            .config
            .module_name
            .split('.')
            .map(String::from)
            .collect::<Vec<_>>();

        for name in entries {
            let mut content = Vec::new();
            {
                let mut emitter = RustEmitter {
                    out: IndentedWriter::new(&mut content, IndentConfig::Space(4)),
                    generator: self,
                    known_sizes: Cow::Borrowed(&known_sizes),
                    current_namespace: current_namespace.clone(),
                };
                let format = &registry[name];
                emitter.output_container(name, format)?;
            }
            known_sizes.insert(name);
            result.insert(
                name.to_string(),
                String::from_utf8_lossy(&content).trim().to_string() + "\n",
            );
        }
        Ok(result)
    }
}

impl<'a, T> RustEmitter<'a, T>
where
    T: std::io::Write,
{
    fn output_comment(&mut self, name: &str) -> std::io::Result<()> {
        let mut path = self.current_namespace.clone();
        path.push(name.to_string());
        if let Some(doc) = self.generator.config.comments.get(&path) {
            let text = textwrap::indent(doc, "/// ").replace("\n\n", "\n///\n");
            write!(self.out, "\n{}", text)?;
        }
        Ok(())
    }

    fn output_custom_code(&mut self, name: &str) -> std::io::Result<()> {
        let mut path = self.current_namespace.clone();
        path.push(name.to_string());
        if let Some(code) = self.generator.config.custom_code.get(&path) {
            write!(self.out, "\n{}", code)?;
        }
        Ok(())
    }

    fn output_preamble(&mut self) -> Result<()> {
        let external_names = self
            .generator
            .config
            .external_definitions
            .values()
            .cloned()
            .flatten()
            .collect::<HashSet<_>>();
        writeln!(self.out, "#![allow(unused_imports)]")?;
        if !external_names.contains("Map") {
            writeln!(self.out, "use std::collections::BTreeMap as Map;")?;
        }
        if self.generator.config.serialization {
            writeln!(self.out, "use serde::{{Serialize, Deserialize}};")?;
        }
        if self.generator.config.serialization && !external_names.contains("Bytes") {
            writeln!(self.out, "use serde_bytes::ByteBuf as Bytes;")?;
        }
        for (module, definitions) in &self.generator.config.external_definitions {
            // Skip the empty module name.
            if !module.is_empty() {
                writeln!(
                    self.out,
                    "use {}::{{{}}};",
                    module,
                    definitions.to_vec().join(", "),
                )?;
            }
        }
        writeln!(self.out)?;
        if !self.generator.config.serialization && !external_names.contains("Bytes") {
            // If we are not going to use Serde derive macros, use plain vectors.
            writeln!(self.out, "type Bytes = Vec<u8>;\n")?;
        }
        Ok(())
    }

    fn quote_type(format: &Format, known_sizes: Option<&HashSet<&str>>) -> String {
        use Format::*;
        match format {
            TypeName(x) => {
                if let Some(set) = known_sizes {
                    if !set.contains(x.as_str()) {
                        return format!("Box<{}>", x);
                    }
                }
                x.to_string()
            }
            Unit => "()".into(),
            Bool => "bool".into(),
            I8 => "i8".into(),
            I16 => "i16".into(),
            I32 => "i32".into(),
            I64 => "i64".into(),
            I128 => "i128".into(),
            U8 => "u8".into(),
            U16 => "u16".into(),
            U32 => "u32".into(),
            U64 => "u64".into(),
            U128 => "u128".into(),
            F32 => "f32".into(),
            F64 => "f64".into(),
            Char => "char".into(),
            Str => "String".into(),
            Bytes => "Bytes".into(),

            Option(format) => format!("Option<{}>", Self::quote_type(format, known_sizes)),
            Seq(format) => format!("Vec<{}>", Self::quote_type(format, None)),
            Map { key, value } => format!(
                "Map<{}, {}>",
                Self::quote_type(key, None),
                Self::quote_type(value, None)
            ),
            Tuple(formats) => format!("({})", Self::quote_types(formats, known_sizes)),
            TupleArray { content, size } => {
                format!("[{}; {}]", Self::quote_type(content, known_sizes), *size)
            }

            Variable(_) => panic!("unexpected value"),
        }
    }

    fn quote_types(formats: &[Format], known_sizes: Option<&HashSet<&str>>) -> String {
        formats
            .iter()
            .map(|x| Self::quote_type(x, known_sizes))
            .collect::<Vec<_>>()
            .join(", ")
    }

    fn output_fields(&mut self, base: &[&str], fields: &[Named<Format>]) -> Result<()> {
        // Do not add 'pub' within variants.
        let prefix = if base.len() <= 1 && self.generator.track_visibility {
            "pub "
        } else {
            ""
        };
        for field in fields {
            self.output_comment(&field.name)?;
            writeln!(
                self.out,
                "{}{}: {},",
                prefix,
                field.name,
                Self::quote_type(&field.value, Some(&self.known_sizes)),
            )?;
        }
        Ok(())
    }

    fn output_variant(&mut self, base: &str, name: &str, variant: &VariantFormat) -> Result<()> {
        self.output_comment(name)?;
        use VariantFormat::*;
        match variant {
            Unit => writeln!(self.out, "{},", name),
            NewType(format) => writeln!(
                self.out,
                "{}({}),",
                name,
                Self::quote_type(format, Some(&self.known_sizes))
            ),
            Tuple(formats) => writeln!(
                self.out,
                "{}({}),",
                name,
                Self::quote_types(formats, Some(&self.known_sizes))
            ),
            Struct(fields) => {
                writeln!(self.out, "{} {{", name)?;
                self.current_namespace.push(name.to_string());
                self.out.indent();
                self.output_fields(&[base, name], fields)?;
                self.out.unindent();
                self.current_namespace.pop();
                writeln!(self.out, "}},")
            }
            Variable(_) => panic!("incorrect value"),
        }
    }

    fn output_variants(
        &mut self,
        base: &str,
        variants: &BTreeMap<u32, Named<VariantFormat>>,
    ) -> Result<()> {
        for (expected_index, (index, variant)) in variants.iter().enumerate() {
            assert_eq!(*index, expected_index as u32);
            self.output_variant(base, &variant.name, &variant.value)?;
        }
        Ok(())
    }

    fn output_container(&mut self, name: &str, format: &ContainerFormat) -> Result<()> {
        self.output_comment(name)?;
        let mut derive_macros = self.generator.derive_macros.clone();
        if self.generator.config.serialization {
            derive_macros.push("Serialize".to_string());
            derive_macros.push("Deserialize".to_string());
        }
        let mut prefix = String::new();
        if !derive_macros.is_empty() {
            prefix.push_str(&format!("#[derive({})]\n", derive_macros.join(", ")));
        }
        if let Some(text) = &self.generator.custom_derive_block {
            prefix.push_str(text);
            prefix.push_str("\n");
        }
        if self.generator.track_visibility {
            prefix.push_str("pub ");
        }

        use ContainerFormat::*;
        match format {
            UnitStruct => writeln!(self.out, "{}struct {};\n", prefix, name)?,
            NewTypeStruct(format) => writeln!(
                self.out,
                "{}struct {}({}{});\n",
                prefix,
                name,
                if self.generator.track_visibility {
                    "pub "
                } else {
                    ""
                },
                Self::quote_type(format, Some(&self.known_sizes))
            )?,
            TupleStruct(formats) => writeln!(
                self.out,
                "{}struct {}({});\n",
                prefix,
                name,
                Self::quote_types(formats, Some(&self.known_sizes))
            )?,
            Struct(fields) => {
                writeln!(self.out, "{}struct {} {{", prefix, name)?;
                self.current_namespace.push(name.to_string());
                self.out.indent();
                self.output_fields(&[name], fields)?;
                self.out.unindent();
                self.current_namespace.pop();
                writeln!(self.out, "}}\n")?;
            }
            Enum(variants) => {
                writeln!(self.out, "{}enum {} {{", prefix, name)?;
                self.current_namespace.push(name.to_string());
                self.out.indent();
                self.output_variants(name, variants)?;
                self.out.unindent();
                self.current_namespace.pop();
                writeln!(self.out, "}}\n")?;
            }
        }
        self.output_custom_code(name)
    }
}

/// Installer for generated source files in Rust.
pub struct Installer {
    install_dir: PathBuf,
}

impl Installer {
    pub fn new(install_dir: PathBuf) -> Self {
        Installer { install_dir }
    }

    fn runtime_installation_message(
        name: &str,
    ) -> std::result::Result<(), Box<dyn std::error::Error>> {
        eprintln!("Not installing sources for published crate {}", name);
        Ok(())
    }
}

impl crate::SourceInstaller for Installer {
    type Error = Box<dyn std::error::Error>;

    fn install_module(
        &self,
        config: &CodeGeneratorConfig,
        registry: &Registry,
    ) -> std::result::Result<(), Self::Error> {
        let generator = CodeGenerator::new(config);
        let (name, version) = {
            let parts = config.module_name.splitn(2, ':').collect::<Vec<_>>();
            if parts.len() >= 2 {
                (parts[0].to_string(), parts[1].to_string())
            } else {
                (parts[0].to_string(), "0.1.0".to_string())
            }
        };
        let dir_path = self.install_dir.join(&name);
        std::fs::create_dir_all(&dir_path)?;
        let mut cargo = std::fs::File::create(&dir_path.join("Cargo.toml"))?;
        write!(
            cargo,
            r#"[package]
name = "{}"
version = "{}"
edition = "2018"

[dependencies]
serde = {{ version = "1.0", features = ["derive"] }}
serde_bytes = "0.11"
"#,
            name, version,
        )?;
        std::fs::create_dir(dir_path.join("src"))?;
        let source_path = dir_path.join("src/lib.rs");
        let mut source = std::fs::File::create(&source_path)?;
        generator.output(&mut source, &registry)
    }

    fn install_serde_runtime(&self) -> std::result::Result<(), Self::Error> {
        Self::runtime_installation_message("serde")
    }

    fn install_bincode_runtime(&self) -> std::result::Result<(), Self::Error> {
        Self::runtime_installation_message("bincode")
    }

    fn install_lcs_runtime(&self) -> std::result::Result<(), Self::Error> {
        Self::runtime_installation_message("libra-canonical-serialization")
    }
}
<commit_msg>rust: clippy<commit_after>// Copyright (c) Facebook, Inc. and its affiliates
// SPDX-License-Identifier: MIT OR Apache-2.0

use crate::{
    analyzer,
    indent::{IndentConfig, IndentedWriter},
    CodeGeneratorConfig,
};
use serde_reflection::{ContainerFormat, Format, Named, Registry, VariantFormat};
use std::borrow::Cow;
use std::collections::{BTreeMap, HashSet};
use std::io::{Result, Write};
use std::path::PathBuf;

/// Main configuration object for code-generation in Rust.
pub struct CodeGenerator<'a> {
    /// Language-independent configuration.
    config: &'a CodeGeneratorConfig,
    /// Which derive macros should be added (independently from serialization).
    derive_macros: Vec<String>,
    /// Additional block of text added before each new container definition.
    custom_derive_block: Option<String>,
    /// Whether definitions and fields should be marked as `pub`.
    track_visibility: bool,
}

/// Shared state for the code generation of a Rust source file.
struct RustEmitter<'a, T> {
    /// Writer.
    out: IndentedWriter<T>,
    /// Generator.
    generator: &'a CodeGenerator<'a>,
    /// Track which definitions have a known size. (Used to add `Box` types.)
    known_sizes: Cow<'a, HashSet<&'a str>>,
    /// Current namespace (e.g. vec!["my_package", "my_module", "MyClass"])
    current_namespace: Vec<String>,
}

impl<'a> CodeGenerator<'a> {
    /// Create a Rust code generator for the given config.
    pub fn new(config: &'a CodeGeneratorConfig) -> Self {
        Self {
            config,
            derive_macros: vec!["Clone", "Debug", "PartialEq", "PartialOrd"]
                .into_iter()
                .map(String::from)
                .collect(),
            custom_derive_block: None,
            track_visibility: true,
        }
    }

    /// Which derive macros should be added (independently from serialization).
    pub fn with_derive_macros(mut self, derive_macros: Vec<String>) -> Self {
        self.derive_macros = derive_macros;
        self
    }

    /// Additional block of text added after `derive_macros` (if any), before each new
    /// container definition.
    pub fn with_custom_derive_block(mut self, custom_derive_block: Option<String>) -> Self {
        self.custom_derive_block = custom_derive_block;
        self
    }

    /// Whether definitions and fields should be marked as `pub`.
    pub fn with_track_visibility(mut self, track_visibility: bool) -> Self {
        self.track_visibility = track_visibility;
        self
    }

    /// Write container definitions in Rust.
    pub fn output(
        &self,
        out: &mut dyn Write,
        registry: &Registry,
    ) -> std::result::Result<(), Box<dyn std::error::Error>> {
        let external_names = self
            .config
            .external_definitions
            .values()
            .cloned()
            .flatten()
            .collect();
        let dependencies =
            analyzer::get_dependency_map_with_external_dependencies(registry, &external_names)?;
        let entries = analyzer::best_effort_topological_sort(&dependencies);

        let known_sizes = external_names
            .iter()
            .map(<String as std::ops::Deref>::deref)
            .collect::<HashSet<_>>();

        let current_namespace = self
            .config
            .module_name
            .split('.')
            .map(String::from)
            .collect();
        let mut emitter = RustEmitter {
            out: IndentedWriter::new(out, IndentConfig::Space(4)),
            generator: self,
            known_sizes: Cow::Owned(known_sizes),
            current_namespace,
        };

        emitter.output_preamble()?;
        for name in entries {
            let format = &registry[name];
            emitter.output_container(name, format)?;
            emitter.known_sizes.to_mut().insert(name);
        }
        Ok(())
    }

    /// For each container, generate a Rust definition.
    pub fn quote_container_definitions(
        &self,
        registry: &Registry,
    ) -> std::result::Result<BTreeMap<String, String>, Box<dyn std::error::Error>> {
        let dependencies = analyzer::get_dependency_map(registry)?;
        let entries = analyzer::best_effort_topological_sort(&dependencies);

        let mut result = BTreeMap::new();
        let mut known_sizes = HashSet::new();
        let current_namespace = self
            .config
            .module_name
            .split('.')
            .map(String::from)
            .collect::<Vec<_>>();

        for name in entries {
            let mut content = Vec::new();
            {
                let mut emitter = RustEmitter {
                    out: IndentedWriter::new(&mut content, IndentConfig::Space(4)),
                    generator: self,
                    known_sizes: Cow::Borrowed(&known_sizes),
                    current_namespace: current_namespace.clone(),
                };
                let format = &registry[name];
                emitter.output_container(name, format)?;
            }
            known_sizes.insert(name);
            result.insert(
                name.to_string(),
                String::from_utf8_lossy(&content).trim().to_string() + "\n",
            );
        }
        Ok(result)
    }
}

impl<'a, T> RustEmitter<'a, T>
where
    T: std::io::Write,
{
    fn output_comment(&mut self, name: &str) -> std::io::Result<()> {
        let mut path = self.current_namespace.clone();
        path.push(name.to_string());
        if let Some(doc) = self.generator.config.comments.get(&path) {
            let text = textwrap::indent(doc, "/// ").replace("\n\n", "\n///\n");
            write!(self.out, "\n{}", text)?;
        }
        Ok(())
    }

    fn output_custom_code(&mut self, name: &str) -> std::io::Result<()> {
        let mut path = self.current_namespace.clone();
        path.push(name.to_string());
        if let Some(code) = self.generator.config.custom_code.get(&path) {
            write!(self.out, "\n{}", code)?;
        }
        Ok(())
    }

    fn output_preamble(&mut self) -> Result<()> {
        let external_names = self
            .generator
            .config
            .external_definitions
            .values()
            .cloned()
            .flatten()
            .collect::<HashSet<_>>();
        writeln!(self.out, "#![allow(unused_imports)]")?;
        if !external_names.contains("Map") {
            writeln!(self.out, "use std::collections::BTreeMap as Map;")?;
        }
        if self.generator.config.serialization {
            writeln!(self.out, "use serde::{{Serialize, Deserialize}};")?;
        }
        if self.generator.config.serialization && !external_names.contains("Bytes") {
            writeln!(self.out, "use serde_bytes::ByteBuf as Bytes;")?;
        }
        for (module, definitions) in &self.generator.config.external_definitions {
            // Skip the empty module name.
            if !module.is_empty() {
                writeln!(
                    self.out,
                    "use {}::{{{}}};",
                    module,
                    definitions.to_vec().join(", "),
                )?;
            }
        }
        writeln!(self.out)?;
        if !self.generator.config.serialization && !external_names.contains("Bytes") {
            // If we are not going to use Serde derive macros, use plain vectors.
            writeln!(self.out, "type Bytes = Vec<u8>;\n")?;
        }
        Ok(())
    }

    fn quote_type(format: &Format, known_sizes: Option<&HashSet<&str>>) -> String {
        use Format::*;
        match format {
            TypeName(x) => {
                if let Some(set) = known_sizes {
                    if !set.contains(x.as_str()) {
                        return format!("Box<{}>", x);
                    }
                }
                x.to_string()
            }
            Unit => "()".into(),
            Bool => "bool".into(),
            I8 => "i8".into(),
            I16 => "i16".into(),
            I32 => "i32".into(),
            I64 => "i64".into(),
            I128 => "i128".into(),
            U8 => "u8".into(),
            U16 => "u16".into(),
            U32 => "u32".into(),
            U64 => "u64".into(),
            U128 => "u128".into(),
            F32 => "f32".into(),
            F64 => "f64".into(),
            Char => "char".into(),
            Str => "String".into(),
            Bytes => "Bytes".into(),

            Option(format) => format!("Option<{}>", Self::quote_type(format, known_sizes)),
            Seq(format) => format!("Vec<{}>", Self::quote_type(format, None)),
            Map { key, value } => format!(
                "Map<{}, {}>",
                Self::quote_type(key, None),
                Self::quote_type(value, None)
            ),
            Tuple(formats) => format!("({})", Self::quote_types(formats, known_sizes)),
            TupleArray { content, size } => {
                format!("[{}; {}]", Self::quote_type(content, known_sizes), *size)
            }

            Variable(_) => panic!("unexpected value"),
        }
    }

    fn quote_types(formats: &[Format], known_sizes: Option<&HashSet<&str>>) -> String {
        formats
            .iter()
            .map(|x| Self::quote_type(x, known_sizes))
            .collect::<Vec<_>>()
            .join(", ")
    }

    fn output_fields(&mut self, base: &[&str], fields: &[Named<Format>]) -> Result<()> {
        // Do not add 'pub' within variants.
        let prefix = if base.len() <= 1 && self.generator.track_visibility {
            "pub "
        } else {
            ""
        };
        for field in fields {
            self.output_comment(&field.name)?;
            writeln!(
                self.out,
                "{}{}: {},",
                prefix,
                field.name,
                Self::quote_type(&field.value, Some(&self.known_sizes)),
            )?;
        }
        Ok(())
    }

    fn output_variant(&mut self, base: &str, name: &str, variant: &VariantFormat) -> Result<()> {
        self.output_comment(name)?;
        use VariantFormat::*;
        match variant {
            Unit => writeln!(self.out, "{},", name),
            NewType(format) => writeln!(
                self.out,
                "{}({}),",
                name,
                Self::quote_type(format, Some(&self.known_sizes))
            ),
            Tuple(formats) => writeln!(
                self.out,
                "{}({}),",
                name,
                Self::quote_types(formats, Some(&self.known_sizes))
            ),
            Struct(fields) => {
                writeln!(self.out, "{} {{", name)?;
                self.current_namespace.push(name.to_string());
                self.out.indent();
                self.output_fields(&[base, name], fields)?;
                self.out.unindent();
                self.current_namespace.pop();
                writeln!(self.out, "}},")
            }
            Variable(_) => panic!("incorrect value"),
        }
    }

    fn output_variants(
        &mut self,
        base: &str,
        variants: &BTreeMap<u32, Named<VariantFormat>>,
    ) -> Result<()> {
        for (expected_index, (index, variant)) in variants.iter().enumerate() {
            assert_eq!(*index, expected_index as u32);
            self.output_variant(base, &variant.name, &variant.value)?;
        }
        Ok(())
    }

    fn output_container(&mut self, name: &str, format: &ContainerFormat) -> Result<()> {
        self.output_comment(name)?;
        let mut derive_macros = self.generator.derive_macros.clone();
        if self.generator.config.serialization {
            derive_macros.push("Serialize".to_string());
            derive_macros.push("Deserialize".to_string());
        }
        let mut prefix = String::new();
        if !derive_macros.is_empty() {
            prefix.push_str(&format!("#[derive({})]\n", derive_macros.join(", ")));
        }
        if let Some(text) = &self.generator.custom_derive_block {
            prefix.push_str(text);
            prefix.push('\n');
        }
        if self.generator.track_visibility {
            prefix.push_str("pub ");
        }

        use ContainerFormat::*;
        match format {
            UnitStruct => writeln!(self.out, "{}struct {};\n", prefix, name)?,
            NewTypeStruct(format) => writeln!(
                self.out,
                "{}struct {}({}{});\n",
                prefix,
                name,
                if self.generator.track_visibility {
                    "pub "
                } else {
                    ""
                },
                Self::quote_type(format, Some(&self.known_sizes))
            )?,
            TupleStruct(formats) => writeln!(
                self.out,
                "{}struct {}({});\n",
                prefix,
                name,
                Self::quote_types(formats, Some(&self.known_sizes))
            )?,
            Struct(fields) => {
                writeln!(self.out, "{}struct {} {{", prefix, name)?;
                self.current_namespace.push(name.to_string());
                self.out.indent();
                self.output_fields(&[name], fields)?;
                self.out.unindent();
                self.current_namespace.pop();
                writeln!(self.out, "}}\n")?;
            }
            Enum(variants) => {
                writeln!(self.out, "{}enum {} {{", prefix, name)?;
                self.current_namespace.push(name.to_string());
                self.out.indent();
                self.output_variants(name, variants)?;
                self.out.unindent();
                self.current_namespace.pop();
                writeln!(self.out, "}}\n")?;
            }
        }
        self.output_custom_code(name)
    }
}

/// Installer for generated source files in Rust.
pub struct Installer {
    install_dir: PathBuf,
}

impl Installer {
    pub fn new(install_dir: PathBuf) -> Self {
        Installer { install_dir }
    }

    fn runtime_installation_message(
        name: &str,
    ) -> std::result::Result<(), Box<dyn std::error::Error>> {
        eprintln!("Not installing sources for published crate {}", name);
        Ok(())
    }
}

impl crate::SourceInstaller for Installer {
    type Error = Box<dyn std::error::Error>;

    fn install_module(
        &self,
        config: &CodeGeneratorConfig,
        registry: &Registry,
    ) -> std::result::Result<(), Self::Error> {
        let generator = CodeGenerator::new(config);
        let (name, version) = {
            let parts = config.module_name.splitn(2, ':').collect::<Vec<_>>();
            if parts.len() >= 2 {
                (parts[0].to_string(), parts[1].to_string())
            } else {
                (parts[0].to_string(), "0.1.0".to_string())
            }
        };
        let dir_path = self.install_dir.join(&name);
        std::fs::create_dir_all(&dir_path)?;
        let mut cargo = std::fs::File::create(&dir_path.join("Cargo.toml"))?;
        write!(
            cargo,
            r#"[package]
name = "{}"
version = "{}"
edition = "2018"

[dependencies]
serde = {{ version = "1.0", features = ["derive"] }}
serde_bytes = "0.11"
"#,
            name, version,
        )?;
        std::fs::create_dir(dir_path.join("src"))?;
        let source_path = dir_path.join("src/lib.rs");
        let mut source = std::fs::File::create(&source_path)?;
        generator.output(&mut source, &registry)
    }

    fn install_serde_runtime(&self) -> std::result::Result<(), Self::Error> {
        Self::runtime_installation_message("serde")
    }

    fn install_bincode_runtime(&self) -> std::result::Result<(), Self::Error> {
        Self::runtime_installation_message("bincode")
    }

    fn install_lcs_runtime(&self) -> std::result::Result<(), Self::Error> {
        Self::runtime_installation_message("libra-canonical-serialization")
    }
}
<|endoftext|>


<commit_before>use crate::QCoreApplication;
use cpp_utils::{MutPtr, MutRef};
use std::os::raw::{c_char, c_int};
use std::process;

/// A struct providing valid `argc` and `argv` values for Qt application
/// objects.
///
/// Constructors of `qt_core::core_application::CoreApplication`,
/// `qt_gui::gui_application::GuiApplication` and `qt_widgets::application::Application`
/// require `argc` and `argv` values that are available in C++'s `main` function but
/// not available in Rust. More importantly, `argc` and `argv` must be valid for the entire
/// life of the application. This struct stores list of arguments in a format compatible with
/// `argc` and `argv`, and can be used to initialize Qt application objects.
/// `CoreApplicationArgs` must live longer than the application object.
///
/// `CoreApplication::init` convenience function
/// and similar functions in the other application types
/// can be used instead of `CoreApplicationArgs`.
pub struct QCoreApplicationArgs {
    _values: Vec<Vec<u8>>,
    argc: Box<c_int>,
    argv: Vec<*mut c_char>,
}

impl QCoreApplicationArgs {
    /// Creates an object containing `args`.
    pub fn new(mut args: Vec<Vec<u8>>) -> QCoreApplicationArgs {
        for arg in &mut args {
            if !arg.ends_with(&[0]) {
                arg.push(0);
            }
        }
        QCoreApplicationArgs {
            argc: Box::new(args.len() as c_int),
            argv: args
                .iter_mut()
                .map(|x| x.as_mut_ptr() as *mut c_char)
                .collect(),
            _values: args,
        }
    }

    /// Returns `(argc, argv)` values in the form accepted by the application objects'
    /// constructors.
    pub fn get(&mut self) -> (*mut c_int, *mut *mut c_char) {
        (self.argc.as_mut(), self.argv.as_mut_ptr())
    }

    #[cfg(unix)]
    /// Creates an object representing real arguments of the application.
    /// On Windows, this function uses empty argument list for performance reasons because
    /// Qt doesn't use `argc` and `argv` on Windows at all.
    pub fn from_real() -> QCoreApplicationArgs {
        use std::os::unix::ffi::OsStringExt;

        let args = std::env::args_os().map(|arg| arg.into_vec()).collect();
        QCoreApplicationArgs::new(args)
    }
    #[cfg(windows)]
    /// Creates an object representing real arguments of the application.
    /// On Windows, this function uses empty argument list for performance reasons because
    /// Qt doesn't use `argc` and `argv` on Windows at all.
    pub fn from_real() -> QCoreApplicationArgs {
        // Qt doesn't use argc and argv on Windows anyway
        // TODO: check this
        QCoreApplicationArgs::new(Vec::new())
    }
}

impl QCoreApplication {
    /// A convenience function for performing proper initialization and de-initialization of
    /// a Qt application.
    ///
    /// This function creates `CoreApplication` with valid `argc` and `argv`, calls the passed
    /// closure `f(app)` with the application object and exist the process with the exit code
    /// returned by the closure. The closure should perform the initialization of the application
    /// and either return immediately or call `CoreApplication::exec()` and return its return value:
    /// ```no_run
    /// use qt_core::QCoreApplication;
    ///
    /// fn main() {
    ///     QCoreApplication::init(|app| {
    ///         unsafe {
    ///             // initialization goes here
    ///             QCoreApplication::exec()
    ///         }
    ///     })
    /// }
    /// ```
    pub fn init<F: FnOnce(::cpp_utils::MutPtr<QCoreApplication>) -> i32>(f: F) -> ! {
        let exit_code = {
            unsafe {
                let mut args = QCoreApplicationArgs::from_real();
                let (argc, argv) = args.get();
                let mut app = QCoreApplication::new_2a(
                    MutRef::from_raw(argc).unwrap(),
                    MutPtr::from_raw(argv),
                );
                f(app.as_mut_ptr())
            }
        }; // drop `app` and `args`
        process::exit(exit_code)
    }
}
<commit_msg>fix QCoreApplicationArgs for Windows<commit_after>use crate::{QCoreApplication, QString};
use cpp_utils::{MutPtr, MutRef};
use std::iter::once;
use std::os::raw::{c_char, c_int};
use std::process;

/// A struct providing valid `argc` and `argv` values for Qt application
/// objects.
///
/// Constructors of `qt_core::core_application::CoreApplication`,
/// `qt_gui::gui_application::GuiApplication` and `qt_widgets::application::Application`
/// require `argc` and `argv` values that are available in C++'s `main` function but
/// not available in Rust. More importantly, `argc` and `argv` must be valid for the entire
/// life of the application. This struct stores list of arguments in a format compatible with
/// `argc` and `argv`, and can be used to initialize Qt application objects.
/// `CoreApplicationArgs` must live longer than the application object.
///
/// `CoreApplication::init` convenience function
/// and similar functions in the other application types
/// can be used instead of `CoreApplicationArgs`.
pub struct QCoreApplicationArgs {
    _values: Vec<Vec<u8>>,
    argc: Box<c_int>,
    argv: Vec<*mut c_char>,
}

impl QCoreApplicationArgs {
    /// Creates an object containing `args`.
    pub fn new() -> QCoreApplicationArgs {
        // Qt uses `QString::fromLocal8Bit()` to decode `argv`,
        // so we use `QString::toLocal8Bit()` to encode it.
        let mut args = std::env::args()
            .map(|arg| unsafe {
                QString::from_std_str(&arg)
                    .to_local8_bit()
                    .as_slice()
                    .iter()
                    .map(|&c| c as u8)
                    .chain(once(0))
                    .collect::<Vec<_>>()
            })
            .collect::<Vec<_>>();

        QCoreApplicationArgs {
            argc: Box::new(args.len() as c_int),
            argv: args
                .iter_mut()
                .map(|x| x.as_mut_ptr() as *mut c_char)
                .collect(),
            _values: args,
        }
    }

    /// Returns `(argc, argv)` values in the form accepted by the application objects'
    /// constructors.
    pub fn get(&mut self) -> (*mut c_int, *mut *mut c_char) {
        (self.argc.as_mut(), self.argv.as_mut_ptr())
    }
}

impl QCoreApplication {
    /// A convenience function for performing proper initialization and de-initialization of
    /// a Qt application.
    ///
    /// This function creates `CoreApplication` with valid `argc` and `argv`, calls the passed
    /// closure `f(app)` with the application object and exist the process with the exit code
    /// returned by the closure. The closure should perform the initialization of the application
    /// and either return immediately or call `CoreApplication::exec()` and return its return value:
    /// ```no_run
    /// use qt_core::QCoreApplication;
    ///
    /// fn main() {
    ///     QCoreApplication::init(|app| {
    ///         unsafe {
    ///             // initialization goes here
    ///             QCoreApplication::exec()
    ///         }
    ///     })
    /// }
    /// ```
    pub fn init<F: FnOnce(::cpp_utils::MutPtr<QCoreApplication>) -> i32>(f: F) -> ! {
        let exit_code = {
            unsafe {
                let mut args = QCoreApplicationArgs::new();
                let (argc, argv) = args.get();
                let mut app = QCoreApplication::new_2a(
                    MutRef::from_raw(argc).unwrap(),
                    MutPtr::from_raw(argv),
                );
                f(app.as_mut_ptr())
            }
        }; // drop `app` and `args`
        process::exit(exit_code)
    }
}
<|endoftext|>


<commit_before>use std::libc::{c_ulonglong};
use std::vec::{with_capacity, append};
use std::vec::raw::{to_mut_ptr, to_ptr};

#[doc(hidden)]
pub fn marshal<T>(buf: &[u8],
                  padbefore: uint,
                  bytestodrop: uint,
                  f: &fn(*mut u8, *u8, c_ulonglong) -> T
                 ) -> (~[u8], T) {
    let mut dst = with_capacity(buf.len() + padbefore);
    for _ in range(0, padbefore) {
        dst.push(0);
    }
    dst = append(dst, buf);
    let pdst = to_mut_ptr(dst);
    let psrc = to_ptr(dst);
    let res = f(pdst, psrc, dst.len() as c_ulonglong);
    for _ in range(0, bytestodrop) {
        dst.shift();
    }
    (dst, res)
}
<commit_msg>Optimized marshal function<commit_after>use std::libc::{c_ulonglong};
use std::vec::{with_capacity, append};
use std::vec::raw::{to_mut_ptr, to_ptr};

#[doc(hidden)]
pub fn marshal<T>(buf: &[u8],
                  padbefore: uint,
                  bytestodrop: uint,
                  f: &fn(*mut u8, *u8, c_ulonglong) -> T
                 ) -> (~[u8], T) {
    let mut dst = with_capacity(buf.len() + padbefore);
    for _ in range(0, padbefore) {
        dst.push(0);
    }
    dst = append(dst, buf);
    let pdst = to_mut_ptr(dst);
    let psrc = to_ptr(dst);
    let res = f(pdst, psrc, dst.len() as c_ulonglong);
    (dst.move_iter().skip(bytestodrop).collect(), res)
}
<|endoftext|>


<commit_before>// Copyright 2015-2019 Benjamin Fry <benjaminfry@me.com>
//
// Licensed under the Apache License, Version 2.0, <LICENSE-APACHE or
// http://apache.org/licenses/LICENSE-2.0> or the MIT license <LICENSE-MIT or
// http://opensource.org/licenses/MIT>, at your option. This file may not be
// copied, modified, or distributed except according to those terms.

use std::pin::Pin;
use std::sync::Arc;
use std::task::Context;

use futures::lock::Mutex;
use futures::{future, Future, Poll, TryFutureExt};
use smallvec::SmallVec;

use proto::error::ProtoError;
use proto::op::ResponseCode;
use proto::xfer::{DnsHandle, DnsRequest, DnsResponse};

use crate::config::{ResolverConfig, ResolverOpts};
#[cfg(feature = "mdns")]
use crate::name_server;
use crate::name_server::{ConnectionHandle, ConnectionProvider, NameServer, StandardConnection};

/// A pool of NameServers
///
/// This is not expected to be used directly, see [`AsyncResolver`].
#[derive(Clone)]
pub struct NameServerPool<C: DnsHandle + 'static, P: ConnectionProvider<ConnHandle = C> + 'static> {
    // TODO: switch to FuturesMutex (Mutex will have some undesireable locking)
    datagram_conns: Arc<Mutex<Vec<NameServer<C, P>>>>, /* All NameServers must be the same type */
    stream_conns: Arc<Mutex<Vec<NameServer<C, P>>>>,   /* All NameServers must be the same type */
    #[cfg(feature = "mdns")]
    mdns_conns: NameServer<C, P>, /* All NameServers must be the same type */
    options: ResolverOpts,
    conn_provider: P,
}

impl NameServerPool<ConnectionHandle, StandardConnection> {
    pub(crate) fn from_config(config: &ResolverConfig, options: &ResolverOpts) -> Self {
        Self::from_config_with_provider(config, options, StandardConnection)
    }
}

impl<C: DnsHandle + 'static, P: ConnectionProvider<ConnHandle = C> + 'static> NameServerPool<C, P> {
    pub(crate) fn from_config_with_provider(
        config: &ResolverConfig,
        options: &ResolverOpts,
        conn_provider: P,
    ) -> NameServerPool<C, P> {
        let datagram_conns: Vec<NameServer<C, P>> = config
            .name_servers()
            .iter()
            .filter(|ns_config| ns_config.protocol.is_datagram())
            .map(|ns_config| {
                #[cfg(feature = "dns-over-rustls")]
                let ns_config = {
                    let mut ns_config = ns_config.clone();
                    ns_config.tls_config = config.client_config().clone();
                    ns_config
                };
                #[cfg(not(feature = "dns-over-rustls"))]
                let ns_config = { ns_config.clone() };

                NameServer::<C, P>::new_with_provider(ns_config, *options, conn_provider.clone())
            })
            .collect();

        let stream_conns: Vec<NameServer<C, P>> = config
            .name_servers()
            .iter()
            .filter(|ns_config| ns_config.protocol.is_stream())
            .map(|ns_config| {
                #[cfg(feature = "dns-over-rustls")]
                let ns_config = {
                    let mut ns_config = ns_config.clone();
                    ns_config.tls_config = config.client_config().clone();
                    ns_config
                };
                #[cfg(not(feature = "dns-over-rustls"))]
                let ns_config = { ns_config.clone() };

                NameServer::<C, P>::new_with_provider(ns_config, *options, conn_provider.clone())
            })
            .collect();

        NameServerPool {
            datagram_conns: Arc::new(Mutex::new(datagram_conns)),
            stream_conns: Arc::new(Mutex::new(stream_conns)),
            #[cfg(feature = "mdns")]
            mdns_conns: name_server::mdns_nameserver(*options, conn_provider.clone()),
            options: *options,
            conn_provider,
        }
    }

    #[doc(hidden)]
    #[cfg(not(feature = "mdns"))]
    pub fn from_nameservers(
        options: &ResolverOpts,
        datagram_conns: Vec<NameServer<C, P>>,
        stream_conns: Vec<NameServer<C, P>>,
        conn_provider: P,
    ) -> Self {
        NameServerPool {
            datagram_conns: Arc::new(Mutex::new(datagram_conns.into_iter().collect())),
            stream_conns: Arc::new(Mutex::new(stream_conns.into_iter().collect())),
            options: *options,
            conn_provider,
        }
    }

    #[doc(hidden)]
    #[cfg(feature = "mdns")]
    pub fn from_nameservers(
        options: &ResolverOpts,
        datagram_conns: Vec<NameServer<C, P>>,
        stream_conns: Vec<NameServer<C, P>>,
        mdns_conns: NameServer<C, P>,
        conn_provider: P,
    ) -> Self {
        NameServerPool {
            datagram_conns: Arc::new(Mutex::new(datagram_conns.into_iter().collect())),
            stream_conns: Arc::new(Mutex::new(stream_conns.into_iter().collect())),
            mdns_conns,
            options: *options,
            conn_provider,
        }
    }

    async fn try_send(
        opts: ResolverOpts,
        conns: Arc<Mutex<Vec<NameServer<C, P>>>>,
        request: DnsRequest,
    ) -> Result<DnsResponse, ProtoError> {
        // pull a lock on the shared connections, lock releases at the end of the method
        let mut conns = conns.lock().await;

        // select the highest priority connection
        //   reorder the connections based on current view...
        //   this reorders the inner set
        conns.sort_unstable();

        // TODO: restrict this size to a maximum # of NameServers to try
        // get a stable view for trying all connections
        //   we split into chunks of the number of parallel requests to issue
        let conns: Vec<NameServer<C, P>> = conns.clone();
        let request_loop = request.clone();

        parallel_conn_loop(conns, request_loop, opts).await
    }
}

impl<C, P> DnsHandle for NameServerPool<C, P>
where
    C: DnsHandle + 'static,
    P: ConnectionProvider<ConnHandle = C> + 'static,
{
    type Response = Pin<Box<dyn Future<Output = Result<DnsResponse, ProtoError>> + Send>>;

    fn send<R: Into<DnsRequest>>(&mut self, request: R) -> Self::Response {
        let opts = self.options;
        let request = request.into();
        let datagram_conns = Arc::clone(&self.datagram_conns);
        let stream_conns1 = Arc::clone(&self.stream_conns);
        let stream_conns2 = Arc::clone(&self.stream_conns);
        // TODO: remove this clone, return the Message in the error?
        let tcp_message1 = request.clone();
        let tcp_message2 = request.clone();

        // if it's a .local. query, then we *only* query mDNS, these should never be sent on to upstream resolvers
        #[cfg(feature = "mdns")]
        let mdns = mdns::maybe_local(&mut self.mdns_conns, request);

        // TODO: limited to only when mDNS is enabled, but this should probably always be enforced?
        #[cfg(not(feature = "mdns"))]
        let mdns = Local::NotMdns(request);

        // local queries are queried through mDNS
        if mdns.is_local() {
            return mdns.take_future();
        }

        // TODO: should we allow mDNS to be used for standard lookups as well?

        // it wasn't a local query, continue with standard lookup path
        let request = mdns.take_request();

        // First try the UDP connections
        Box::pin(
            Self::try_send(opts, datagram_conns, request)
                .and_then(move |response| {
                    // handling promotion from datagram to stream base on truncation in message
                    if ResponseCode::NoError == response.response_code() && response.truncated() {
                        // TCP connections should not truncate
                        future::Either::Left(Self::try_send(opts, stream_conns1, tcp_message1))
                    } else {
                        // Return the result from the UDP connection
                        future::Either::Right(future::ok(response))
                    }
                })
                // if UDP fails, try TCP
                .or_else(move |_| Self::try_send(opts, stream_conns2, tcp_message2)),
        )
    }
}

// TODO: we should be able to have a self-referential future here with Pin and not require cloned conns
/// An async function that will loop over all the conns with a max parallel request count of ops.num_concurrent_req
async fn parallel_conn_loop<C, P>(
    mut conns: Vec<NameServer<C, P>>,
    request: DnsRequest,
    opts: ResolverOpts,
) -> Result<DnsResponse, ProtoError>
where
    C: DnsHandle + 'static,
    P: ConnectionProvider<ConnHandle = C> + 'static,
{
    let mut err = ProtoError::from("No connections available");

    loop {
        let request_cont = request.clone();

        // construct the parallel requests, 2 is the default
        let mut par_conns = SmallVec::<[NameServer<C, P>; 2]>::new();
        let count = conns.len().min(opts.num_concurrent_reqs.max(1));
        for conn in conns.drain(..count) {
            par_conns.push(conn);
        }

        // construct the requests to send
        let requests = if par_conns.is_empty() {
            return Err(err);
        } else {
            par_conns
                .into_iter()
                .map(move |mut conn| conn.send(request_cont.clone()))
        };

        match future::select_ok(requests).await {
            Ok((sent, _)) => return Ok(sent),
            // consider a debug msg here
            Err(e) => {
                err = e;
                continue;
            }
        };
    }
}

#[cfg(feature = "mdns")]
mod mdns {
    use super::*;

    use proto::rr::domain::usage;
    use proto::DnsHandle;

    /// Returns true
    pub fn maybe_local<C, P>(name_server: &mut NameServer<C, P>, request: DnsRequest) -> Local
    where
        C: DnsHandle + 'static,
        P: ConnectionProvider<ConnHandle = C> + 'static,
    {
        if request
            .queries()
            .iter()
            .any(|query| usage::LOCAL.name().zone_of(query.name()))
        {
            Local::ResolveFuture(name_server.send(request))
        } else {
            Local::NotMdns(request)
        }
    }
}

pub enum Local {
    #[allow(dead_code)]
    ResolveFuture(Pin<Box<dyn Future<Output = Result<DnsResponse, ProtoError>> + Send + Unpin>>),
    NotMdns(DnsRequest),
}

impl Local {
    fn is_local(&self) -> bool {
        if let Local::ResolveFuture(..) = *self {
            true
        } else {
            false
        }
    }

    /// Takes the future
    ///
    /// # Panics
    ///
    /// Panics if this is in fact a Local::NotMdns
    fn take_future(
        self,
    ) -> Pin<Box<dyn Future<Output = Result<DnsResponse, ProtoError>> + Send + Unpin>> {
        match self {
            Local::ResolveFuture(future) => future,
            _ => panic!("non Local queries have no future, see take_message()"),
        }
    }

    /// Takes the message
    ///
    /// # Panics
    ///
    /// Panics if this is in fact a Local::ResolveFuture
    fn take_request(self) -> DnsRequest {
        match self {
            Local::NotMdns(request) => request,
            _ => panic!("Local queries must be polled, see take_future()"),
        }
    }
}

impl Future for Local {
    type Output = Result<DnsResponse, ProtoError>;

    fn poll(mut self: Pin<&mut Self>, cx: &mut Context) -> Poll<Self::Output> {
        match *self {
            Local::ResolveFuture(ref mut ns) => ns.as_mut().poll(cx),
            // TODO: making this a panic for now
            Local::NotMdns(..) => panic!("Local queries that are not mDNS should not be polled"), //Local::NotMdns(message) => return Err(ResolveErrorKind::Message("not mDNS")),
        }
    }
}

#[cfg(test)]
mod tests {
    extern crate env_logger;

    use std::net::{IpAddr, Ipv4Addr, SocketAddr};

    use tokio::runtime::current_thread::Runtime;

    use proto::op::Query;
    use proto::rr::{Name, RecordType};
    use proto::xfer::{DnsHandle, DnsRequestOptions};

    use super::*;
    use crate::config::NameServerConfig;
    use crate::config::Protocol;

    #[ignore]
    // because of there is a real connection that needs a reasonable timeout
    #[test]
    fn test_failed_then_success_pool() {
        let config1 = NameServerConfig {
            socket_addr: SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 0, 252)), 253),
            protocol: Protocol::Udp,
            tls_dns_name: None,
            #[cfg(feature = "dns-over-rustls")]
            tls_config: None,
        };

        let config2 = NameServerConfig {
            socket_addr: SocketAddr::new(IpAddr::V4(Ipv4Addr::new(8, 8, 8, 8)), 53),
            protocol: Protocol::Udp,
            tls_dns_name: None,
            #[cfg(feature = "dns-over-rustls")]
            tls_config: None,
        };

        let mut resolver_config = ResolverConfig::new();
        resolver_config.add_name_server(config1);
        resolver_config.add_name_server(config2);

        let mut io_loop = Runtime::new().unwrap();
        let mut pool = NameServerPool::<_, StandardConnection>::from_config(
            &resolver_config,
            &ResolverOpts::default(),
        );

        let name = Name::parse("www.example.com.", None).unwrap();

        // TODO: it's not clear why there are two failures before the success
        for i in 0..2 {
            assert!(
                io_loop
                    .block_on(pool.lookup(
                        Query::query(name.clone(), RecordType::A),
                        DnsRequestOptions::default()
                    ))
                    .is_err(),
                "iter: {}",
                i
            );
        }

        for i in 0..10 {
            assert!(
                io_loop
                    .block_on(pool.lookup(
                        Query::query(name.clone(), RecordType::A),
                        DnsRequestOptions::default()
                    ))
                    .is_ok(),
                "iter: {}",
                i
            );
        }
    }
}
<commit_msg>lock prevent the nameserver to do concurrent request<commit_after>// Copyright 2015-2019 Benjamin Fry <benjaminfry@me.com>
//
// Licensed under the Apache License, Version 2.0, <LICENSE-APACHE or
// http://apache.org/licenses/LICENSE-2.0> or the MIT license <LICENSE-MIT or
// http://opensource.org/licenses/MIT>, at your option. This file may not be
// copied, modified, or distributed except according to those terms.

use std::pin::Pin;
use std::sync::Arc;
use std::task::Context;

use futures::lock::Mutex;
use futures::{future, Future, Poll, TryFutureExt};
use smallvec::SmallVec;

use proto::error::ProtoError;
use proto::op::ResponseCode;
use proto::xfer::{DnsHandle, DnsRequest, DnsResponse};

use crate::config::{ResolverConfig, ResolverOpts};
#[cfg(feature = "mdns")]
use crate::name_server;
use crate::name_server::{ConnectionHandle, ConnectionProvider, NameServer, StandardConnection};

/// A pool of NameServers
///
/// This is not expected to be used directly, see [`AsyncResolver`].
#[derive(Clone)]
pub struct NameServerPool<C: DnsHandle + 'static, P: ConnectionProvider<ConnHandle = C> + 'static> {
    // TODO: switch to FuturesMutex (Mutex will have some undesireable locking)
    datagram_conns: Arc<Mutex<Vec<NameServer<C, P>>>>, /* All NameServers must be the same type */
    stream_conns: Arc<Mutex<Vec<NameServer<C, P>>>>,   /* All NameServers must be the same type */
    #[cfg(feature = "mdns")]
    mdns_conns: NameServer<C, P>, /* All NameServers must be the same type */
    options: ResolverOpts,
    conn_provider: P,
}

impl NameServerPool<ConnectionHandle, StandardConnection> {
    pub(crate) fn from_config(config: &ResolverConfig, options: &ResolverOpts) -> Self {
        Self::from_config_with_provider(config, options, StandardConnection)
    }
}

impl<C: DnsHandle + 'static, P: ConnectionProvider<ConnHandle = C> + 'static> NameServerPool<C, P> {
    pub(crate) fn from_config_with_provider(
        config: &ResolverConfig,
        options: &ResolverOpts,
        conn_provider: P,
    ) -> NameServerPool<C, P> {
        let datagram_conns: Vec<NameServer<C, P>> = config
            .name_servers()
            .iter()
            .filter(|ns_config| ns_config.protocol.is_datagram())
            .map(|ns_config| {
                #[cfg(feature = "dns-over-rustls")]
                let ns_config = {
                    let mut ns_config = ns_config.clone();
                    ns_config.tls_config = config.client_config().clone();
                    ns_config
                };
                #[cfg(not(feature = "dns-over-rustls"))]
                let ns_config = { ns_config.clone() };

                NameServer::<C, P>::new_with_provider(ns_config, *options, conn_provider.clone())
            })
            .collect();

        let stream_conns: Vec<NameServer<C, P>> = config
            .name_servers()
            .iter()
            .filter(|ns_config| ns_config.protocol.is_stream())
            .map(|ns_config| {
                #[cfg(feature = "dns-over-rustls")]
                let ns_config = {
                    let mut ns_config = ns_config.clone();
                    ns_config.tls_config = config.client_config().clone();
                    ns_config
                };
                #[cfg(not(feature = "dns-over-rustls"))]
                let ns_config = { ns_config.clone() };

                NameServer::<C, P>::new_with_provider(ns_config, *options, conn_provider.clone())
            })
            .collect();

        NameServerPool {
            datagram_conns: Arc::new(Mutex::new(datagram_conns)),
            stream_conns: Arc::new(Mutex::new(stream_conns)),
            #[cfg(feature = "mdns")]
            mdns_conns: name_server::mdns_nameserver(*options, conn_provider.clone()),
            options: *options,
            conn_provider,
        }
    }

    #[doc(hidden)]
    #[cfg(not(feature = "mdns"))]
    pub fn from_nameservers(
        options: &ResolverOpts,
        datagram_conns: Vec<NameServer<C, P>>,
        stream_conns: Vec<NameServer<C, P>>,
        conn_provider: P,
    ) -> Self {
        NameServerPool {
            datagram_conns: Arc::new(Mutex::new(datagram_conns.into_iter().collect())),
            stream_conns: Arc::new(Mutex::new(stream_conns.into_iter().collect())),
            options: *options,
            conn_provider,
        }
    }

    #[doc(hidden)]
    #[cfg(feature = "mdns")]
    pub fn from_nameservers(
        options: &ResolverOpts,
        datagram_conns: Vec<NameServer<C, P>>,
        stream_conns: Vec<NameServer<C, P>>,
        mdns_conns: NameServer<C, P>,
        conn_provider: P,
    ) -> Self {
        NameServerPool {
            datagram_conns: Arc::new(Mutex::new(datagram_conns.into_iter().collect())),
            stream_conns: Arc::new(Mutex::new(stream_conns.into_iter().collect())),
            mdns_conns,
            options: *options,
            conn_provider,
        }
    }

    async fn try_send(
        opts: ResolverOpts,
        conns: Arc<Mutex<Vec<NameServer<C, P>>>>,
        request: DnsRequest,
    ) -> Result<DnsResponse, ProtoError> {
        let conns: Vec<NameServer<C, P>> = {
            // pull a lock on the shared connections, lock releases at the end of this scope
            let mut conns = conns.lock().await;

            // select the highest priority connection
            //   reorder the connections based on current view...
            //   this reorders the inner set
            conns.sort_unstable();

            // TODO: restrict this size to a maximum # of NameServers to try
            // get a stable view for trying all connections
            //   we split into chunks of the number of parallel requests to issue
            conns.clone()
         };

        let request_loop = request.clone();

        parallel_conn_loop(conns, request_loop, opts).await
    }
}

impl<C, P> DnsHandle for NameServerPool<C, P>
where
    C: DnsHandle + 'static,
    P: ConnectionProvider<ConnHandle = C> + 'static,
{
    type Response = Pin<Box<dyn Future<Output = Result<DnsResponse, ProtoError>> + Send>>;

    fn send<R: Into<DnsRequest>>(&mut self, request: R) -> Self::Response {
        let opts = self.options;
        let request = request.into();
        let datagram_conns = Arc::clone(&self.datagram_conns);
        let stream_conns1 = Arc::clone(&self.stream_conns);
        let stream_conns2 = Arc::clone(&self.stream_conns);
        // TODO: remove this clone, return the Message in the error?
        let tcp_message1 = request.clone();
        let tcp_message2 = request.clone();

        // if it's a .local. query, then we *only* query mDNS, these should never be sent on to upstream resolvers
        #[cfg(feature = "mdns")]
        let mdns = mdns::maybe_local(&mut self.mdns_conns, request);

        // TODO: limited to only when mDNS is enabled, but this should probably always be enforced?
        #[cfg(not(feature = "mdns"))]
        let mdns = Local::NotMdns(request);

        // local queries are queried through mDNS
        if mdns.is_local() {
            return mdns.take_future();
        }

        // TODO: should we allow mDNS to be used for standard lookups as well?

        // it wasn't a local query, continue with standard lookup path
        let request = mdns.take_request();

        // First try the UDP connections
        Box::pin(
            Self::try_send(opts, datagram_conns, request)
                .and_then(move |response| {
                    // handling promotion from datagram to stream base on truncation in message
                    if ResponseCode::NoError == response.response_code() && response.truncated() {
                        // TCP connections should not truncate
                        future::Either::Left(Self::try_send(opts, stream_conns1, tcp_message1))
                    } else {
                        // Return the result from the UDP connection
                        future::Either::Right(future::ok(response))
                    }
                })
                // if UDP fails, try TCP
                .or_else(move |_| Self::try_send(opts, stream_conns2, tcp_message2)),
        )
    }
}

// TODO: we should be able to have a self-referential future here with Pin and not require cloned conns
/// An async function that will loop over all the conns with a max parallel request count of ops.num_concurrent_req
async fn parallel_conn_loop<C, P>(
    mut conns: Vec<NameServer<C, P>>,
    request: DnsRequest,
    opts: ResolverOpts,
) -> Result<DnsResponse, ProtoError>
where
    C: DnsHandle + 'static,
    P: ConnectionProvider<ConnHandle = C> + 'static,
{
    let mut err = ProtoError::from("No connections available");

    loop {
        let request_cont = request.clone();

        // construct the parallel requests, 2 is the default
        let mut par_conns = SmallVec::<[NameServer<C, P>; 2]>::new();
        let count = conns.len().min(opts.num_concurrent_reqs.max(1));
        for conn in conns.drain(..count) {
            par_conns.push(conn);
        }

        // construct the requests to send
        let requests = if par_conns.is_empty() {
            return Err(err);
        } else {
            par_conns
                .into_iter()
                .map(move |mut conn| conn.send(request_cont.clone()))
        };

        match future::select_ok(requests).await {
            Ok((sent, _)) => return Ok(sent),
            // consider a debug msg here
            Err(e) => {
                err = e;
                continue;
            }
        };
    }
}

#[cfg(feature = "mdns")]
mod mdns {
    use super::*;

    use proto::rr::domain::usage;
    use proto::DnsHandle;

    /// Returns true
    pub fn maybe_local<C, P>(name_server: &mut NameServer<C, P>, request: DnsRequest) -> Local
    where
        C: DnsHandle + 'static,
        P: ConnectionProvider<ConnHandle = C> + 'static,
    {
        if request
            .queries()
            .iter()
            .any(|query| usage::LOCAL.name().zone_of(query.name()))
        {
            Local::ResolveFuture(name_server.send(request))
        } else {
            Local::NotMdns(request)
        }
    }
}

pub enum Local {
    #[allow(dead_code)]
    ResolveFuture(Pin<Box<dyn Future<Output = Result<DnsResponse, ProtoError>> + Send + Unpin>>),
    NotMdns(DnsRequest),
}

impl Local {
    fn is_local(&self) -> bool {
        if let Local::ResolveFuture(..) = *self {
            true
        } else {
            false
        }
    }

    /// Takes the future
    ///
    /// # Panics
    ///
    /// Panics if this is in fact a Local::NotMdns
    fn take_future(
        self,
    ) -> Pin<Box<dyn Future<Output = Result<DnsResponse, ProtoError>> + Send + Unpin>> {
        match self {
            Local::ResolveFuture(future) => future,
            _ => panic!("non Local queries have no future, see take_message()"),
        }
    }

    /// Takes the message
    ///
    /// # Panics
    ///
    /// Panics if this is in fact a Local::ResolveFuture
    fn take_request(self) -> DnsRequest {
        match self {
            Local::NotMdns(request) => request,
            _ => panic!("Local queries must be polled, see take_future()"),
        }
    }
}

impl Future for Local {
    type Output = Result<DnsResponse, ProtoError>;

    fn poll(mut self: Pin<&mut Self>, cx: &mut Context) -> Poll<Self::Output> {
        match *self {
            Local::ResolveFuture(ref mut ns) => ns.as_mut().poll(cx),
            // TODO: making this a panic for now
            Local::NotMdns(..) => panic!("Local queries that are not mDNS should not be polled"), //Local::NotMdns(message) => return Err(ResolveErrorKind::Message("not mDNS")),
        }
    }
}

#[cfg(test)]
mod tests {
    extern crate env_logger;

    use std::net::{IpAddr, Ipv4Addr, SocketAddr};

    use tokio::runtime::current_thread::Runtime;

    use proto::op::Query;
    use proto::rr::{Name, RecordType};
    use proto::xfer::{DnsHandle, DnsRequestOptions};

    use super::*;
    use crate::config::NameServerConfig;
    use crate::config::Protocol;

    #[ignore]
    // because of there is a real connection that needs a reasonable timeout
    #[test]
    fn test_failed_then_success_pool() {
        let config1 = NameServerConfig {
            socket_addr: SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 0, 252)), 253),
            protocol: Protocol::Udp,
            tls_dns_name: None,
            #[cfg(feature = "dns-over-rustls")]
            tls_config: None,
        };

        let config2 = NameServerConfig {
            socket_addr: SocketAddr::new(IpAddr::V4(Ipv4Addr::new(8, 8, 8, 8)), 53),
            protocol: Protocol::Udp,
            tls_dns_name: None,
            #[cfg(feature = "dns-over-rustls")]
            tls_config: None,
        };

        let mut resolver_config = ResolverConfig::new();
        resolver_config.add_name_server(config1);
        resolver_config.add_name_server(config2);

        let mut io_loop = Runtime::new().unwrap();
        let mut pool = NameServerPool::<_, StandardConnection>::from_config(
            &resolver_config,
            &ResolverOpts::default(),
        );

        let name = Name::parse("www.example.com.", None).unwrap();

        // TODO: it's not clear why there are two failures before the success
        for i in 0..2 {
            assert!(
                io_loop
                    .block_on(pool.lookup(
                        Query::query(name.clone(), RecordType::A),
                        DnsRequestOptions::default()
                    ))
                    .is_err(),
                "iter: {}",
                i
            );
        }

        for i in 0..10 {
            assert!(
                io_loop
                    .block_on(pool.lookup(
                        Query::query(name.clone(), RecordType::A),
                        DnsRequestOptions::default()
                    ))
                    .is_ok(),
                "iter: {}",
                i
            );
        }
    }
}
<|endoftext|>


<commit_before>

use std::mem;

use std::io::prelude::*;

use std::fs;
use std::fs::{OpenOptions, create_dir, remove_dir_all};

use byteorder::{WriteBytesExt, ReadBytesExt, BigEndian};

use bincode::SizeLimit;
use bincode::rustc_serialize::{encode_into, decode_from};

use super::SqlType;

use super::Engine;
use super::Error;
use super::engine::FlatFile;
use super::types::Column;
use super::data::Rows;
use super::super::parse::ast::DataSrc;

/// constants
const MAGIC_NUMBER: u64 = 0x6561742073686974; // secret
const VERSION_NO: u8 = 1;

//---------------------------------------------------------------
// DataType
//---------------------------------------------------------------
/// A Enum for Datatypes (will be removed later)
#[repr(u8)]
#[derive(Clone, Copy, Debug, RustcDecodable, RustcEncodable)]
pub enum DataType { Integer = 1, Float = 2, }

impl DataType {
    pub fn value(&self) -> u8 {
       *self as u8
    }
}

//---------------------------------------------------------------
// Database
//---------------------------------------------------------------
#[derive(Debug)]
pub struct Database {
    name: String,
}

impl Database {
    /// Starts the process of creating a new Database
    /// Returns database or on fail Error
    pub fn create(name: &str) -> Result<Database, Error> {
        let d = Database{ name: name.to_string() };
        try!(d.save());
        info!("created new database {:?}", d);
        Ok(d)
    }

    /// Loads already existing Database
    /// returns DataBase Error when database does not exist else the loaded DB
    pub fn load(name: &str) -> Result<Database, Error> {
        if try!(fs::metadata(name)).is_dir() {
            info!("loaded Database {:?}", name.to_string());
            Ok(Database{ name: name.to_string() })
        } else {
            warn!("could not load database {:?}", name.to_string());
            return Err(Error::LoadDataBase)
        }
    }

    /// Creates a folder for the database
    fn save(&self) -> Result<(), Error> {
        info!("trying to create dir!");
        try!(create_dir(&self.name));
        info!("created dir");
        Ok(())
    }
    /// Deletes the database folder and all its contents
    /// do not use RANDOM!!
    pub fn delete(&self) -> Result<(), Error> {
        info!("deleting Database and all its tables");
        try!(remove_dir_all(&self.name));
        Ok(())
    }
    /// Creates a new table in the DB folder
    /// Returns with Error on fail else Table
    pub fn create_table(&self, name: &str, columns: Vec<Column>, engine_id: u8)
        -> Result<Table, Error>
    {

        let t = Table::new(&self, name, columns, engine_id);
        try!(t.save());
        info!("created new table {:?}", t);
        Ok(t)
    }

    /// calls load for table with the database path
    /// Returns with Error on fail else Table
    pub fn load_table(&self, name: &str) -> Result<Table, Error> {
        Table::load(&self, name)
    }
}

//---------------------------------------------------------------
// TableMetaData
//---------------------------------------------------------------

#[derive(Debug, RustcDecodable, RustcEncodable)]
pub struct TableMetaData {
    version_nmbr: u8,
    engine_id: u8,
    columns: Vec<Column>,
    //primary_key: String
}

//---------------------------------------------------------------
// Table
//---------------------------------------------------------------

/// Table struct that contains the table information
#[derive(Debug)]
pub struct Table<'a> {
    database: &'a Database,
    pub name: String,
    meta_data: TableMetaData,
}

impl<'a> Table<'a> {
    /// Creates new table object
    /// Returns Table
    pub fn new<'b>(database: &'b Database, name: &str,
                   columns: Vec<Column>, engine_id: u8)
        -> Table<'b>
    {

        let meta_data = TableMetaData {
            version_nmbr: VERSION_NO,
            engine_id: engine_id,
            columns: columns,
        };
        info!("created meta data: {:?}", meta_data);

        Table {
            name: name.to_string(),
            database: database,
            meta_data: meta_data,
        }
    }

    /// Loads the table from the DB
    /// Returns with Error on fail else Table
    fn load<'b>(database: &'b Database, name: &str)
        -> Result<Table<'b>, Error>
    {
        // TODO: Read the .tbl file from disk and parse it

        let path_to_table = Table::get_path(&database.name, name, "tbl");
        info!("getting path and opening file: {:?}", path_to_table);
        let mut file = try!(OpenOptions::new()
            .read(true)
            .open(path_to_table));
        info!("reading file: {:?}", file);
        let ma_nmbr = try!(file.read_uint::<BigEndian>(mem::size_of_val(&MAGIC_NUMBER)));

        info!("checking magic number: {:?}", ma_nmbr);
        if ma_nmbr != MAGIC_NUMBER {
            println!("Magic Number not correct");
            return Err(Error::WrongMagicNmbr)
        }
        let meta_data: TableMetaData = try!(decode_from(&mut file, SizeLimit::Infinite));
        info!("getting meta data{:?}", meta_data);

        let table = Table::new(database, name, meta_data.columns, meta_data.engine_id);
        info!("returning table: {:?}", table);
        Ok(table)
    }

    /// Saves the table with a identification number in table file
    /// Returns Error on fail else Nothing
    pub fn save(&self) -> Result<(), Error> {
        // call for open file
        info!("opening file to write");
        let mut file = try!(OpenOptions::new()
            .write(true)
            .create(true)
            .open(self.get_table_metadata_path()));
        info!("writing magic number in file: {:?}", file);
        try!(file.write_u64::<BigEndian>(MAGIC_NUMBER));//MAGIC_NUMBER
        info!("writing meta data in file: {:?}", file);
        try!(encode_into(&self.meta_data, &mut file, SizeLimit::Infinite));

        // debug message all okay
        info!("I Wrote my File");
        Ok(())
    }

    /// Deletes the .tbl files
    /// Returns Error on fail if path points to a directory,
    /// if the user lacks permissions to remove the file,
    /// or if some other filesystem-level error occurs.
    pub fn delete(&self) -> Result<(), Error> {

        info!("remove meta file: {:?}", self.get_table_metadata_path());
        try!(fs::remove_file(self.get_table_metadata_path()));

        info!("remove data file: {:?}", self.get_table_data_path());
        try!(fs::remove_file(self.get_table_data_path()));

        Ok(())
    }

    /// Returns columns of table as array
    pub fn columns(&self) -> &[Column] {
        &self.meta_data.columns
    }

    /// Adds a column to the tabel
    /// Returns name of Column or on fail Error
    pub fn add_column(
        &mut self,
        name: &str,
        sql_type: SqlType,
        allow_null: bool,
        description: &str,
        is_primary_key: bool
        ) -> Result<(), Error> {

        match self.meta_data.columns.iter().find(|x| x.name == name) {
            Some(_) => {
                warn!("Column {:?} already exists", name);
                return Err(Error::AddColumn)
            },
            None => {
                info!("Column {:?} was added", name);
            },
        }

        self.meta_data.columns.push(Column::new(
            name,
            sql_type,
            allow_null,
            description,
            is_primary_key)
        );
        Ok(())
    }

    /// Removes a column from the table
    /// Returns name of Column or on fail Error
    pub fn remove_column(&mut self, name: &str) -> Result<(), Error> {
        let index = match self.meta_data.columns.iter().position(|x| x.name == name) {
            Some(x) => {
                info!("Column {:?} was removed", self.name);
                x
            },
            None => {
                warn!("Column {:?} could not be found", self.name);
                return Err(Error::RemoveColumn)
            },
        };
        self.meta_data.columns.swap_remove(index);
        Ok(())
    }

    /// Creates an engine for Table
    /// Returns Box<Engine>
    pub fn create_engine(self) -> Box<Engine + 'a> {
        Box::new(FlatFile::new(self))
    }

    /// Returns the path for the metadata files
    fn get_table_metadata_path(&self) -> String {
        Self::get_path(&self.database.name, &self.name, "tbl")
    }

    /// Returns the path for the data files
    pub fn get_table_data_path(&self) -> String {
        Self::get_path(&self.database.name, &self.name, "dat")
    }

    /// Returns the path of the table
    fn get_path(database: &str, name: &str, ext: &str) -> String {
         format!("{}/{}.{}", database, name, ext)
    }

    pub fn get_description(&self) -> Result<Rows, Error> {
        let mut data = Vec::<u8>::new();

        let mut columns = Vec::<Column>::new();
        columns.push(Column::new("Name", SqlType::VarChar(255), false, "description", false));
        columns.push(Column::new("SqlType", SqlType::VarChar(255), false, "SqlType", false));
        columns.push(Column::new("allow_null", SqlType::Bool, false, "description", false));
        columns.push(Column::new("is_primary_key", SqlType::Bool, false, "description", false));
        columns.push(Column::new("description", SqlType::VarChar(255), false, "description", false));

        for c in &self.columns()[..] {
            let mut datasrc = DataSrc::String(c.name.clone());
            SqlType::VarChar(255).encode_into(&mut data, &datasrc);

            datasrc = DataSrc::String(format!("{:?}", c.sql_type));
            SqlType::VarChar(255).encode_into(&mut data, &datasrc);

            if c.is_primary_key == true {
                datasrc = DataSrc::Bool(1);
            }
            else {
                datasrc = DataSrc::Bool(0);
            }

            SqlType::Bool.encode_into(&mut data, &datasrc);

            if c.allow_null == true {
                datasrc = DataSrc::Bool(1);
            }
            else {
                datasrc = DataSrc::Bool(0);
            }

            SqlType::Bool.encode_into(&mut data, &datasrc);

            datasrc = DataSrc::String(c.description.clone());
            SqlType::VarChar(255).encode_into(&mut data, &datasrc);
        }
        Ok(Rows {data: data, columns: columns})
    }

/*
    pub name: String, // name of column
    pub sql_type: SqlType, // name of the data type that is contained in this column
    pub is_primary_key: bool, // defines if column is PK
    pub allow_null: bool, // defines if cloumn allows null
    pub description: String //Displays text describing this column.
*/

 /*   fn get_primary_key(&self) -> Column {
        match self.meta_data.columns.iter().find(|x| x.name == self.meta_data.primary_key) {
            Some(x) => {
                info!("Primaty Key: {:?} was found", self.meta_data.primary_key);
                Ok(x)
            },
            None => {
                warn!("Primary Key: {:?} was not found", self.meta_data.primary_key);
                Err(Error::PrimaryKey)
            },
        }
    }*/
}
<commit_msg>Fix style check<commit_after>

use std::mem;

use std::io::prelude::*;

use std::fs;
use std::fs::{OpenOptions, create_dir, remove_dir_all};

use byteorder::{WriteBytesExt, ReadBytesExt, BigEndian};

use bincode::SizeLimit;
use bincode::rustc_serialize::{encode_into, decode_from};

use super::SqlType;

use super::Engine;
use super::Error;
use super::engine::FlatFile;
use super::types::Column;
use super::data::Rows;
use super::super::parse::ast::DataSrc;

/// constants
const MAGIC_NUMBER: u64 = 0x6561742073686974; // secret
const VERSION_NO: u8 = 1;

//---------------------------------------------------------------
// DataType
//---------------------------------------------------------------
/// A Enum for Datatypes (will be removed later)
#[repr(u8)]
#[derive(Clone, Copy, Debug, RustcDecodable, RustcEncodable)]
pub enum DataType { Integer = 1, Float = 2, }

impl DataType {
    pub fn value(&self) -> u8 {
       *self as u8
    }
}

//---------------------------------------------------------------
// Database
//---------------------------------------------------------------
#[derive(Debug)]
pub struct Database {
    name: String,
}

impl Database {
    /// Starts the process of creating a new Database
    /// Returns database or on fail Error
    pub fn create(name: &str) -> Result<Database, Error> {
        let d = Database{ name: name.to_string() };
        try!(d.save());
        info!("created new database {:?}", d);
        Ok(d)
    }

    /// Loads already existing Database
    /// returns DataBase Error when database does not exist else the loaded DB
    pub fn load(name: &str) -> Result<Database, Error> {
        if try!(fs::metadata(name)).is_dir() {
            info!("loaded Database {:?}", name.to_string());
            Ok(Database{ name: name.to_string() })
        } else {
            warn!("could not load database {:?}", name.to_string());
            return Err(Error::LoadDataBase)
        }
    }

    /// Creates a folder for the database
    fn save(&self) -> Result<(), Error> {
        info!("trying to create dir!");
        try!(create_dir(&self.name));
        info!("created dir");
        Ok(())
    }
    /// Deletes the database folder and all its contents
    /// do not use RANDOM!!
    pub fn delete(&self) -> Result<(), Error> {
        info!("deleting Database and all its tables");
        try!(remove_dir_all(&self.name));
        Ok(())
    }
    /// Creates a new table in the DB folder
    /// Returns with Error on fail else Table
    pub fn create_table(&self, name: &str, columns: Vec<Column>, engine_id: u8)
        -> Result<Table, Error>
    {

        let t = Table::new(&self, name, columns, engine_id);
        try!(t.save());
        info!("created new table {:?}", t);
        Ok(t)
    }

    /// calls load for table with the database path
    /// Returns with Error on fail else Table
    pub fn load_table(&self, name: &str) -> Result<Table, Error> {
        Table::load(&self, name)
    }
}

//---------------------------------------------------------------
// TableMetaData
//---------------------------------------------------------------

#[derive(Debug, RustcDecodable, RustcEncodable)]
pub struct TableMetaData {
    version_nmbr: u8,
    engine_id: u8,
    columns: Vec<Column>,
    //primary_key: String
}

//---------------------------------------------------------------
// Table
//---------------------------------------------------------------

/// Table struct that contains the table information
#[derive(Debug)]
pub struct Table<'a> {
    database: &'a Database,
    pub name: String,
    meta_data: TableMetaData,
}

impl<'a> Table<'a> {
    /// Creates new table object
    /// Returns Table
    pub fn new<'b>(database: &'b Database, name: &str,
                   columns: Vec<Column>, engine_id: u8)
        -> Table<'b>
    {

        let meta_data = TableMetaData {
            version_nmbr: VERSION_NO,
            engine_id: engine_id,
            columns: columns,
        };
        info!("created meta data: {:?}", meta_data);

        Table {
            name: name.to_string(),
            database: database,
            meta_data: meta_data,
        }
    }

    /// Loads the table from the DB
    /// Returns with Error on fail else Table
    fn load<'b>(database: &'b Database, name: &str)
        -> Result<Table<'b>, Error>
    {
        // TODO: Read the .tbl file from disk and parse it

        let path_to_table = Table::get_path(&database.name, name, "tbl");
        info!("getting path and opening file: {:?}", path_to_table);
        let mut file = try!(OpenOptions::new()
            .read(true)
            .open(path_to_table));
        info!("reading file: {:?}", file);
        let ma_nmbr = try!(file.read_uint::<BigEndian>(mem::size_of_val(&MAGIC_NUMBER)));

        info!("checking magic number: {:?}", ma_nmbr);
        if ma_nmbr != MAGIC_NUMBER {
            println!("Magic Number not correct");
            return Err(Error::WrongMagicNmbr)
        }
        let meta_data: TableMetaData = try!(decode_from(&mut file, SizeLimit::Infinite));
        info!("getting meta data{:?}", meta_data);

        let table = Table::new(database, name, meta_data.columns, meta_data.engine_id);
        info!("returning table: {:?}", table);
        Ok(table)
    }

    /// Saves the table with a identification number in table file
    /// Returns Error on fail else Nothing
    pub fn save(&self) -> Result<(), Error> {
        // call for open file
        info!("opening file to write");
        let mut file = try!(OpenOptions::new()
            .write(true)
            .create(true)
            .open(self.get_table_metadata_path()));
        info!("writing magic number in file: {:?}", file);
        try!(file.write_u64::<BigEndian>(MAGIC_NUMBER));//MAGIC_NUMBER
        info!("writing meta data in file: {:?}", file);
        try!(encode_into(&self.meta_data, &mut file, SizeLimit::Infinite));

        // debug message all okay
        info!("I Wrote my File");
        Ok(())
    }

    /// Deletes the .tbl files
    /// Returns Error on fail if path points to a directory,
    /// if the user lacks permissions to remove the file,
    /// or if some other filesystem-level error occurs.
    pub fn delete(&self) -> Result<(), Error> {

        info!("remove meta file: {:?}", self.get_table_metadata_path());
        try!(fs::remove_file(self.get_table_metadata_path()));

        info!("remove data file: {:?}", self.get_table_data_path());
        try!(fs::remove_file(self.get_table_data_path()));

        Ok(())
    }

    /// Returns columns of table as array
    pub fn columns(&self) -> &[Column] {
        &self.meta_data.columns
    }

    /// Adds a column to the tabel
    /// Returns name of Column or on fail Error
    pub fn add_column(
        &mut self,
        name: &str,
        sql_type: SqlType,
        allow_null: bool,
        description: &str,
        is_primary_key: bool
        ) -> Result<(), Error> {

        match self.meta_data.columns.iter().find(|x| x.name == name) {
            Some(_) => {
                warn!("Column {:?} already exists", name);
                return Err(Error::AddColumn)
            },
            None => {
                info!("Column {:?} was added", name);
            },
        }

        self.meta_data.columns.push(Column::new(
            name,
            sql_type,
            allow_null,
            description,
            is_primary_key)
        );
        Ok(())
    }

    /// Removes a column from the table
    /// Returns name of Column or on fail Error
    pub fn remove_column(&mut self, name: &str) -> Result<(), Error> {
        let index = match self.meta_data.columns.iter().position(|x| x.name == name) {
            Some(x) => {
                info!("Column {:?} was removed", self.name);
                x
            },
            None => {
                warn!("Column {:?} could not be found", self.name);
                return Err(Error::RemoveColumn)
            },
        };
        self.meta_data.columns.swap_remove(index);
        Ok(())
    }

    /// Creates an engine for Table
    /// Returns Box<Engine>
    pub fn create_engine(self) -> Box<Engine + 'a> {
        Box::new(FlatFile::new(self))
    }

    /// Returns the path for the metadata files
    fn get_table_metadata_path(&self) -> String {
        Self::get_path(&self.database.name, &self.name, "tbl")
    }

    /// Returns the path for the data files
    pub fn get_table_data_path(&self) -> String {
        Self::get_path(&self.database.name, &self.name, "dat")
    }

    /// Returns the path of the table
    fn get_path(database: &str, name: &str, ext: &str) -> String {
         format!("{}/{}.{}", database, name, ext)
    }

    pub fn get_description(&self) -> Result<Rows, Error> {
        let mut data = Vec::<u8>::new();

        let mut columns = Vec::<Column>::new();
        columns.push(Column::new("Name",
                                 SqlType::VarChar(255),
                                 false,
                                 "description",
                                 false));
        columns.push(Column::new("SqlType",
                                 SqlType::VarChar(255),
                                 false,
                                 "SqlType",
                                 false));
        columns.push(Column::new("allow_null",
                                 SqlType::Bool,
                                 false,
                                 "description",
                                 false));
        columns.push(Column::new("is_primary_key",
                                 SqlType::Bool,
                                 false,
                                 "description",
                                 false));
        columns.push(Column::new("description",
                                 SqlType::VarChar(255),
                                 false,
                                 "description",
                                 false));

        for c in &self.columns()[..] {
            let mut datasrc = DataSrc::String(c.name.clone());
            SqlType::VarChar(255).encode_into(&mut data, &datasrc);

            datasrc = DataSrc::String(format!("{:?}", c.sql_type));
            SqlType::VarChar(255).encode_into(&mut data, &datasrc);

            if c.is_primary_key == true {
                datasrc = DataSrc::Bool(1);
            }
            else {
                datasrc = DataSrc::Bool(0);
            }

            SqlType::Bool.encode_into(&mut data, &datasrc);

            if c.allow_null == true {
                datasrc = DataSrc::Bool(1);
            }
            else {
                datasrc = DataSrc::Bool(0);
            }

            SqlType::Bool.encode_into(&mut data, &datasrc);

            datasrc = DataSrc::String(c.description.clone());
            SqlType::VarChar(255).encode_into(&mut data, &datasrc);
        }
        Ok(Rows {data: data, columns: columns})
    }

/*
    pub name: String, // name of column
    pub sql_type: SqlType, // name of the data type that is contained in this column
    pub is_primary_key: bool, // defines if column is PK
    pub allow_null: bool, // defines if cloumn allows null
    pub description: String //Displays text describing this column.
*/

 /*   fn get_primary_key(&self) -> Column {
        match self.meta_data.columns.iter().find(|x| x.name == self.meta_data.primary_key) {
            Some(x) => {
                info!("Primaty Key: {:?} was found", self.meta_data.primary_key);
                Ok(x)
            },
            None => {
                warn!("Primary Key: {:?} was not found", self.meta_data.primary_key);
                Err(Error::PrimaryKey)
            },
        }
    }*/
}
<|endoftext|>