<commit_before>use base64;
use crate::buffer;
use crate::image;
use std::{fs, io};

use image_crate::ImageFormat::{Jpeg, Png};
use std::path::Path;
use crate::{Document, Error, Gltf, Result};

/// Return type of `import`.
type Import = (Document, Vec<buffer::Data>, Vec<image::Data>);

/// Represents the set of URI schemes the importer supports.
#[derive(Clone, Debug, Eq, Hash, PartialEq)]
enum Scheme<'a> {
    /// `data:[<media type>];base64,<data>`.
    Data(Option<&'a str>, &'a str),

    /// `file:[//]<absolute file path>`.
    ///
    /// Note: The file scheme does not implement authority.
    File(&'a str),

    /// `../foo`, etc.
    Relative,

    /// Placeholder for an unsupported URI scheme identifier.
    Unsupported,
}

impl<'a> Scheme<'a> {
    fn parse<'s>(uri: &'s str) -> Scheme<'s> {
        if uri.contains(":") {
            if uri.starts_with("data:") {
                let match0 = &uri["data:".len()..].split(";base64,").nth(0);
                let match1 = &uri["data:".len()..].split(";base64,").nth(1);
                if match1.is_some() {
                    Scheme::Data(Some(match0.unwrap()), match1.unwrap())
                } else if match0.is_some() {
                    Scheme::Data(None, match0.unwrap())
                } else {
                    Scheme::Unsupported
                }
            } else if uri.starts_with("file://") {
                Scheme::File(&uri["file://".len()..])
            } else if uri.starts_with("file:") {
                Scheme::File(&uri["file:".len()..])
            } else {
                Scheme::Unsupported
            }
        } else {
            Scheme::Relative
        }
    }

    fn read(base: &Path, uri: &str) -> Result<Vec<u8>> {
        match Scheme::parse(uri) {
            Scheme::Data(_, base64) => base64::decode(&base64).map_err(Error::Base64),
            Scheme::File(path) => read_to_end(path),
            Scheme::Relative => read_to_end(base.join(uri)),
            Scheme::Unsupported => Err(Error::UnsupportedScheme),
        }
    }
}

fn read_to_end<P>(path: P) -> Result<Vec<u8>>
where P: AsRef<Path>
{
    use io::Read;
    let file = fs::File::open(path.as_ref()).map_err(Error::Io)?;
    // Allocate one extra byte so the buffer doesn't need to grow before the
    // final `read` call at the end of the file.  Don't worry about `usize`
    // overflow because reading will fail regardless in that case.
    let length = file.metadata().map(|x| x.len() + 1).unwrap_or(0);
    let mut reader = io::BufReader::new(file);
    let mut data = Vec::with_capacity(length as usize);
    reader.read_to_end(&mut data).map_err(Error::Io)?;
    Ok(data)
}

/// Import the buffer data referenced by a glTF document.
pub fn import_buffer_data(
    document: &Document,
    base: Option<&Path>,
    mut blob: Option<Vec<u8>>,
) -> Result<Vec<buffer::Data>> {
    let mut buffers = Vec::new();
    for buffer in document.buffers() {
        let mut data = match buffer.source() {
            buffer::Source::Uri(uri) if base.is_some() => Scheme::read(base.unwrap(), uri),
            buffer::Source::Bin => blob.take().ok_or(Error::MissingBlob),
            _ => Err(Error::ExternalReferenceInSliceImport)
        }?;
        if data.len() < buffer.length() {
            return Err(
                Error::BufferLength {
                    buffer: buffer.index(),
                    expected: buffer.length(),
                    actual: data.len(),
                }
            );
        }
        while data.len() % 4 != 0 {
            data.push(0);
        }
        buffers.push(buffer::Data(data));
    }
    Ok(buffers)
}

/// Import the image data referenced by a glTF document.
pub fn import_image_data(
    document: &Document,
    base: Option<&Path>,
    buffer_data: &[buffer::Data],
) -> Result<Vec<image::Data>> {
    let mut images = Vec::new();
    #[cfg(feature = "guess_mime_type")]
    let guess_format = |encoded_image: &[u8]| {
        match image_crate::guess_format(encoded_image) {
            Ok(image_crate::ImageFormat::Png) => Some(Png),
            Ok(image_crate::ImageFormat::Jpeg) => Some(Jpeg),
            _ => None,
        }
    };
    #[cfg(not(feature = "guess_mime_type"))]
    let guess_format = |_encoded_image: &[u8]| {
        None
    };
    for image in document.images() {
        match image.source() {
            image::Source::Uri { uri, mime_type } if base.is_some() => {
                match Scheme::parse(uri) {
                    Scheme::Data(Some(annoying_case), base64) => {
                        let encoded_image = base64::decode(&base64).map_err(Error::Base64)?;
                        let encoded_format = match annoying_case.as_ref() {
                            "image/png" => Png,
                            "image/jpeg" => Jpeg,
                            _ => match guess_format(&encoded_image) {
                                Some(format) => format,
                                None => return Err(Error::UnsupportedImageEncoding),
                            },
                        };
                        let decoded_image = image_crate::load_from_memory_with_format(&encoded_image, encoded_format)?;
                        images.push(image::Data::new(decoded_image));
                        continue;
                    },
                    Scheme::Unsupported => return Err(Error::UnsupportedScheme),
                    _ => {},
                }
                let encoded_image = Scheme::read(base.unwrap(), uri)?;
                let encoded_format =  match mime_type {
                    Some("image/png") => Png,
                    Some("image/jpeg") => Jpeg,
                    Some(_) => match guess_format(&encoded_image) {
                        Some(format) => format,
                        None => return Err(Error::UnsupportedImageEncoding),
                    },
                    None => match uri.rsplit(".").next() {
                        Some("png") => Png,
                        Some("jpg") | Some("jpeg") => Jpeg,
                        _ => match guess_format(&encoded_image) {
                            Some(format) => format,
                            None => return Err(Error::UnsupportedImageEncoding),
                        },
                    },
                };
                let decoded_image = image_crate::load_from_memory_with_format(&encoded_image, encoded_format)?;
                images.push(image::Data::new(decoded_image));
            },
            image::Source::View { view, mime_type } => {
                let parent_buffer_data = &buffer_data[view.buffer().index()].0;
                let begin = view.offset();
                let end = begin + view.length();
                let encoded_image = &parent_buffer_data[begin..end];
                let encoded_format = match mime_type {
                    "image/png" => Png,
                    "image/jpeg" => Jpeg,
                    _ => match guess_format(encoded_image) {
                        Some(format) => format,
                        None => return Err(Error::UnsupportedImageEncoding),
                    },
                };
                let decoded_image = image_crate::load_from_memory_with_format(encoded_image, encoded_format)?;
                images.push(image::Data::new(decoded_image));
            },
            _ => return Err(Error::ExternalReferenceInSliceImport)
        }
    }

    Ok(images)
}

fn import_impl(Gltf { document, blob }: Gltf, base: Option<&Path>) -> Result<Import> {
    let buffer_data = import_buffer_data(&document, base, blob)?;
    let image_data = import_image_data(&document, base, &buffer_data)?;
    let import = (document, buffer_data, image_data);
    Ok(import)
}

fn import_path(path: &Path) -> Result<Import> {
    let base = path.parent().unwrap_or(Path::new("./"));
    let file = fs::File::open(path).map_err(Error::Io)?;
    let reader = io::BufReader::new(file);
    import_impl(Gltf::from_reader(reader)?, Some(base))
}

/// Import some glTF 2.0 from the file system.
///
/// ```
/// # fn run() -> Result<(), gltf::Error> {
/// # let path = "examples/Box.gltf";
/// # #[allow(unused)]
/// let (document, buffers, images) = gltf::import(path)?;
/// # Ok(())
/// # }
/// # fn main() {
/// #     run().expect("test failure");
/// # }
/// ```
///
/// ### Note
///
/// This function is provided as a convenience for loading glTF and associated
/// resources from the file system. It is suitable for real world use but may
/// not be suitable for all real world use cases. More complex import scenarios
/// such downloading from web URLs are not handled by this function. These
/// scenarios are delegated to the user.
///
/// You can read glTF without loading resources by constructing the [`Gltf`]
/// (standard glTF) or [`Glb`] (binary glTF) data structures explicitly.
///
/// [`Gltf`]: struct.Gltf.html
/// [`Glb`]: struct.Glb.html
pub fn import<P>(path: P) -> Result<Import>
    where P: AsRef<Path>
{
    import_path(path.as_ref())
}

pub fn import_slice_impl(slice: &[u8]) -> Result<Import> {
    import_impl(Gltf::from_slice(slice)?, None)
}

/// Import some glTF 2.0 from a slice
///
/// ```
/// # extern crate gltf;
/// # use std::fs;
/// # use std::io::Read;
/// # fn run() -> Result<(), gltf::Error> {
/// # let path = "examples/Box.glb";
/// # let mut file = fs::File::open(path).map_err(gltf::Error::Io)?;
/// # let mut bytes = Vec::new();
/// # file.read_to_end(&mut bytes).map_err(gltf::Error::Io)?;
/// # #[allow(unused)]
/// let (document, buffers, images) = gltf::import_slice(bytes.as_slice())?;
/// # Ok(())
/// # }
/// # fn main() {
/// #     run().expect("test failure");
/// # }
/// ```
pub fn import_slice<S>(slice: S) -> Result<Import>
    where S: AsRef<[u8]>
{
    import_slice_impl(slice.as_ref())
}
<commit_msg>Support URIs with embedded data when importing from a slice<commit_after>use base64;
use crate::buffer;
use crate::image;
use std::{fs, io};

use image_crate::ImageFormat::{Jpeg, Png};
use std::path::Path;
use crate::{Document, Error, Gltf, Result};

/// Return type of `import`.
type Import = (Document, Vec<buffer::Data>, Vec<image::Data>);

/// Represents the set of URI schemes the importer supports.
#[derive(Clone, Debug, Eq, Hash, PartialEq)]
enum Scheme<'a> {
    /// `data:[<media type>];base64,<data>`.
    Data(Option<&'a str>, &'a str),

    /// `file:[//]<absolute file path>`.
    ///
    /// Note: The file scheme does not implement authority.
    File(&'a str),

    /// `../foo`, etc.
    Relative,

    /// Placeholder for an unsupported URI scheme identifier.
    Unsupported,
}

impl<'a> Scheme<'a> {
    fn parse<'s>(uri: &'s str) -> Scheme<'s> {
        if uri.contains(":") {
            if uri.starts_with("data:") {
                let match0 = &uri["data:".len()..].split(";base64,").nth(0);
                let match1 = &uri["data:".len()..].split(";base64,").nth(1);
                if match1.is_some() {
                    Scheme::Data(Some(match0.unwrap()), match1.unwrap())
                } else if match0.is_some() {
                    Scheme::Data(None, match0.unwrap())
                } else {
                    Scheme::Unsupported
                }
            } else if uri.starts_with("file://") {
                Scheme::File(&uri["file://".len()..])
            } else if uri.starts_with("file:") {
                Scheme::File(&uri["file:".len()..])
            } else {
                Scheme::Unsupported
            }
        } else {
            Scheme::Relative
        }
    }

    fn read(base: Option<&Path>, uri: &str) -> Result<Vec<u8>> {
        match Scheme::parse(uri) {
            // The path may be unused in the Scheme::Data case
            // Example: "uri" : "data:application/octet-stream;base64,wsVHPgA...."
            Scheme::Data(_, base64) => base64::decode(&base64).map_err(Error::Base64),
            Scheme::File(path) if base.is_some() => read_to_end(path),
            Scheme::Relative if base.is_some() => read_to_end(base.unwrap().join(uri)),
            Scheme::Unsupported => Err(Error::UnsupportedScheme),
            _ => Err(Error::ExternalReferenceInSliceImport)
        }
    }
}

