<commit_before>//! A macro for writing HTML templates.
//!
//! This crate, along with its sister `maud_macros`, lets you generate
//! HTML markup from within Rust. It exposes a single macro, `html!`,
//! which compiles your markup to specialized Rust code.
//!
//! # Dependencies
//!
//! To get started, add `maud` and `maud_macros` to your `Cargo.toml`:
//!
//! ```toml
//! [dependencies]
//! # ...
//! maud = "*"
//! maud_macros = "*"
//! ```
//!
//! # Example
//!
//! ```rust
//! #![feature(plugin)]
//! #![allow(unstable)]
//!
//! extern crate maud;
//! #[plugin] #[no_link] extern crate maud_macros;
//!
//! fn main() {
//!     let name = "Lyra";
//!     let markup = html! {
//!         p { "Hi, " $name "!" }
//!     };
//!     assert_eq!(markup.render(), "<p>Hi, Lyra!</p>");
//! }
//! ```
//!
//! # Syntax
//!
//! Maud uses its own syntax, loosely inspired by HTML and Rust.
//!
//! **Note:** The markup you see below has been cleaned up a bit. In
//! reality, Maud doesn't add extra whitespace to the HTML it generates.
//!
//! ## Literals
//!
//! ```
//! html! {
//!     "Oatmeal, are you crazy?\n"
//!     "<script>alert(\"XSS\")</script>\n"
//!     $$"<script>alert(\"XSS\")</script>\n"
//! }
//! ```
//!
//! ```html
//! Oatmeal, are you crazy?
//! <script>alert("XSS")</script>
//! <script>alert("XSS")</script>
//! ```
//!
//! Literal strings use the same syntax as Rust. Wrap them in
//! double quotes, and use a backslash for escapes.
//!
//! By default, HTML special characters are escaped automatically. Add a
//! `$$` prefix to disable this escaping.
//!
//! ## Elements
//!
//! ```
//! html! {
//!     h1 "Pinkie's Brew"
//!     p {
//!         "Watch as I work my gypsy magic"
//!         br /
//!         "Eye of a newt and cinnamon"
//!     }
//!     p small em "squee"
//! }
//! ```
//!
//! ```html
//! <h1>Pinkie's Brew</h1>
//! <p>
//!   Watch as I work my gypsy magic
//!   <br>
//!   Eye of a newt and cinnamon
//! </p>
//! <p><small><em>squee</em></small></p>
//! ```
//!
//! Write an element using curly braces (`p {}`).
//!
//! Terminate a void element using a slash (`br /`).
//!
//! If the element has only a single child, you can omit the brackets
//! (`h1 "Pinkie's Brew"`). This shorthand works with nested elements
//! too.
//!
//! ## Attributes
//!
//! ```
//! html! {
//!     form method="POST" {
//!         label for="waffles" "Do you like waffles?"
//!         input name="waffles" type="checkbox" checked=! /
//!     }
//! }
//! ```
//!
//! ```html
//! <form method="POST">
//!   <label for="waffles">Do you like waffles?</label>
//!   <input name="waffles" type="checkbox" checked>
//! </form>
//! ```
//!
//! Add attributes using the syntax `attr="value"`. Attributes must be
//! quoted: they are parsed as string literals.
//!
//! To declare an empty attribute, use `!` for the value: `checked=!`.
//!
//! ## Splices
//!
//! ```
//! let best_pony = "Pinkie Pie";
//! let numbers = [1, 2, 3, 4];
//! let secret_message = "Surprise!";
//! let pre_escaped = "<p>Pre-escaped</p>";
//! html! {
//!     h1 { $best_pony " says:" }
//!     p {
//!         "I have " $numbers.len() " numbers, "
//!         "and the first one is " $numbers[0]
//!     }
//!     p title=$secret_message {
//!         "1 + 1 = " $(1 + 1)
//!     }
//!     $$pre_escaped
//! }
//! ```
//!
//! ```html
//! <h1>Pinkie Pie says:</h1>
//! <p>I have 4 numbers, and the first one is 1</p>
//! <p title="Surprise!">1 + 1 = 2</p>
//! <p>Pre-escaped</p>
//! ```
//!
//! Use `$(expr)` syntax to splice a Rust expression into the output.
//! You can omit the brackets if it's just a variable (`$var`), indexing
//! operation (`$var[0]`), method call (`$var.method()`), or property
//! lookup (`$var.property`).
//!
//! As with literals, expression values are escaped by default. Use a
//! `$$` prefix to disable this behavior.

#![allow(unstable)]

use std::fmt;
use std::io::{IoError, IoErrorKind, IoResult};

/// Escape an HTML value.
pub fn escape(s: &str) -> String {
    use std::fmt::Writer;
    let mut buf = String::new();
    rt::Escaper { inner: &mut buf }.write_str(s).unwrap();
    buf
}

/// A block of HTML markup, as returned by the `html!` macro.
///
/// Use `.render()` to convert it to a `String`, or `.render_to()` to
/// write it directly to a handle.
pub struct Markup<'a> {
    callback: &'a (Fn(&mut fmt::Writer) -> fmt::Result + 'a),
}

impl<'a> Markup<'a> {
    /// Render the markup to a `String`.
    pub fn render(&self) -> String {
        let mut buf = String::new();
        self.render_fmt(&mut buf).unwrap();
        buf
    }

    /// Render the markup to a `std::io::Writer`.
    ///
    /// If you're writing to, say, a file or a socket, this may be more
    /// efficient than calling `.render()`.
    pub fn render_to(&self, w: &mut Writer) -> IoResult<()> {
        struct WriterWrapper<'a, 'b: 'a> {
            inner: &'a mut (Writer + 'b),
        }
        impl<'a, 'b> fmt::Writer for WriterWrapper<'a, 'b> {
            fn write_str(&mut self, s: &str) -> fmt::Result {
                self.inner.write_str(s).map_err(|_| fmt::Error)
            }
        }
        self.render_fmt(&mut WriterWrapper { inner: w })
            .map_err(|_| IoError {
                kind: IoErrorKind::OtherIoError,
                desc: "formatting error",
                detail: None,
            })
    }

    /// Render the markup to a `std::fmt::Writer`.
    pub fn render_fmt(&self, w: &mut fmt::Writer) -> fmt::Result {
        (self.callback)(w)
    }
}

/// Internal functions used by the `maud_macros` package. You should
/// never need to call these directly.
#[doc(hidden)]
pub mod rt {
    use std::fmt;
    use super::Markup;

    #[inline]
    pub fn make_markup<'a>(f: &'a (Fn(&mut fmt::Writer) -> fmt::Result + 'a)) -> Markup<'a> {
        Markup { callback: f }
    }

    /// rustc is a butt and doesn't let us quote macro invocations
    /// directly. So we factor the `write!` call into a separate
    /// function and use that instead.
    ///
    /// See <https://github.com/rust-lang/rust/issues/16617>
    #[inline]
    pub fn write_fmt<T: fmt::Display>(w: &mut fmt::Writer, value: T) -> fmt::Result {
        write!(w, "{}", value)
    }

    pub struct Escaper<'a, 'b: 'a> {
        pub inner: &'a mut (fmt::Writer + 'b),
    }

    impl<'a, 'b> fmt::Writer for Escaper<'a, 'b> {
        fn write_str(&mut self, s: &str) -> fmt::Result {
            for c in s.chars() {
                try!(match c {
                    '&' => self.inner.write_str("&"),
                    '<' => self.inner.write_str("<"),
                    '>' => self.inner.write_str(">"),
                    '"' => self.inner.write_str("""),
                    '\'' => self.inner.write_str("'"),
                    _ => write!(self.inner, "{}", c),
                });
            }
            Ok(())
        }
    }
}
<commit_msg>Mention Display trait in docs<commit_after>//! A macro for writing HTML templates.
//!
//! This crate, along with its sister `maud_macros`, lets you generate
//! HTML markup from within Rust. It exposes a single macro, `html!`,
//! which compiles your markup to specialized Rust code.
//!
//! # Dependencies
//!
//! To get started, add `maud` and `maud_macros` to your `Cargo.toml`:
//!
//! ```toml
//! [dependencies]
//! # ...
//! maud = "*"
//! maud_macros = "*"
//! ```
//!
//! # Example
//!
//! ```rust
//! #![feature(plugin)]
//! #![allow(unstable)]
//!
//! extern crate maud;
//! #[plugin] #[no_link] extern crate maud_macros;
//!
//! fn main() {
//!     let name = "Lyra";
//!     let markup = html! {
//!         p { "Hi, " $name "!" }
//!     };
//!     assert_eq!(markup.render(), "<p>Hi, Lyra!</p>");
//! }
//! ```
//!
//! # Syntax
//!
//! Maud uses its own syntax, loosely inspired by HTML and Rust.
//!
//! **Note:** The markup you see below has been cleaned up a bit. In
//! reality, Maud doesn't add extra whitespace to the HTML it generates.
//!
//! ## Literals
//!
//! ```
//! html! {
//!     "Oatmeal, are you crazy?\n"
//!     "<script>alert(\"XSS\")</script>\n"
//!     $$"<script>alert(\"XSS\")</script>\n"
//! }
//! ```
//!
//! ```html
//! Oatmeal, are you crazy?
//! <script>alert("XSS")</script>
//! <script>alert("XSS")</script>
//! ```
//!
//! Literal strings use the same syntax as Rust. Wrap them in
//! double quotes, and use a backslash for escapes.
//!
//! By default, HTML special characters are escaped automatically. Add a
//! `$$` prefix to disable this escaping.
//!
//! ## Elements
//!
//! ```
//! html! {
//!     h1 "Pinkie's Brew"
//!     p {
//!         "Watch as I work my gypsy magic"
//!         br /
//!         "Eye of a newt and cinnamon"
//!     }
//!     p small em "squee"
//! }
//! ```
//!
//! ```html
//! <h1>Pinkie's Brew</h1>
//! <p>
//!   Watch as I work my gypsy magic
//!   <br>
//!   Eye of a newt and cinnamon
//! </p>
//! <p><small><em>squee</em></small></p>
//! ```
//!
//! Write an element using curly braces (`p {}`).
//!
//! Terminate a void element using a slash (`br /`).
//!
//! If the element has only a single child, you can omit the brackets
//! (`h1 "Pinkie's Brew"`). This shorthand works with nested elements
//! too.
//!
//! ## Attributes
//!
//! ```
//! html! {
//!     form method="POST" {
//!         label for="waffles" "Do you like waffles?"
//!         input name="waffles" type="checkbox" checked=! /
//!     }
//! }
//! ```
//!
//! ```html
//! <form method="POST">
//!   <label for="waffles">Do you like waffles?</label>
//!   <input name="waffles" type="checkbox" checked>
//! </form>
//! ```
//!
//! Add attributes using the syntax `attr="value"`. Attributes must be
//! quoted: they are parsed as string literals.
//!
//! To declare an empty attribute, use `!` for the value: `checked=!`.
//!
//! ## Splices
//!
//! ```
//! let best_pony = "Pinkie Pie";
//! let numbers = [1, 2, 3, 4];
//! let secret_message = "Surprise!";
//! let pre_escaped = "<p>Pre-escaped</p>";
//! html! {
//!     h1 { $best_pony " says:" }
//!     p {
//!         "I have " $numbers.len() " numbers, "
//!         "and the first one is " $numbers[0]
//!     }
//!     p title=$secret_message {
//!         "1 + 1 = " $(1 + 1)
//!     }
//!     $$pre_escaped
//! }
//! ```
//!
//! ```html
//! <h1>Pinkie Pie says:</h1>
//! <p>I have 4 numbers, and the first one is 1</p>
//! <p title="Surprise!">1 + 1 = 2</p>
//! <p>Pre-escaped</p>
//! ```
//!
//! Use `$(expr)` syntax to splice a Rust expression into the output.
//! The expression may evaluate to anything that implements
//! `std::fmt::Display`.
//!
//! You can omit the brackets if it's just a variable (`$var`), indexing
//! operation (`$var[0]`), method call (`$var.method()`), or property
//! lookup (`$var.property`).
//!
//! As with literals, expression values are escaped by default. Use a
//! `$$` prefix to disable this behavior.

#![allow(unstable)]

use std::fmt;
use std::io::{IoError, IoErrorKind, IoResult};

/// Escape an HTML value.
pub fn escape(s: &str) -> String {
    use std::fmt::Writer;
    let mut buf = String::new();
    rt::Escaper { inner: &mut buf }.write_str(s).unwrap();
    buf
}

/// A block of HTML markup, as returned by the `html!` macro.
///
/// Use `.render()` to convert it to a `String`, or `.render_to()` to
/// write it directly to a handle.
pub struct Markup<'a> {
    callback: &'a (Fn(&mut fmt::Writer) -> fmt::Result + 'a),
}

impl<'a> Markup<'a> {
    /// Render the markup to a `String`.
    pub fn render(&self) -> String {
        let mut buf = String::new();
        self.render_fmt(&mut buf).unwrap();
        buf
    }

    /// Render the markup to a `std::io::Writer`.
    ///
    /// If you're writing to, say, a file or a socket, this may be more
    /// efficient than calling `.render()`.
    pub fn render_to(&self, w: &mut Writer) -> IoResult<()> {
        struct WriterWrapper<'a, 'b: 'a> {
            inner: &'a mut (Writer + 'b),
        }
        impl<'a, 'b> fmt::Writer for WriterWrapper<'a, 'b> {
            fn write_str(&mut self, s: &str) -> fmt::Result {
                self.inner.write_str(s).map_err(|_| fmt::Error)
            }
        }
        self.render_fmt(&mut WriterWrapper { inner: w })
            .map_err(|_| IoError {
                kind: IoErrorKind::OtherIoError,
                desc: "formatting error",
                detail: None,
            })
    }

    /// Render the markup to a `std::fmt::Writer`.
    pub fn render_fmt(&self, w: &mut fmt::Writer) -> fmt::Result {
        (self.callback)(w)
    }
}

/// Internal functions used by the `maud_macros` package. You should
/// never need to call these directly.
#[doc(hidden)]
pub mod rt {
    use std::fmt;
    use super::Markup;

    #[inline]
    pub fn make_markup<'a>(f: &'a (Fn(&mut fmt::Writer) -> fmt::Result + 'a)) -> Markup<'a> {
        Markup { callback: f }
    }

    /// rustc is a butt and doesn't let us quote macro invocations
    /// directly. So we factor the `write!` call into a separate
    /// function and use that instead.
    ///
    /// See <https://github.com/rust-lang/rust/issues/16617>
    #[inline]
    pub fn write_fmt<T: fmt::Display>(w: &mut fmt::Writer, value: T) -> fmt::Result {
        write!(w, "{}", value)
    }

    pub struct Escaper<'a, 'b: 'a> {
        pub inner: &'a mut (fmt::Writer + 'b),
    }

    impl<'a, 'b> fmt::Writer for Escaper<'a, 'b> {
        fn write_str(&mut self, s: &str) -> fmt::Result {
            for c in s.chars() {
                try!(match c {
                    '&' => self.inner.write_str("&"),
                    '<' => self.inner.write_str("<"),
                    '>' => self.inner.write_str(">"),
                    '"' => self.inner.write_str("""),
                    '\'' => self.inner.write_str("'"),
                    _ => write!(self.inner, "{}", c),
                });
            }
            Ok(())
        }
    }
}
<|endoftext|>


<commit_before>#![allow(missing_copy_implementations)]

#[cfg(test)]
mod tests;

use crate::fmt;
use crate::io::{self, BufRead, ErrorKind, Initializer, IoSlice, IoSliceMut, Read, Write};
use crate::mem::MaybeUninit;

/// Copies the entire contents of a reader into a writer.
///
/// This function will continuously read data from `reader` and then
/// write it into `writer` in a streaming fashion until `reader`
/// returns EOF.
///
/// On success, the total number of bytes that were copied from
/// `reader` to `writer` is returned.
///
/// If you’re wanting to copy the contents of one file to another and you’re
/// working with filesystem paths, see the [`fs::copy`] function.
///
/// [`fs::copy`]: crate::fs::copy
///
/// # Errors
///
/// This function will return an error immediately if any call to [`read`] or
/// [`write`] returns an error. All instances of [`ErrorKind::Interrupted`] are
/// handled by this function and the underlying operation is retried.
///
/// [`read`]: Read::read
/// [`write`]: Write::write
///
/// # Examples
///
/// ```
/// use std::io;
///
/// fn main() -> io::Result<()> {
///     let mut reader: &[u8] = b"hello";
///     let mut writer: Vec<u8> = vec![];
///
///     io::copy(&mut reader, &mut writer)?;
///
///     assert_eq!(&b"hello"[..], &writer[..]);
///     Ok(())
/// }
/// ```
#[stable(feature = "rust1", since = "1.0.0")]
pub fn copy<R: ?Sized, W: ?Sized>(reader: &mut R, writer: &mut W) -> io::Result<u64>
where
    R: Read,
    W: Write,
{
    let mut buf = MaybeUninit::<[u8; super::DEFAULT_BUF_SIZE]>::uninit();
    // FIXME(#76092): This is calling `get_mut` and `get_ref` on an uninitialized
    // `MaybeUninit`. Revisit this once we decided whether that is valid or not.
    // This is still technically undefined behavior due to creating a reference
    // to uninitialized data, but within libstd we can rely on more guarantees
    // than if this code were in an external lib.
    unsafe {
        reader.initializer().initialize(buf.assume_init_mut());
    }

    let mut written = 0;
    loop {
        let len = match reader.read(unsafe { buf.assume_init_mut() }) {
            Ok(0) => return Ok(written),
            Ok(len) => len,
            Err(ref e) if e.kind() == ErrorKind::Interrupted => continue,
            Err(e) => return Err(e),
        };
        writer.write_all(unsafe { &buf.assume_init_ref()[..len] })?;
        written += len as u64;
    }
}

/// A reader which is always at EOF.
///
/// This struct is generally created by calling [`empty()`]. Please see
/// the documentation of [`empty()`] for more details.
#[stable(feature = "rust1", since = "1.0.0")]
pub struct Empty {
    _priv: (),
}

/// Constructs a new handle to an empty reader.
///
/// All reads from the returned reader will return [`Ok`]`(0)`.
///
/// # Examples
///
/// A slightly sad example of not reading anything into a buffer:
///
/// ```
/// use std::io::{self, Read};
///
/// let mut buffer = String::new();
/// io::empty().read_to_string(&mut buffer).unwrap();
/// assert!(buffer.is_empty());
/// ```
#[stable(feature = "rust1", since = "1.0.0")]
pub fn empty() -> Empty {
    Empty { _priv: () }
}

#[stable(feature = "rust1", since = "1.0.0")]
impl Read for Empty {
    #[inline]
    fn read(&mut self, _buf: &mut [u8]) -> io::Result<usize> {
        Ok(0)
    }

    #[inline]
    unsafe fn initializer(&self) -> Initializer {
        Initializer::nop()
    }
}
#[stable(feature = "rust1", since = "1.0.0")]
impl BufRead for Empty {
    #[inline]
    fn fill_buf(&mut self) -> io::Result<&[u8]> {
        Ok(&[])
    }
    #[inline]
    fn consume(&mut self, _n: usize) {}
}

#[stable(feature = "std_debug", since = "1.16.0")]
impl fmt::Debug for Empty {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.pad("Empty { .. }")
    }
}

/// A reader which yields one byte over and over and over and over and over and...
///
/// This struct is generally created by calling [`repeat()`]. Please
/// see the documentation of [`repeat()`] for more details.
#[stable(feature = "rust1", since = "1.0.0")]
pub struct Repeat {
    byte: u8,
}

/// Creates an instance of a reader that infinitely repeats one byte.
///
/// All reads from this reader will succeed by filling the specified buffer with
/// the given byte.
///
/// # Examples
///
/// ```
/// use std::io::{self, Read};
///
/// let mut buffer = [0; 3];
/// io::repeat(0b101).read_exact(&mut buffer).unwrap();
/// assert_eq!(buffer, [0b101, 0b101, 0b101]);
/// ```
#[stable(feature = "rust1", since = "1.0.0")]
pub fn repeat(byte: u8) -> Repeat {
    Repeat { byte }
}

#[stable(feature = "rust1", since = "1.0.0")]
impl Read for Repeat {
    #[inline]
    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
        for slot in &mut *buf {
            *slot = self.byte;
        }
        Ok(buf.len())
    }

    #[inline]
    fn read_vectored(&mut self, bufs: &mut [IoSliceMut<'_>]) -> io::Result<usize> {
        let mut nwritten = 0;
        for buf in bufs {
            nwritten += self.read(buf)?;
        }
        Ok(nwritten)
    }

    #[inline]
    fn is_read_vectored(&self) -> bool {
        true
    }

    #[inline]
    unsafe fn initializer(&self) -> Initializer {
        Initializer::nop()
    }
}

#[stable(feature = "std_debug", since = "1.16.0")]
impl fmt::Debug for Repeat {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.pad("Repeat { .. }")
    }
}

/// A writer which will move data into the void.
///
/// This struct is generally created by calling [`sink`]. Please
/// see the documentation of [`sink()`] for more details.
#[stable(feature = "rust1", since = "1.0.0")]
pub struct Sink {
    _priv: (),
}

/// Creates an instance of a writer which will successfully consume all data.
///
/// All calls to [`write`] on the returned instance will return `Ok(buf.len())`
/// and the contents of the buffer will not be inspected.
///
/// [`write`]: Write::write
///
/// # Examples
///
/// ```rust
/// use std::io::{self, Write};
///
/// let buffer = vec![1, 2, 3, 5, 8];
/// let num_bytes = io::sink().write(&buffer).unwrap();
/// assert_eq!(num_bytes, 5);
/// ```
#[stable(feature = "rust1", since = "1.0.0")]
pub fn sink() -> Sink {
    Sink { _priv: () }
}