fn read_to_end<P>(path: P) -> Result<Vec<u8>>
where P: AsRef<Path>
{
    use io::Read;
    let file = fs::File::open(path.as_ref()).map_err(Error::Io)?;
    // Allocate one extra byte so the buffer doesn't need to grow before the
    // final `read` call at the end of the file.  Don't worry about `usize`
    // overflow because reading will fail regardless in that case.
    let length = file.metadata().map(|x| x.len() + 1).unwrap_or(0);
    let mut reader = io::BufReader::new(file);
    let mut data = Vec::with_capacity(length as usize);
    reader.read_to_end(&mut data).map_err(Error::Io)?;
    Ok(data)
}

/// Import the buffer data referenced by a glTF document.
pub fn import_buffer_data(
    document: &Document,
    base: Option<&Path>,
    mut blob: Option<Vec<u8>>,
) -> Result<Vec<buffer::Data>> {
    let mut buffers = Vec::new();
    for buffer in document.buffers() {
        let mut data = match buffer.source() {
            buffer::Source::Uri(uri) => Scheme::read(base, uri),
            buffer::Source::Bin => blob.take().ok_or(Error::MissingBlob),
        }?;
        if data.len() < buffer.length() {
            return Err(
                Error::BufferLength {
                    buffer: buffer.index(),
                    expected: buffer.length(),
                    actual: data.len(),
                }
            );
        }
        while data.len() % 4 != 0 {
            data.push(0);
        }
        buffers.push(buffer::Data(data));
    }
    Ok(buffers)
}

/// Import the image data referenced by a glTF document.
pub fn import_image_data(
    document: &Document,
    base: Option<&Path>,
    buffer_data: &[buffer::Data],
) -> Result<Vec<image::Data>> {
    let mut images = Vec::new();
    #[cfg(feature = "guess_mime_type")]
    let guess_format = |encoded_image: &[u8]| {
        match image_crate::guess_format(encoded_image) {
            Ok(image_crate::ImageFormat::Png) => Some(Png),
            Ok(image_crate::ImageFormat::Jpeg) => Some(Jpeg),
            _ => None,
        }
    };
    #[cfg(not(feature = "guess_mime_type"))]
    let guess_format = |_encoded_image: &[u8]| {
        None
    };
    for image in document.images() {
        match image.source() {
            image::Source::Uri { uri, mime_type } if base.is_some() => {
                match Scheme::parse(uri) {
                    Scheme::Data(Some(annoying_case), base64) => {
                        let encoded_image = base64::decode(&base64).map_err(Error::Base64)?;
                        let encoded_format = match annoying_case.as_ref() {
                            "image/png" => Png,
                            "image/jpeg" => Jpeg,
                            _ => match guess_format(&encoded_image) {
                                Some(format) => format,
                                None => return Err(Error::UnsupportedImageEncoding),
                            },
                        };
                        let decoded_image = image_crate::load_from_memory_with_format(&encoded_image, encoded_format)?;
                        images.push(image::Data::new(decoded_image));
                        continue;
                    },
                    Scheme::Unsupported => return Err(Error::UnsupportedScheme),
                    _ => {},
                }
                let encoded_image = Scheme::read(base, uri)?;
                let encoded_format =  match mime_type {
                    Some("image/png") => Png,
                    Some("image/jpeg") => Jpeg,
                    Some(_) => match guess_format(&encoded_image) {
                        Some(format) => format,
                        None => return Err(Error::UnsupportedImageEncoding),
                    },
                    None => match uri.rsplit(".").next() {
                        Some("png") => Png,
                        Some("jpg") | Some("jpeg") => Jpeg,
                        _ => match guess_format(&encoded_image) {
                            Some(format) => format,
                            None => return Err(Error::UnsupportedImageEncoding),
                        },
                    },
                };
                let decoded_image = image_crate::load_from_memory_with_format(&encoded_image, encoded_format)?;
                images.push(image::Data::new(decoded_image));
            },
            image::Source::View { view, mime_type } => {
                let parent_buffer_data = &buffer_data[view.buffer().index()].0;
                let begin = view.offset();
                let end = begin + view.length();
                let encoded_image = &parent_buffer_data[begin..end];
                let encoded_format = match mime_type {
                    "image/png" => Png,
                    "image/jpeg" => Jpeg,
                    _ => match guess_format(encoded_image) {
                        Some(format) => format,
                        None => return Err(Error::UnsupportedImageEncoding),
                    },
                };
                let decoded_image = image_crate::load_from_memory_with_format(encoded_image, encoded_format)?;
                images.push(image::Data::new(decoded_image));
            },
            _ => return Err(Error::ExternalReferenceInSliceImport)
        }
    }

    Ok(images)
}

fn import_impl(Gltf { document, blob }: Gltf, base: Option<&Path>) -> Result<Import> {
    let buffer_data = import_buffer_data(&document, base, blob)?;
    let image_data = import_image_data(&document, base, &buffer_data)?;
    let import = (document, buffer_data, image_data);
    Ok(import)
}

fn import_path(path: &Path) -> Result<Import> {
    let base = path.parent().unwrap_or(Path::new("./"));
    let file = fs::File::open(path).map_err(Error::Io)?;
    let reader = io::BufReader::new(file);
    import_impl(Gltf::from_reader(reader)?, Some(base))
}

/// Import some glTF 2.0 from the file system.
///
/// ```
/// # fn run() -> Result<(), gltf::Error> {
/// # let path = "examples/Box.gltf";
/// # #[allow(unused)]
/// let (document, buffers, images) = gltf::import(path)?;
/// # Ok(())
/// # }
/// # fn main() {
/// #     run().expect("test failure");
/// # }
/// ```
///
/// ### Note
///
/// This function is provided as a convenience for loading glTF and associated
/// resources from the file system. It is suitable for real world use but may
/// not be suitable for all real world use cases. More complex import scenarios
/// such downloading from web URLs are not handled by this function. These
/// scenarios are delegated to the user.
///
/// You can read glTF without loading resources by constructing the [`Gltf`]
/// (standard glTF) or [`Glb`] (binary glTF) data structures explicitly.
///
/// [`Gltf`]: struct.Gltf.html
/// [`Glb`]: struct.Glb.html
pub fn import<P>(path: P) -> Result<Import>
    where P: AsRef<Path>
{
    import_path(path.as_ref())
}

pub fn import_slice_impl(slice: &[u8]) -> Result<Import> {
    import_impl(Gltf::from_slice(slice)?, None)
}

/// Import some glTF 2.0 from a slice
///
/// ```
/// # extern crate gltf;
/// # use std::fs;
/// # use std::io::Read;
/// # fn run() -> Result<(), gltf::Error> {
/// # let path = "examples/Box.glb";
/// # let mut file = fs::File::open(path).map_err(gltf::Error::Io)?;
/// # let mut bytes = Vec::new();
/// # file.read_to_end(&mut bytes).map_err(gltf::Error::Io)?;
/// # #[allow(unused)]
/// let (document, buffers, images) = gltf::import_slice(bytes.as_slice())?;
/// # Ok(())
/// # }
/// # fn main() {
/// #     run().expect("test failure");
/// # }
/// ```
pub fn import_slice<S>(slice: S) -> Result<Import>
    where S: AsRef<[u8]>
{
    import_slice_impl(slice.as_ref())
}
<|endoftext|>


<commit_before>/// The Toml Insert extensions

use toml::Value;

use tokenizer::Token;
use tokenizer::tokenize_with_seperator;
use error::*;

pub trait TomlValueInsertExt {

    /// Extension function for inserting a value in the current toml::Value document
    /// using a custom seperator.
    ///
    /// For difference to TomlSetExt::set() and friends, read [#semantics].
    ///
    /// # Semantics
    ///
    /// The function automatically creates intermediate data structures based on the query string.
    /// That means, if the query string is `"a.b.c.[0]"`, but only a table `"a"` exists in the
    /// document, the function automatically creates a table `"b"` inside `"a"` and `"c"` inside
    /// `"b"`, and an array in `"c"`. The array index is ignored if the array is created.
    ///
    /// If an Array exists, but the specified index is larger than the last index, the array will
    /// be expanded by one element: If the array has a length of 3, but the query string specifies
    /// that the element should be put at 1000, the function ignores the large index and simply
    /// appends the value to the index.
    ///
    /// # Return value
    ///
    /// If the insert operation worked correctly, `Ok(None)` is returned.
    /// If the insert operation replaced an existing value `Ok(Some(old_value))` is returned
    /// On failure, `Err(e)` is returned
    ///
    fn insert_with_seperator(&mut self, query: &String, sep: char, value: Value) -> Result<Option<Value>>;

    /// Extension function for inserting a value from the current toml::Value document
    ///
    /// See documentation of `TomlValueinsertExt::insert_with_seperator`
    fn insert(&mut self, query: &String, value: Value) -> Result<Option<Value>> {
        self.insert_with_seperator(query, '.', value)
    }

}

impl TomlValueInsertExt for Value {

    fn insert_with_seperator(&mut self, query: &String, sep: char, value: Value) -> Result<Option<Value>> {
        use resolver::mut_resolver::resolve;

        let mut tokens = try!(tokenize_with_seperator(query, sep));
        let last       = tokens.pop_last().unwrap();
        let mut val    = try!(resolve(self, &tokens));

        match *last {
            Token::Identifier { ident, .. } => {
                match val {
                    &mut Value::Table(ref mut t) => {
                        Ok(t.insert(ident, value))
                    },
                    _ => unimplemented!()
                }
            },

            Token::Index { idx , .. } => {
                match val {
                    &mut Value::Array(ref mut a) => {
                        if a.len() > idx {
                            unimplemented!()
                        } else {
                            a.push(value);
                            Ok(None)
                        }
                    },
                    _ => unimplemented!()
                }
            },
        }
    }

}

#[cfg(test)]
mod test {
    use super::*;
    use error::*;
    use toml::Value;
    use toml::from_str as toml_from_str;