#[stable(feature = "rust1", since = "1.0.0")]
impl Write for Sink {
    #[inline]
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
        Ok(buf.len())
    }

    #[inline]
    fn write_vectored(&mut self, bufs: &[IoSlice<'_>]) -> io::Result<usize> {
        let total_len = bufs.iter().map(|b| b.len()).sum();
        Ok(total_len)
    }

    #[inline]
    fn is_write_vectored(&self) -> bool {
        true
    }

    #[inline]
    fn flush(&mut self) -> io::Result<()> {
        Ok(())
    }
}

#[stable(feature = "std_debug", since = "1.16.0")]
impl fmt::Debug for Sink {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.pad("Sink { .. }")
    }
}
<commit_msg>Rollup merge of #76242 - RalfJung:read-fixme, r=Dylan-DPC<commit_after>#![allow(missing_copy_implementations)]

#[cfg(test)]
mod tests;

use crate::fmt;
use crate::io::{self, BufRead, ErrorKind, Initializer, IoSlice, IoSliceMut, Read, Write};
use crate::mem::MaybeUninit;

/// Copies the entire contents of a reader into a writer.
///
/// This function will continuously read data from `reader` and then
/// write it into `writer` in a streaming fashion until `reader`
/// returns EOF.
///
/// On success, the total number of bytes that were copied from
/// `reader` to `writer` is returned.
///
/// If you’re wanting to copy the contents of one file to another and you’re
/// working with filesystem paths, see the [`fs::copy`] function.
///
/// [`fs::copy`]: crate::fs::copy
///
/// # Errors
///
/// This function will return an error immediately if any call to [`read`] or
/// [`write`] returns an error. All instances of [`ErrorKind::Interrupted`] are
/// handled by this function and the underlying operation is retried.
///
/// [`read`]: Read::read
/// [`write`]: Write::write
///
/// # Examples
///
/// ```
/// use std::io;
///
/// fn main() -> io::Result<()> {
///     let mut reader: &[u8] = b"hello";
///     let mut writer: Vec<u8> = vec![];
///
///     io::copy(&mut reader, &mut writer)?;
///
///     assert_eq!(&b"hello"[..], &writer[..]);
///     Ok(())
/// }
/// ```
#[stable(feature = "rust1", since = "1.0.0")]
pub fn copy<R: ?Sized, W: ?Sized>(reader: &mut R, writer: &mut W) -> io::Result<u64>
where
    R: Read,
    W: Write,
{
    let mut buf = MaybeUninit::<[u8; super::DEFAULT_BUF_SIZE]>::uninit();
    // FIXME: #42788
    //
    //   - This creates a (mut) reference to a slice of
    //     _uninitialized_ integers, which is **undefined behavior**
    //
    //   - Only the standard library gets to soundly "ignore" this,
    //     based on its privileged knowledge of unstable rustc
    //     internals;
    unsafe {
        reader.initializer().initialize(buf.assume_init_mut());
    }

    let mut written = 0;
    loop {
        let len = match reader.read(unsafe { buf.assume_init_mut() }) {
            Ok(0) => return Ok(written),
            Ok(len) => len,
            Err(ref e) if e.kind() == ErrorKind::Interrupted => continue,
            Err(e) => return Err(e),
        };
        writer.write_all(unsafe { &buf.assume_init_ref()[..len] })?;
        written += len as u64;
    }
}

/// A reader which is always at EOF.
///
/// This struct is generally created by calling [`empty()`]. Please see
/// the documentation of [`empty()`] for more details.
#[stable(feature = "rust1", since = "1.0.0")]
pub struct Empty {
    _priv: (),
}

/// Constructs a new handle to an empty reader.
///
/// All reads from the returned reader will return [`Ok`]`(0)`.
///
/// # Examples
///
/// A slightly sad example of not reading anything into a buffer:
///
/// ```
/// use std::io::{self, Read};
///
/// let mut buffer = String::new();
/// io::empty().read_to_string(&mut buffer).unwrap();
/// assert!(buffer.is_empty());
/// ```
#[stable(feature = "rust1", since = "1.0.0")]
pub fn empty() -> Empty {
    Empty { _priv: () }
}

#[stable(feature = "rust1", since = "1.0.0")]
impl Read for Empty {
    #[inline]
    fn read(&mut self, _buf: &mut [u8]) -> io::Result<usize> {
        Ok(0)
    }

    #[inline]
    unsafe fn initializer(&self) -> Initializer {
        Initializer::nop()
    }
}
#[stable(feature = "rust1", since = "1.0.0")]
impl BufRead for Empty {
    #[inline]
    fn fill_buf(&mut self) -> io::Result<&[u8]> {
        Ok(&[])
    }
    #[inline]
    fn consume(&mut self, _n: usize) {}
}

#[stable(feature = "std_debug", since = "1.16.0")]
impl fmt::Debug for Empty {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.pad("Empty { .. }")
    }
}

/// A reader which yields one byte over and over and over and over and over and...
///
/// This struct is generally created by calling [`repeat()`]. Please
/// see the documentation of [`repeat()`] for more details.
#[stable(feature = "rust1", since = "1.0.0")]
pub struct Repeat {
    byte: u8,
}

/// Creates an instance of a reader that infinitely repeats one byte.
///
/// All reads from this reader will succeed by filling the specified buffer with
/// the given byte.
///
/// # Examples
///
/// ```
/// use std::io::{self, Read};
///
/// let mut buffer = [0; 3];
/// io::repeat(0b101).read_exact(&mut buffer).unwrap();
/// assert_eq!(buffer, [0b101, 0b101, 0b101]);
/// ```
#[stable(feature = "rust1", since = "1.0.0")]
pub fn repeat(byte: u8) -> Repeat {
    Repeat { byte }
}

#[stable(feature = "rust1", since = "1.0.0")]
impl Read for Repeat {
    #[inline]
    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
        for slot in &mut *buf {
            *slot = self.byte;
        }
        Ok(buf.len())
    }

    #[inline]
    fn read_vectored(&mut self, bufs: &mut [IoSliceMut<'_>]) -> io::Result<usize> {
        let mut nwritten = 0;
        for buf in bufs {
            nwritten += self.read(buf)?;
        }
        Ok(nwritten)
    }

    #[inline]
    fn is_read_vectored(&self) -> bool {
        true
    }

    #[inline]
    unsafe fn initializer(&self) -> Initializer {
        Initializer::nop()
    }
}

#[stable(feature = "std_debug", since = "1.16.0")]
impl fmt::Debug for Repeat {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.pad("Repeat { .. }")
    }
}

/// A writer which will move data into the void.
///
/// This struct is generally created by calling [`sink`]. Please
/// see the documentation of [`sink()`] for more details.
#[stable(feature = "rust1", since = "1.0.0")]
pub struct Sink {
    _priv: (),
}

/// Creates an instance of a writer which will successfully consume all data.
///
/// All calls to [`write`] on the returned instance will return `Ok(buf.len())`
/// and the contents of the buffer will not be inspected.
///
/// [`write`]: Write::write
///
/// # Examples
///
/// ```rust
/// use std::io::{self, Write};
///
/// let buffer = vec![1, 2, 3, 5, 8];
/// let num_bytes = io::sink().write(&buffer).unwrap();
/// assert_eq!(num_bytes, 5);
/// ```
#[stable(feature = "rust1", since = "1.0.0")]
pub fn sink() -> Sink {
    Sink { _priv: () }
}

#[stable(feature = "rust1", since = "1.0.0")]
impl Write for Sink {
    #[inline]
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
        Ok(buf.len())
    }

    #[inline]
    fn write_vectored(&mut self, bufs: &[IoSlice<'_>]) -> io::Result<usize> {
        let total_len = bufs.iter().map(|b| b.len()).sum();
        Ok(total_len)
    }

    #[inline]
    fn is_write_vectored(&self) -> bool {
        true
    }

    #[inline]
    fn flush(&mut self) -> io::Result<()> {
        Ok(())
    }
}

#[stable(feature = "std_debug", since = "1.16.0")]
impl fmt::Debug for Sink {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.pad("Sink { .. }")
    }
}
<|endoftext|>


<commit_before>use core::cmp;
use core::mem;
use core::ops::{Deref, DerefMut};

use io::{Read, Write, Result};

use thread::{Scheduler, WaitQueue};

use hal::xen::sched::yield_cpu;
use hal::xen::event::{dispatcher, send};
use hal::xen::defs::{ConsoleInterface, ConsRingIdx, EvtchnPort};

use hal::intrinsics::wmb;

static mut CONSOLE: Option<Console> = None;

pub fn console<'a>() -> &'a mut Console {
    unsafe {
        CONSOLE.as_mut().expect("Using console before initialization")
    }
}

pub unsafe fn init(interface: *mut ConsoleInterface, port: EvtchnPort) {
    CONSOLE = Some((Console::new(interface, port)));
}

pub struct Console {
    interface: *mut ConsoleInterface,
    port: EvtchnPort,
    queue: WaitQueue,
}

impl Console {
    pub fn console_callback(port: EvtchnPort, data: *mut u8) {
        let console = unsafe { &mut *(data as *mut Console) };
        let cons = console.in_cons;
        let prod = console.in_prod;

        if prod - cons > 0 {
            console.queue.unblock();
        }

        send(port);
    }

    pub unsafe fn new(interface: *mut ConsoleInterface,
                      port: EvtchnPort) -> Self {
        Console {
            interface: interface,
            port: port,
            queue: WaitQueue::new(),
        }
    }

    pub unsafe fn init_input(&mut self) {
        dispatcher().bind_port(self.port, Console::console_callback,
                               self as *mut Self as *mut u8);
        dispatcher().unmask_event(self.port);
        send(self.port);
    }

    fn is_output_full(&self) -> bool {
        let data: ConsRingIdx;

        data = self.out_prod - self.out_cons;

        // size_of output field
        (data as usize) >= mem::size_of::<[u8; 2048]>()
    }

    fn out_idx(&self) -> usize {
        let size_of_output = (mem::size_of::<[u8; 2048]>() - 1) as u32;

        (self.out_prod & size_of_output) as usize
    }

    fn in_idx(&self) -> usize {
        let size_of_input = (mem::size_of::<[u8; 1024]>() - 1) as u32;

        (self.in_cons & size_of_input) as usize
    }
}

impl Deref for Console {
    type Target = ConsoleInterface;

    #[inline]
    fn deref(&self) -> &ConsoleInterface {
        // This isn't totally safe. If the interface has not been initialized
        // this would cause some problems. For this to be safe, the console
        // MUST be initialized before being used
        unsafe {
            &(*self.interface)
        }
    }
}

impl DerefMut for Console {
    #[inline]
    fn deref_mut(&mut self) -> &mut ConsoleInterface {
        // See deref
        unsafe {
            &mut (*self.interface)
        }
    }
}

impl ::hal::Console for Console {}

impl Read for Console {
    fn read(&mut self, buf: &mut [u8]) -> Result<usize> {
        // More locked or atomic checks of the condition are not necessary
        // here. Indeed this Console is a "raw" console which means that it is
        // not thread safe. This safety will be guaranteed by the stdio
        // functions such as stdin, stdout, ...
        wait_event!(self.queue, self.in_cons - self.in_prod > 0);

        let i = 0;

        let size = cmp::min((self.in_cons - self.in_prod) as usize, buf.len());

        while i < buf.len() && self.in_cons < self.in_prod {
            let index = self.in_idx();

            buf[i] = self.input[index];

            self.in_cons += 1;

            wmb();
        }

        Ok(size)
    }
}

impl Write for Console {
    fn write(&mut self, buf: &[u8]) -> Result<usize> {
        for c in buf {
            while self.is_output_full() {
                send(self.port);
            }

            let index = self.out_idx();

            self.output[index] = *c;
            wmb();

            self.out_prod += 1;
        }

        send(self.port);

        Ok(buf.len())
    }

    fn flush(&mut self) -> Result<()> {
        while self.out_cons < self.out_prod {
            yield_cpu();
        }

        Ok(())
    }
}
<commit_msg>xen: Console::console_callback() is now private<commit_after>use core::cmp;
use core::mem;
use core::ops::{Deref, DerefMut};

use io::{Read, Write, Result};

use thread::WaitQueue;

use hal::xen::sched::yield_cpu;
use hal::xen::event::{dispatcher, send};
use hal::xen::defs::{ConsoleInterface, ConsRingIdx, EvtchnPort};

use hal::intrinsics::wmb;

static mut CONSOLE: Option<Console> = None;

pub fn console<'a>() -> &'a mut Console {
    unsafe {
        CONSOLE.as_mut().expect("Using console before initialization")
    }
}

pub unsafe fn init(interface: *mut ConsoleInterface, port: EvtchnPort) {
    CONSOLE = Some((Console::new(interface, port)));
}

pub struct Console {
    interface: *mut ConsoleInterface,
    port: EvtchnPort,
    queue: WaitQueue,
}

impl Console {
    pub unsafe fn new(interface: *mut ConsoleInterface,
                      port: EvtchnPort) -> Self {
        Console {
            interface: interface,
            port: port,
            queue: WaitQueue::new(),
        }
    }

    fn console_callback(port: EvtchnPort, data: *mut u8) {
        let console = unsafe { &mut *(data as *mut Console) };
        let cons = console.in_cons;
        let prod = console.in_prod;

        if prod - cons > 0 {
            console.queue.unblock();
        }

        send(port);
    }

    pub unsafe fn init_input(&mut self) {
        dispatcher().bind_port(self.port, Self::console_callback,
                               self as *mut Self as *mut u8);
        dispatcher().unmask_event(self.port);
        send(self.port);
    }

    fn is_output_full(&self) -> bool {
        let data: ConsRingIdx;

        data = self.out_prod - self.out_cons;

        // size_of output field
        (data as usize) >= mem::size_of::<[u8; 2048]>()
    }

    fn out_idx(&self) -> usize {
        let size_of_output = (mem::size_of::<[u8; 2048]>() - 1) as u32;

        (self.out_prod & size_of_output) as usize
    }

    fn in_idx(&self) -> usize {
        let size_of_input = (mem::size_of::<[u8; 1024]>() - 1) as u32;

        (self.in_cons & size_of_input) as usize
    }
}

impl Deref for Console {
    type Target = ConsoleInterface;

    #[inline]
    fn deref(&self) -> &ConsoleInterface {
        // This isn't totally safe. If the interface has not been initialized
        // this would cause some problems. For this to be safe, the console
        // MUST be initialized before being used
        unsafe {
            &(*self.interface)
        }
    }
}

impl DerefMut for Console {
    #[inline]
    fn deref_mut(&mut self) -> &mut ConsoleInterface {
        // See deref
        unsafe {
            &mut (*self.interface)
        }
    }
}

impl ::hal::Console for Console {}

impl Read for Console {
    fn read(&mut self, buf: &mut [u8]) -> Result<usize> {
        // More locked or atomic checks of the condition are not necessary
        // here. Indeed this Console is a "raw" console which means that it is
        // not thread safe. This safety will be guaranteed by the stdio
        // functions such as stdin, stdout, ...
        wait_event!(self.queue, self.in_cons - self.in_prod > 0);

        let i = 0;

        let size = cmp::min((self.in_cons - self.in_prod) as usize, buf.len());

        while i < buf.len() && self.in_cons < self.in_prod {
            let index = self.in_idx();

            buf[i] = self.input[index];

            self.in_cons += 1;

            wmb();
        }

        Ok(size)
    }
}

impl Write for Console {
    fn write(&mut self, buf: &[u8]) -> Result<usize> {
        for c in buf {
            while self.is_output_full() {
                send(self.port);
            }

            let index = self.out_idx();

            self.output[index] = *c;
            wmb();

            self.out_prod += 1;
        }

        send(self.port);

        Ok(buf.len())
    }

    fn flush(&mut self) -> Result<()> {
        while self.out_cons < self.out_prod {
            yield_cpu();
        }

        Ok(())
    }
}
<|endoftext|>


<commit_before>//! Defines the Path integrator which implements path tracing with
//! explicit light sampling

use std::vec::Vec;
use std::iter;
use std::num::Float;
use rand::{Rng, StdRng};

use scene::Scene;
use linalg;
use linalg::Ray;
use geometry::Intersection;
use film::Colorf;
use integrator::Integrator;
use bxdf::BxDFType;
use sampler::{Sampler, Sample};

/// The path integrator implementing Path tracing with explicit light sampling
/// See [Kajiya, The Rendering Equation](http://dl.acm.org/citation.cfm?id=15902)
#[derive(Copy, Debug)]
pub struct Path {
    min_depth: usize,
    max_depth: usize,
}

impl Path {
    /// Create a new path integrator with the min and max length desired for paths
    pub fn new(min_depth: u32, max_depth: u32) -> Path {
        Path { min_depth: min_depth as usize, max_depth: max_depth as usize }
    }
}

impl Integrator for Path {
    fn illumination(&self, scene: &Scene, r: &Ray, hit: &Intersection, sampler: &mut Sampler,
                    rng: &mut StdRng) -> Colorf {
        // TODO: We really need the memory pool now
        let mut l_samples: Vec<_> = iter::repeat((0.0, 0.0)).take(self.max_depth as usize + 1us).collect();
        let mut l_samples_comp: Vec<_> = iter::repeat(0.0).take(self.max_depth as usize + 1us).collect();
        let mut bsdf_samples: Vec<_> = iter::repeat((0.0, 0.0)).take(self.max_depth as usize + 1us).collect();
        let mut bsdf_samples_comp: Vec<_> = iter::repeat(0.0).take(self.max_depth as usize + 1us).collect();
        let mut path_samples: Vec<_> = iter::repeat((0.0, 0.0)).take(self.max_depth as usize + 1us).collect();
        let mut path_samples_comp: Vec<_> = iter::repeat(0.0).take(self.max_depth as usize + 1us).collect();
        sampler.get_samples_2d(&mut l_samples[], rng);
        sampler.get_samples_2d(&mut bsdf_samples[], rng);
        sampler.get_samples_2d(&mut path_samples[], rng);
        sampler.get_samples_1d(&mut l_samples_comp[], rng);
        sampler.get_samples_1d(&mut bsdf_samples_comp[], rng);
        sampler.get_samples_1d(&mut path_samples_comp[], rng);

        let mut illum = Colorf::black();
        let mut path_throughput = Colorf::broadcast(1.0);
        // Track if the previous bounce was a specular one
        //let mut specular_bounce = false;
        let mut current_hit = *hit;
        let mut ray = *r;
        let mut bounce = 0us;
        loop {
            // TODO: Sample emissive objects on first bounce and specular bounces
            // when emissive objects are added
            /*
            if bounce == 0 || specular_bounce {
            }
            */
            let bsdf = current_hit.instance.material.bsdf(&current_hit);
            let w_o = -ray.d;
            let light_sample = Sample::new(&l_samples[bounce], l_samples_comp[bounce]);
            let bsdf_sample = Sample::new(&bsdf_samples[bounce], bsdf_samples_comp[bounce]);
            let li = self.sample_one_light(scene, &w_o, &bsdf, &light_sample, &bsdf_sample);
            illum = illum + path_throughput * li;

            // Determine the next direction to take the path by sampling the BSDF
            let path_sample = Sample::new(&path_samples[bounce], path_samples_comp[bounce]);
            let (f, w_i, pdf, sampled_type) = bsdf.sample(&w_o, BxDFType::all(), &path_sample);
            if f.is_black() || pdf == 0.0 {
                break;
            }
            //specular_bounce = sampled_type.contains(&BxDFType::Specular);
            path_throughput = path_throughput * f * Float::abs(linalg::dot(&w_i, &bsdf.n)) / pdf;

            // Check if we're beyond the min depth at which point we start trying to
            // terminate rays using Russian Roulette
            if bounce > self.min_depth {
                let cont_prob = Float::min(0.5, path_throughput.luminance());
                if rng.next_f32() > cont_prob {
                    break;
                }
                // Re-weight the sum terms accordingly with the Russian roulette weight
                path_throughput = path_throughput / cont_prob;
            }
            if bounce == self.max_depth {
                break;
            }

            ray = ray.child(&bsdf.p, &w_i.normalized());
            ray.min_t = 0.001;
            // Find the next vertex on the path
            match scene.intersect(&mut ray) {
                Some(h) => current_hit = h,
                None => break,
            }
            bounce += 1;
        }
        illum
    }
}

<commit_msg>Tweak the roulette to always give 50% chance to continue path<commit_after>//! Defines the Path integrator which implements path tracing with
//! explicit light sampling

use std::vec::Vec;
use std::iter;
use std::num::Float;
use rand::{Rng, StdRng};

use scene::Scene;
use linalg;
use linalg::Ray;
use geometry::Intersection;
use film::Colorf;
use integrator::Integrator;
use bxdf::BxDFType;
use sampler::{Sampler, Sample};

/// The path integrator implementing Path tracing with explicit light sampling
/// See [Kajiya, The Rendering Equation](http://dl.acm.org/citation.cfm?id=15902)
#[derive(Copy, Debug)]
pub struct Path {
    min_depth: usize,
    max_depth: usize,
}

impl Path {
    /// Create a new path integrator with the min and max length desired for paths
    pub fn new(min_depth: u32, max_depth: u32) -> Path {
        Path { min_depth: min_depth as usize, max_depth: max_depth as usize }
    }
}