    #[test]
    fn test_insert_with_seperator_into_table() {
        let mut toml : Value = toml_from_str(r#"
        [table]
        "#).unwrap();

        let res = toml.insert_with_seperator(&String::from("table.a"), '.', Value::Integer(1));

        assert!(res.is_ok());

        let res = res.unwrap();
        assert!(res.is_none());

        assert!(is_match!(toml, Value::Table(_)));
        match toml {
            Value::Table(ref t) => {
                assert!(!t.is_empty());

                let table = t.get("table");
                assert!(table.is_some());

                let table = table.unwrap();
                assert!(is_match!(table, &Value::Table(_)));
                match table {
                    &Value::Table(ref t) => {
                        assert!(!t.is_empty());

                        let a = t.get("a");
                        assert!(a.is_some());

                        let a = a.unwrap();
                        assert!(is_match!(a, &Value::Integer(1)));
                    },
                    _ => panic!("What just happenend?"),
                }
            },
            _ => panic!("What just happenend?"),
        }
    }

    #[test]
    fn test_insert_with_seperator_into_array() {
        use std::ops::Index;

        let mut toml : Value = toml_from_str(r#"
        array = []
        "#).unwrap();

        let res = toml.insert_with_seperator(&String::from("array.[0]"), '.', Value::Integer(1));

        assert!(res.is_ok());

        let res = res.unwrap();
        assert!(res.is_none());

        assert!(is_match!(toml, Value::Table(_)));
        match toml {
            Value::Table(ref t) => {
                assert!(!t.is_empty());

                let array = t.get("array");
                assert!(array.is_some());

                let array = array.unwrap();
                assert!(is_match!(array, &Value::Array(_)));
                match array {
                    &Value::Array(ref a) => {
                        assert!(!a.is_empty());
                        assert!(is_match!(a.index(0), &Value::Integer(1)));
                    },
                    _ => panic!("What just happenend?"),
                }
            },
            _ => panic!("What just happenend?"),
        }
    }

    #[test]
    fn test_insert_with_seperator_into_nested_table() {
        let mut toml : Value = toml_from_str(r#"
        [a.b.c]
        "#).unwrap();

        let res = toml.insert_with_seperator(&String::from("a.b.c.d"), '.', Value::Integer(1));

        assert!(res.is_ok());

        let res = res.unwrap();
        assert!(res.is_none());

        assert!(is_match!(toml, Value::Table(_)));
        match toml {
            Value::Table(ref outer) => {
                assert!(!outer.is_empty());
                let a_tab = outer.get("a");
                assert!(a_tab.is_some());

                let a_tab = a_tab.unwrap();
                assert!(is_match!(a_tab, &Value::Table(_)));
                match a_tab {
                    &Value::Table(ref a) => {
                        assert!(!a.is_empty());

                        let b_tab = a.get("b");
                        assert!(b_tab.is_some());

                        let b_tab = b_tab.unwrap();
                        assert!(is_match!(b_tab, &Value::Table(_)));
                        match b_tab {
                            &Value::Table(ref b) => {
                                assert!(!b.is_empty());

                                let c_tab = b.get("c");
                                assert!(c_tab.is_some());

                                let c_tab = c_tab.unwrap();
                                assert!(is_match!(c_tab, &Value::Table(_)));
                                match c_tab {
                                    &Value::Table(ref c) => {
                                        assert!(!c.is_empty());

                                        let d = c.get("d");
                                        assert!(d.is_some());

                                        let d = d.unwrap();
                                        assert!(is_match!(d, &Value::Integer(1)));
                                    },
                                    _ => panic!("What just happenend?"),
                                }
                            },
                            _ => panic!("What just happenend?"),
                        }
                    },
                    _ => panic!("What just happenend?"),
                }
            },
            _ => panic!("What just happened?"),
        }
    }

}

<commit_msg>Add tests for error values<commit_after>/// The Toml Insert extensions

use toml::Value;

use tokenizer::Token;
use tokenizer::tokenize_with_seperator;
use error::*;

pub trait TomlValueInsertExt {

    /// Extension function for inserting a value in the current toml::Value document
    /// using a custom seperator.
    ///
    /// For difference to TomlSetExt::set() and friends, read [#semantics].
    ///
    /// # Semantics
    ///
    /// The function automatically creates intermediate data structures based on the query string.
    /// That means, if the query string is `"a.b.c.[0]"`, but only a table `"a"` exists in the
    /// document, the function automatically creates a table `"b"` inside `"a"` and `"c"` inside
    /// `"b"`, and an array in `"c"`. The array index is ignored if the array is created.
    ///
    /// If an Array exists, but the specified index is larger than the last index, the array will
    /// be expanded by one element: If the array has a length of 3, but the query string specifies
    /// that the element should be put at 1000, the function ignores the large index and simply
    /// appends the value to the index.
    ///
    /// # Return value
    ///
    /// If the insert operation worked correctly, `Ok(None)` is returned.
    /// If the insert operation replaced an existing value `Ok(Some(old_value))` is returned
    /// On failure, `Err(e)` is returned
    ///
    fn insert_with_seperator(&mut self, query: &String, sep: char, value: Value) -> Result<Option<Value>>;

    /// Extension function for inserting a value from the current toml::Value document
    ///
    /// See documentation of `TomlValueinsertExt::insert_with_seperator`
    fn insert(&mut self, query: &String, value: Value) -> Result<Option<Value>> {
        self.insert_with_seperator(query, '.', value)
    }

}

impl TomlValueInsertExt for Value {

    fn insert_with_seperator(&mut self, query: &String, sep: char, value: Value) -> Result<Option<Value>> {
        use resolver::mut_resolver::resolve;

        let mut tokens = try!(tokenize_with_seperator(query, sep));
        let last       = tokens.pop_last().unwrap();
        let mut val    = try!(resolve(self, &tokens));

        match *last {
            Token::Identifier { ident, .. } => {
                match val {
                    &mut Value::Table(ref mut t) => {
                        Ok(t.insert(ident, value))
                    },
                    _ => unimplemented!()
                }
            },

            Token::Index { idx , .. } => {
                match val {
                    &mut Value::Array(ref mut a) => {
                        if a.len() > idx {
                            unimplemented!()
                        } else {
                            a.push(value);
                            Ok(None)
                        }
                    },
                    _ => unimplemented!()
                }
            },
        }
    }

}

#[cfg(test)]
mod test {
    use super::*;
    use error::*;
    use toml::Value;
    use toml::from_str as toml_from_str;

    #[test]
    fn test_insert_with_seperator_into_table() {
        let mut toml : Value = toml_from_str(r#"
        [table]
        "#).unwrap();

        let res = toml.insert_with_seperator(&String::from("table.a"), '.', Value::Integer(1));

        assert!(res.is_ok());

        let res = res.unwrap();
        assert!(res.is_none());

        assert!(is_match!(toml, Value::Table(_)));
        match toml {
            Value::Table(ref t) => {
                assert!(!t.is_empty());

                let table = t.get("table");
                assert!(table.is_some());

                let table = table.unwrap();
                assert!(is_match!(table, &Value::Table(_)));
                match table {
                    &Value::Table(ref t) => {
                        assert!(!t.is_empty());

                        let a = t.get("a");
                        assert!(a.is_some());

                        let a = a.unwrap();
                        assert!(is_match!(a, &Value::Integer(1)));
                    },
                    _ => panic!("What just happenend?"),
                }
            },
            _ => panic!("What just happenend?"),
        }
    }

    #[test]
    fn test_insert_with_seperator_into_array() {
        use std::ops::Index;

        let mut toml : Value = toml_from_str(r#"
        array = []
        "#).unwrap();

        let res = toml.insert_with_seperator(&String::from("array.[0]"), '.', Value::Integer(1));

        assert!(res.is_ok());

        let res = res.unwrap();
        assert!(res.is_none());

        assert!(is_match!(toml, Value::Table(_)));
        match toml {
            Value::Table(ref t) => {
                assert!(!t.is_empty());

                let array = t.get("array");
                assert!(array.is_some());

                let array = array.unwrap();
                assert!(is_match!(array, &Value::Array(_)));
                match array {
                    &Value::Array(ref a) => {
                        assert!(!a.is_empty());
                        assert!(is_match!(a.index(0), &Value::Integer(1)));
                    },
                    _ => panic!("What just happenend?"),
                }
            },
            _ => panic!("What just happenend?"),
        }
    }

    #[test]
    fn test_insert_with_seperator_into_nested_table() {
        let mut toml : Value = toml_from_str(r#"
        [a.b.c]
        "#).unwrap();

        let res = toml.insert_with_seperator(&String::from("a.b.c.d"), '.', Value::Integer(1));

        assert!(res.is_ok());

        let res = res.unwrap();
        assert!(res.is_none());

        assert!(is_match!(toml, Value::Table(_)));
        match toml {
            Value::Table(ref outer) => {
                assert!(!outer.is_empty());
                let a_tab = outer.get("a");
                assert!(a_tab.is_some());

                let a_tab = a_tab.unwrap();
                assert!(is_match!(a_tab, &Value::Table(_)));
                match a_tab {
                    &Value::Table(ref a) => {
                        assert!(!a.is_empty());

                        let b_tab = a.get("b");
                        assert!(b_tab.is_some());

                        let b_tab = b_tab.unwrap();
                        assert!(is_match!(b_tab, &Value::Table(_)));
                        match b_tab {
                            &Value::Table(ref b) => {
                                assert!(!b.is_empty());

                                let c_tab = b.get("c");
                                assert!(c_tab.is_some());

                                let c_tab = c_tab.unwrap();
                                assert!(is_match!(c_tab, &Value::Table(_)));
                                match c_tab {
                                    &Value::Table(ref c) => {
                                        assert!(!c.is_empty());

                                        let d = c.get("d");
                                        assert!(d.is_some());

                                        let d = d.unwrap();
                                        assert!(is_match!(d, &Value::Integer(1)));
                                    },
                                    _ => panic!("What just happenend?"),
                                }
                            },
                            _ => panic!("What just happenend?"),
                        }
                    },
                    _ => panic!("What just happenend?"),
                }
            },
            _ => panic!("What just happened?"),
        }
    }

    #[test]
    fn test_insert_with_seperator_into_table_where_array_is() {
        let mut toml : Value = toml_from_str(r#"
        table = []
        "#).unwrap();

        let res = toml.insert_with_seperator(&String::from("table.a"), '.', Value::Integer(1));

        assert!(res.is_err());

        let err = res.unwrap_err();
        assert!(is_match!(err.kind(), &ErrorKind::NoIdentifierInArray(_)));
    }

    #[test]
    fn test_insert_with_seperator_into_array_where_table_is() {
        let mut toml : Value = toml_from_str(r#"
        [table]
        "#).unwrap();

        let res = toml.insert_with_seperator(&String::from("table.[0]"), '.', Value::Integer(1));

        assert!(res.is_err());

        let err = res.unwrap_err();
        assert!(is_match!(err.kind(), &ErrorKind::NoIndexInTable(_)));
    }

}

<|endoftext|>


<commit_before>// Copyright (c) 2016 Anatoly Ikorsky
//
// Licensed under the Apache License, Version 2.0
// <LICENSE-APACHE or http://www.apache.org/licenses/LICENSE-2.0> or the MIT
// license <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. All files in the project carrying such notice may not be copied,
// modified, or distributed except according to those terms.

use BoxFuture;
use errors::*;
use std::fmt;
#[cfg(feature = "ssl")]
use std::fs::File;
use lib_futures::Async::NotReady;
use lib_futures::Async::Ready;
#[cfg(feature = "ssl")]
use lib_futures::{Future, IntoFuture};
use lib_futures::Poll;
use lib_futures::stream;
use io::futures::ConnectingStream;
use io::futures::new_connecting_stream;
use io::futures::new_write_packet;
use io::futures::WritePacket;
#[cfg(feature = "ssl")]
use native_tls::{Certificate, Pkcs12, TlsConnector};
use myc::packets::{PacketParser, ParseResult, RawPacket};
use opts::SslOpts;
use std::io;
use std::net::ToSocketAddrs;
use std::time::Duration;
use tokio::net::TcpStream;
use tokio::reactor::Handle;
use tokio_io::AsyncRead;
use tokio_io::AsyncWrite;
#[cfg(feature = "ssl")]
use tokio_tls::{TlsConnectorExt, TlsStream};

pub mod futures;

#[derive(Debug)]
enum Endpoint {
    Plain(TcpStream),
    #[cfg(feature = "ssl")]
    Secure(TlsStream<TcpStream>),
}

impl Endpoint {
    pub fn set_keepalive_ms(&self, ms: Option<u32>) -> Result<()> {
        let ms = ms.map(|val| Duration::from_millis(val as u64));
        match *self {
            Endpoint::Plain(ref stream) => stream.set_keepalive(ms)?,
            #[cfg(feature = "ssl")]
            Endpoint::Secure(ref stream) => stream.get_ref().get_ref().set_keepalive(ms)?,
        }
        Ok(())
    }

    pub fn set_tcp_nodelay(&self, val: bool) -> Result<()> {
        match *self {
            Endpoint::Plain(ref stream) => stream.set_nodelay(val)?,
            #[cfg(feature = "ssl")]
            Endpoint::Secure(ref stream) => stream.get_ref().get_ref().set_nodelay(val)?,
        }
        Ok(())
    }

    #[cfg(feature = "ssl")]
    pub fn make_secure(self, domain: String, ssl_opts: SslOpts) -> BoxFuture<Self> {
        let fut = ::std::fs::File::open(ssl_opts.pkcs12_path())
            .map_err(Error::from)
            .and_then(|mut file| {
                let mut der = vec![];
                file.read_to_end(&mut der)?;
                let identity = Pkcs12::from_der(&*der, ssl_opts.password().unwrap_or(""))
                    .chain_err(|| "Can't parse der")?;
                let mut builder = TlsConnector::builder().chain_err(
                    || "Can't create TlsConnectorBuilder)",
                )?;
                match ssl_opts.root_cert_path() {
                    Some(root_cert_path) => {
                        let mut root_cert_der = vec![];
                        let mut root_cert_file = File::open(root_cert_path)?;
                        root_cert_file.read_to_end(&mut root_cert_der)?;
                        let root_cert = Certificate::from_der(&*root_cert_der)
                            .chain_err(|| "Can't parse root certificate")?;
                        builder.add_root_certificate(root_cert)
                            .chain_err(|| "Can't add root certificate")?;
                    },
                    None => (),
                }
                builder.identity(identity).chain_err(
                    || "Can't set identity for TlsConnectorBuilder",
                )?;
                builder.build().chain_err(
                    || "Can't build TlsConnectorBuilder",
                )
            })
            .into_future()
            .and_then(move |tls_connector| match self {
                Endpoint::Plain(stream) => {
                    let fut = if ssl_opts.skip_domain_validation() {
                        TlsConnectorExt::danger_connect_async_without_providing_domain_for_certificate_verification_and_server_name_indication(&tls_connector, stream)
                    } else {
                        TlsConnectorExt::connect_async(&tls_connector, &*domain, stream)
                    };
                    fut.then(|result| result.chain_err(|| "Can't connect TlsConnector"))
                }
                Endpoint::Secure(_) => unreachable!(),
            })
            .map(|tls_stream| Endpoint::Secure(tls_stream));
        Box::new(fut)
    }
}

impl From<TcpStream> for Endpoint {
    fn from(stream: TcpStream) -> Self {
        Endpoint::Plain(stream)
    }
}

#[cfg(feature = "ssl")]
impl From<TlsStream<TcpStream>> for Endpoint {
    fn from(stream: TlsStream<TcpStream>) -> Self {
        Endpoint::Secure(stream)
    }
}

impl io::Read for Endpoint {
    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
        match *self {
            Endpoint::Plain(ref mut stream) => stream.read(buf),
            #[cfg(feature = "ssl")]
            Endpoint::Secure(ref mut stream) => stream.read(buf),
        }
    }
}

impl io::Write for Endpoint {
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
        match *self {
            Endpoint::Plain(ref mut stream) => stream.write(buf),
            #[cfg(feature = "ssl")]
            Endpoint::Secure(ref mut stream) => stream.write(buf),
        }
    }

    fn flush(&mut self) -> io::Result<()> {
        match *self {
            Endpoint::Plain(ref mut stream) => stream.flush(),
            #[cfg(feature = "ssl")]
            Endpoint::Secure(ref mut stream) => stream.flush(),
        }
    }
}

impl AsyncRead for Endpoint {
    unsafe fn prepare_uninitialized_buffer(&self, _buf: &mut [u8]) -> bool {
        match *self {
            Endpoint::Plain(ref stream) => stream.prepare_uninitialized_buffer(_buf),
            #[cfg(feature = "ssl")]
            Endpoint::Secure(ref stream) => stream.prepare_uninitialized_buffer(_buf),
        }
    }
}

impl AsyncWrite for Endpoint {
    fn shutdown(&mut self) -> Poll<(), io::Error> {
        match *self {
            Endpoint::Plain(ref mut stream) => stream.shutdown(),
            #[cfg(feature = "ssl")]
            Endpoint::Secure(ref mut stream) => stream.shutdown(),
        }
    }
}

/// Stream connected to MySql server.
pub struct Stream {
    endpoint: Option<Endpoint>,
    closed: bool,
    parser: Option<PacketParser>,
    packets: ::std::collections::VecDeque<(RawPacket, u8)>,
    buf: Vec<u8>,
}

impl fmt::Debug for Stream {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "Stream (endpoint={:?})", self.endpoint)
    }
}

impl Stream {
    pub fn connect<S>(addr: S, handle: &Handle) -> ConnectingStream
    where
        S: ToSocketAddrs,
    {
        new_connecting_stream(addr, handle)
    }

    pub fn write_packet(self, data: Vec<u8>, seq_id: u8) -> WritePacket {
        new_write_packet(self, data, seq_id)
    }

    pub fn set_keepalive_ms(&self, ms: Option<u32>) -> Result<()> {
        match self.endpoint {
            Some(ref endpoint) => endpoint.set_keepalive_ms(ms),
            None => unreachable!(),
        }
    }

    pub fn set_tcp_nodelay(&self, val: bool) -> Result<()> {
        match self.endpoint {
            Some(ref endpoint) => endpoint.set_tcp_nodelay(val),
            None => unreachable!(),
        }
    }

    #[cfg(not(feature = "ssl"))]
    #[allow(unused)]
    pub fn make_secure(self, domain: String, ssl_opts: SslOpts) -> BoxFuture<Self> {
        panic!("Ssl connection requires `ssl` feature");
    }

    #[cfg(feature = "ssl")]
    pub fn make_secure(mut self, domain: String, ssl_opts: SslOpts) -> BoxFuture<Self> {
        match self.endpoint.take() {
            Some(endpoint) => {
                let fut = endpoint.make_secure(domain, ssl_opts).map(|endpoint| {
                    self.endpoint = Some(endpoint);
                    self
                });
                Box::new(fut)
            }
            None => unreachable!(),
        }
    }
}

impl io::Read for Stream {
    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
        match self.endpoint {
            Some(ref mut endpoint) => endpoint.read(buf),
            None => unreachable!(),
        }
    }
}

impl io::Write for Stream {
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
        match self.endpoint {
            Some(ref mut endpoint) => endpoint.write(buf),
            None => unreachable!(),
        }
    }

    fn flush(&mut self) -> io::Result<()> {
        match self.endpoint {
            Some(ref mut endpoint) => endpoint.flush(),
            None => unreachable!(),
        }
    }
}

impl AsyncRead for Stream {
    unsafe fn prepare_uninitialized_buffer(&self, _buf: &mut [u8]) -> bool {
        match self.endpoint {
            Some(ref endpoint) => endpoint.prepare_uninitialized_buffer(_buf),
            None => unreachable!(),
        }
    }
}

impl AsyncWrite for Stream {
    fn shutdown(&mut self) -> Poll<(), io::Error> {
        match self.endpoint {
            Some(ref mut endpoint) => endpoint.shutdown(),
            None => unreachable!(),
        }
    }
}

impl stream::Stream for Stream {
    type Item = (RawPacket, u8);
    type Error = Error;

    fn poll(&mut self) -> Poll<Option<(RawPacket, u8)>, Error> {
        // emit pending packets
        if let Some(packet) = self.packets.pop_front() {
            return Ok(Ready(Some(packet)));
        }

        // should read everything from self.endpoint
        let mut buf = [0u8; 4096];
        while !self.closed {
            match self.poll_read(&mut buf[..]) {
                Err(error) => {
                    self.closed = true;
                    bail!(error)
                }
                Ok(Ready(0)) => {
                    self.closed = true;
                }
                Ok(Ready(size)) => {
                    self.buf.extend_from_slice(&buf[..size]);
                }
                Ok(NotReady) => break,
            };
        }

        // parse buffer into packets
        let (packets, parser, buf) = parse_packet(self.parser.take().unwrap(), &self.buf);
        self.packets = packets.into();
        self.parser = Some(parser);
        self.buf = buf;

        if let Some(packet) = self.packets.pop_front() {
            return Ok(Ready(Some(packet)));
        }

        if self.closed {
            bail!(ErrorKind::ConnectionClosed);
        }
        return Ok(NotReady);
    }
}

fn parse_packet(
    mut parser: PacketParser,
    mut buf: &[u8],
) -> (Vec<(RawPacket, u8)>, PacketParser, Vec<u8>) {
    let mut packets = Vec::new();

    loop {
        parser = match parser.parse() {
            ParseResult::Done(packet, seq_id) => {
                packets.push((packet, seq_id));
                PacketParser::empty()
            }
            ParseResult::Incomplete(mut parser, needed) => {
                if buf.len() < needed {
                    return (packets, parser, Vec::from(buf));
                }

                parser.extend_from_slice(&buf[..needed]);
                buf = &buf[needed..];
                parser
            }
        };
    }
}
<commit_msg>Fix build with ssl feature<commit_after>// Copyright (c) 2016 Anatoly Ikorsky
//
// Licensed under the Apache License, Version 2.0
// <LICENSE-APACHE or http://www.apache.org/licenses/LICENSE-2.0> or the MIT
// license <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. All files in the project carrying such notice may not be copied,
// modified, or distributed except according to those terms.

use BoxFuture;
use errors::*;
use std::fmt;
#[cfg(feature = "ssl")]
use std::fs::File;
#[cfg(feature = "ssl")]
use std::io::Read;
use lib_futures::Async::NotReady;
use lib_futures::Async::Ready;
#[cfg(feature = "ssl")]
use lib_futures::{Future, IntoFuture};
use lib_futures::Poll;
use lib_futures::stream;
use io::futures::ConnectingStream;
use io::futures::new_connecting_stream;
use io::futures::new_write_packet;
use io::futures::WritePacket;
#[cfg(feature = "ssl")]
use native_tls::{Certificate, Pkcs12, TlsConnector};
use myc::packets::{PacketParser, ParseResult, RawPacket};
use opts::SslOpts;
use std::io;
use std::net::ToSocketAddrs;
use std::time::Duration;
use tokio::net::TcpStream;
use tokio::reactor::Handle;
use tokio_io::AsyncRead;
use tokio_io::AsyncWrite;
#[cfg(feature = "ssl")]
use tokio_tls::{TlsConnectorExt, TlsStream};

pub mod futures;

#[derive(Debug)]
enum Endpoint {
    Plain(TcpStream),
    #[cfg(feature = "ssl")]
    Secure(TlsStream<TcpStream>),
}

impl Endpoint {
    pub fn set_keepalive_ms(&self, ms: Option<u32>) -> Result<()> {
        let ms = ms.map(|val| Duration::from_millis(val as u64));
        match *self {
            Endpoint::Plain(ref stream) => stream.set_keepalive(ms)?,
            #[cfg(feature = "ssl")]
            Endpoint::Secure(ref stream) => stream.get_ref().get_ref().set_keepalive(ms)?,
        }
        Ok(())
    }

    pub fn set_tcp_nodelay(&self, val: bool) -> Result<()> {
        match *self {
            Endpoint::Plain(ref stream) => stream.set_nodelay(val)?,
            #[cfg(feature = "ssl")]
            Endpoint::Secure(ref stream) => stream.get_ref().get_ref().set_nodelay(val)?,
        }
        Ok(())
    }

    #[cfg(feature = "ssl")]
    pub fn make_secure(self, domain: String, ssl_opts: SslOpts) -> BoxFuture<Self> {
        let fut = ::std::fs::File::open(ssl_opts.pkcs12_path())
            .map_err(Error::from)
            .and_then(|mut file| {
                let mut der = vec![];
                file.read_to_end(&mut der)?;
                let identity = Pkcs12::from_der(&*der, ssl_opts.password().unwrap_or(""))
                    .chain_err(|| "Can't parse der")?;
                let mut builder = TlsConnector::builder().chain_err(
                    || "Can't create TlsConnectorBuilder)",
                )?;
                match ssl_opts.root_cert_path() {
                    Some(root_cert_path) => {
                        let mut root_cert_der = vec![];
                        let mut root_cert_file = File::open(root_cert_path)?;
                        root_cert_file.read_to_end(&mut root_cert_der)?;
                        let root_cert = Certificate::from_der(&*root_cert_der)
                            .chain_err(|| "Can't parse root certificate")?;
                        builder.add_root_certificate(root_cert)
                            .chain_err(|| "Can't add root certificate")?;
                    },
                    None => (),
                }
                builder.identity(identity).chain_err(
                    || "Can't set identity for TlsConnectorBuilder",
                )?;
                builder.build().chain_err(
                    || "Can't build TlsConnectorBuilder",
                )
            })
            .into_future()
            .and_then(move |tls_connector| match self {
                Endpoint::Plain(stream) => {
                    let fut = if ssl_opts.skip_domain_validation() {
                        TlsConnectorExt::danger_connect_async_without_providing_domain_for_certificate_verification_and_server_name_indication(&tls_connector, stream)
                    } else {
                        TlsConnectorExt::connect_async(&tls_connector, &*domain, stream)
                    };
                    fut.then(|result| result.chain_err(|| "Can't connect TlsConnector"))
                }
                Endpoint::Secure(_) => unreachable!(),
            })
            .map(|tls_stream| Endpoint::Secure(tls_stream));
        Box::new(fut)
    }
}

impl From<TcpStream> for Endpoint {
    fn from(stream: TcpStream) -> Self {
        Endpoint::Plain(stream)
    }
}

#[cfg(feature = "ssl")]
impl From<TlsStream<TcpStream>> for Endpoint {
    fn from(stream: TlsStream<TcpStream>) -> Self {
        Endpoint::Secure(stream)
    }
}

impl io::Read for Endpoint {
    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
        match *self {
            Endpoint::Plain(ref mut stream) => stream.read(buf),
            #[cfg(feature = "ssl")]
            Endpoint::Secure(ref mut stream) => stream.read(buf),
        }
    }
}

impl io::Write for Endpoint {
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
        match *self {
            Endpoint::Plain(ref mut stream) => stream.write(buf),
            #[cfg(feature = "ssl")]
            Endpoint::Secure(ref mut stream) => stream.write(buf),
        }
    }

    fn flush(&mut self) -> io::Result<()> {
        match *self {
            Endpoint::Plain(ref mut stream) => stream.flush(),
            #[cfg(feature = "ssl")]
            Endpoint::Secure(ref mut stream) => stream.flush(),
        }
    }
}

impl AsyncRead for Endpoint {
    unsafe fn prepare_uninitialized_buffer(&self, _buf: &mut [u8]) -> bool {
        match *self {
            Endpoint::Plain(ref stream) => stream.prepare_uninitialized_buffer(_buf),
            #[cfg(feature = "ssl")]
            Endpoint::Secure(ref stream) => stream.prepare_uninitialized_buffer(_buf),
        }
    }
}

impl AsyncWrite for Endpoint {
    fn shutdown(&mut self) -> Poll<(), io::Error> {
        match *self {
            Endpoint::Plain(ref mut stream) => stream.shutdown(),
            #[cfg(feature = "ssl")]
            Endpoint::Secure(ref mut stream) => stream.shutdown(),
        }
    }
}

/// Stream connected to MySql server.
pub struct Stream {
    endpoint: Option<Endpoint>,
    closed: bool,
    parser: Option<PacketParser>,
    packets: ::std::collections::VecDeque<(RawPacket, u8)>,
    buf: Vec<u8>,
}