impl Integrator for Path {
    fn illumination(&self, scene: &Scene, r: &Ray, hit: &Intersection, sampler: &mut Sampler,
                    rng: &mut StdRng) -> Colorf {
        // TODO: We really need the memory pool now
        let mut l_samples: Vec<_> = iter::repeat((0.0, 0.0)).take(self.max_depth as usize + 1us).collect();
        let mut l_samples_comp: Vec<_> = iter::repeat(0.0).take(self.max_depth as usize + 1us).collect();
        let mut bsdf_samples: Vec<_> = iter::repeat((0.0, 0.0)).take(self.max_depth as usize + 1us).collect();
        let mut bsdf_samples_comp: Vec<_> = iter::repeat(0.0).take(self.max_depth as usize + 1us).collect();
        let mut path_samples: Vec<_> = iter::repeat((0.0, 0.0)).take(self.max_depth as usize + 1us).collect();
        let mut path_samples_comp: Vec<_> = iter::repeat(0.0).take(self.max_depth as usize + 1us).collect();
        sampler.get_samples_2d(&mut l_samples[], rng);
        sampler.get_samples_2d(&mut bsdf_samples[], rng);
        sampler.get_samples_2d(&mut path_samples[], rng);
        sampler.get_samples_1d(&mut l_samples_comp[], rng);
        sampler.get_samples_1d(&mut bsdf_samples_comp[], rng);
        sampler.get_samples_1d(&mut path_samples_comp[], rng);

        let mut illum = Colorf::black();
        let mut path_throughput = Colorf::broadcast(1.0);
        // Track if the previous bounce was a specular one
        //let mut specular_bounce = false;
        let mut current_hit = *hit;
        let mut ray = *r;
        let mut bounce = 0us;
        loop {
            // TODO: Sample emissive objects on first bounce and specular bounces
            // when emissive objects are added
            /*
            if bounce == 0 || specular_bounce {
            }
            */
            let bsdf = current_hit.instance.material.bsdf(&current_hit);
            let w_o = -ray.d;
            let light_sample = Sample::new(&l_samples[bounce], l_samples_comp[bounce]);
            let bsdf_sample = Sample::new(&bsdf_samples[bounce], bsdf_samples_comp[bounce]);
            let li = self.sample_one_light(scene, &w_o, &bsdf, &light_sample, &bsdf_sample);
            illum = illum + path_throughput * li;

            // Determine the next direction to take the path by sampling the BSDF
            let path_sample = Sample::new(&path_samples[bounce], path_samples_comp[bounce]);
            let (f, w_i, pdf, sampled_type) = bsdf.sample(&w_o, BxDFType::all(), &path_sample);
            if f.is_black() || pdf == 0.0 {
                break;
            }
            //specular_bounce = sampled_type.contains(&BxDFType::Specular);
            path_throughput = path_throughput * f * Float::abs(linalg::dot(&w_i, &bsdf.n)) / pdf;

            // Check if we're beyond the min depth at which point we start trying to
            // terminate rays using Russian Roulette
            if bounce > self.min_depth {
                let cont_prob = Float::max(0.5, path_throughput.luminance());
                if rng.next_f32() > cont_prob {
                    break;
                }
                // Re-weight the sum terms accordingly with the Russian roulette weight
                path_throughput = path_throughput / cont_prob;
            }
            if bounce == self.max_depth {
                break;
            }

            ray = ray.child(&bsdf.p, &w_i.normalized());
            ray.min_t = 0.001;
            // Find the next vertex on the path
            match scene.intersect(&mut ray) {
                Some(h) => current_hit = h,
                None => break,
            }
            bounce += 1;
        }
        illum
    }
}

<|endoftext|>