impl fmt::Debug for Stream {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "Stream (endpoint={:?})", self.endpoint)
    }
}

impl Stream {
    pub fn connect<S>(addr: S, handle: &Handle) -> ConnectingStream
    where
        S: ToSocketAddrs,
    {
        new_connecting_stream(addr, handle)
    }

    pub fn write_packet(self, data: Vec<u8>, seq_id: u8) -> WritePacket {
        new_write_packet(self, data, seq_id)
    }

    pub fn set_keepalive_ms(&self, ms: Option<u32>) -> Result<()> {
        match self.endpoint {
            Some(ref endpoint) => endpoint.set_keepalive_ms(ms),
            None => unreachable!(),
        }
    }

    pub fn set_tcp_nodelay(&self, val: bool) -> Result<()> {
        match self.endpoint {
            Some(ref endpoint) => endpoint.set_tcp_nodelay(val),
            None => unreachable!(),
        }
    }

    #[cfg(not(feature = "ssl"))]
    #[allow(unused)]
    pub fn make_secure(self, domain: String, ssl_opts: SslOpts) -> BoxFuture<Self> {
        panic!("Ssl connection requires `ssl` feature");
    }

    #[cfg(feature = "ssl")]
    pub fn make_secure(mut self, domain: String, ssl_opts: SslOpts) -> BoxFuture<Self> {
        match self.endpoint.take() {
            Some(endpoint) => {
                let fut = endpoint.make_secure(domain, ssl_opts).map(|endpoint| {
                    self.endpoint = Some(endpoint);
                    self
                });
                Box::new(fut)
            }
            None => unreachable!(),
        }
    }
}

impl io::Read for Stream {
    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
        match self.endpoint {
            Some(ref mut endpoint) => endpoint.read(buf),
            None => unreachable!(),
        }
    }
}

impl io::Write for Stream {
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
        match self.endpoint {
            Some(ref mut endpoint) => endpoint.write(buf),
            None => unreachable!(),
        }
    }

    fn flush(&mut self) -> io::Result<()> {
        match self.endpoint {
            Some(ref mut endpoint) => endpoint.flush(),
            None => unreachable!(),
        }
    }
}

impl AsyncRead for Stream {
    unsafe fn prepare_uninitialized_buffer(&self, _buf: &mut [u8]) -> bool {
        match self.endpoint {
            Some(ref endpoint) => endpoint.prepare_uninitialized_buffer(_buf),
            None => unreachable!(),
        }
    }
}

impl AsyncWrite for Stream {
    fn shutdown(&mut self) -> Poll<(), io::Error> {
        match self.endpoint {
            Some(ref mut endpoint) => endpoint.shutdown(),
            None => unreachable!(),
        }
    }
}

impl stream::Stream for Stream {
    type Item = (RawPacket, u8);
    type Error = Error;

    fn poll(&mut self) -> Poll<Option<(RawPacket, u8)>, Error> {
        // emit pending packets
        if let Some(packet) = self.packets.pop_front() {
            return Ok(Ready(Some(packet)));
        }

        // should read everything from self.endpoint
        let mut buf = [0u8; 4096];
        while !self.closed {
            match self.poll_read(&mut buf[..]) {
                Err(error) => {
                    self.closed = true;
                    bail!(error)
                }
                Ok(Ready(0)) => {
                    self.closed = true;
                }
                Ok(Ready(size)) => {
                    self.buf.extend_from_slice(&buf[..size]);
                }
                Ok(NotReady) => break,
            };
        }

        // parse buffer into packets
        let (packets, parser, buf) = parse_packet(self.parser.take().unwrap(), &self.buf);
        self.packets = packets.into();
        self.parser = Some(parser);
        self.buf = buf;

        if let Some(packet) = self.packets.pop_front() {
            return Ok(Ready(Some(packet)));
        }

        if self.closed {
            bail!(ErrorKind::ConnectionClosed);
        }
        return Ok(NotReady);
    }
}

fn parse_packet(
    mut parser: PacketParser,
    mut buf: &[u8],
) -> (Vec<(RawPacket, u8)>, PacketParser, Vec<u8>) {
    let mut packets = Vec::new();

    loop {
        parser = match parser.parse() {
            ParseResult::Done(packet, seq_id) => {
                packets.push((packet, seq_id));
                PacketParser::empty()
            }
            ParseResult::Incomplete(mut parser, needed) => {
                if buf.len() < needed {
                    return (packets, parser, Vec::from(buf));
                }

                parser.extend_from_slice(&buf[..needed]);
                buf = &buf[needed..];
                parser
            }
        };
    }
}
<|endoftext|>


<commit_before>use std;
use std::io::prelude::*;
use std::net::{UdpSocket, ToSocketAddrs};
use std::mem;
use std::hash::{Hash, Hasher, SipHasher};

use log::{Log, LogRecord, LogLevel, LogMetadata};
use flate2::Compression;
use flate2::write::GzEncoder;
use chrono::{DateTime, UTC};
use serde_json;
use byteorder::{BigEndian, WriteBytesExt};

use GraylogError;

const MAX_PACKET_SIZE: usize = 8192;
const MAGIC_BYTES: u16 = 0x1E0F;

struct GraylogLogger {
    server: A,
    sock: UdpSocket,
    level: LogLevel,
    hostname: String,
}

#[derive(Serialize, Deserialize)]
struct Gelf {
    version: String,
    host: String,
    short_message: String,
    full_message: Option<String>,
    timestamp: Option<i64>,
    level: Option<u8>,
}

struct GelfChunks<'a> {
    count: usize,
    total: usize,
    #[allow(dead_code)]
    data: &'a [u8],
    chunks: Vec<&'a [u8]>,
    message_id: u64,
}

struct Chunk<'a> {
    header: ChunkHeader,
    data: &'a [u8],
}

impl<'a> Chunk<'a> {
    fn to_binary(&self) -> Result<Vec<u8>, GraylogError> {
        let mut wrt = vec![];
        try!(wrt.write_u16::<BigEndian>(self.header.magic_bytes));
        try!(wrt.write_u64::<BigEndian>(self.header.message_id));
        try!(wrt.write_u8(self.header.seq_number));
        try!(wrt.write_u8(self.header.seq_count));
         wrt.extend(self.data.iter());
        Ok(wrt)
    }
}

struct ChunkHeader {
    magic_bytes: u16,
    message_id: u64,
    seq_number: u8,
    seq_count: u8,
}

impl<'a> GelfChunks<'a> {
    fn new(data: &'a [u8], message_id: u64) -> GelfChunks<'a> {
        let chunk_size = MAX_PACKET_SIZE - mem::size_of::<ChunkHeader>();
        let mut chunks = data.chunks(chunk_size).collect::<Vec<&[u8]>>();
        chunks.reverse();

        GelfChunks {
            count: 0,
            total: chunks.len(),
            chunks: chunks,
            data: data,
            message_id: message_id,
        }
    }
}

impl<'a> Iterator for GelfChunks<'a> {
    type Item = Chunk<'a>;

    fn next(&mut self) -> Option<Chunk<'a>> {
        self.chunks.pop().map(|data| {
            let header = ChunkHeader {
                magic_bytes: MAGIC_BYTES,
                message_id: self.message_id,
                seq_number: self.count as u8,
                seq_count: self.total as u8,
            };

            self.count += 1;

            Chunk {
                header: header,
                data: data,
            }
        })
    }
}

pub fn init(addr: A, level: LogLevel) -> Result<(), GraylogError>
    where A: std::marker::Send + std::marker::Sync
{
    use log;
    use std::process::Command;
    let sock = try!(UdpSocket::bind("0.0.0.0:0"));

    try!(log::set_logger(|max_log_level| {
        max_log_level.set(level.to_log_level_filter());
        let hostname = Command::new("hostname")
            .arg("-f")
            .output()
            .map(|out| String::from_utf8_lossy(&out.stdout).into_owned())
            .unwrap_or("unknown".to_string());

        Box::new(GraylogLogger {
            level: level,
            server: addr,
            sock: sock,
            hostname: hostname,
        })
    }));
    Ok(())
}

impl Gelf {
    fn message_id(&self) -> u64 {
        use chrono::Timelike;
        let mut s = SipHasher::new();
        self.host.hash(&mut s);
        self.timestamp.hash(&mut s);
        UTC::now().nanosecond().hash(&mut s);
        s.finish()
    }
}

impl GraylogLogger {
    fn send_as_gelf(&self, record: &LogRecord) -> Result<usize, GraylogError> {
        let mut e = GzEncoder::new(Vec::new(), Compression::Default);
        let utc: DateTime<UTC> = UTC::now();

        let gelf = Gelf {
            version: "1.1".to_string(),
            short_message: format!("{}", record.args()),
            full_message: None,
            timestamp: Some(utc.timestamp()),
            level: None,
            host: self.hostname.clone(),
        };

        let json = try!(serde_json::to_string(&gelf));
        try!(e.write(json.as_bytes()));

        let compressed_bytes = try!(e.finish());
        if compressed_bytes.len() > MAX_PACKET_SIZE {
            self.send_chunked_gelf(&compressed_bytes, gelf.message_id())
        } else {
            self.send(&compressed_bytes)
        }
    }

    fn send_chunked_gelf(&self, buffer: &[u8], message_id: u64) -> Result<usize, GraylogError> {
        let chunks = GelfChunks::new(buffer, message_id);
        println!("{:x}", message_id);
        for chunk in chunks {
            let data = try!(chunk.to_binary());
            try!(self.send(&data));
        }
        Ok(1)
    }

    fn send(&self, buffer: &[u8]) -> Result<usize, GraylogError> {
        let s = try!(self.sock.send_to(buffer, &self.server).or_else(|e| {
            println!("Error writing to Graylog: {}", e);
            Err(e)
        }));
        Ok(s)
    }
}

impl Log for GraylogLogger
    where A: std::marker::Send + std::marker::Sync
{
    fn enabled(&self, metadata: &LogMetadata) -> bool {
        metadata.level() <= self.level
    }

    fn log(&self, record: &LogRecord) {
        if self.enabled(record.metadata()) {
            if self.send_as_gelf(record).is_err() {
                println!("{} {} - {}", UTC::now(), record.level(), record.args());
            };
        }
    }
}
<commit_msg>print the chunk header for debugging purpose<commit_after>use std;
use std::io::prelude::*;
use std::net::{UdpSocket, ToSocketAddrs};
use std::mem;
use std::hash::{Hash, Hasher, SipHasher};

use log::{Log, LogRecord, LogLevel, LogMetadata};
use flate2::Compression;
use flate2::write::GzEncoder;
use chrono::{DateTime, UTC};
use serde_json;
use byteorder::{BigEndian, WriteBytesExt};

use GraylogError;

const MAX_PACKET_SIZE: usize = 8192;
const MAGIC_BYTES: u16 = 0x1E0F;

struct GraylogLogger {
    server: A,
    sock: UdpSocket,
    level: LogLevel,
    hostname: String,
}

#[derive(Serialize, Deserialize)]
struct Gelf {
    version: String,
    host: String,
    short_message: String,
    full_message: Option<String>,
    timestamp: Option<i64>,
    level: Option<u8>,
}

struct GelfChunks<'a> {
    count: usize,
    total: usize,
    #[allow(dead_code)]
    data: &'a [u8],
    chunks: Vec<&'a [u8]>,
    message_id: u64,
}

struct Chunk<'a> {
    header: ChunkHeader,
    data: &'a [u8],
}

impl<'a> Chunk<'a> {
    fn to_binary(&self) -> Result<Vec<u8>, GraylogError> {
        let mut wrt = vec![];
        try!(wrt.write_u16::<BigEndian>(self.header.magic_bytes));
        try!(wrt.write_u64::<BigEndian>(self.header.message_id));
        try!(wrt.write_u8(self.header.seq_number));
        try!(wrt.write_u8(self.header.seq_count));
         wrt.extend(self.data.iter());
        Ok(wrt)
    }
}

#[derive(Debug)]
struct ChunkHeader {
    magic_bytes: u16,
    message_id: u64,
    seq_number: u8,
    seq_count: u8,
}

impl<'a> GelfChunks<'a> {
    fn new(data: &'a [u8], message_id: u64) -> GelfChunks<'a> {
        let chunk_size = MAX_PACKET_SIZE - mem::size_of::<ChunkHeader>();
        let mut chunks = data.chunks(chunk_size).collect::<Vec<&[u8]>>();
        chunks.reverse();

        GelfChunks {
            count: 0,
            total: chunks.len(),
            chunks: chunks,
            data: data,
            message_id: message_id,
        }
    }
}

impl<'a> Iterator for GelfChunks<'a> {
    type Item = Chunk<'a>;

    fn next(&mut self) -> Option<Chunk<'a>> {
        self.chunks.pop().map(|data| {
            let header = ChunkHeader {
                magic_bytes: MAGIC_BYTES,
                message_id: self.message_id,
                seq_number: self.count as u8,
                seq_count: self.total as u8,
            };
            println!("{:?}", header);

            self.count += 1;

            Chunk {
                header: header,
                data: data,
            }
        })
    }
}

pub fn init(addr: A, level: LogLevel) -> Result<(), GraylogError>
    where A: std::marker::Send + std::marker::Sync
{
    use log;
    use std::process::Command;
    let sock = try!(UdpSocket::bind("0.0.0.0:0"));

    try!(log::set_logger(|max_log_level| {
        max_log_level.set(level.to_log_level_filter());
        let hostname = Command::new("hostname")
            .arg("-f")
            .output()
            .map(|out| String::from_utf8_lossy(&out.stdout).into_owned())
            .unwrap_or("unknown".to_string());

        Box::new(GraylogLogger {
            level: level,
            server: addr,
            sock: sock,
            hostname: hostname,
        })
    }));
    Ok(())
}

impl Gelf {
    fn message_id(&self) -> u64 {
        use chrono::Timelike;
        let mut s = SipHasher::new();
        self.host.hash(&mut s);
        self.timestamp.hash(&mut s);
        UTC::now().nanosecond().hash(&mut s);
        s.finish()
    }
}

impl GraylogLogger {
    fn send_as_gelf(&self, record: &LogRecord) -> Result<usize, GraylogError> {
        let mut e = GzEncoder::new(Vec::new(), Compression::Default);
        let utc: DateTime<UTC> = UTC::now();

        let gelf = Gelf {
            version: "1.1".to_string(),
            short_message: format!("{}", record.args()),
            full_message: None,
            timestamp: Some(utc.timestamp()),
            level: None,
            host: self.hostname.clone(),
        };

        let json = try!(serde_json::to_string(&gelf));
        try!(e.write(json.as_bytes()));

        let compressed_bytes = try!(e.finish());
        if compressed_bytes.len() > MAX_PACKET_SIZE {
            self.send_chunked_gelf(&compressed_bytes, gelf.message_id())
        } else {
            self.send(&compressed_bytes)
        }
    }

    fn send_chunked_gelf(&self, buffer: &[u8], message_id: u64) -> Result<usize, GraylogError> {
        let chunks = GelfChunks::new(buffer, message_id);
        println!("{:x}", message_id);
        for chunk in chunks {
            let data = try!(chunk.to_binary());
            try!(self.send(&data));
        }
        Ok(1)
    }

    fn send(&self, buffer: &[u8]) -> Result<usize, GraylogError> {
        let s = try!(self.sock.send_to(buffer, &self.server).or_else(|e| {
            println!("Error writing to Graylog: {}", e);
            Err(e)
        }));
        Ok(s)
    }
}

impl Log for GraylogLogger
    where A: std::marker::Send + std::marker::Sync
{
    fn enabled(&self, metadata: &LogMetadata) -> bool {
        metadata.level() <= self.level
    }

    fn log(&self, record: &LogRecord) {
        if self.enabled(record.metadata()) {
            if self.send_as_gelf(record).is_err() {
                println!("{} {} - {}", UTC::now(), record.level(), record.args());
            };
        }
    }
}
<|endoftext|>


<commit_before>use std::collections::HashMap;

use log;
use hostname;

use backends::Backend;
use message::{Message, WireMessage};
use errors::{Result, ErrorKind};

/// Logger for sending log-messages in [GELF-format](http://docs.graylog.org/en/2.2/pages/gelf.html)
/// to a GELF server, e.g. [Graylog](https://www.graylog.org/)
///
/// A `Logger` instance can be either used as a standalone object to log directly
/// to a log-server or it can installed as `log`-crate log-handler.
///
/// # Standalone usage
///
/// Note that the simple constructor (`Logger::new`) is used. It determines the
/// local hostname by itself.
///
/// ``` no_run
/// use gelf::{Logger, UdpBackend, Message, Level};
///
/// let backend = UdpBackend::new("127.0.0.1:12201");
/// let logger = Logger::new(Box::new(backend)).expect("Failed to determine hostname");
///
/// logger.log_message(Message::new("Test log message!", Some(Level::Debug)));
/// ```
///
/// # Usage with the `log`-crate
///
/// We assume the `log`-crate has been imported with macros enabled, like
/// this `#[macro_use] extern crate log;`.
///
/// ``` no_run
/// use gelf::{Logger, UdpBackend, Message, Level};
/// use log::LogLevelFilter;
///
/// let backend = UdpBackend::new("127.0.0.1:12201");
/// let logger = Logger::new(Box::new(backend)).expect("Failed to determine hostname");
///
/// logger.install(LogLevelFilter::Trace).expect("Failed to install logger");
///
/// debug!("Debug message");
/// warn!("Warning");
/// ```
///
/// # Further information
///
/// For more information check out the `examples/` directory.
pub struct Logger {
    hostname: String,
    backend: Box<Backend>,
    default_metadata: HashMap<String, String>,
}

impl Logger {
    /// Construct a new `Logger` instance
    ///
    /// The backend needs to be boxed for usage as a logger with the `log`-crate.
    /// This constructor tries to determine the local hostname (required by GELF)
    /// with the help of the `hostname`-crate. If you want to set a custom hostname
    /// check out the `Logger::new_with_hostname` constructor
    pub fn new(backend: Box<Backend>) -> Result<Self> {
        hostname::get_hostname()
            .map(|hostname| Logger::new_with_hostname(backend, &hostname))
            .ok_or(ErrorKind::LoggerCreateFailed("Failed to determine local hostname").into())
    }

    /// Construct a new `Logger` instance with predetermined hostname
    ///
    /// The backend needs to be boxed for usage as a logger with the `log`-crate. It
    /// uses the passed hostname for the GELF `host` field
    pub fn new_with_hostname(backend: Box<Backend>, hostname: &str) -> Logger {
        Logger {
            hostname: String::from(hostname),
            backend: backend,
            default_metadata: HashMap::new(),
        }
    }

    /// Install a logger instance as a `log`-Logger
    ///
    /// This method wraps `log::set_logger` as a convenience function as required
    /// by the `log`-crate. `log_level` defines the maximum log-level the logger
    /// should log.
    ///
    /// Note that installing the logger consumes it. To uninstall you need to call
    /// `log::shutdown_logger` which returns the boxed, original `Logger` instance.
    pub fn install<T: Into<log::LogLevelFilter>>(self, log_level: T) -> Result<()> {
        log::set_logger(|max_level| {
            max_level.set(Into::into(log_level));
            Box::new(self)
        })?;

        Ok(())
    }

    /// Log a message via the logger's transport to a GELF server.
    ///
    /// The logger will automatically all `default_metadata` fields to the message
    /// which are missing in the passed `Message`.
    pub fn log_message(&self, msg: Message) {
        self.backend.log(WireMessage::new(msg, &self));
    }

    /// Return the hostname used for GELF's `host`-field
    pub fn hostname(&self) -> &String {
        &self.hostname
    }

    /// Set the hostname used for GELF's `host`-field
    pub fn set_hostname(&mut self, hostname: String) -> &mut Self {
        self.hostname = hostname;
        self
    }

    /// Return all default metadata
    pub fn default_metadata(&self) -> &HashMap<String, String> {
        &self.default_metadata
    }

    /// Add a default metadata field
    ///
    /// Every logged `Message` is checked for every default_metadata field.
    /// If it contains an entry with the key, the default is ignored. But if
    /// there is no additional information present, the default is added to the message.
    ///
    /// This can be used for example to add a `facility` to every message:
    ///
    /// ``` no_run
    /// let mut logger = Logger::new(..);
    /// logger.add_default_metadata(String::from("facility"), String::from("my_awesome_rust_service"))
    ///
    /// logger.log_message(Message::new("This is important information"), None);
    /// // -> The message will contain an additional field "_facility" with the value "my_awesome_rust_service"
    /// ```
    pub fn add_default_metadata(&mut self, key: String, value: String) -> &mut Self {
        self.default_metadata.insert(key, value);
        self
    }
}

impl log::Log for Logger {
    /// Determines if a log message with the specified metadata would be logged.
    ///
    /// See [docs](https://doc.rust-lang.org/log/log/trait.Log.html#tymethod.enabled)
    /// for more details
    fn enabled(&self, _: &log::LogMetadata) -> bool {
        // The logger does not dicard any log-level by itself, therefore it is
        // always enabled
        true
    }

    /// Logs the `LogRecord`.
    /// See [docs](https://doc.rust-lang.org/log/log/trait.Log.html#tymethod.log)
    /// for more details
    fn log(&self, record: &log::LogRecord) {
        if !self.enabled(record.metadata()) {
            ()
        }

        self.log_message(From::from(record))
    }
}
<commit_msg>Improve Logger tests with runnable/testable docs<commit_after>use std::collections::HashMap;