<commit_before>/*
 * Copyright (C) 2015 Benjamin Fry <benjaminfry@me.com>
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
// TODO, I've implemented this as a seperate entity from the cache, but I wonder if the cache
//  should be the only "front-end" for lookups, where if that misses, then we go to the catalog
//  then, if requested, do a recursive lookup... i.e. the catalog would only point to files.
use std::collections::HashMap;
use std::io;
use std::sync::RwLock;

use server::{Request, RequestHandler, ResponseHandler};
use trust_dns::op::{Edns, Header, LowerQuery, MessageType, OpCode, ResponseCode};
use trust_dns::rr::dnssec::{Algorithm, SupportedAlgorithms};
use trust_dns::rr::rdata::opt::{EdnsCode, EdnsOption};
use trust_dns::rr::{LowerName, RecordType};

use authority::{
    AuthLookup, Authority, LookupRecords, MessageRequest, MessageResponse, MessageResponseBuilder,
    ZoneType,
};

/// Set of authorities, zones, available to this server.
pub struct Catalog {
    authorities: HashMap<LowerName, RwLock<Authority>>,
}

fn send_response<R: ResponseHandler + 'static>(
    response_edns: Option<Edns>,
    mut response: MessageResponse,
    response_handle: R,
) -> io::Result<()> {
    if let Some(mut resp_edns) = response_edns {
        // set edns DAU and DHU
        // send along the algorithms which are supported by this authority
        let mut algorithms = SupportedAlgorithms::new();
        algorithms.set(Algorithm::RSASHA256);
        algorithms.set(Algorithm::ECDSAP256SHA256);
        algorithms.set(Algorithm::ECDSAP384SHA384);
        algorithms.set(Algorithm::ED25519);

        let dau = EdnsOption::DAU(algorithms);
        let dhu = EdnsOption::DHU(algorithms);

        resp_edns.set_option(dau);
        resp_edns.set_option(dhu);

        response.set_edns(resp_edns);
    }

    response_handle.send(response)
}

impl RequestHandler for Catalog {
    /// Determine's what needs to happen given the type of request, i.e. Query or Update.
    ///
    /// # Arguments
    ///
    /// * `request` - the requested action to perform.
    /// * `response_handle` - sink for the response message to be sent
    fn handle_request<'q, 'a, R: ResponseHandler + 'static>(
        &'a self,
        request: &'q Request,
        response_handle: R,
    ) -> io::Result<()> {
        let request_message = &request.message;
        trace!("request: {:?}", request_message);

        let response_edns: Option<Edns>;

        // check if it's edns
        if let Some(req_edns) = request_message.edns() {
            let mut response = MessageResponseBuilder::new(Some(request_message.raw_queries()));
            let mut response_header = Header::default();
            response_header.set_id(request_message.id());

            let mut resp_edns: Edns = Edns::new();

            // check our version against the request
            // TODO: what version are we?
            let our_version = 0;
            resp_edns.set_dnssec_ok(true);
            resp_edns.set_max_payload(req_edns.max_payload().max(512));
            resp_edns.set_version(our_version);

            if req_edns.version() > our_version {
                warn!(
                    "request edns version greater than {}: {}",
                    our_version,
                    req_edns.version()
                );
                response_header.set_response_code(ResponseCode::BADVERS);
                response.edns(resp_edns);

                // TODO: should ResponseHandle consume self?
                return response_handle.send(response.build(response_header));
            }

            response_edns = Some(resp_edns);
        } else {
            response_edns = None;
        }

        match request_message.message_type() {
            // TODO think about threading query lookups for multiple lookups, this could be a huge improvement
            //  especially for recursive lookups
            MessageType::Query => match request_message.op_code() {
                OpCode::Query => {
                    return self.lookup(request_message, response_edns, response_handle)
                }
                OpCode::Update => {
                    return self.update(request_message, response_edns, response_handle)
                }
                c @ _ => {
                    error!("unimplemented op_code: {:?}", c);
                    let response = MessageResponseBuilder::new(Some(request_message.raw_queries()));
                    return response_handle.send(response.error_msg(
                        request_message.id(),
                        request_message.op_code(),
                        ResponseCode::NotImp,
                    ));
                }
            },
            MessageType::Response => {
                warn!(
                    "got a response as a request from id: {}",
                    request_message.id()
                );
                let response = MessageResponseBuilder::new(Some(request_message.raw_queries()));

                return response_handle.send(response.error_msg(
                    request_message.id(),
                    request_message.op_code(),
                    ResponseCode::FormErr,
                ));
            }
        };
    }
}

impl Catalog {
    /// Constructs a new Catalog
    pub fn new() -> Self {
        Catalog {
            authorities: HashMap::new(),
        }
    }

    /// Insert or update a zone authority
    ///
    /// # Arguments
    ///
    /// * `name` - zone name, e.g. example.com.
    /// * `authority` - the zone data
    pub fn upsert(&mut self, name: LowerName, authority: Authority) {
        self.authorities.insert(name, RwLock::new(authority));
    }

    /// Update the zone given the Update request.
    ///
    /// [RFC 2136](https://tools.ietf.org/html/rfc2136), DNS Update, April 1997
    ///
    /// ```text
    /// 3.1 - Process Zone Section
    ///
    ///   3.1.1. The Zone Section is checked to see that there is exactly one
    ///   RR therein and that the RR's ZTYPE is SOA, else signal FORMERR to the
    ///   requestor.  Next, the ZNAME and ZCLASS are checked to see if the zone
    ///   so named is one of this server's authority zones, else signal NOTAUTH
    ///   to the requestor.  If the server is a zone slave, the request will be
    ///   forwarded toward the primary master.
    ///
    ///   3.1.2 - Pseudocode For Zone Section Processing
    ///
    ///      if (zcount != 1 || ztype != SOA)
    ///           return (FORMERR)
    ///      if (zone_type(zname, zclass) == SLAVE)
    ///           return forward()
    ///      if (zone_type(zname, zclass) == MASTER)
    ///           return update()
    ///      return (NOTAUTH)
    ///
    ///   Sections 3.2 through 3.8 describe the primary master's behaviour,
    ///   whereas Section 6 describes a forwarder's behaviour.
    ///
    /// 3.8 - Response
    ///
    ///   At the end of UPDATE processing, a response code will be known.  A
    ///   response message is generated by copying the ID and Opcode fields
    ///   from the request, and either copying the ZOCOUNT, PRCOUNT, UPCOUNT,
    ///   and ADCOUNT fields and associated sections, or placing zeros (0) in
    ///   the these "count" fields and not including any part of the original
    ///   update.  The QR bit is set to one (1), and the response is sent back
    ///   to the requestor.  If the requestor used UDP, then the response will
    ///   be sent to the requestor's source UDP port.  If the requestor used
    ///   TCP, then the response will be sent back on the requestor's open TCP
    ///   connection.
    /// ```
    ///
    /// The "request" should be an update formatted message.
    ///  The response will be in the alternate, all 0's format described in RFC 2136 section 3.8
    ///  as this is more efficient.
    ///
    /// # Arguments
    ///
    /// * `request` - an update message
    /// * `response_handle` - sink for the response message to be sent
    pub fn update<'q, R: ResponseHandler + 'static>(
        &self,
        update: &'q MessageRequest,
        response_edns: Option<Edns>,
        response_handle: R,
    ) -> io::Result<()> {
        let response = MessageResponseBuilder::new(None);
        let mut response_header = Header::default();
        response_header.set_id(update.id());
        response_header.set_op_code(OpCode::Update);
        response_header.set_message_type(MessageType::Response);

        let zones: &[LowerQuery] = update.queries();

        // 2.3 - Zone Section
        //
        //  All records to be updated must be in the same zone, and
        //  therefore the Zone Section is allowed to contain exactly one record.
        //  The ZNAME is the zone name, the ZTYPE must be SOA, and the ZCLASS is
        //  the zone's class.
        if zones.len() != 1 || zones[0].query_type() != RecordType::SOA {
            response_header.set_response_code(ResponseCode::FormErr);

            return send_response(
                response_edns,
                response.build(response_header),
                response_handle,
            );
        }

        if let Some(authority) = self.find_auth_recurse(zones[0].name()) {
            let mut authority = authority.write().unwrap(); // poison errors should panic...
            match authority.zone_type() {
                ZoneType::Slave => {
                    error!("slave forwarding for update not yet implemented");
                    response_header.set_response_code(ResponseCode::NotImp);

                    return send_response(
                        response_edns,
                        response.build(response_header),
                        response_handle,
                    );
                }
                ZoneType::Master => {
                    let update_result = authority.update(update);
                    match update_result {
                        // successful update
                        Ok(..) => {
                            response_header.set_response_code(ResponseCode::NoError);
                        }
                        Err(response_code) => {
                            response_header.set_response_code(response_code);
                        }
                    }

                    return send_response(
                        response_edns,
                        response.build(response_header),
                        response_handle,
                    );
                }
                _ => {
                    response_header.set_response_code(ResponseCode::NotAuth);

                    return send_response(
                        response_edns,
                        response.build(response_header),
                        response_handle,
                    );
                }
            }
        } else {
            response_header.set_response_code(ResponseCode::NXDomain);

            return send_response(
                response_edns,
                response.build(response_header),
                response_handle,
            );
        }
    }

    /// Checks whether the `Catalog` contains DNS records for `name`
    ///
    /// Use this when you know the exact `LowerName` that was used when
    /// adding an authority and you don't care about the authority it
    /// contains. For public domain names, `LowerName` is usually the
    /// top level domain name like `example.com.`.
    ///
    /// If you do not know the exact domain name to use or you actually
    /// want to use the authority it contains, use `find` instead.
    pub fn contains(&self, name: &LowerName) -> bool {
        self.authorities.contains_key(name)
    }

    /// Given the requested query, lookup and return any matching results.
    ///
    /// # Arguments
    ///
    /// * `request` - the query message.
    /// * `response_handle` - sink for the response message to be sent
    pub fn lookup<'q, R: ResponseHandler + 'static>(
        &self,
        request: &'q MessageRequest,
        response_edns: Option<Edns>,
        response_handle: R,
    ) -> io::Result<()> {
        // TODO: the spec is very unclear on what to do with multiple queries
        //  we will search for each, in the future, maybe make this threaded to respond even faster.
        for query in request.queries() {
            if let Some(ref_authority) = self.find_auth_recurse(query.name()) {
                let authority = &ref_authority.read().unwrap(); // poison errors should panic
                info!(
                    "request: {} found authority: {}",
                    request.id(),
                    authority.origin()
                );

                let mut response = MessageResponseBuilder::new(Some(request.raw_queries()));
                let mut response_header = Header::new();
                response_header.set_id(request.id());
                response_header.set_op_code(OpCode::Query);
                response_header.set_message_type(MessageType::Response);

                let (is_dnssec, supported_algorithms) = request.edns().map_or(
                    (false, SupportedAlgorithms::new()),
                    |edns| {
                        let supported_algorithms =
                            if let Some(&EdnsOption::DAU(algs)) = edns.option(&EdnsCode::DAU) {
                                algs
                            } else {
                                debug!("no DAU in request, used default SupportAlgorithms");
                                Default::default()
                            };

                        (edns.dnssec_ok(), supported_algorithms)
                    },
                );

                // log algorithms being requested
                if is_dnssec {
                    info!(
                        "request: {} supported_algs: {}",
                        request.id(),
                        supported_algorithms
                    );
                }

                let records = authority.search(query, is_dnssec, supported_algorithms);

                // setup headers
                //  and add records
                if !records.is_empty() {
                    response_header.set_response_code(ResponseCode::NoError);
                    response_header.set_authoritative(true);
                    response.answers(records);

                    // get the NS records
                    let ns = authority.ns(is_dnssec, supported_algorithms);
                    // chain here to match type below...
                    response.name_servers(ns.chain(LookupRecords::NxDomain));
                } else {
                    // in the not found case it's standard to return the SOA in the authority section
                    //   if the name is in this zone, etc.
                    // see https://tools.ietf.org/html/rfc2308 for proper response construct
                    match records {
                        AuthLookup::NxDomain => {
                            response_header.set_response_code(ResponseCode::NXDomain)
                        }
                        AuthLookup::NameExists => {
                            response_header.set_response_code(ResponseCode::NoError)
                        }
                        _ => panic!(
                            "programming error, should have return NoError with records above"
                        ),
                    };

                    // in the dnssec case, nsec records should exist, we return NoError + NoData + NSec...
                    let ns = if is_dnssec {
                        // get NSEC records
                        let mut nsecs = authority.get_nsec_records(
                            query.name(),
                            is_dnssec,
                            supported_algorithms,
                        );
                        debug!("request: {} non-existent adding nsecs", request.id(),);

                        response_header.set_response_code(ResponseCode::NoError);
                        nsecs
                    } else {
                        // place holder...
                        debug!("request: {} non-existent", request.id());
                        LookupRecords::NxDomain
                    };

                    let soa = authority.soa_secure(is_dnssec, supported_algorithms);
                    let ns = ns.chain(soa);

                    response.name_servers(ns);
                }

                return send_response(
                    response_edns,
                    response.build(response_header),
                    response_handle,
                );
            }
        }

        let response = MessageResponseBuilder::new(Some(request.raw_queries()));
        return send_response(
            response_edns,
            response.error_msg(request.id(), request.op_code(), ResponseCode::NXDomain),
            response_handle,
        );
    }

    /// recursively searches the catalog for a matching auhtority.
    fn find_auth_recurse(&self, name: &LowerName) -> Option<&RwLock<Authority>> {
        let authority = self.authorities.get(name);
        if authority.is_some() {
            return authority;
        } else {
            let name = name.base_name();
            if !name.is_root() {
                return self.find_auth_recurse(&name);
            }
        }

        None
    }
}
<commit_msg>add `Catalog::find`<commit_after>/*
 * Copyright (C) 2015 Benjamin Fry <benjaminfry@me.com>
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
// TODO, I've implemented this as a seperate entity from the cache, but I wonder if the cache
//  should be the only "front-end" for lookups, where if that misses, then we go to the catalog
//  then, if requested, do a recursive lookup... i.e. the catalog would only point to files.
use std::collections::HashMap;
use std::io;
use std::sync::RwLock;

use server::{Request, RequestHandler, ResponseHandler};
use trust_dns::op::{Edns, Header, LowerQuery, MessageType, OpCode, ResponseCode};
use trust_dns::rr::dnssec::{Algorithm, SupportedAlgorithms};
use trust_dns::rr::rdata::opt::{EdnsCode, EdnsOption};
use trust_dns::rr::{LowerName, RecordType};

use authority::{
    AuthLookup, Authority, LookupRecords, MessageRequest, MessageResponse, MessageResponseBuilder,
    ZoneType,
};

/// Set of authorities, zones, available to this server.
pub struct Catalog {
    authorities: HashMap<LowerName, RwLock<Authority>>,
}

fn send_response<R: ResponseHandler + 'static>(
    response_edns: Option<Edns>,
    mut response: MessageResponse,
    response_handle: R,
) -> io::Result<()> {
    if let Some(mut resp_edns) = response_edns {
        // set edns DAU and DHU
        // send along the algorithms which are supported by this authority
        let mut algorithms = SupportedAlgorithms::new();
        algorithms.set(Algorithm::RSASHA256);
        algorithms.set(Algorithm::ECDSAP256SHA256);
        algorithms.set(Algorithm::ECDSAP384SHA384);
        algorithms.set(Algorithm::ED25519);

        let dau = EdnsOption::DAU(algorithms);
        let dhu = EdnsOption::DHU(algorithms);

        resp_edns.set_option(dau);
        resp_edns.set_option(dhu);

        response.set_edns(resp_edns);
    }

    response_handle.send(response)
}

impl RequestHandler for Catalog {
    /// Determine's what needs to happen given the type of request, i.e. Query or Update.
    ///
    /// # Arguments
    ///
    /// * `request` - the requested action to perform.
    /// * `response_handle` - sink for the response message to be sent
    fn handle_request<'q, 'a, R: ResponseHandler + 'static>(
        &'a self,
        request: &'q Request,
        response_handle: R,
    ) -> io::Result<()> {
        let request_message = &request.message;
        trace!("request: {:?}", request_message);

        let response_edns: Option<Edns>;

        // check if it's edns
        if let Some(req_edns) = request_message.edns() {
            let mut response = MessageResponseBuilder::new(Some(request_message.raw_queries()));
            let mut response_header = Header::default();
            response_header.set_id(request_message.id());

            let mut resp_edns: Edns = Edns::new();

            // check our version against the request
            // TODO: what version are we?
            let our_version = 0;
            resp_edns.set_dnssec_ok(true);
            resp_edns.set_max_payload(req_edns.max_payload().max(512));
            resp_edns.set_version(our_version);

            if req_edns.version() > our_version {
                warn!(
                    "request edns version greater than {}: {}",
                    our_version,
                    req_edns.version()
                );
                response_header.set_response_code(ResponseCode::BADVERS);
                response.edns(resp_edns);

                // TODO: should ResponseHandle consume self?
                return response_handle.send(response.build(response_header));
            }

            response_edns = Some(resp_edns);
        } else {
            response_edns = None;
        }

        match request_message.message_type() {
            // TODO think about threading query lookups for multiple lookups, this could be a huge improvement
            //  especially for recursive lookups
            MessageType::Query => match request_message.op_code() {
                OpCode::Query => {
                    return self.lookup(request_message, response_edns, response_handle)
                }
                OpCode::Update => {
                    return self.update(request_message, response_edns, response_handle)
                }
                c @ _ => {
                    error!("unimplemented op_code: {:?}", c);
                    let response = MessageResponseBuilder::new(Some(request_message.raw_queries()));
                    return response_handle.send(response.error_msg(
                        request_message.id(),
                        request_message.op_code(),
                        ResponseCode::NotImp,
                    ));
                }
            },
            MessageType::Response => {
                warn!(
                    "got a response as a request from id: {}",
                    request_message.id()
                );
                let response = MessageResponseBuilder::new(Some(request_message.raw_queries()));

                return response_handle.send(response.error_msg(
                    request_message.id(),
                    request_message.op_code(),
                    ResponseCode::FormErr,
                ));
            }
        };
    }
}

impl Catalog {
    /// Constructs a new Catalog
    pub fn new() -> Self {
        Catalog {
            authorities: HashMap::new(),
        }
    }

    /// Insert or update a zone authority
    ///
    /// # Arguments
    ///
    /// * `name` - zone name, e.g. example.com.
    /// * `authority` - the zone data
    pub fn upsert(&mut self, name: LowerName, authority: Authority) {
        self.authorities.insert(name, RwLock::new(authority));
    }

    /// Update the zone given the Update request.
    ///
    /// [RFC 2136](https://tools.ietf.org/html/rfc2136), DNS Update, April 1997
    ///
    /// ```text
    /// 3.1 - Process Zone Section
    ///
    ///   3.1.1. The Zone Section is checked to see that there is exactly one
    ///   RR therein and that the RR's ZTYPE is SOA, else signal FORMERR to the
    ///   requestor.  Next, the ZNAME and ZCLASS are checked to see if the zone
    ///   so named is one of this server's authority zones, else signal NOTAUTH
    ///   to the requestor.  If the server is a zone slave, the request will be
    ///   forwarded toward the primary master.
    ///
    ///   3.1.2 - Pseudocode For Zone Section Processing
    ///
    ///      if (zcount != 1 || ztype != SOA)
    ///           return (FORMERR)
    ///      if (zone_type(zname, zclass) == SLAVE)
    ///           return forward()
    ///      if (zone_type(zname, zclass) == MASTER)
    ///           return update()
    ///      return (NOTAUTH)
    ///
    ///   Sections 3.2 through 3.8 describe the primary master's behaviour,
    ///   whereas Section 6 describes a forwarder's behaviour.
    ///
    /// 3.8 - Response
    ///
    ///   At the end of UPDATE processing, a response code will be known.  A
    ///   response message is generated by copying the ID and Opcode fields
    ///   from the request, and either copying the ZOCOUNT, PRCOUNT, UPCOUNT,
    ///   and ADCOUNT fields and associated sections, or placing zeros (0) in
    ///   the these "count" fields and not including any part of the original
    ///   update.  The QR bit is set to one (1), and the response is sent back
    ///   to the requestor.  If the requestor used UDP, then the response will
    ///   be sent to the requestor's source UDP port.  If the requestor used
    ///   TCP, then the response will be sent back on the requestor's open TCP
    ///   connection.
    /// ```
    ///
    /// The "request" should be an update formatted message.
    ///  The response will be in the alternate, all 0's format described in RFC 2136 section 3.8
    ///  as this is more efficient.
    ///
    /// # Arguments
    ///
    /// * `request` - an update message
    /// * `response_handle` - sink for the response message to be sent
    pub fn update<'q, R: ResponseHandler + 'static>(
        &self,
        update: &'q MessageRequest,
        response_edns: Option<Edns>,
        response_handle: R,
    ) -> io::Result<()> {
        let response = MessageResponseBuilder::new(None);
        let mut response_header = Header::default();
        response_header.set_id(update.id());
        response_header.set_op_code(OpCode::Update);
        response_header.set_message_type(MessageType::Response);

        let zones: &[LowerQuery] = update.queries();

        // 2.3 - Zone Section
        //
        //  All records to be updated must be in the same zone, and
        //  therefore the Zone Section is allowed to contain exactly one record.
        //  The ZNAME is the zone name, the ZTYPE must be SOA, and the ZCLASS is
        //  the zone's class.
        if zones.len() != 1 || zones[0].query_type() != RecordType::SOA {
            response_header.set_response_code(ResponseCode::FormErr);

            return send_response(
                response_edns,
                response.build(response_header),
                response_handle,
            );
        }

        if let Some(authority) = self.find(zones[0].name()) {
            let mut authority = authority.write().unwrap(); // poison errors should panic...
            match authority.zone_type() {
                ZoneType::Slave => {
                    error!("slave forwarding for update not yet implemented");
                    response_header.set_response_code(ResponseCode::NotImp);

                    return send_response(
                        response_edns,
                        response.build(response_header),
                        response_handle,
                    );
                }
                ZoneType::Master => {
                    let update_result = authority.update(update);
                    match update_result {
                        // successful update
                        Ok(..) => {
                            response_header.set_response_code(ResponseCode::NoError);
                        }
                        Err(response_code) => {
                            response_header.set_response_code(response_code);
                        }
                    }

                    return send_response(
                        response_edns,
                        response.build(response_header),
                        response_handle,
                    );
                }
                _ => {
                    response_header.set_response_code(ResponseCode::NotAuth);

                    return send_response(
                        response_edns,
                        response.build(response_header),
                        response_handle,
                    );
                }
            }
        } else {
            response_header.set_response_code(ResponseCode::NXDomain);

            return send_response(
                response_edns,
                response.build(response_header),
                response_handle,
            );
        }
    }

    /// Checks whether the `Catalog` contains DNS records for `name`
    ///
    /// Use this when you know the exact `LowerName` that was used when
    /// adding an authority and you don't care about the authority it
    /// contains. For public domain names, `LowerName` is usually the
    /// top level domain name like `example.com.`.
    ///
    /// If you do not know the exact domain name to use or you actually
    /// want to use the authority it contains, use `find` instead.
    pub fn contains(&self, name: &LowerName) -> bool {
        self.authorities.contains_key(name)
    }

    /// Given the requested query, lookup and return any matching results.
    ///
    /// # Arguments
    ///
    /// * `request` - the query message.
    /// * `response_handle` - sink for the response message to be sent
    pub fn lookup<'q, R: ResponseHandler + 'static>(
        &self,
        request: &'q MessageRequest,
        response_edns: Option<Edns>,
        response_handle: R,
    ) -> io::Result<()> {
        // TODO: the spec is very unclear on what to do with multiple queries
        //  we will search for each, in the future, maybe make this threaded to respond even faster.
        for query in request.queries() {
            if let Some(ref_authority) = self.find(query.name()) {
                let authority = &ref_authority.read().unwrap(); // poison errors should panic
                info!(
                    "request: {} found authority: {}",
                    request.id(),
                    authority.origin()
                );

                let mut response = MessageResponseBuilder::new(Some(request.raw_queries()));
                let mut response_header = Header::new();
                response_header.set_id(request.id());
                response_header.set_op_code(OpCode::Query);
                response_header.set_message_type(MessageType::Response);

                let (is_dnssec, supported_algorithms) = request.edns().map_or(
                    (false, SupportedAlgorithms::new()),
                    |edns| {
                        let supported_algorithms =
                            if let Some(&EdnsOption::DAU(algs)) = edns.option(&EdnsCode::DAU) {
                                algs
                            } else {
                                debug!("no DAU in request, used default SupportAlgorithms");
                                Default::default()
                            };

                        (edns.dnssec_ok(), supported_algorithms)
                    },
                );

                // log algorithms being requested
                if is_dnssec {
                    info!(
                        "request: {} supported_algs: {}",
                        request.id(),
                        supported_algorithms
                    );
                }

                let records = authority.search(query, is_dnssec, supported_algorithms);

                // setup headers
                //  and add records
                if !records.is_empty() {
                    response_header.set_response_code(ResponseCode::NoError);
                    response_header.set_authoritative(true);
                    response.answers(records);

                    // get the NS records
                    let ns = authority.ns(is_dnssec, supported_algorithms);
                    // chain here to match type below...
                    response.name_servers(ns.chain(LookupRecords::NxDomain));
                } else {
                    // in the not found case it's standard to return the SOA in the authority section
                    //   if the name is in this zone, etc.
                    // see https://tools.ietf.org/html/rfc2308 for proper response construct
                    match records {
                        AuthLookup::NxDomain => {
                            response_header.set_response_code(ResponseCode::NXDomain)
                        }
                        AuthLookup::NameExists => {
                            response_header.set_response_code(ResponseCode::NoError)
                        }
                        _ => panic!(
                            "programming error, should have return NoError with records above"
                        ),
                    };

                    // in the dnssec case, nsec records should exist, we return NoError + NoData + NSec...
                    let ns = if is_dnssec {
                        // get NSEC records
                        let mut nsecs = authority.get_nsec_records(
                            query.name(),
                            is_dnssec,
                            supported_algorithms,
                        );
                        debug!("request: {} non-existent adding nsecs", request.id(),);

                        response_header.set_response_code(ResponseCode::NoError);
                        nsecs
                    } else {
                        // place holder...
                        debug!("request: {} non-existent", request.id());
                        LookupRecords::NxDomain
                    };

                    let soa = authority.soa_secure(is_dnssec, supported_algorithms);
                    let ns = ns.chain(soa);

                    response.name_servers(ns);
                }

                return send_response(
                    response_edns,
                    response.build(response_header),
                    response_handle,
                );
            }
        }

        let response = MessageResponseBuilder::new(Some(request.raw_queries()));
        return send_response(
            response_edns,
            response.error_msg(request.id(), request.op_code(), ResponseCode::NXDomain),
            response_handle,
        );
    }

    /// Recursively searches the catalog for a matching authority
    pub fn find(&self, name: &LowerName) -> Option<&RwLock<Authority>> {
        self.authorities.get(name)
            .or_else(|| {
                let name = name.base_name();
                if !name.is_root() {
                    self.find(&name)
                } else {
                    None
                }
            })
    }
}
<|endoftext|>


<commit_before>use errors::*;
use std::fmt;
use std::io::{Read, Seek, SeekFrom};
use unterflow_protocol::convert::*;
use unterflow_protocol::protocol::client::*;
use unterflow_protocol::protocol::gossip::*;
use unterflow_protocol::protocol::management::*;
use unterflow_protocol::protocol::raft::*;
use unterflow_protocol::protocol::transport::*;

#[derive(Default, Debug)]
pub struct Protocol {
    pub frame: Option<FrameHeader>,
    transport: Option<TransportHeader>,
    protocol: Option<Box<fmt::Debug>>,
    message: Option<Box<fmt::Debug>>,
    pretty: bool,
}

impl fmt::Display for Protocol {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        if let Some(ref frame) = self.frame {
            writeln!(f, "{:?}", frame)?;
        }
        if let Some(ref transport) = self.transport {
            writeln!(f, "{:?}", transport)?;
        }
        if let Some(ref protocol) = self.protocol {
            writeln!(f, "{:?}", protocol)?;
        }
        if let Some(ref message) = self.message {
            if self.pretty {
                writeln!(f, "{:#?}", message)?;
            } else {
                writeln!(f, "{:?}", message)?;
            }
        }

        Ok(())
    }
}

impl Protocol {
    pub fn parse<T: Read + Seek>(mut payload: &mut T, offset: u64, pretty: bool) -> Result<Self> {
        payload.seek(SeekFrom::Start(offset))?;
        let mut protocol = Protocol::default();
        protocol.pretty = pretty;

        if let Ok(frame) = FrameHeader::from_bytes(&mut payload) {

            if frame.type_id == FrameType::Unknown {
                bail!("Unknown frame type: {:?}", frame);
            }

            if frame.type_id == FrameType::Message {

                let transport = TransportHeader::from_bytes(&mut payload)?;
                protocol.protocol = match transport.protocol_id {
                    TransportProtocol::RequestResponse => {
                        Some(Box::new(RequestResponseHeader::from_bytes(&mut payload)?))
                    }
                    TransportProtocol::FullDuplexSingleMessage => {
                        Some(Box::new(SingleMessageHeader::from_bytes(&mut payload)?))
                    }
                    TransportProtocol::Unknown => bail!("Unknown transport protocol: {:?}", transport),
                };
                protocol.transport = Some(transport);

                let header = MessageHeader::from_bytes(&mut payload)?;

                // Client Protocol
                if header == ErrorResponse::message_header() {
                    protocol.message = Some(Box::new(ErrorResponse::from_bytes(&mut payload)?));
                } else if header == ControlMessageRequest::message_header() {
                    protocol.message = Some(Box::new(ControlMessageRequest::from_bytes(&mut payload)?));
                } else if header == ControlMessageResponse::message_header() {
                    protocol.message = Some(Box::new(ControlMessageResponse::from_bytes(&mut payload)?));
                } else if header == ExecuteCommandRequest::message_header() {
                    protocol.message = Some(Box::new(ExecuteCommandRequest::from_bytes(&mut payload)?));
                } else if header == ExecuteCommandResponse::message_header() {
                    protocol.message = Some(Box::new(ExecuteCommandResponse::from_bytes(&mut payload)?));
                } else if header == SubscribedEvent::message_header() {
                    protocol.message = Some(Box::new(SubscribedEvent::from_bytes(&mut payload)?));
                } else if header == BrokerEventMetadata::message_header() {
                    protocol.message = Some(Box::new(BrokerEventMetadata::from_bytes(&mut payload)?));
                }
                // Management Protocol
                else if header == InvitationRequest::message_header() {
                    protocol.message = Some(Box::new(InvitationRequest::from_bytes(&mut payload)?));
                } else if header == InvitationResponse::message_header() {
                    protocol.message = Some(Box::new(InvitationResponse::from_bytes(&mut payload)?));
                }
                // Raft Protocol
                else if header == JoinRequest::message_header() {
                    protocol.message = Some(Box::new(JoinRequest::from_bytes(&mut payload)?));
                } else if header == JoinResponse::message_header() {
                    protocol.message = Some(Box::new(JoinResponse::from_bytes(&mut payload)?));
                } else if header == LeaveRequest::message_header() {
                    protocol.message = Some(Box::new(LeaveRequest::from_bytes(&mut payload)?));
                } else if header == LeaveResponse::message_header() {
                    protocol.message = Some(Box::new(LeaveResponse::from_bytes(&mut payload)?));
                } else if header == ConfigurationRequest::message_header() {
                    protocol.message = Some(Box::new(ConfigurationRequest::from_bytes(&mut payload)?));
                } else if header == ConfigurationResponse::message_header() {
                    protocol.message = Some(Box::new(ConfigurationResponse::from_bytes(&mut payload)?));
                } else if header == PollRequest::message_header() {
                    protocol.message = Some(Box::new(PollRequest::from_bytes(&mut payload)?));
                } else if header == PollResponse::message_header() {
                    protocol.message = Some(Box::new(PollResponse::from_bytes(&mut payload)?));
                } else if header == VoteRequest::message_header() {
                    protocol.message = Some(Box::new(VoteRequest::from_bytes(&mut payload)?));
                } else if header == VoteResponse::message_header() {
                    protocol.message = Some(Box::new(VoteResponse::from_bytes(&mut payload)?));
                } else if header == AppendRequest::message_header() {
                    protocol.message = Some(Box::new(AppendRequest::from_bytes(&mut payload)?));
                } else if header == AppendResponse::message_header() {
                    protocol.message = Some(Box::new(AppendResponse::from_bytes(&mut payload)?));
                }
                // Gossip Protocol
                else if header == Gossip::message_header() {
                    protocol.message = Some(Box::new(Gossip::from_bytes(&mut payload)?));
                } else if header == Probe::message_header() {
                    protocol.message = Some(Box::new(Probe::from_bytes(&mut payload)?));
                } else if header == PeerDescriptor::message_header() {
                    protocol.message = Some(Box::new(PeerDescriptor::from_bytes(&mut payload)?));
                } else {
                    bail!("Unknown message header: {:?}", header);
                }
            }

            protocol.frame = Some(frame);
        }

        Ok(protocol)
    }
}
<commit_msg>feat(unterflow-dump): print also message header<commit_after>use errors::*;
use std::fmt;
use std::io::{Read, Seek, SeekFrom};
use unterflow_protocol::convert::*;
use unterflow_protocol::protocol::client::*;
use unterflow_protocol::protocol::gossip::*;
use unterflow_protocol::protocol::management::*;
use unterflow_protocol::protocol::raft::*;
use unterflow_protocol::protocol::transport::*;

#[derive(Default, Debug)]
pub struct Protocol {
    pub frame: Option<FrameHeader>,
    transport: Option<TransportHeader>,
    protocol: Option<Box<fmt::Debug>>,
    header: Option<MessageHeader>,
    message: Option<Box<fmt::Debug>>,
    pretty: bool,
}

impl fmt::Display for Protocol {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        if let Some(ref frame) = self.frame {
            writeln!(f, "{:?}", frame)?;
        }
        if let Some(ref transport) = self.transport {
            writeln!(f, "{:?}", transport)?;
        }
        if let Some(ref protocol) = self.protocol {
            writeln!(f, "{:?}", protocol)?;
        }
        if let Some(ref header) = self.header {
            writeln!(f, "{:?}", header)?;
        }
        if let Some(ref message) = self.message {
            if self.pretty {
                writeln!(f, "{:#?}", message)?;
            } else {
                writeln!(f, "{:?}", message)?;
            }
        }

        Ok(())
    }
}

impl Protocol {
    pub fn parse<T: Read + Seek>(mut payload: &mut T, offset: u64, pretty: bool) -> Result<Self> {
        payload.seek(SeekFrom::Start(offset))?;
        let mut protocol = Protocol::default();
        protocol.pretty = pretty;

        if let Ok(frame) = FrameHeader::from_bytes(&mut payload) {

            if frame.type_id == FrameType::Unknown {
                bail!("Unknown frame type: {:?}", frame);
            }

            if frame.type_id == FrameType::Message {

                let transport = TransportHeader::from_bytes(&mut payload)?;
                protocol.protocol = match transport.protocol_id {
                    TransportProtocol::RequestResponse => {
                        Some(Box::new(RequestResponseHeader::from_bytes(&mut payload)?))
                    }
                    TransportProtocol::FullDuplexSingleMessage => {
                        Some(Box::new(SingleMessageHeader::from_bytes(&mut payload)?))
                    }
                    TransportProtocol::Unknown => bail!("Unknown transport protocol: {:?}", transport),
                };
                protocol.transport = Some(transport);

                let header = MessageHeader::from_bytes(&mut payload)?;

                // Client Protocol
                if header == ErrorResponse::message_header() {
                    protocol.message = Some(Box::new(ErrorResponse::from_bytes(&mut payload)?));
                } else if header == ControlMessageRequest::message_header() {
                    protocol.message = Some(Box::new(ControlMessageRequest::from_bytes(&mut payload)?));
                } else if header == ControlMessageResponse::message_header() {
                    protocol.message = Some(Box::new(ControlMessageResponse::from_bytes(&mut payload)?));
                } else if header == ExecuteCommandRequest::message_header() {
                    protocol.message = Some(Box::new(ExecuteCommandRequest::from_bytes(&mut payload)?));
                } else if header == ExecuteCommandResponse::message_header() {
                    protocol.message = Some(Box::new(ExecuteCommandResponse::from_bytes(&mut payload)?));
                } else if header == SubscribedEvent::message_header() {
                    protocol.message = Some(Box::new(SubscribedEvent::from_bytes(&mut payload)?));
                } else if header == BrokerEventMetadata::message_header() {
                    protocol.message = Some(Box::new(BrokerEventMetadata::from_bytes(&mut payload)?));
                }
                // Management Protocol
                else if header == InvitationRequest::message_header() {
                    protocol.message = Some(Box::new(InvitationRequest::from_bytes(&mut payload)?));
                } else if header == InvitationResponse::message_header() {
                    protocol.message = Some(Box::new(InvitationResponse::from_bytes(&mut payload)?));
                }
                // Raft Protocol
                else if header == JoinRequest::message_header() {
                    protocol.message = Some(Box::new(JoinRequest::from_bytes(&mut payload)?));
                } else if header == JoinResponse::message_header() {
                    protocol.message = Some(Box::new(JoinResponse::from_bytes(&mut payload)?));
                } else if header == LeaveRequest::message_header() {
                    protocol.message = Some(Box::new(LeaveRequest::from_bytes(&mut payload)?));
                } else if header == LeaveResponse::message_header() {
                    protocol.message = Some(Box::new(LeaveResponse::from_bytes(&mut payload)?));
                } else if header == ConfigurationRequest::message_header() {
                    protocol.message = Some(Box::new(ConfigurationRequest::from_bytes(&mut payload)?));
                } else if header == ConfigurationResponse::message_header() {
                    protocol.message = Some(Box::new(ConfigurationResponse::from_bytes(&mut payload)?));
                } else if header == PollRequest::message_header() {
                    protocol.message = Some(Box::new(PollRequest::from_bytes(&mut payload)?));
                } else if header == PollResponse::message_header() {
                    protocol.message = Some(Box::new(PollResponse::from_bytes(&mut payload)?));
                } else if header == VoteRequest::message_header() {
                    protocol.message = Some(Box::new(VoteRequest::from_bytes(&mut payload)?));
                } else if header == VoteResponse::message_header() {
                    protocol.message = Some(Box::new(VoteResponse::from_bytes(&mut payload)?));
                } else if header == AppendRequest::message_header() {
                    protocol.message = Some(Box::new(AppendRequest::from_bytes(&mut payload)?));
                } else if header == AppendResponse::message_header() {
                    protocol.message = Some(Box::new(AppendResponse::from_bytes(&mut payload)?));
                }
                // Gossip Protocol
                else if header == Gossip::message_header() {
                    protocol.message = Some(Box::new(Gossip::from_bytes(&mut payload)?));
                } else if header == Probe::message_header() {
                    protocol.message = Some(Box::new(Probe::from_bytes(&mut payload)?));
                } else if header == PeerDescriptor::message_header() {
                    protocol.message = Some(Box::new(PeerDescriptor::from_bytes(&mut payload)?));
                } else {
                    bail!("Unknown message header: {:?}", header);
                }

                protocol.header = Some(header);
            }

            protocol.frame = Some(frame);
        }

        Ok(protocol)
    }
}
<|endoftext|>