use log;
use hostname;

use backends::Backend;
use message::{Message, WireMessage};
use errors::{Result, ErrorKind};

/// Logger for sending log-messages in [GELF-format](http://docs.graylog.org/en/2.2/pages/gelf.html)
/// to a GELF server, e.g. [Graylog](https://www.graylog.org/)
///
/// A `Logger` instance can be either used as a standalone object to log directly
/// to a log-server or it can installed as `log`-crate log-handler.
///
/// # Standalone usage
///
/// Note that the simple constructor (`Logger::new`) is used. It determines the
/// local hostname by itself.
///
/// ```
/// # use gelf::NullBackend;
/// use gelf::{Logger, UdpBackend, Message, Level};
///
/// let backend = UdpBackend::new("127.0.0.1:12201").expect("Failed to create UDP backend");
/// # let backend = NullBackend::new();
/// let logger = Logger::new(Box::new(backend)).expect("Failed to determine hostname");
///
/// logger.log_message(Message::new(String::from("Test log message!"), Some(Level::Debug)));
/// ```
///
/// # Usage with the `log`-crate
///
/// We assume the `log`-crate has been imported with macros enabled, like
/// this `#[macro_use] extern crate log;`.
///
/// ```
/// # #[macro_use] extern crate log;
/// # extern crate gelf;
/// # use gelf::NullBackend;
/// use gelf::{Logger, UdpBackend, Message, Level};
/// use log::LogLevelFilter;
///
/// # pub fn main() {
/// let backend = UdpBackend::new("127.0.0.1:12201").expect("Failed to create UDP backend");
/// # let backend = NullBackend::new();
/// let logger = Logger::new(Box::new(backend)).expect("Failed to determine hostname");
///
/// logger.install(LogLevelFilter::Trace).expect("Failed to install logger");
///
/// debug!("Debug message");
/// warn!("Warning");
/// # }
/// ```
///
/// # Further information
///
/// For more information check out the `examples/` directory.
pub struct Logger {
    hostname: String,
    backend: Box<Backend>,
    default_metadata: HashMap<String, String>,
}

impl Logger {
    /// Construct a new `Logger` instance
    ///
    /// The backend needs to be boxed for usage as a logger with the `log`-crate.
    /// This constructor tries to determine the local hostname (required by GELF)
    /// with the help of the `hostname`-crate. If you want to set a custom hostname
    /// check out the `Logger::new_with_hostname` constructor
    pub fn new(backend: Box<Backend>) -> Result<Self> {
        hostname::get_hostname()
            .map(|hostname| Logger::new_with_hostname(backend, &hostname))
            .ok_or(ErrorKind::LoggerCreateFailed("Failed to determine local hostname").into())
    }

    /// Construct a new `Logger` instance with predetermined hostname
    ///
    /// The backend needs to be boxed for usage as a logger with the `log`-crate. It
    /// uses the passed hostname for the GELF `host` field
    pub fn new_with_hostname(backend: Box<Backend>, hostname: &str) -> Logger {
        Logger {
            hostname: String::from(hostname),
            backend: backend,
            default_metadata: HashMap::new(),
        }
    }

    /// Install a logger instance as a `log`-Logger
    ///
    /// This method wraps `log::set_logger` as a convenience function as required
    /// by the `log`-crate. `log_level` defines the maximum log-level the logger
    /// should log.
    ///
    /// Note that installing the logger consumes it. To uninstall you need to call
    /// `log::shutdown_logger` which returns the boxed, original `Logger` instance.
    pub fn install<T: Into<log::LogLevelFilter>>(self, log_level: T) -> Result<()> {
        log::set_logger(|max_level| {
            max_level.set(Into::into(log_level));
            Box::new(self)
        })?;

        Ok(())
    }

    /// Log a message via the logger's transport to a GELF server.
    ///
    /// The logger will automatically all `default_metadata` fields to the message
    /// which are missing in the passed `Message`.
    pub fn log_message(&self, msg: Message) {
        self.backend.log(WireMessage::new(msg, &self));
    }

    /// Return the hostname used for GELF's `host`-field
    pub fn hostname(&self) -> &String {
        &self.hostname
    }

    /// Set the hostname used for GELF's `host`-field
    pub fn set_hostname(&mut self, hostname: String) -> &mut Self {
        self.hostname = hostname;
        self
    }

    /// Return all default metadata
    pub fn default_metadata(&self) -> &HashMap<String, String> {
        &self.default_metadata
    }

    /// Add a default metadata field
    ///
    /// Every logged `Message` is checked for every default_metadata field.
    /// If it contains an entry with the key, the default is ignored. But if
    /// there is no additional information present, the default is added to the message.
    ///
    /// This can be used for example to add a `facility` to every message:
    ///
    /// ```
    /// # use gelf::{Logger, NullBackend, Message};
    /// # let backend = NullBackend::new();
    /// # let mut logger = Logger::new(Box::new(backend)).unwrap();
    /// logger.add_default_metadata(String::from("facility"), String::from("my_awesome_rust_service"));
    ///
    /// logger.log_message(Message::new(String::from("This is important information"), None));
    /// // -> The message will contain an additional field "_facility" with the value "my_awesome_rust_service"
    /// ```
    pub fn add_default_metadata(&mut self, key: String, value: String) -> &mut Self {
        self.default_metadata.insert(key, value);
        self
    }
}

impl log::Log for Logger {
    /// Determines if a log message with the specified metadata would be logged.
    ///
    /// See [docs](https://doc.rust-lang.org/log/log/trait.Log.html#tymethod.enabled)
    /// for more details
    fn enabled(&self, _: &log::LogMetadata) -> bool {
        // The logger does not dicard any log-level by itself, therefore it is
        // always enabled
        true
    }

    /// Logs the `LogRecord`.
    /// See [docs](https://doc.rust-lang.org/log/log/trait.Log.html#tymethod.log)
    /// for more details
    fn log(&self, record: &log::LogRecord) {
        if !self.enabled(record.metadata()) {
            ()
        }

        self.log_message(From::from(record))
    }
}
<|endoftext|>


<commit_before>//! A filesystem logger for use (primarily) with the DirectIP server.
//!
//! This logger is very dumb, as it was hacked together just to provide *some* sort of logging for
//! the sever. In particular, it always it set to `Debug`, and has a bunch of panics in its `log`
//! method. A more robust implementation would be better about this stuff, but as I don't need more
//! functionality yet I'm going to keep it simple until I find out otherwise.

use std::fs::OpenOptions;
use std::io::Write;
use std::path::Path;

use chrono::UTC;

use log::{Log, LogLevel, LogLevelFilter, LogMetadata, LogRecord, set_logger, SetLoggerError};

pub fn init<P: 'static + AsRef<Path> + Send + Sync>(path: P) -> Result<(), SetLoggerError> {
    set_logger(|max_log_level| {
        max_log_level.set(LogLevelFilter::Debug);
        Box::new(Logger { path: path })
    })
}

struct Logger<P: AsRef<Path>> {
    path: P,
}

impl<P: AsRef<Path> + Send + Sync> Log for Logger<P> {
    fn enabled(&self, metadata: &LogMetadata) -> bool {
        metadata.level() <= LogLevel::Debug
    }

    /// Log a message.
    ///
    /// This function has some panics in it. I'm not sure of the "right" way to handle exceptional
    /// situaions in this logging module. Part of me wants to ignore everything, since logging
    /// should not interfere with the functioning of the program as a whole. However, since I'm in
    /// dev mode for the whole system, silent logs might be worse than a crashing program. For now,
    /// I'll keep the panics, but with the idea that I need to fix this in the future.
    fn log(&self, record: &LogRecord) {
        if self.enabled(record.metadata()) {
            let mut file = OpenOptions::new()
                .create(true)
                .write(true)
                .append(true)
                .open(&self.path).unwrap();
            file.write_all(format!("({}) {}: {}\n", UTC::now().format("%Y-%m-%d %H:%M:%S"),
                record.level(), record.args()).as_bytes()).unwrap();
        }
    }
}
<commit_msg>Document logger::init<commit_after>//! A filesystem logger for use (primarily) with the DirectIP server.
//!
//! This logger is very dumb, as it was hacked together just to provide *some* sort of logging for
//! the sever. In particular, it always it set to `Debug`, and has a bunch of panics in its `log`
//! method. A more robust implementation would be better about this stuff, but as I don't need more
//! functionality yet I'm going to keep it simple until I find out otherwise.

use std::fs::OpenOptions;
use std::io::Write;
use std::path::Path;

use chrono::UTC;

use log::{Log, LogLevel, LogLevelFilter, LogMetadata, LogRecord, set_logger, SetLoggerError};

/// Initialize a new logger.
///
/// This logger will write all messages with level Debug or less to the file specified by `path`.
pub fn init<P: 'static + AsRef<Path> + Send + Sync>(path: P) -> Result<(), SetLoggerError> {
    set_logger(|max_log_level| {
        max_log_level.set(LogLevelFilter::Debug);
        Box::new(Logger { path: path })
    })
}

struct Logger<P: AsRef<Path>> {
    path: P,
}

impl<P: AsRef<Path> + Send + Sync> Log for Logger<P> {
    fn enabled(&self, metadata: &LogMetadata) -> bool {
        metadata.level() <= LogLevel::Debug
    }

    /// Log a message.
    ///
    /// This function has some panics in it. I'm not sure of the "right" way to handle exceptional
    /// situaions in this logging module. Part of me wants to ignore everything, since logging
    /// should not interfere with the functioning of the program as a whole. However, since I'm in
    /// dev mode for the whole system, silent logs might be worse than a crashing program. For now,
    /// I'll keep the panics, but with the idea that I need to fix this in the future.
    fn log(&self, record: &LogRecord) {
        if self.enabled(record.metadata()) {
            let mut file = OpenOptions::new()
                .create(true)
                .write(true)
                .append(true)
                .open(&self.path).unwrap();
            file.write_all(format!("({}) {}: {}\n", UTC::now().format("%Y-%m-%d %H:%M:%S"),
                record.level(), record.args()).as_bytes()).unwrap();
        }
    }
}
<|endoftext|>


<commit_before>use crate::common::util::*;

#[test]
#[cfg(unix)]
fn test_groups() {
    if !is_ci() {
        new_ucmd!().succeeds().stdout_is(expected_result(&[]));
    } else {
        // TODO: investigate how this could be tested in CI
        // stderr = groups: cannot find name for group ID 116
        println!("test skipped:");
    }
}

#[test]
#[cfg(unix)]
#[ignore = "fixme: 'groups USERNAME' needs more debugging"]
fn test_groups_username() {
    let scene = TestScenario::new(util_name!());
    let whoami_result = scene.cmd("whoami").run();

    let username = if whoami_result.succeeded() {
        whoami_result.stdout_move_str()
    } else if is_ci() {
        String::from("docker")
    } else {
        println!("test skipped:");
        return;
    };

    // TODO: stdout should be in the form: "username : group1 group2 group3"

    scene
        .ucmd()
        .arg(&username)
        .succeeds()
        .stdout_is(expected_result(&[&username]));
}

#[cfg(unix)]
fn expected_result(args: &[&str]) -> String {
    // We want to use GNU id. On most linux systems, this is "id", but on
    // bsd-like systems (e.g. FreeBSD, MacOS), it is commonly "gid".
    #[cfg(any(target_os = "linux"))]
    let util_name = "id";
    #[cfg(not(target_os = "linux"))]
    let util_name = "gid";

    TestScenario::new(&util_name)
        .cmd_keepenv(util_name)
        .env("LANGUAGE", "C")
        .args(args)
        .args(&["-Gn"])
        .succeeds()
        .stdout_move_str()
}
<commit_msg>groups: fix clippy lint in test<commit_after>use crate::common::util::*;

#[test]
#[cfg(unix)]
fn test_groups() {
    if !is_ci() {
        new_ucmd!().succeeds().stdout_is(expected_result(&[]));
    } else {
        // TODO: investigate how this could be tested in CI
        // stderr = groups: cannot find name for group ID 116
        println!("test skipped:");
    }
}

#[test]
#[cfg(unix)]
#[ignore = "fixme: 'groups USERNAME' needs more debugging"]
fn test_groups_username() {
    let scene = TestScenario::new(util_name!());
    let whoami_result = scene.cmd("whoami").run();

    let username = if whoami_result.succeeded() {
        whoami_result.stdout_move_str()
    } else if is_ci() {
        String::from("docker")
    } else {
        println!("test skipped:");
        return;
    };

    // TODO: stdout should be in the form: "username : group1 group2 group3"

    scene
        .ucmd()
        .arg(&username)
        .succeeds()
        .stdout_is(expected_result(&[&username]));
}

#[cfg(unix)]
fn expected_result(args: &[&str]) -> String {
    // We want to use GNU id. On most linux systems, this is "id", but on
    // bsd-like systems (e.g. FreeBSD, MacOS), it is commonly "gid".
    #[cfg(any(target_os = "linux"))]
    let util_name = "id";
    #[cfg(not(target_os = "linux"))]
    let util_name = "gid";

    TestScenario::new(util_name)
        .cmd_keepenv(util_name)
        .env("LANGUAGE", "C")
        .args(args)
        .args(&["-Gn"])
        .succeeds()
        .stdout_move_str()
}
<|endoftext|>


<commit_before>#![feature(globs)]

#![experimental]
#![crate_id = "lrustc"]
#![desc = "lrustc, the Literate Rust extension"]
#![license = "MIT/ASL2"]

#![feature(globs)]
#![feature(phase)]

extern crate getopts;
extern crate regex;
extern crate debug;
#[phase(plugin, link)] extern crate log;
use std::os;

mod liblrustc;

pub fn opts() -> Vec<getopts::OptGroup> {
    use getopts::optopt;
    vec!(optopt("d", "dump", "Dump a regular rust file", "PATH"),)
}

fn main() {
    use getopts::getopts;
    let args: Vec<String> = os::args();
    let matches = match getopts(args.tail(), opts().as_slice()) {
        Ok(m) => { m }
        Err(f) => { fail!(f.to_str()) }
    };
    if matches.opt_present("d") {

    }
}
<commit_msg>remove redundant feature gate<commit_after>#![experimental]
#![crate_id = "lrustc"]
#![desc = "lrustc, the Literate Rust extension"]
#![license = "MIT/ASL2"]

#![feature(globs)]
#![feature(phase)]

extern crate getopts;
extern crate regex;
extern crate debug;
#[phase(plugin, link)] extern crate log;
use std::os;

mod liblrustc;

pub fn opts() -> Vec<getopts::OptGroup> {
    use getopts::optopt;
    vec!(optopt("d", "dump", "Dump a regular rust file", "PATH"),)
}

fn main() {
    use getopts::getopts;
    let args: Vec<String> = os::args();
    let matches = match getopts(args.tail(), opts().as_slice()) {
        Ok(m) => { m }
        Err(f) => { fail!(f.to_str()) }
    };
    if matches.opt_present("d") {

    }
}
<|endoftext|>


<commit_before>/// Macro to print the current stats of ArrayFire's memory manager.
///
/// `mem_info!` print 4 values:
///
///  Name                    | Description
/// -------------------------|-------------------------
///  Allocated Bytes         | Total number of bytes allocated by the memory manager
///  Allocated Buffers       | Total number of buffers allocated
///  Locked (In Use) Bytes   | Number of bytes that are in use by active arrays
///  Locked (In Use) Buffers | Number of buffers that are in use by active arrays
///
///  The `Allocated Bytes` is always a multiple of the memory step size. The
///  default step size is 1024 bytes. This means when a buffer is to be
///  allocated, the size is always rounded up to a multiple of the step size.
///  You can use [get_mem_step_size](./fn.get_mem_step_size.html) to check the
///  current step size and [set_mem_step_size](./fn.set_mem_step_size.html) to
///  set a custom resolution size.
///
///  The `Allocated Buffers` is the number of buffers that use up the allocated
///  bytes. This includes buffers currently in scope, as well as buffers marked
///  as free, ie, from arrays gone out of scope. The free buffers are available
///  for use by new arrays that might be created.
///
///  The `Locked Bytes` is the number of bytes in use that cannot be
///  reallocated at the moment. The difference of Allocated Bytes and Locked
///  Bytes is the total bytes available for reallocation.
///
///  The `Locked Buffers` is the number of buffer in use that cannot be
///  reallocated at the moment. The difference of Allocated Buffers and Locked
///  Buffers is the number of buffers available for reallocation.
///
/// # Parameters
///
/// - `msg` is the message that is printed to screen before printing stats
///
/// # Examples
///
/// ```rust
/// use arrayfire::{Dim4, device_mem_info, print, randu, mem_info};
///
/// let dims = Dim4::new(&[5, 5, 1, 1]);
/// let a = randu::<f32>(dims);
/// print(&a);
/// mem_info!("Hello!");
/// ```
///
/// Sample Output:
///
/// ```text
/// AF Memory: Here
/// Allocated [ Bytes | Buffers ] = [ 4096 | 4 ]
/// In Use    [ Bytes | Buffers ] = [ 2048 | 2 ]
/// ```
#[macro_export]
macro_rules! mem_info {
    [$msg: expr] => {
        {
            let (abytes, abuffs, lbytes, lbuffs) = device_mem_info();
            println!("AF Memory: {:?}", $msg);
            println!("Allocated [Bytes | Buffers] = [ {} | {} ]", abytes, abuffs);
            println!("In Use    [Bytes | Buffers] = [ {} | {} ]", lbytes, lbuffs);
        }
    };
}

/// Join multiple Arrays along a given dimension
///
/// All the Arrays provided to this macro should be of type `&Array`
///
/// # Examples
///
/// ```rust
/// use arrayfire::{Dim4, join_many, print, randu};
///
/// let a = &randu::<f32>(Dim4::new(&[5, 3, 1, 1]));
/// let b = &randu::<f32>(Dim4::new(&[5, 3, 1, 1]));
/// let c = &randu::<f32>(Dim4::new(&[5, 3, 1, 1]));
/// let d = join_many![2; a, b, c];
/// print(&d);
/// ```
///
/// # Panics
///
/// This macro just calls [join_many](./fn.join_many.html) function after collecting all
/// the input arrays into a vector.
// Using macro to implement join many wrapper
#[macro_export]
macro_rules! join_many {
    [$dim: expr; $($x:ident),+] => {
        {
            let mut temp_vec = Vec::new();
            $(
                temp_vec.push($x);
             )*
            join_many($dim, temp_vec)
        }
    };
}

/// Print given message before printing out the Array to standard output
///
/// # Examples
///
/// ```rust
/// use arrayfire::{Dim4, print_gen, randu, af_print};
/// let dims = Dim4::new(&[3, 1, 1, 1]);
/// let a = randu::<f32>(dims);
/// af_print!("Create a 5-by-3 matrix of random floats on the GPU", a);
/// ```
///
#[macro_export]
macro_rules! af_print {
    [$msg: expr, $x: ident] => {
        {
            print_gen(String::from($msg), &$x, Some(4));
        }
    };
}

/// Evaluate arbitrary number of arrays
#[macro_export]
macro_rules! eval {
    [$($x:ident),+] => {
        {
            let mut temp_vec = Vec::new();
            $(
                temp_vec.push($x);
             )*
            eval_multiple(temp_vec)
        }
    };
}
<commit_msg>Update macros.rs<commit_after>/// Macro to print the current stats of ArrayFire's memory manager.
///
/// `mem_info!` print 4 values:
///
///  Name                    | Description
/// -------------------------|-------------------------
///  Allocated Bytes         | Total number of bytes allocated by the memory manager
///  Allocated Buffers       | Total number of buffers allocated
///  Locked (In Use) Bytes   | Number of bytes that are in use by active arrays
///  Locked (In Use) Buffers | Number of buffers that are in use by active arrays
///
///  The `Allocated Bytes` is always a multiple of the memory step size. The
///  default step size is 1024 bytes. This means when a buffer is to be
///  allocated, the size is always rounded up to a multiple of the step size.
///  You can use [get_mem_step_size](./fn.get_mem_step_size.html) to check the
///  current step size and [set_mem_step_size](./fn.set_mem_step_size.html) to
///  set a custom resolution size.
///
///  The `Allocated Buffers` is the number of buffers that use up the allocated
///  bytes. This includes buffers currently in scope, as well as buffers marked
///  as free, ie, from arrays gone out of scope. The free buffers are available
///  for use by new arrays that might be created.
///
///  The `Locked Bytes` is the number of bytes in use that cannot be
///  reallocated at the moment. The difference of Allocated Bytes and Locked
///  Bytes is the total bytes available for reallocation.
///
///  The `Locked Buffers` is the number of buffer in use that cannot be
///  reallocated at the moment. The difference of Allocated Buffers and Locked
///  Buffers is the number of buffers available for reallocation.
///
/// # Parameters
///
/// - `msg` is the message that is printed to screen before printing stats
///
/// # Examples
///
/// ```rust
/// use arrayfire::{Dim4, device_mem_info, print, randu, mem_info};
///
/// let dims = Dim4::new(&[5, 5, 1, 1]);
/// let a = randu::<f32>(dims);
/// print(&a);
/// mem_info!("Hello!");
/// ```
///
/// Sample Output:
///
/// ```text
/// AF Memory: Here
/// Allocated [ Bytes | Buffers ] = [ 4096 | 4 ]
/// In Use    [ Bytes | Buffers ] = [ 2048 | 2 ]
/// ```
#[macro_export]
macro_rules! mem_info {
    [$msg: expr] => {
        {
            let (abytes, abuffs, lbytes, lbuffs) = device_mem_info();
            println!("AF Memory: {:?}", $msg);
            println!("Allocated [Bytes | Buffers] = [ {} | {} ]", abytes, abuffs);
            println!("In Use    [Bytes | Buffers] = [ {} | {} ]", lbytes, lbuffs);
        }
    };
}

/// Join multiple Arrays along a given dimension
///
/// All the Arrays provided to this macro should be of type `&Array`
///
/// # Examples
///
/// ```rust
/// use arrayfire::{Dim4, join_many, print, randu};
///
/// let a = &randu::<f32>(Dim4::new(&[5, 3, 1, 1]));
/// let b = &randu::<f32>(Dim4::new(&[5, 3, 1, 1]));
/// let c = &randu::<f32>(Dim4::new(&[5, 3, 1, 1]));
/// let d = join_many![2; a, b, c];
/// print(&d);
/// ```
///
/// # Panics
///
/// This macro just calls [join_many](./fn.join_many.html) function after collecting all
/// the input arrays into a vector.
// Using macro to implement join many wrapper
#[macro_export]
macro_rules! join_many {
    [$dim: expr; $($x:expr),+] => {
        {
            let mut temp_vec = Vec::new();
            $(
                temp_vec.push($x);
             )*
            join_many($dim, temp_vec)
        }
    };
}

/// Print given message before printing out the Array to standard output
///
/// # Examples
///
/// ```rust
/// use arrayfire::{Dim4, print_gen, randu, af_print};
/// let dims = Dim4::new(&[3, 1, 1, 1]);
/// let a = randu::<f32>(dims);
/// af_print!("Create a 5-by-3 matrix of random floats on the GPU", a);
/// ```
///
#[macro_export]
macro_rules! af_print {
    [$msg: expr, $x: expr] => {
        {
            print_gen(String::from($msg), &$x, Some(4));
        }
    };
}

/// Evaluate arbitrary number of arrays
#[macro_export]
macro_rules! eval {
    [$($x:expr),+] => {
        {
            let mut temp_vec = Vec::new();
            $(
                temp_vec.push($x);
             )*
            eval_multiple(temp_vec)
        }
    };
}
<|endoftext|>