<commit_before>// Copyright ⓒ 2015-2018 Kevin B. Knapp
//
// `clap_generate` is distributed under the terms of both the MIT license and the Apache License
// (Version 2.0).
// See the [LICENSE-APACHE](LICENSE-APACHE) and [LICENSE-MIT](LICENSE-MIT) files in this repository
// for more information.

//! Generates completions (and other things) for [`clap`](https://github.com/clap-rs/clap) based CLIs
//!
//! - For generating at compile-time, see [`generate_to`]
//! - For generating at runtime, see [`generate`]

#![doc(html_logo_url = "https://clap.rs/images/media/clap.png")]
#![doc(html_root_url = "https://docs.rs/clap_generate/3.0.0-beta.4")]
#![deny(missing_docs, trivial_casts, unused_allocation, trivial_numeric_casts)]
#![forbid(unsafe_code)]
#![allow(clippy::needless_doctest_main)]

const INTERNAL_ERROR_MSG: &str = "Fatal internal error. Please consider filing a bug \
                                  report at https://github.com/clap-rs/clap/issues";

#[macro_use]
#[allow(missing_docs)]
mod macros;

/// Contains some popular generators
pub mod generators;
/// Contains supported shells for auto-completion scripts
mod shell;

use std::ffi::OsString;
use std::fs::File;
use std::io::Error;
use std::io::Write;
use std::path::PathBuf;

#[doc(inline)]
pub use generators::Generator;
#[doc(inline)]
pub use shell::Shell;

/// Generate a completions file for a specified shell at compile-time.
///
/// **NOTE:** to generate the file at compile time you must use a `build.rs` "Build Script" or a
/// [`cargo-xtask`](https://github.com/matklad/cargo-xtask)
///
/// # Examples
///
/// The following example generates a bash completion script via a `build.rs` script. In this
/// simple example, we'll demo a very small application with only a single subcommand and two
/// args. Real applications could be many multiple levels deep in subcommands, and have tens or
/// potentially hundreds of arguments.
///
/// First, it helps if we separate out our `App` definition into a separate file. Whether you
/// do this as a function, or bare App definition is a matter of personal preference.
///
/// ```
/// // src/cli.rs
///
/// use clap::{App, Arg};
///
/// pub fn build_cli() -> App<'static> {
///     App::new("compl")
///         .about("Tests completions")
///         .arg(Arg::new("file")
///             .about("some input file"))
///         .subcommand(App::new("test")
///             .about("tests things")
///             .arg(Arg::new("case")
///                 .long("case")
///                 .takes_value(true)
///                 .about("the case to test")))
/// }
/// ```
///
/// In our regular code, we can simply call this `build_cli()` function, then call
/// `get_matches()`, or any of the other normal methods directly after. For example:
///
/// ```ignore
/// // src/main.rs
///
/// mod cli;
///
/// fn main() {
///     let _m = cli::build_cli().get_matches();
///
///     // normal logic continues...
/// }
/// ```
///
/// Next, we set up our `Cargo.toml` to use a `build.rs` build script.
///
/// ```toml
/// # Cargo.toml
/// build = "build.rs"
///
/// [dependencies]
/// clap = "*"
///
/// [build-dependencies]
/// clap = "*"
/// clap_generate = "*"
/// ```
///
/// Next, we place a `build.rs` in our project root.
///
/// ```ignore
/// use clap_generate::{generate_to, generators::Bash};
/// use std::env;
/// use std::io::Error;
///
/// include!("src/cli.rs");
///
/// fn main() -> Result<(), Error> {
///     let outdir = match env::var_os("OUT_DIR") {
///         None => return Ok(()),
///         Some(outdir) => outdir,
///     };
///
///     let mut app = build_cli();
///     let path = generate_to(
///         Bash,
///         &mut app, // We need to specify what generator to use
///         "myapp",  // We need to specify the bin name manually
///         outdir,   // We need to specify where to write to
///     )?;
///
///     println!("cargo:warning=completion file is generated: {:?}", path);
///
///     Ok(())
/// }
/// ```
///
/// Now, once we compile there will be a `{bin_name}.bash` file in the directory.
/// Assuming we compiled with debug mode, it would be somewhere similar to
/// `<project>/target/debug/build/myapp-<hash>/out/myapp.bash`.
///
/// **NOTE:** Please look at the individual [generators]
/// to see the name of the files generated.
pub fn generate_to<G, S, T>(
    gen: G,
    app: &mut clap::App,
    bin_name: S,
    out_dir: T,
) -> Result<PathBuf, Error>
where
    G: Generator,
    S: Into<String>,
    T: Into<OsString>,
{
    app.set_bin_name(bin_name);

    let out_dir = PathBuf::from(out_dir.into());
    let file_name = gen.file_name(app.get_bin_name().unwrap());

    let path = out_dir.join(file_name);
    let mut file = File::create(&path)?;

    _generate::<G, S>(gen, app, &mut file);
    Ok(path)
}

/// Generate a completions file for a specified shell at runtime.
///
/// Until `cargo install` can install extra files like a completion script, this may be
/// used e.g. in a command that outputs the contents of the completion script, to be
/// redirected into a file by the user.
///
/// # Examples
///
/// Assuming a separate `cli.rs` like the [example above](generate_to()),
/// we can let users generate a completion script using a command:
///
/// ```ignore
/// // src/main.rs
///
/// mod cli;
/// use std::io;
/// use clap_generate::{generate, generators::Bash};
///
/// fn main() {
///     let matches = cli::build_cli().get_matches();
///
///     if matches.is_present("generate-bash-completions") {
///         generate(Bash, &mut cli::build_cli(), "myapp", &mut io::stdout());
///     }
///
///     // normal logic continues...
/// }
///
/// ```
///
/// Usage:
///
/// ```shell
/// $ myapp generate-bash-completions > /usr/share/bash-completion/completions/myapp.bash
/// ```
pub fn generate<G, S>(gen: G, app: &mut clap::App, bin_name: S, buf: &mut dyn Write)
where
    G: Generator,
    S: Into<String>,
{
    app.set_bin_name(bin_name);
    _generate::<G, S>(gen, app, buf)
}

fn _generate<G, S>(gen: G, app: &mut clap::App, buf: &mut dyn Write)
where
    G: Generator,
    S: Into<String>,
{
    // TODO: All the subcommands need to be built instead of just the top one
    app._build();
    app._build_bin_names();

    gen.generate(app, buf)
}
<commit_msg>docs: Bridge the user to 'Shell'<commit_after>// Copyright ⓒ 2015-2018 Kevin B. Knapp
//
// `clap_generate` is distributed under the terms of both the MIT license and the Apache License
// (Version 2.0).
// See the [LICENSE-APACHE](LICENSE-APACHE) and [LICENSE-MIT](LICENSE-MIT) files in this repository
// for more information.

//! Generates completions (and other things) for [`clap`](https://github.com/clap-rs/clap) based CLIs
//!
//! - For generating at compile-time, see [`generate_to`]
//! - For generating at runtime, see [`generate`]
//!
//! [`Shell`] is a convenience `enum` for an argument value type that implements `Generator`
//! for each natively-supported shell type.
//!
//! # Example
//!
//! ```rust,no_run
//! use clap::{App, AppSettings, Arg, ValueHint};
//! use clap_generate::{generate, Generator, Shell};
//! use std::io;
//!
//! fn build_cli() -> App<'static> {
//!     App::new("example")
//!          .arg(Arg::new("file")
//!              .about("some input file")
//!                 .value_hint(ValueHint::AnyPath),
//!         )
//!        .arg(
//!            Arg::new("generator")
//!                .long("generate")
//!                .possible_values(Shell::arg_values()),
//!        )
//! }
//!
//! fn print_completions<G: Generator>(gen: G, app: &mut App) {
//!     generate(gen, app, app.get_name().to_string(), &mut io::stdout());
//! }
//!
//! fn main() {
//!     let matches = build_cli().get_matches();
//!
//!     if let Ok(generator) = matches.value_of_t::<Shell>("generator") {
//!         let mut app = build_cli();
//!         eprintln!("Generating completion file for {}...", generator);
//!         print_completions(generator, &mut app);
//!     }
//! }
//! ```

#![doc(html_logo_url = "https://clap.rs/images/media/clap.png")]
#![doc(html_root_url = "https://docs.rs/clap_generate/3.0.0-beta.4")]
#![deny(missing_docs, trivial_casts, unused_allocation, trivial_numeric_casts)]
#![forbid(unsafe_code)]
#![allow(clippy::needless_doctest_main)]

const INTERNAL_ERROR_MSG: &str = "Fatal internal error. Please consider filing a bug \
                                  report at https://github.com/clap-rs/clap/issues";

#[macro_use]
#[allow(missing_docs)]
mod macros;

/// Contains some popular generators
pub mod generators;
/// Contains supported shells for auto-completion scripts
mod shell;

use std::ffi::OsString;
use std::fs::File;
use std::io::Error;
use std::io::Write;
use std::path::PathBuf;

#[doc(inline)]
pub use generators::Generator;
#[doc(inline)]
pub use shell::Shell;

/// Generate a completions file for a specified shell at compile-time.
///
/// **NOTE:** to generate the file at compile time you must use a `build.rs` "Build Script" or a
/// [`cargo-xtask`](https://github.com/matklad/cargo-xtask)
///
/// # Examples
///
/// The following example generates a bash completion script via a `build.rs` script. In this
/// simple example, we'll demo a very small application with only a single subcommand and two
/// args. Real applications could be many multiple levels deep in subcommands, and have tens or
/// potentially hundreds of arguments.
///
/// First, it helps if we separate out our `App` definition into a separate file. Whether you
/// do this as a function, or bare App definition is a matter of personal preference.
///
/// ```
/// // src/cli.rs
///
/// use clap::{App, Arg};
///
/// pub fn build_cli() -> App<'static> {
///     App::new("compl")
///         .about("Tests completions")
///         .arg(Arg::new("file")
///             .about("some input file"))
///         .subcommand(App::new("test")
///             .about("tests things")
///             .arg(Arg::new("case")
///                 .long("case")
///                 .takes_value(true)
///                 .about("the case to test")))
/// }
/// ```
///
/// In our regular code, we can simply call this `build_cli()` function, then call
/// `get_matches()`, or any of the other normal methods directly after. For example:
///
/// ```ignore
/// // src/main.rs
///
/// mod cli;
///
/// fn main() {
///     let _m = cli::build_cli().get_matches();
///
///     // normal logic continues...
/// }
/// ```
///
/// Next, we set up our `Cargo.toml` to use a `build.rs` build script.
///
/// ```toml
/// # Cargo.toml
/// build = "build.rs"
///
/// [dependencies]
/// clap = "*"
///
/// [build-dependencies]
/// clap = "*"
/// clap_generate = "*"
/// ```
///
/// Next, we place a `build.rs` in our project root.
///
/// ```ignore
/// use clap_generate::{generate_to, generators::Bash};
/// use std::env;
/// use std::io::Error;
///
/// include!("src/cli.rs");
///
/// fn main() -> Result<(), Error> {
///     let outdir = match env::var_os("OUT_DIR") {
///         None => return Ok(()),
///         Some(outdir) => outdir,
///     };
///
///     let mut app = build_cli();
///     let path = generate_to(
///         Bash,
///         &mut app, // We need to specify what generator to use
///         "myapp",  // We need to specify the bin name manually
///         outdir,   // We need to specify where to write to
///     )?;
///
///     println!("cargo:warning=completion file is generated: {:?}", path);
///
///     Ok(())
/// }
/// ```
///
/// Now, once we compile there will be a `{bin_name}.bash` file in the directory.
/// Assuming we compiled with debug mode, it would be somewhere similar to
/// `<project>/target/debug/build/myapp-<hash>/out/myapp.bash`.
///
/// **NOTE:** Please look at the individual [generators]
/// to see the name of the files generated.
pub fn generate_to<G, S, T>(
    gen: G,
    app: &mut clap::App,
    bin_name: S,
    out_dir: T,
) -> Result<PathBuf, Error>
where
    G: Generator,
    S: Into<String>,
    T: Into<OsString>,
{
    app.set_bin_name(bin_name);

    let out_dir = PathBuf::from(out_dir.into());
    let file_name = gen.file_name(app.get_bin_name().unwrap());

    let path = out_dir.join(file_name);
    let mut file = File::create(&path)?;

    _generate::<G, S>(gen, app, &mut file);
    Ok(path)
}

/// Generate a completions file for a specified shell at runtime.
///
/// Until `cargo install` can install extra files like a completion script, this may be
/// used e.g. in a command that outputs the contents of the completion script, to be
/// redirected into a file by the user.
///
/// # Examples
///
/// Assuming a separate `cli.rs` like the [example above](generate_to()),
/// we can let users generate a completion script using a command:
///
/// ```ignore
/// // src/main.rs
///
/// mod cli;
/// use std::io;
/// use clap_generate::{generate, generators::Bash};
///
/// fn main() {
///     let matches = cli::build_cli().get_matches();
///
///     if matches.is_present("generate-bash-completions") {
///         generate(Bash, &mut cli::build_cli(), "myapp", &mut io::stdout());
///     }
///
///     // normal logic continues...
/// }
///
/// ```
///
/// Usage:
///
/// ```shell
/// $ myapp generate-bash-completions > /usr/share/bash-completion/completions/myapp.bash
/// ```
pub fn generate<G, S>(gen: G, app: &mut clap::App, bin_name: S, buf: &mut dyn Write)
where
    G: Generator,
    S: Into<String>,
{
    app.set_bin_name(bin_name);
    _generate::<G, S>(gen, app, buf)
}

fn _generate<G, S>(gen: G, app: &mut clap::App, buf: &mut dyn Write)
where
    G: Generator,
    S: Into<String>,
{
    // TODO: All the subcommands need to be built instead of just the top one
    app._build();
    app._build_bin_names();

    gen.generate(app, buf)
}
<|endoftext|>


<commit_before>/*
 * Copyright (c) 2013-2014, David Renshaw (dwrenshaw@gmail.com)
 *
 * See the LICENSE file in the capnproto-rust root directory.
 */

use std;
use std::vec::Vec;
use common::*;
use endian::*;
use message::*;
use arena;
use io;

pub struct OwnedSpaceMessageReader {
    options : ReaderOptions,
    arena : Box<arena::ReaderArena>,
    segment_slices : Vec<(uint, uint)>,
    owned_space : Vec<Word>,
}

impl MessageReader for OwnedSpaceMessageReader {
    fn get_segment<'b>(&'b self, id : uint) -> &'b [Word] {
        let (a,b) = self.segment_slices.as_slice()[id];
        self.owned_space.slice(a, b)
    }

    fn arena<'b>(&'b self) -> &'b arena::ReaderArena { &*self.arena }
    fn mut_arena<'b>(&'b mut self) -> &'b mut arena::ReaderArena { &mut *self.arena }

    fn get_options<'b>(&'b self) -> &'b ReaderOptions {
        return &self.options;
    }
}

fn invalid_input<T>(desc : &'static str) -> std::io::IoResult<T> {
    return Err(std::io::IoError{ kind : std::io::InvalidInput,
                                 desc : desc,
                                 detail : None});
}

pub fn new_reader<U : std::io::Reader>(inputStream : &mut U,
                                       options : ReaderOptions)
                                       -> std::io::IoResult<OwnedSpaceMessageReader> {

    let firstWord = try!(inputStream.read_exact(8));

    let segmentCount : u32 =
        unsafe {let p : *WireValue<u32> = std::mem::transmute(firstWord.as_ptr());
                (*p).get() + 1
    };

    let segment0Size =
        if segmentCount == 0 { 0 } else {
        unsafe {let p : *WireValue<u32> = std::mem::transmute(firstWord.as_slice().unsafe_ref(4));
                (*p).get()
        }
    };

    let mut totalWords = segment0Size;

    if segmentCount >= 512 {
        return invalid_input("too many segments");
    }

    let mut moreSizes : Vec<u32> = Vec::with_capacity((segmentCount & !1) as uint);

    if segmentCount > 1 {
        let moreSizesRaw = try!(inputStream.read_exact((4 * (segmentCount & !1)) as uint));
        for ii in range(0, segmentCount as uint - 1) {
            let size = unsafe {
                let p : *WireValue<u32> =
                    std::mem::transmute(moreSizesRaw.as_slice().unsafe_ref(ii * 4));
                (*p).get()
            };
            moreSizes.push(size);
            totalWords += size;
        }
    }

    //# Don't accept a message which the receiver couldn't possibly
    //# traverse without hitting the traversal limit. Without this
    //# check, a malicious client could transmit a very large
    //# segment size to make the receiver allocate excessive space
    //# and possibly crash.
    if ! (totalWords as u64 <= options.traversal_limit_in_words)  {
        return invalid_input("Message is too large. To increase the limit on the \
                              receiving end, see capnp::ReaderOptions.");
    }

    let mut ownedSpace : Vec<Word> = allocate_zeroed_words(totalWords as uint);
    let bufLen = totalWords as uint * BYTES_PER_WORD;

    unsafe {
        let ptr : *mut u8 = std::mem::transmute(ownedSpace.as_mut_slice().as_mut_ptr());
        try!(std::slice::raw::mut_buf_as_slice::<u8,std::io::IoResult<uint>>(ptr, bufLen, |buf| {
                    io::read_at_least(inputStream, buf, bufLen)
                }));
    }

    // TODO(maybe someday) lazy reading like in capnp-c++?

    let mut segment_slices : Vec<(uint, uint)> = vec!((0, segment0Size as uint));

    let arena = {
        let segment0 : &[Word] = ownedSpace.slice(0, segment0Size as uint);
        let mut segments : Vec<&[Word]> = vec!(segment0);

        if segmentCount > 1 {
            let mut offset = segment0Size;

            for ii in range(0, segmentCount as uint - 1) {
                segments.push(ownedSpace.slice(offset as uint,
                                               (offset + moreSizes.as_slice()[ii]) as uint));
                segment_slices.push((offset as uint,
                                     (offset + moreSizes.as_slice()[ii]) as uint));
                offset += moreSizes.as_slice()[ii];
            }
        }
        arena::ReaderArena::new(segments.as_slice(), options)
    };

    Ok(OwnedSpaceMessageReader {
        segment_slices : segment_slices,
        owned_space : ownedSpace,
        arena : arena,
        options : options,
    })
}

pub fn write_message<T : std::io::Writer, U : MessageBuilder>(
    outputStream : &mut T,
    message : &U) -> std::io::IoResult<()> {

    try!(message.get_segments_for_output(
        |segments| {

            let tableSize : uint = (segments.len() + 2) & (!1);

            let mut table : Vec<WireValue<u32>> = Vec::with_capacity(tableSize);
            unsafe { table.set_len(tableSize) }

            table.as_mut_slice()[0].set((segments.len() - 1) as u32);

            for i in range(0, segments.len()) {
                table.as_mut_slice()[i + 1].set(segments[i].len() as u32);
            }
            if segments.len() % 2 == 0 {
                // Set padding.
                table.as_mut_slice()[segments.len() + 1].set( 0 );
            }

            unsafe {
                let ptr : *u8 = std::mem::transmute(table.as_ptr());
                try!(std::slice::raw::buf_as_slice::<u8,std::io::IoResult<()>>(ptr, table.len() * 4, |buf| {
                        outputStream.write(buf)
                    }));
            }

            for i in range(0, segments.len()) {
                unsafe {
                    let ptr : *u8 = std::mem::transmute(segments[i].as_ptr());
                    try!(std::slice::raw::buf_as_slice::<u8,std::io::IoResult<()>>(
                        ptr,
                        segments[i].len() * BYTES_PER_WORD,
                        |buf| { outputStream.write(buf) }));
                }
            }
            Ok(())
        }));

    outputStream.flush()
}
<commit_msg>more snake_case<commit_after>/*
 * Copyright (c) 2013-2014, David Renshaw (dwrenshaw@gmail.com)
 *
 * See the LICENSE file in the capnproto-rust root directory.
 */

use std;
use std::vec::Vec;
use common::*;
use endian::*;
use message::*;
use arena;
use io;

pub struct OwnedSpaceMessageReader {
    options : ReaderOptions,
    arena : Box<arena::ReaderArena>,
    segment_slices : Vec<(uint, uint)>,
    owned_space : Vec<Word>,
}

impl MessageReader for OwnedSpaceMessageReader {
    fn get_segment<'b>(&'b self, id : uint) -> &'b [Word] {
        let (a,b) = self.segment_slices.as_slice()[id];
        self.owned_space.slice(a, b)
    }

    fn arena<'b>(&'b self) -> &'b arena::ReaderArena { &*self.arena }
    fn mut_arena<'b>(&'b mut self) -> &'b mut arena::ReaderArena { &mut *self.arena }

    fn get_options<'b>(&'b self) -> &'b ReaderOptions {
        return &self.options;
    }
}

fn invalid_input<T>(desc : &'static str) -> std::io::IoResult<T> {
    return Err(std::io::IoError{ kind : std::io::InvalidInput,
                                 desc : desc,
                                 detail : None});
}

pub fn new_reader<U : std::io::Reader>(inputStream : &mut U,
                                       options : ReaderOptions)
                                       -> std::io::IoResult<OwnedSpaceMessageReader> {

    let firstWord = try!(inputStream.read_exact(8));

    let segmentCount : u32 =
        unsafe {let p : *WireValue<u32> = std::mem::transmute(firstWord.as_ptr());
                (*p).get() + 1
    };

    let segment0Size =
        if segmentCount == 0 { 0 } else {
        unsafe {let p : *WireValue<u32> = std::mem::transmute(firstWord.as_slice().unsafe_ref(4));
                (*p).get()
        }
    };

    let mut totalWords = segment0Size;

    if segmentCount >= 512 {
        return invalid_input("too many segments");
    }

    let mut moreSizes : Vec<u32> = Vec::with_capacity((segmentCount & !1) as uint);

    if segmentCount > 1 {
        let moreSizesRaw = try!(inputStream.read_exact((4 * (segmentCount & !1)) as uint));
        for ii in range(0, segmentCount as uint - 1) {
            let size = unsafe {
                let p : *WireValue<u32> =
                    std::mem::transmute(moreSizesRaw.as_slice().unsafe_ref(ii * 4));
                (*p).get()
            };
            moreSizes.push(size);
            totalWords += size;
        }
    }

    //# Don't accept a message which the receiver couldn't possibly
    //# traverse without hitting the traversal limit. Without this
    //# check, a malicious client could transmit a very large
    //# segment size to make the receiver allocate excessive space
    //# and possibly crash.
    if ! (totalWords as u64 <= options.traversal_limit_in_words)  {
        return invalid_input("Message is too large. To increase the limit on the \
                              receiving end, see capnp::ReaderOptions.");
    }

    let mut ownedSpace : Vec<Word> = allocate_zeroed_words(totalWords as uint);
    let bufLen = totalWords as uint * BYTES_PER_WORD;

    unsafe {
        let ptr : *mut u8 = std::mem::transmute(ownedSpace.as_mut_slice().as_mut_ptr());
        try!(std::slice::raw::mut_buf_as_slice::<u8,std::io::IoResult<uint>>(ptr, bufLen, |buf| {
                    io::read_at_least(inputStream, buf, bufLen)
                }));
    }

    // TODO(maybe someday) lazy reading like in capnp-c++?

    let mut segment_slices : Vec<(uint, uint)> = vec!((0, segment0Size as uint));

    let arena = {
        let segment0 : &[Word] = ownedSpace.slice(0, segment0Size as uint);
        let mut segments : Vec<&[Word]> = vec!(segment0);

        if segmentCount > 1 {
            let mut offset = segment0Size;

            for ii in range(0, segmentCount as uint - 1) {
                segments.push(ownedSpace.slice(offset as uint,
                                               (offset + moreSizes.as_slice()[ii]) as uint));
                segment_slices.push((offset as uint,
                                     (offset + moreSizes.as_slice()[ii]) as uint));
                offset += moreSizes.as_slice()[ii];
            }
        }
        arena::ReaderArena::new(segments.as_slice(), options)
    };

    Ok(OwnedSpaceMessageReader {
        segment_slices : segment_slices,
        owned_space : ownedSpace,
        arena : arena,
        options : options,
    })
}

pub fn write_message<T : std::io::Writer, U : MessageBuilder>(
    output_stream : &mut T,
    message : &U) -> std::io::IoResult<()> {

    try!(message.get_segments_for_output(
        |segments| {

            let table_size : uint = (segments.len() + 2) & (!1);

            let mut table : Vec<WireValue<u32>> = Vec::with_capacity(table_size);
            unsafe { table.set_len(table_size) }

            table.as_mut_slice()[0].set((segments.len() - 1) as u32);

            for i in range(0, segments.len()) {
                table.as_mut_slice()[i + 1].set(segments[i].len() as u32);
            }
            if segments.len() % 2 == 0 {
                // Set padding.
                table.as_mut_slice()[segments.len() + 1].set( 0 );
            }

            unsafe {
                let ptr : *u8 = std::mem::transmute(table.as_ptr());
                try!(std::slice::raw::buf_as_slice::<u8,std::io::IoResult<()>>(ptr, table.len() * 4, |buf| {
                        output_stream.write(buf)
                    }));
            }

            for i in range(0, segments.len()) {
                unsafe {
                    let ptr : *u8 = std::mem::transmute(segments[i].as_ptr());
                    try!(std::slice::raw::buf_as_slice::<u8,std::io::IoResult<()>>(
                        ptr,
                        segments[i].len() * BYTES_PER_WORD,
                        |buf| { output_stream.write(buf) }));
                }
            }
            Ok(())
        }));

    output_stream.flush()
}
<|endoftext|>


<commit_before>// Copyright 2018 Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0
use std::convert::TryInto;

use seccomp::{
    allow_syscall, allow_syscall_if, BpfProgram, Error, SeccompAction, SeccompCmpArgLen as ArgLen,
    SeccompCmpOp::Eq, SeccompCondition as Cond, SeccompError, SeccompFilter, SeccompLevel,
    SeccompRule,
};
use utils::signal::sigrtmin;

/// The default filter containing the white listed syscall rules required by `Firecracker` to
/// function.
/// Any non-trivial modification to this allow list needs a proper comment to specify its source
/// or why the sycall/condition is needed.
pub fn default_filter() -> Result<SeccompFilter, Error> {
    Ok(SeccompFilter::new(
        vec![
            // Called by the api thread to receive data on socket
            allow_syscall_if(
                libc::SYS_accept4,
                or![and![Cond::new(
                    3,
                    ArgLen::DWORD,
                    Eq,
                    libc::SOCK_CLOEXEC as u64
                )?],],
            ),
            // Called for expanding the heap
            allow_syscall(libc::SYS_brk),
            // Used for metrics, via the helpers in utils/src/time.rs
            allow_syscall_if(
                libc::SYS_clock_gettime,
                or![and![Cond::new(
                    0,
                    ArgLen::DWORD,
                    Eq,
                    libc::CLOCK_PROCESS_CPUTIME_ID as u64
                )?],],
            ),
            allow_syscall(libc::SYS_close),
            // Needed for vsock
            allow_syscall(libc::SYS_connect),
            allow_syscall(libc::SYS_epoll_ctl),
            allow_syscall(libc::SYS_epoll_pwait),
            #[cfg(all(target_env = "gnu", target_arch = "x86_64"))]
            allow_syscall(libc::SYS_epoll_wait),
            allow_syscall(libc::SYS_exit),
            allow_syscall(libc::SYS_exit_group),
            // Used by snapshotting, drive patching and rescanning
            allow_syscall_if(
                libc::SYS_fcntl,
                or![and![
                    Cond::new(1, ArgLen::DWORD, Eq, super::FCNTL_F_SETFD)?,
                    Cond::new(2, ArgLen::DWORD, Eq, super::FCNTL_FD_CLOEXEC)?,
                ],],
            ),
            // Used for drive patching & rescanning, for reading the local timezone
            allow_syscall(libc::SYS_fstat),
            // Used for snapshotting
            allow_syscall(libc::SYS_ftruncate),
            // Used for synchronization
            allow_syscall_if(
                libc::SYS_futex,
                or![
                    and![Cond::new(1, ArgLen::DWORD, Eq, super::FUTEX_WAIT_PRIVATE)?],
                    and![Cond::new(1, ArgLen::DWORD, Eq, super::FUTEX_WAKE_PRIVATE)?],
                    #[cfg(target_env = "gnu")]
                    and![Cond::new(
                        1,
                        ArgLen::DWORD,
                        Eq,
                        super::FUTEX_CMP_REQUEUE_PRIVATE
                    )?],
                ],
            ),
            // Used by glibc's tgkill
            #[cfg(target_env = "gnu")]
            allow_syscall(libc::SYS_getpid),
            allow_syscall_if(libc::SYS_ioctl, super::create_ioctl_seccomp_rule()?),
            // Used by the block device
            allow_syscall(libc::SYS_lseek),
            // Triggered by musl for some customer workloads
            #[cfg(target_env = "musl")]
            allow_syscall_if(
                libc::SYS_madvise,
                or![and![Cond::new(
                    2,
                    ArgLen::DWORD,
                    Eq,
                    libc::MADV_DONTNEED as u64
                )?],],
            ),
            // Used for re-allocating large memory regions, for example vectors
            allow_syscall(libc::SYS_mremap),
            // Used for freeing memory
            allow_syscall(libc::SYS_munmap),
            // Used for reading the timezone in LocalTime::now()
            allow_syscall_if(
                libc::SYS_mmap,
                or![
                    and![Cond::new(3, ArgLen::DWORD, Eq, libc::MAP_SHARED as u64)?],
                    and![Cond::new(
                        3,
                        ArgLen::DWORD,
                        Eq,
                        (libc::MAP_FIXED | libc::MAP_ANONYMOUS | libc::MAP_PRIVATE) as u64
                    )?],
                ],
            ),
            #[cfg(target_arch = "x86_64")]
            allow_syscall(libc::SYS_open),
            #[cfg(target_arch = "aarch64")]
            allow_syscall(libc::SYS_openat),
            allow_syscall(libc::SYS_read),
            // Used by the API thread and vsock
            allow_syscall(libc::SYS_recvfrom),
            // SYS_rt_sigreturn is needed in case a fault does occur, so that the signal handler
            // can return. Otherwise we get stuck in a fault loop.
            allow_syscall(libc::SYS_rt_sigreturn),
            // Used by the API thread and vsock
            allow_syscall_if(
                libc::SYS_socket,
                or![and![
                    Cond::new(0, ArgLen::DWORD, Eq, libc::AF_UNIX as u64)?,
                    Cond::new(
                        1,
                        ArgLen::DWORD,
                        Eq,
                        (libc::SOCK_STREAM as u64) | (libc::SOCK_CLOEXEC as u64)
                    )?,
                    Cond::new(2, ArgLen::DWORD, Eq, 0u64)?
                ],],
            ),
            // Used to kick vcpus
            allow_syscall_if(
                libc::SYS_tkill,
                or![and![Cond::new(
                    1,
                    ArgLen::DWORD,
                    Eq,
                    (sigrtmin() + super::super::vstate::vcpu::VCPU_RTSIG_OFFSET) as u64
                )?]],
            ),
            // Used to kick vcpus, on gnu
            #[cfg(target_env = "gnu")]
            allow_syscall(libc::SYS_tgkill),
            // Needed for rate limiting
            allow_syscall_if(
                libc::SYS_timerfd_create,
                or![and![
                    Cond::new(0, ArgLen::DWORD, Eq, libc::CLOCK_MONOTONIC as u64)?,
                    Cond::new(
                        1,
                        ArgLen::DWORD,
                        Eq,
                        (libc::TFD_CLOEXEC as u64) | (libc::TFD_NONBLOCK as u64)
                    )?,
                ],],
            ),
            // Needed for rate limiting
            allow_syscall_if(
                libc::SYS_timerfd_settime,
                or![and![Cond::new(1, ArgLen::DWORD, Eq, 0u64)?],],
            ),
            allow_syscall(libc::SYS_write),
        ]
        .into_iter()
        .collect(),
        SeccompAction::Trap,
    )?)
}

/// Generate a BPF program based on a seccomp level value.
pub fn get_seccomp_filter(seccomp_level: SeccompLevel) -> Result<BpfProgram, SeccompError> {
    match seccomp_level {
        SeccompLevel::None => Ok(vec![]),
        SeccompLevel::Basic => default_filter()
            .map(|filter| filter.allow_all())
            .and_then(|filter| filter.try_into())
            .map_err(SeccompError::SeccompFilter),
        SeccompLevel::Advanced => default_filter()
            .and_then(|filter| filter.try_into())
            .map_err(SeccompError::SeccompFilter),
    }
}

#[cfg(test)]
mod tests {
    use super::get_seccomp_filter;
    use seccomp::SeccompLevel;

    #[test]
    fn test_get_seccomp_filter() {
        assert!(get_seccomp_filter(SeccompLevel::None).is_ok());
        assert!(get_seccomp_filter(SeccompLevel::Basic).is_ok());
        assert!(get_seccomp_filter(SeccompLevel::Advanced).is_ok());
    }
}
<commit_msg>Remove seccomp advanced rule for clock_gettime<commit_after>// Copyright 2018 Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0
use std::convert::TryInto;

use seccomp::{
    allow_syscall, allow_syscall_if, BpfProgram, Error, SeccompAction, SeccompCmpArgLen as ArgLen,
    SeccompCmpOp::Eq, SeccompCondition as Cond, SeccompError, SeccompFilter, SeccompLevel,
    SeccompRule,
};
use utils::signal::sigrtmin;

/// The default filter containing the white listed syscall rules required by `Firecracker` to
/// function.
/// Any non-trivial modification to this allow list needs a proper comment to specify its source
/// or why the sycall/condition is needed.
pub fn default_filter() -> Result<SeccompFilter, Error> {
    Ok(SeccompFilter::new(
        vec![
            // Called by the api thread to receive data on socket
            allow_syscall_if(
                libc::SYS_accept4,
                or![and![Cond::new(
                    3,
                    ArgLen::DWORD,
                    Eq,
                    libc::SOCK_CLOEXEC as u64
                )?],],
            ),
            // Called for expanding the heap
            allow_syscall(libc::SYS_brk),
            // Used for metrics and logging, via the helpers in utils/src/time.rs
            // It's not called on some platforms, because of vdso optimisations. In those cases,
            // musl falls back to the regular syscall.
            allow_syscall(libc::SYS_clock_gettime),
            allow_syscall(libc::SYS_close),
            // Needed for vsock
            allow_syscall(libc::SYS_connect),
            allow_syscall(libc::SYS_epoll_ctl),
            allow_syscall(libc::SYS_epoll_pwait),
            #[cfg(all(target_env = "gnu", target_arch = "x86_64"))]
            allow_syscall(libc::SYS_epoll_wait),
            allow_syscall(libc::SYS_exit),
            allow_syscall(libc::SYS_exit_group),
            // Used by snapshotting, drive patching and rescanning
            allow_syscall_if(
                libc::SYS_fcntl,
                or![and![
                    Cond::new(1, ArgLen::DWORD, Eq, super::FCNTL_F_SETFD)?,
                    Cond::new(2, ArgLen::DWORD, Eq, super::FCNTL_FD_CLOEXEC)?,
                ],],
            ),
            // Used for drive patching & rescanning, for reading the local timezone
            allow_syscall(libc::SYS_fstat),
            // Used for snapshotting
            allow_syscall(libc::SYS_ftruncate),
            // Used for synchronization
            allow_syscall_if(
                libc::SYS_futex,
                or![
                    and![Cond::new(1, ArgLen::DWORD, Eq, super::FUTEX_WAIT_PRIVATE)?],
                    and![Cond::new(1, ArgLen::DWORD, Eq, super::FUTEX_WAKE_PRIVATE)?],
                    #[cfg(target_env = "gnu")]
                    and![Cond::new(
                        1,
                        ArgLen::DWORD,
                        Eq,
                        super::FUTEX_CMP_REQUEUE_PRIVATE
                    )?],
                ],
            ),
            // Used by glibc's tgkill
            #[cfg(target_env = "gnu")]
            allow_syscall(libc::SYS_getpid),
            allow_syscall_if(libc::SYS_ioctl, super::create_ioctl_seccomp_rule()?),
            // Used by the block device
            allow_syscall(libc::SYS_lseek),
            // Triggered by musl for some customer workloads
            #[cfg(target_env = "musl")]
            allow_syscall_if(
                libc::SYS_madvise,
                or![and![Cond::new(
                    2,
                    ArgLen::DWORD,
                    Eq,
                    libc::MADV_DONTNEED as u64
                )?],],
            ),
            // Used for re-allocating large memory regions, for example vectors
            allow_syscall(libc::SYS_mremap),
            // Used for freeing memory
            allow_syscall(libc::SYS_munmap),
            // Used for reading the timezone in LocalTime::now()
            allow_syscall_if(
                libc::SYS_mmap,
                or![
                    and![Cond::new(3, ArgLen::DWORD, Eq, libc::MAP_SHARED as u64)?],
                    and![Cond::new(
                        3,
                        ArgLen::DWORD,
                        Eq,
                        (libc::MAP_FIXED | libc::MAP_ANONYMOUS | libc::MAP_PRIVATE) as u64
                    )?],
                ],
            ),
            #[cfg(target_arch = "x86_64")]
            allow_syscall(libc::SYS_open),
            #[cfg(target_arch = "aarch64")]
            allow_syscall(libc::SYS_openat),
            allow_syscall(libc::SYS_read),
            // Used by the API thread and vsock
            allow_syscall(libc::SYS_recvfrom),
            // SYS_rt_sigreturn is needed in case a fault does occur, so that the signal handler
            // can return. Otherwise we get stuck in a fault loop.
            allow_syscall(libc::SYS_rt_sigreturn),
            // Used by the API thread and vsock
            allow_syscall_if(
                libc::SYS_socket,
                or![and![
                    Cond::new(0, ArgLen::DWORD, Eq, libc::AF_UNIX as u64)?,
                    Cond::new(
                        1,
                        ArgLen::DWORD,
                        Eq,
                        (libc::SOCK_STREAM as u64) | (libc::SOCK_CLOEXEC as u64)
                    )?,
                    Cond::new(2, ArgLen::DWORD, Eq, 0u64)?
                ],],
            ),
            // Used to kick vcpus
            allow_syscall_if(
                libc::SYS_tkill,
                or![and![Cond::new(
                    1,
                    ArgLen::DWORD,
                    Eq,
                    (sigrtmin() + super::super::vstate::vcpu::VCPU_RTSIG_OFFSET) as u64
                )?]],
            ),
            // Used to kick vcpus, on gnu
            #[cfg(target_env = "gnu")]
            allow_syscall(libc::SYS_tgkill),
            // Needed for rate limiting
            allow_syscall_if(
                libc::SYS_timerfd_create,
                or![and![
                    Cond::new(0, ArgLen::DWORD, Eq, libc::CLOCK_MONOTONIC as u64)?,
                    Cond::new(
                        1,
                        ArgLen::DWORD,
                        Eq,
                        (libc::TFD_CLOEXEC as u64) | (libc::TFD_NONBLOCK as u64)
                    )?,
                ],],
            ),
            // Needed for rate limiting
            allow_syscall_if(
                libc::SYS_timerfd_settime,
                or![and![Cond::new(1, ArgLen::DWORD, Eq, 0u64)?],],
            ),
            allow_syscall(libc::SYS_write),
        ]
        .into_iter()
        .collect(),
        SeccompAction::Trap,
    )?)
}

/// Generate a BPF program based on a seccomp level value.
pub fn get_seccomp_filter(seccomp_level: SeccompLevel) -> Result<BpfProgram, SeccompError> {
    match seccomp_level {
        SeccompLevel::None => Ok(vec![]),
        SeccompLevel::Basic => default_filter()
            .map(|filter| filter.allow_all())
            .and_then(|filter| filter.try_into())
            .map_err(SeccompError::SeccompFilter),
        SeccompLevel::Advanced => default_filter()
            .and_then(|filter| filter.try_into())
            .map_err(SeccompError::SeccompFilter),
    }
}

#[cfg(test)]
mod tests {
    use super::get_seccomp_filter;
    use seccomp::SeccompLevel;

    #[test]
    fn test_get_seccomp_filter() {
        assert!(get_seccomp_filter(SeccompLevel::None).is_ok());
        assert!(get_seccomp_filter(SeccompLevel::Basic).is_ok());
        assert!(get_seccomp_filter(SeccompLevel::Advanced).is_ok());
    }
}
<|endoftext|>


<commit_before>use std::fmt::Show;
use test::Bencher;

use serialize::{Decoder, Decodable};

use de::{Deserializer, Deserializable};

//////////////////////////////////////////////////////////////////////////////

#[deriving(Show)]
enum Error {
    EndOfStream,
    SyntaxError,
}

//////////////////////////////////////////////////////////////////////////////

mod decoder {
    use std::vec;
    use serialize::Decoder;

    use super::{Error, EndOfStream, SyntaxError};

    pub struct IntsDecoder {
        len: uint,
        iter: vec::MoveItems<int>,
    }

    impl IntsDecoder {
        #[inline]
        pub fn new(values: Vec<int>) -> IntsDecoder {
            IntsDecoder {
                len: values.len(),
                iter: values.move_iter(),
            }
        }
    }

    impl Decoder<Error> for IntsDecoder {
        // Primitive types:
        fn read_nil(&mut self) -> Result<(), Error> { Err(SyntaxError) }
        fn read_uint(&mut self) -> Result<uint, Error> { Err(SyntaxError) }
        fn read_u64(&mut self) -> Result<u64, Error> { Err(SyntaxError) }
        fn read_u32(&mut self) -> Result<u32, Error> { Err(SyntaxError) }
        fn read_u16(&mut self) -> Result<u16, Error> { Err(SyntaxError) }
        fn read_u8(&mut self) -> Result<u8, Error> { Err(SyntaxError) }
        #[inline]
        fn read_int(&mut self) -> Result<int, Error> {
            match self.iter.next() {
                Some(value) => Ok(value),
                None => Err(EndOfStream),
            }
        }
        fn read_i64(&mut self) -> Result<i64, Error> { Err(SyntaxError) }
        fn read_i32(&mut self) -> Result<i32, Error> { Err(SyntaxError) }
        fn read_i16(&mut self) -> Result<i16, Error> { Err(SyntaxError) }
        fn read_i8(&mut self) -> Result<i8, Error> { Err(SyntaxError) }
        fn read_bool(&mut self) -> Result<bool, Error> { Err(SyntaxError) }
        fn read_f64(&mut self) -> Result<f64, Error> { Err(SyntaxError) }
        fn read_f32(&mut self) -> Result<f32, Error> { Err(SyntaxError) }
        fn read_char(&mut self) -> Result<char, Error> { Err(SyntaxError) }
        fn read_str(&mut self) -> Result<StrBuf, Error> { Err(SyntaxError) }

        // Compound types:
        fn read_enum<T>(&mut self, _name: &str, _f: |&mut IntsDecoder| -> Result<T, Error>) -> Result<T, Error> { Err(SyntaxError) }

        fn read_enum_variant<T>(&mut self,
                                _names: &[&str],
                                _f: |&mut IntsDecoder, uint| -> Result<T, Error>)
                                -> Result<T, Error> { Err(SyntaxError) }
        fn read_enum_variant_arg<T>(&mut self,
                                    _a_idx: uint,
                                    _f: |&mut IntsDecoder| -> Result<T, Error>)
                                    -> Result<T, Error> { Err(SyntaxError) }

        fn read_enum_struct_variant<T>(&mut self,
                                       _names: &[&str],
                                       _f: |&mut IntsDecoder, uint| -> Result<T, Error>)
                                       -> Result<T, Error> { Err(SyntaxError) }
        fn read_enum_struct_variant_field<T>(&mut self,
                                             _f_name: &str,
                                             _f_idx: uint,
                                             _f: |&mut IntsDecoder| -> Result<T, Error>)
                                             -> Result<T, Error> { Err(SyntaxError) }

        fn read_struct<T>(&mut self, _s_name: &str, _len: uint, _f: |&mut IntsDecoder| -> Result<T, Error>)
                          -> Result<T, Error> { Err(SyntaxError) }
        fn read_struct_field<T>(&mut self,
                                _f_name: &str,
                                _f_idx: uint,
                                _f: |&mut IntsDecoder| -> Result<T, Error>)
                                -> Result<T, Error> { Err(SyntaxError) }

        fn read_tuple<T>(&mut self, _f: |&mut IntsDecoder, uint| -> Result<T, Error>) -> Result<T, Error> { Err(SyntaxError) }
        fn read_tuple_arg<T>(&mut self, _a_idx: uint, _f: |&mut IntsDecoder| -> Result<T, Error>) -> Result<T, Error> { Err(SyntaxError) }

        fn read_tuple_struct<T>(&mut self,
                                _s_name: &str,
                                _f: |&mut IntsDecoder, uint| -> Result<T, Error>)
                                -> Result<T, Error> { Err(SyntaxError) }
        fn read_tuple_struct_arg<T>(&mut self,
                                    _a_idx: uint,
                                    _f: |&mut IntsDecoder| -> Result<T, Error>)
                                    -> Result<T, Error> { Err(SyntaxError) }

        // Specialized types:
        fn read_option<T>(&mut self, _f: |&mut IntsDecoder, bool| -> Result<T, Error>) -> Result<T, Error> { Err(SyntaxError) }

        #[inline]
        fn read_seq<T>(&mut self, f: |&mut IntsDecoder, uint| -> Result<T, Error>) -> Result<T, Error> {
            f(self, self.len)
        }
        #[inline]
        fn read_seq_elt<T>(&mut self, _idx: uint, f: |&mut IntsDecoder| -> Result<T, Error>) -> Result<T, Error> {
            f(self)
        }

        fn read_map<T>(&mut self, _f: |&mut IntsDecoder, uint| -> Result<T, Error>) -> Result<T, Error> { Err(SyntaxError) }
        fn read_map_elt_key<T>(&mut self, _idx: uint, _f: |&mut IntsDecoder| -> Result<T, Error>) -> Result<T, Error> { Err(SyntaxError) }
        fn read_map_elt_val<T>(&mut self, _idx: uint, _f: |&mut IntsDecoder| -> Result<T, Error>) -> Result<T, Error> { Err(SyntaxError) }
    }
}

//////////////////////////////////////////////////////////////////////////////

mod deserializer {
    use std::num;
    use std::vec;

    use super::{Error, EndOfStream, SyntaxError};

    use de::Deserializer;
    use de::{Token, Int, SeqStart, Sep, End};

    #[deriving(Eq, Show)]
    enum State {
        StartState,
        SepOrEndState,
        EndState,
    }

    pub struct IntsDeserializer {
        state: State,
        len: uint,
        iter: vec::MoveItems<int>,
    }

    impl IntsDeserializer {
        #[inline]
        pub fn new(values: Vec<int>) -> IntsDeserializer {
            IntsDeserializer {
                state: StartState,
                len: values.len(),
                iter: values.move_iter(),
            }
        }
    }

    impl Iterator<Result<Token, Error>> for IntsDeserializer {
        #[inline]
        fn next(&mut self) -> Option<Result<Token, Error>> {
            match self.state {
                StartState => {
                    self.state = SepOrEndState;
                    Some(Ok(SeqStart(self.len)))
                }
                SepOrEndState => {
                    match self.iter.next() {
                        Some(value) => {
                            Some(Ok(Int(value)))
                        }
                        None => {
                            self.state = EndState;
                            Some(Ok(End))
                        }
                    }
                }
                EndState => {
                    None
                }
            }
        }
    }

    impl Deserializer<Error> for IntsDeserializer {
        #[inline]
        fn end_of_stream_error(&self) -> Error {
            EndOfStream
        }

        #[inline]
        fn syntax_error(&self) -> Error {
            SyntaxError
        }
    }
}

//////////////////////////////////////////////////////////////////////////////

fn run_decoder<
    E: Show,
    D: Decoder<E>,
    T: Clone + Eq + Show + Decodable<D, E>
>(mut d: D, value: T) {
    let v: T = Decodable::decode(&mut d).unwrap();

    assert_eq!(value, v);
}

fn run_deserializer<
    E: Show,
    D: Deserializer<E>,
    T: Clone + Eq + Show + Deserializable<E, D>
>(mut d: D, value: T) {
    let v: T = Deserializable::deserialize(&mut d).unwrap();

    assert_eq!(value, v);
}

#[bench]
fn bench_decoder_vec_int_000(b: &mut Bencher) {
    b.iter(|| {
        let v: Vec<int> = vec!();
        run_decoder(decoder::IntsDecoder::new(v.clone()), v)
    })
}

#[bench]
fn bench_deserializer_vec_int_000(b: &mut Bencher) {
    b.iter(|| {
        let v: Vec<int> = vec!();
        run_deserializer(deserializer::IntsDeserializer::new(v.clone()), v)
    })
}

#[bench]
fn bench_decoder_vec_int_003(b: &mut Bencher) {
    b.iter(|| {
        let v: Vec<int> = vec!(1, 2, 3);
        run_decoder(decoder::IntsDecoder::new(v.clone()), v)
    })
}

#[bench]
fn bench_deserializer_vec_int_003(b: &mut Bencher) {
    b.iter(|| {
        let v: Vec<int> = vec!(1, 2, 3);
        run_deserializer(deserializer::IntsDeserializer::new(v.clone()), v)
    })
}

#[bench]
fn bench_decoder_vec_int_100(b: &mut Bencher) {
    b.iter(|| {
        let v: Vec<int> = range(0, 100).collect();
        run_decoder(decoder::IntsDecoder::new(v.clone()), v)
    })
}

#[bench]
fn bench_deserializer_vec_int_100(b: &mut Bencher) {
    b.iter(|| {
        let v: Vec<int> = range(0, 100).collect();
        run_deserializer(deserializer::IntsDeserializer::new(v.clone()), v)
    })
}
<commit_msg>minor cleanup<commit_after>use std::fmt::Show;
use test::Bencher;

use serialize::{Decoder, Decodable};

use de::{Deserializer, Deserializable};

//////////////////////////////////////////////////////////////////////////////

#[deriving(Show)]
enum Error {
    EndOfStream,
    SyntaxError,
}

//////////////////////////////////////////////////////////////////////////////

mod decoder {
    use std::vec;
    use serialize;

    use super::{Error, EndOfStream, SyntaxError};

    pub struct Decoder {
        len: uint,
        iter: vec::MoveItems<int>,
    }

    impl Decoder {
        #[inline]
        pub fn new(values: Vec<int>) -> Decoder {
            Decoder {
                len: values.len(),
                iter: values.move_iter(),
            }
        }
    }

    impl serialize::Decoder<Error> for Decoder {
        // Primitive types:
        fn read_nil(&mut self) -> Result<(), Error> { Err(SyntaxError) }
        fn read_uint(&mut self) -> Result<uint, Error> { Err(SyntaxError) }
        fn read_u64(&mut self) -> Result<u64, Error> { Err(SyntaxError) }
        fn read_u32(&mut self) -> Result<u32, Error> { Err(SyntaxError) }
        fn read_u16(&mut self) -> Result<u16, Error> { Err(SyntaxError) }
        fn read_u8(&mut self) -> Result<u8, Error> { Err(SyntaxError) }
        #[inline]
        fn read_int(&mut self) -> Result<int, Error> {
            match self.iter.next() {
                Some(value) => Ok(value),
                None => Err(EndOfStream),
            }
        }
        fn read_i64(&mut self) -> Result<i64, Error> { Err(SyntaxError) }
        fn read_i32(&mut self) -> Result<i32, Error> { Err(SyntaxError) }
        fn read_i16(&mut self) -> Result<i16, Error> { Err(SyntaxError) }
        fn read_i8(&mut self) -> Result<i8, Error> { Err(SyntaxError) }
        fn read_bool(&mut self) -> Result<bool, Error> { Err(SyntaxError) }
        fn read_f64(&mut self) -> Result<f64, Error> { Err(SyntaxError) }
        fn read_f32(&mut self) -> Result<f32, Error> { Err(SyntaxError) }
        fn read_char(&mut self) -> Result<char, Error> { Err(SyntaxError) }
        fn read_str(&mut self) -> Result<StrBuf, Error> { Err(SyntaxError) }

        // Compound types:
        fn read_enum<T>(&mut self, _name: &str, _f: |&mut Decoder| -> Result<T, Error>) -> Result<T, Error> { Err(SyntaxError) }

        fn read_enum_variant<T>(&mut self,
                                _names: &[&str],
                                _f: |&mut Decoder, uint| -> Result<T, Error>)
                                -> Result<T, Error> { Err(SyntaxError) }
        fn read_enum_variant_arg<T>(&mut self,
                                    _a_idx: uint,
                                    _f: |&mut Decoder| -> Result<T, Error>)
                                    -> Result<T, Error> { Err(SyntaxError) }

        fn read_enum_struct_variant<T>(&mut self,
                                       _names: &[&str],
                                       _f: |&mut Decoder, uint| -> Result<T, Error>)
                                       -> Result<T, Error> { Err(SyntaxError) }
        fn read_enum_struct_variant_field<T>(&mut self,
                                             _f_name: &str,
                                             _f_idx: uint,
                                             _f: |&mut Decoder| -> Result<T, Error>)
                                             -> Result<T, Error> { Err(SyntaxError) }

        fn read_struct<T>(&mut self, _s_name: &str, _len: uint, _f: |&mut Decoder| -> Result<T, Error>)
                          -> Result<T, Error> { Err(SyntaxError) }
        fn read_struct_field<T>(&mut self,
                                _f_name: &str,
                                _f_idx: uint,
                                _f: |&mut Decoder| -> Result<T, Error>)
                                -> Result<T, Error> { Err(SyntaxError) }

        fn read_tuple<T>(&mut self, _f: |&mut Decoder, uint| -> Result<T, Error>) -> Result<T, Error> { Err(SyntaxError) }
        fn read_tuple_arg<T>(&mut self, _a_idx: uint, _f: |&mut Decoder| -> Result<T, Error>) -> Result<T, Error> { Err(SyntaxError) }

        fn read_tuple_struct<T>(&mut self,
                                _s_name: &str,
                                _f: |&mut Decoder, uint| -> Result<T, Error>)
                                -> Result<T, Error> { Err(SyntaxError) }
        fn read_tuple_struct_arg<T>(&mut self,
                                    _a_idx: uint,
                                    _f: |&mut Decoder| -> Result<T, Error>)
                                    -> Result<T, Error> { Err(SyntaxError) }

        // Specialized types:
        fn read_option<T>(&mut self, _f: |&mut Decoder, bool| -> Result<T, Error>) -> Result<T, Error> { Err(SyntaxError) }

        #[inline]
        fn read_seq<T>(&mut self, f: |&mut Decoder, uint| -> Result<T, Error>) -> Result<T, Error> {
            f(self, self.len)
        }
        #[inline]
        fn read_seq_elt<T>(&mut self, _idx: uint, f: |&mut Decoder| -> Result<T, Error>) -> Result<T, Error> {
            f(self)
        }

        fn read_map<T>(&mut self, _f: |&mut Decoder, uint| -> Result<T, Error>) -> Result<T, Error> { Err(SyntaxError) }
        fn read_map_elt_key<T>(&mut self, _idx: uint, _f: |&mut Decoder| -> Result<T, Error>) -> Result<T, Error> { Err(SyntaxError) }
        fn read_map_elt_val<T>(&mut self, _idx: uint, _f: |&mut Decoder| -> Result<T, Error>) -> Result<T, Error> { Err(SyntaxError) }
    }
}

//////////////////////////////////////////////////////////////////////////////

mod deserializer {
    //use std::num;
    use std::vec;

    use super::{Error, EndOfStream, SyntaxError};

    use de;

    #[deriving(Eq, Show)]
    enum State {
        StartState,
        SepOrEndState,
        EndState,
    }

    pub struct Deserializer {
        state: State,
        len: uint,
        iter: vec::MoveItems<int>,
    }

    impl Deserializer {
        #[inline]
        pub fn new(values: Vec<int>) -> Deserializer {
            Deserializer {
                state: StartState,
                len: values.len(),
                iter: values.move_iter(),
            }
        }
    }

    impl Iterator<Result<de::Token, Error>> for Deserializer {
        #[inline]
        fn next(&mut self) -> Option<Result<de::Token, Error>> {
            match self.state {
                StartState => {
                    self.state = SepOrEndState;
                    Some(Ok(de::SeqStart(self.len)))
                }
                SepOrEndState => {
                    match self.iter.next() {
                        Some(value) => {
                            Some(Ok(de::Int(value)))
                        }
                        None => {
                            self.state = EndState;
                            Some(Ok(de::End))
                        }
                    }
                }
                EndState => {
                    None
                }
            }
        }
    }

    impl de::Deserializer<Error> for Deserializer {
        #[inline]
        fn end_of_stream_error(&self) -> Error {
            EndOfStream
        }

        #[inline]
        fn syntax_error(&self) -> Error {
            SyntaxError
        }
    }
}

//////////////////////////////////////////////////////////////////////////////

fn run_decoder<
    E: Show,
    D: Decoder<E>,
    T: Clone + Eq + Show + Decodable<D, E>
>(mut d: D, value: T) {
    let v: T = Decodable::decode(&mut d).unwrap();

    assert_eq!(value, v);
}

fn run_deserializer<
    E: Show,
    D: Deserializer<E>,
    T: Clone + Eq + Show + Deserializable<E, D>
>(mut d: D, value: T) {
    let v: T = Deserializable::deserialize(&mut d).unwrap();

    assert_eq!(value, v);
}

#[bench]
fn bench_decoder_vec_int_000(b: &mut Bencher) {
    b.iter(|| {
        let v: Vec<int> = vec!();
        run_decoder(decoder::Decoder::new(v.clone()), v)
    })
}

#[bench]
fn bench_deserializer_vec_int_000(b: &mut Bencher) {
    b.iter(|| {
        let v: Vec<int> = vec!();
        run_deserializer(deserializer::Deserializer::new(v.clone()), v)
    })
}

#[bench]
fn bench_decoder_vec_int_003(b: &mut Bencher) {
    b.iter(|| {
        let v: Vec<int> = vec!(1, 2, 3);
        run_decoder(decoder::Decoder::new(v.clone()), v)
    })
}

#[bench]
fn bench_deserializer_vec_int_003(b: &mut Bencher) {
    b.iter(|| {
        let v: Vec<int> = vec!(1, 2, 3);
        run_deserializer(deserializer::Deserializer::new(v.clone()), v)
    })
}

#[bench]
fn bench_decoder_vec_int_100(b: &mut Bencher) {
    b.iter(|| {
        let v: Vec<int> = range(0, 100).collect();
        run_decoder(decoder::Decoder::new(v.clone()), v)
    })
}

#[bench]
fn bench_deserializer_vec_int_100(b: &mut Bencher) {
    b.iter(|| {
        let v: Vec<int> = range(0, 100).collect();
        run_deserializer(deserializer::Deserializer::new(v.clone()), v)
    })
}
<|endoftext|>


<commit_before>//! This transport sends emails using the SMTP protocol

use EmailTransport;
use SendableEmail;
use openssl::ssl::{SslContext, SslMethod};
use smtp::authentication::Mechanism;
use smtp::client::Client;
use smtp::error::{Error, SmtpResult};
use smtp::extension::{Extension, ServerInfo};
use std::net::{SocketAddr, ToSocketAddrs};
use std::string::String;
use std::time::Duration;

pub mod extension;
pub mod authentication;
pub mod response;
pub mod client;
pub mod error;

// Registrated port numbers:
// https://www.iana.
// org/assignments/service-names-port-numbers/service-names-port-numbers.xhtml

/// Default smtp port
pub const SMTP_PORT: u16 = 25;

/// Default submission port
pub const SUBMISSION_PORT: u16 = 587;

// Useful strings and characters

/// The word separator for SMTP transactions
pub const SP: &'static str = " ";

/// The line ending for SMTP transactions (carriage return + line feed)
pub const CRLF: &'static str = "\r\n";

/// Colon
pub const COLON: &'static str = ":";

/// The ending of message content
pub const MESSAGE_ENDING: &'static str = "\r\n.\r\n";

/// NUL unicode character
pub const NUL: &'static str = "\0";

/// TLS security level
#[derive(PartialEq,Eq,Copy,Clone,Debug)]
pub enum SecurityLevel {
    /// Use a TLS wrapped connection
    ///
    /// Non RFC-compliant, should only be used if the server does not support STARTTLS.
    EncryptedWrapper,
    /// Only send an email on encrypted connection (with STARTTLS)
    ///
    /// Recommended mode, prevents MITM when used with verified certificates.
    AlwaysEncrypt,
    /// Use TLS when available (with STARTTLS)
    ///
    /// Default mode.
    Opportunistic,
    /// Never use TLS
    NeverEncrypt,
}

/// Contains client configuration
#[derive(Debug)]
pub struct SmtpTransportBuilder {
    /// Maximum connection reuse
    ///
    /// Zero means no limitation
    connection_reuse_count_limit: u16,
    /// Enable connection reuse
    connection_reuse: bool,
    /// Name sent during HELO or EHLO
    hello_name: String,
    /// Credentials
    credentials: Option<(String, String)>,
    /// Socket we are connecting to
    server_addr: SocketAddr,
    /// SSL context to use
    ssl_context: SslContext,
    /// TLS security level
    security_level: SecurityLevel,
    /// Enable UTF8 mailboxes in envelope or headers
    smtp_utf8: bool,
    /// Optional enforced authentication mechanism
    authentication_mechanism: Option<Mechanism>,
    /// Define network timeout
    /// It can be changed later for specific needs (like a different timeout for each SMTP command)
    timeout: Option<Duration>,
}

/// Builder for the SMTP `SmtpTransport`
impl SmtpTransportBuilder {
    /// Creates a new local SMTP client
    pub fn new<A: ToSocketAddrs>(addr: A) -> Result<SmtpTransportBuilder, Error> {
        let mut addresses = try!(addr.to_socket_addrs());

        match addresses.next() {
            Some(addr) => {
                Ok(SmtpTransportBuilder {
                       server_addr: addr,
                       ssl_context: SslContext::builder(SslMethod::tls()).unwrap().build(),
                       security_level: SecurityLevel::AlwaysEncrypt,
                       smtp_utf8: false,
                       credentials: None,
                       connection_reuse_count_limit: 100,
                       connection_reuse: false,
                       hello_name: "localhost".to_string(),
                       authentication_mechanism: None,
                       timeout: Some(Duration::new(60, 0)),
                   })
            }
            None => Err(From::from("Could nor resolve hostname")),
        }
    }

    /// Creates a new local SMTP client to port 25
    pub fn localhost() -> Result<SmtpTransportBuilder, Error> {
        SmtpTransportBuilder::new(("localhost", SMTP_PORT))
    }

    /// Use STARTTLS with a specific context
    pub fn ssl_context(mut self, ssl_context: SslContext) -> SmtpTransportBuilder {
        self.ssl_context = ssl_context;
        self
    }

    /// Set the security level for SSL/TLS
    pub fn security_level(mut self, level: SecurityLevel) -> SmtpTransportBuilder {
        self.security_level = level;
        self
    }

    /// Require SSL/TLS using STARTTLS
    ///
    /// Incompatible with `ssl_wrapper()``
    pub fn encrypt(mut self) -> SmtpTransportBuilder {
        self.security_level = SecurityLevel::AlwaysEncrypt;
        self
    }

    /// Require SSL/TLS using SMTPS
    ///
    /// Incompatible with `encrypt()`
    pub fn ssl_wrapper(mut self) -> SmtpTransportBuilder {
        self.security_level = SecurityLevel::EncryptedWrapper;
        self
    }

    /// Enable SMTPUTF8 if the server supports it
    pub fn smtp_utf8(mut self, enabled: bool) -> SmtpTransportBuilder {
        self.smtp_utf8 = enabled;
        self
    }

    /// Set the name used during HELO or EHLO
    pub fn hello_name<S: Into<String>>(mut self, name: S) -> SmtpTransportBuilder {
        self.hello_name = name.into();
        self
    }

    /// Enable connection reuse
    pub fn connection_reuse(mut self, enable: bool) -> SmtpTransportBuilder {
        self.connection_reuse = enable;
        self
    }

    /// Set the maximum number of emails sent using one connection
    pub fn connection_reuse_count_limit(mut self, limit: u16) -> SmtpTransportBuilder {
        self.connection_reuse_count_limit = limit;
        self
    }

    /// Set the client credentials
    pub fn credentials<S: Into<String>>(mut self,
                                        username: S,
                                        password: S)
                                        -> SmtpTransportBuilder {
        self.credentials = Some((username.into(), password.into()));
        self
    }

    /// Set the authentication mechanisms
    pub fn authentication_mechanism(mut self, mechanism: Mechanism) -> SmtpTransportBuilder {
        self.authentication_mechanism = Some(mechanism);
        self
    }

    /// Set the timeout duration
    pub fn timeout(mut self, timeout: Option<Duration>) -> SmtpTransportBuilder {
        self.timeout = timeout;
        self
    }

    /// Build the SMTP client
    ///
    /// It does not connect to the server, but only creates the `SmtpTransport`
    pub fn build(self) -> SmtpTransport {
        SmtpTransport::new(self)
    }
}

/// Represents the state of a client
#[derive(Debug)]
struct State {
    /// Panic state
    pub panic: bool,
    /// Connection reuse counter
    pub connection_reuse_count: u16,
}

/// Structure that implements the high level SMTP client
#[derive(Debug)]
pub struct SmtpTransport {
    /// Information about the server
    /// Value is None before HELO/EHLO
    server_info: Option<ServerInfo>,
    /// SmtpTransport variable states
    state: State,
    /// Information about the client
    client_info: SmtpTransportBuilder,
    /// Low level client
    client: Client,
}

macro_rules! try_smtp (
    ($err: expr, $client: ident) => ({
        match $err {
            Ok(val) => val,
            Err(err) => {
                if !$client.state.panic {
                    $client.state.panic = true;
                    $client.reset();
                }
                return Err(From::from(err))
            },
        }
    })
);

impl SmtpTransport {
    /// Creates a new SMTP client
    ///
    /// It does not connect to the server, but only creates the `SmtpTransport`
    pub fn new(builder: SmtpTransportBuilder) -> SmtpTransport {

        let client = Client::new();

        SmtpTransport {
            client: client,
            server_info: None,
            client_info: builder,
            state: State {
                panic: false,
                connection_reuse_count: 0,
            },
        }
    }

    /// Reset the client state
    fn reset(&mut self) {
        // Close the SMTP transaction if needed
        self.close();

        // Reset the client state
        self.server_info = None;
        self.state.panic = false;
        self.state.connection_reuse_count = 0;
    }

    /// Gets the EHLO response and updates server information
    pub fn get_ehlo(&mut self) -> SmtpResult {
        // Extended Hello
        let ehlo_response = try_smtp!(self.client.ehlo(&self.client_info.hello_name), self);

        self.server_info = Some(try_smtp!(ServerInfo::from_response(&ehlo_response), self));

        // Print server information
        debug!("server {}", self.server_info.as_ref().unwrap());

        Ok(ehlo_response)
    }
}

impl EmailTransport<SmtpResult> for SmtpTransport {
    /// Sends an email
    #[cfg_attr(feature = "cargo-clippy", allow(match_same_arms, cyclomatic_complexity))]
    fn send<T: SendableEmail>(&mut self, email: T) -> SmtpResult {

        // Extract email information
        let message_id = email.message_id();

        // Check if the connection is still available
        if (self.state.connection_reuse_count > 0) && (!self.client.is_connected()) {
            self.reset();
        }

        if self.state.connection_reuse_count == 0 {
            try!(self.client
                     .connect(&self.client_info.server_addr,
                              match self.client_info.security_level {
                                  SecurityLevel::EncryptedWrapper => {
                                      Some(&self.client_info.ssl_context)
                                  }
                                  _ => None,
                              }));

            try!(self.client.set_timeout(self.client_info.timeout));

            // Log the connection
            info!("connection established to {}", self.client_info.server_addr);

            try!(self.get_ehlo());

            match (&self.client_info.security_level,
                   self.server_info
                       .as_ref()
                       .unwrap()
                       .supports_feature(&Extension::StartTls)) {
                (&SecurityLevel::AlwaysEncrypt, false) => {
                    return Err(From::from("Could not encrypt connection, aborting"))
                }
                (&SecurityLevel::Opportunistic, false) => (),
                (&SecurityLevel::NeverEncrypt, _) => (),
                (&SecurityLevel::EncryptedWrapper, _) => (),
                (_, true) => {
                    try_smtp!(self.client.starttls(), self);
                    try_smtp!(self.client
                                  .upgrade_tls_stream(&self.client_info.ssl_context),
                              self);

                    debug!("connection encrypted");

                    // Send EHLO again
                    try!(self.get_ehlo());
                }
            }

            if self.client_info.credentials.is_some() {
                let (username, password) = self.client_info.credentials.clone().unwrap();

                let mut found = false;

                // Compute accepted mechanism
                let accepted_mechanisms = match self.client_info.authentication_mechanism {
                    Some(mechanism) => vec![mechanism],
                    None => {
                        if self.client.is_encrypted() {
                            // If encrypted, allow all mechanisms, with a preference for the
                            // simplest
                            // Login is obsolete so try it last
                            vec![Mechanism::Plain, Mechanism::CramMd5, Mechanism::Login]
                        } else {
                            // If not encrypted, do not allow clear-text passwords
                            vec![Mechanism::CramMd5]
                        }
                    }
                };

                for mechanism in accepted_mechanisms {
                    if self.server_info
                           .as_ref()
                           .unwrap()
                           .supports_auth_mechanism(mechanism) {
                        found = true;
                        try_smtp!(self.client.auth(mechanism, &username, &password), self);
                        break;
                    }
                }

                if !found {
                    info!("No supported authentication mechanisms available");
                }
            }
        }

        // Mail
        let mail_options = match (self.server_info
                                      .as_ref()
                                      .unwrap()
                                      .supports_feature(&Extension::EightBitMime),
                                  self.server_info
                                      .as_ref()
                                      .unwrap()
                                      .supports_feature(&Extension::SmtpUtfEight)) {
            (true, true) => Some("BODY=8BITMIME SMTPUTF8"),
            (true, false) => Some("BODY=8BITMIME"),
            (false, _) => None,
        };

        try_smtp!(self.client.mail(&email.from(), mail_options), self);

        // Log the mail command
        info!("{}: from=<{}>", message_id, email.from());

        // Recipient
        for to_address in &email.to() {
            try_smtp!(self.client.rcpt(to_address), self);
            // Log the rcpt command
            info!("{}: to=<{}>", message_id, to_address);
        }

        // Data
        try_smtp!(self.client.data(), self);

        // Message content
        let message = email.message();
        let result = self.client.message(&message);

        if result.is_ok() {
            // Increment the connection reuse counter
            self.state.connection_reuse_count += 1;

            // Log the message
            info!("{}: conn_use={}, size={}, status=sent ({})",
                  message_id,
                  self.state.connection_reuse_count,
                  message.len(),
                  result
                      .as_ref()
                      .ok()
                      .unwrap()
                      .message()
                      .iter()
                      .next()
                      .unwrap_or(&"no response".to_string()));
        }

        // Test if we can reuse the existing connection
        if (!self.client_info.connection_reuse) ||
           (self.state.connection_reuse_count >= self.client_info.connection_reuse_count_limit) {
            self.reset();
        }

        result
    }

    /// Closes the inner connection
    fn close(&mut self) {
        self.client.close();
    }
}
<commit_msg>Fix typo in error string<commit_after>//! This transport sends emails using the SMTP protocol

use EmailTransport;
use SendableEmail;
use openssl::ssl::{SslContext, SslMethod};
use smtp::authentication::Mechanism;
use smtp::client::Client;
use smtp::error::{Error, SmtpResult};
use smtp::extension::{Extension, ServerInfo};
use std::net::{SocketAddr, ToSocketAddrs};
use std::string::String;
use std::time::Duration;

pub mod extension;
pub mod authentication;
pub mod response;
pub mod client;
pub mod error;

// Registrated port numbers:
// https://www.iana.
// org/assignments/service-names-port-numbers/service-names-port-numbers.xhtml

/// Default smtp port
pub const SMTP_PORT: u16 = 25;

/// Default submission port
pub const SUBMISSION_PORT: u16 = 587;

// Useful strings and characters

/// The word separator for SMTP transactions
pub const SP: &'static str = " ";

/// The line ending for SMTP transactions (carriage return + line feed)
pub const CRLF: &'static str = "\r\n";

/// Colon
pub const COLON: &'static str = ":";

/// The ending of message content
pub const MESSAGE_ENDING: &'static str = "\r\n.\r\n";

/// NUL unicode character
pub const NUL: &'static str = "\0";

/// TLS security level
#[derive(PartialEq,Eq,Copy,Clone,Debug)]
pub enum SecurityLevel {
    /// Use a TLS wrapped connection
    ///
    /// Non RFC-compliant, should only be used if the server does not support STARTTLS.
    EncryptedWrapper,
    /// Only send an email on encrypted connection (with STARTTLS)
    ///
    /// Recommended mode, prevents MITM when used with verified certificates.
    AlwaysEncrypt,
    /// Use TLS when available (with STARTTLS)
    ///
    /// Default mode.
    Opportunistic,
    /// Never use TLS
    NeverEncrypt,
}

/// Contains client configuration
#[derive(Debug)]
pub struct SmtpTransportBuilder {
    /// Maximum connection reuse
    ///
    /// Zero means no limitation
    connection_reuse_count_limit: u16,
    /// Enable connection reuse
    connection_reuse: bool,
    /// Name sent during HELO or EHLO
    hello_name: String,
    /// Credentials
    credentials: Option<(String, String)>,
    /// Socket we are connecting to
    server_addr: SocketAddr,
    /// SSL context to use
    ssl_context: SslContext,
    /// TLS security level
    security_level: SecurityLevel,
    /// Enable UTF8 mailboxes in envelope or headers
    smtp_utf8: bool,
    /// Optional enforced authentication mechanism
    authentication_mechanism: Option<Mechanism>,
    /// Define network timeout
    /// It can be changed later for specific needs (like a different timeout for each SMTP command)
    timeout: Option<Duration>,
}

/// Builder for the SMTP `SmtpTransport`
impl SmtpTransportBuilder {
    /// Creates a new local SMTP client
    pub fn new<A: ToSocketAddrs>(addr: A) -> Result<SmtpTransportBuilder, Error> {
        let mut addresses = try!(addr.to_socket_addrs());

        match addresses.next() {
            Some(addr) => {
                Ok(SmtpTransportBuilder {
                       server_addr: addr,
                       ssl_context: SslContext::builder(SslMethod::tls()).unwrap().build(),
                       security_level: SecurityLevel::AlwaysEncrypt,
                       smtp_utf8: false,
                       credentials: None,
                       connection_reuse_count_limit: 100,
                       connection_reuse: false,
                       hello_name: "localhost".to_string(),
                       authentication_mechanism: None,
                       timeout: Some(Duration::new(60, 0)),
                   })
            }
            None => Err(From::from("Could not resolve hostname")),
        }
    }

    /// Creates a new local SMTP client to port 25
    pub fn localhost() -> Result<SmtpTransportBuilder, Error> {
        SmtpTransportBuilder::new(("localhost", SMTP_PORT))
    }

    /// Use STARTTLS with a specific context
    pub fn ssl_context(mut self, ssl_context: SslContext) -> SmtpTransportBuilder {
        self.ssl_context = ssl_context;
        self
    }

    /// Set the security level for SSL/TLS
    pub fn security_level(mut self, level: SecurityLevel) -> SmtpTransportBuilder {
        self.security_level = level;
        self
    }

    /// Require SSL/TLS using STARTTLS
    ///
    /// Incompatible with `ssl_wrapper()``
    pub fn encrypt(mut self) -> SmtpTransportBuilder {
        self.security_level = SecurityLevel::AlwaysEncrypt;
        self
    }

    /// Require SSL/TLS using SMTPS
    ///
    /// Incompatible with `encrypt()`
    pub fn ssl_wrapper(mut self) -> SmtpTransportBuilder {
        self.security_level = SecurityLevel::EncryptedWrapper;
        self
    }

    /// Enable SMTPUTF8 if the server supports it
    pub fn smtp_utf8(mut self, enabled: bool) -> SmtpTransportBuilder {
        self.smtp_utf8 = enabled;
        self
    }

    /// Set the name used during HELO or EHLO
    pub fn hello_name<S: Into<String>>(mut self, name: S) -> SmtpTransportBuilder {
        self.hello_name = name.into();
        self
    }

    /// Enable connection reuse
    pub fn connection_reuse(mut self, enable: bool) -> SmtpTransportBuilder {
        self.connection_reuse = enable;
        self
    }

    /// Set the maximum number of emails sent using one connection
    pub fn connection_reuse_count_limit(mut self, limit: u16) -> SmtpTransportBuilder {
        self.connection_reuse_count_limit = limit;
        self
    }

    /// Set the client credentials
    pub fn credentials<S: Into<String>>(mut self,
                                        username: S,
                                        password: S)
                                        -> SmtpTransportBuilder {
        self.credentials = Some((username.into(), password.into()));
        self
    }

    /// Set the authentication mechanisms
    pub fn authentication_mechanism(mut self, mechanism: Mechanism) -> SmtpTransportBuilder {
        self.authentication_mechanism = Some(mechanism);
        self
    }

    /// Set the timeout duration
    pub fn timeout(mut self, timeout: Option<Duration>) -> SmtpTransportBuilder {
        self.timeout = timeout;
        self
    }

    /// Build the SMTP client
    ///
    /// It does not connect to the server, but only creates the `SmtpTransport`
    pub fn build(self) -> SmtpTransport {
        SmtpTransport::new(self)
    }
}

/// Represents the state of a client
#[derive(Debug)]
struct State {
    /// Panic state
    pub panic: bool,
    /// Connection reuse counter
    pub connection_reuse_count: u16,
}

/// Structure that implements the high level SMTP client
#[derive(Debug)]
pub struct SmtpTransport {
    /// Information about the server
    /// Value is None before HELO/EHLO
    server_info: Option<ServerInfo>,
    /// SmtpTransport variable states
    state: State,
    /// Information about the client
    client_info: SmtpTransportBuilder,
    /// Low level client
    client: Client,
}

macro_rules! try_smtp (
    ($err: expr, $client: ident) => ({
        match $err {
            Ok(val) => val,
            Err(err) => {
                if !$client.state.panic {
                    $client.state.panic = true;
                    $client.reset();
                }
                return Err(From::from(err))
            },
        }
    })
);

impl SmtpTransport {
    /// Creates a new SMTP client
    ///
    /// It does not connect to the server, but only creates the `SmtpTransport`
    pub fn new(builder: SmtpTransportBuilder) -> SmtpTransport {

        let client = Client::new();

        SmtpTransport {
            client: client,
            server_info: None,
            client_info: builder,
            state: State {
                panic: false,
                connection_reuse_count: 0,
            },
        }
    }

    /// Reset the client state
    fn reset(&mut self) {
        // Close the SMTP transaction if needed
        self.close();

        // Reset the client state
        self.server_info = None;
        self.state.panic = false;
        self.state.connection_reuse_count = 0;
    }

    /// Gets the EHLO response and updates server information
    pub fn get_ehlo(&mut self) -> SmtpResult {
        // Extended Hello
        let ehlo_response = try_smtp!(self.client.ehlo(&self.client_info.hello_name), self);

        self.server_info = Some(try_smtp!(ServerInfo::from_response(&ehlo_response), self));

        // Print server information
        debug!("server {}", self.server_info.as_ref().unwrap());

        Ok(ehlo_response)
    }
}

impl EmailTransport<SmtpResult> for SmtpTransport {
    /// Sends an email
    #[cfg_attr(feature = "cargo-clippy", allow(match_same_arms, cyclomatic_complexity))]
    fn send<T: SendableEmail>(&mut self, email: T) -> SmtpResult {

        // Extract email information
        let message_id = email.message_id();

        // Check if the connection is still available
        if (self.state.connection_reuse_count > 0) && (!self.client.is_connected()) {
            self.reset();
        }

        if self.state.connection_reuse_count == 0 {
            try!(self.client
                     .connect(&self.client_info.server_addr,
                              match self.client_info.security_level {
                                  SecurityLevel::EncryptedWrapper => {
                                      Some(&self.client_info.ssl_context)
                                  }
                                  _ => None,
                              }));

            try!(self.client.set_timeout(self.client_info.timeout));

            // Log the connection
            info!("connection established to {}", self.client_info.server_addr);

            try!(self.get_ehlo());

            match (&self.client_info.security_level,
                   self.server_info
                       .as_ref()
                       .unwrap()
                       .supports_feature(&Extension::StartTls)) {
                (&SecurityLevel::AlwaysEncrypt, false) => {
                    return Err(From::from("Could not encrypt connection, aborting"))
                }
                (&SecurityLevel::Opportunistic, false) => (),
                (&SecurityLevel::NeverEncrypt, _) => (),
                (&SecurityLevel::EncryptedWrapper, _) => (),
                (_, true) => {
                    try_smtp!(self.client.starttls(), self);
                    try_smtp!(self.client
                                  .upgrade_tls_stream(&self.client_info.ssl_context),
                              self);

                    debug!("connection encrypted");

                    // Send EHLO again
                    try!(self.get_ehlo());
                }
            }

            if self.client_info.credentials.is_some() {
                let (username, password) = self.client_info.credentials.clone().unwrap();

                let mut found = false;

                // Compute accepted mechanism
                let accepted_mechanisms = match self.client_info.authentication_mechanism {
                    Some(mechanism) => vec![mechanism],
                    None => {
                        if self.client.is_encrypted() {
                            // If encrypted, allow all mechanisms, with a preference for the
                            // simplest
                            // Login is obsolete so try it last
                            vec![Mechanism::Plain, Mechanism::CramMd5, Mechanism::Login]
                        } else {
                            // If not encrypted, do not allow clear-text passwords
                            vec![Mechanism::CramMd5]
                        }
                    }
                };

                for mechanism in accepted_mechanisms {
                    if self.server_info
                           .as_ref()
                           .unwrap()
                           .supports_auth_mechanism(mechanism) {
                        found = true;
                        try_smtp!(self.client.auth(mechanism, &username, &password), self);
                        break;
                    }
                }

                if !found {
                    info!("No supported authentication mechanisms available");
                }
            }
        }

        // Mail
        let mail_options = match (self.server_info
                                      .as_ref()
                                      .unwrap()
                                      .supports_feature(&Extension::EightBitMime),
                                  self.server_info
                                      .as_ref()
                                      .unwrap()
                                      .supports_feature(&Extension::SmtpUtfEight)) {
            (true, true) => Some("BODY=8BITMIME SMTPUTF8"),
            (true, false) => Some("BODY=8BITMIME"),
            (false, _) => None,
        };

        try_smtp!(self.client.mail(&email.from(), mail_options), self);

        // Log the mail command
        info!("{}: from=<{}>", message_id, email.from());

        // Recipient
        for to_address in &email.to() {
            try_smtp!(self.client.rcpt(to_address), self);
            // Log the rcpt command
            info!("{}: to=<{}>", message_id, to_address);
        }

        // Data
        try_smtp!(self.client.data(), self);

        // Message content
        let message = email.message();
        let result = self.client.message(&message);

        if result.is_ok() {
            // Increment the connection reuse counter
            self.state.connection_reuse_count += 1;

            // Log the message
            info!("{}: conn_use={}, size={}, status=sent ({})",
                  message_id,
                  self.state.connection_reuse_count,
                  message.len(),
                  result
                      .as_ref()
                      .ok()
                      .unwrap()
                      .message()
                      .iter()
                      .next()
                      .unwrap_or(&"no response".to_string()));
        }

        // Test if we can reuse the existing connection
        if (!self.client_info.connection_reuse) ||
           (self.state.connection_reuse_count >= self.client_info.connection_reuse_count_limit) {
            self.reset();
        }

        result
    }

    /// Closes the inner connection
    fn close(&mut self) {
        self.client.close();
    }
}
<|endoftext|>


<commit_before>// Copyright (c) 2014, 2015 Robert Clipsham <robert@octarineparrot.com>
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

extern crate pnet;
extern crate time;

use pnet::datalink::{self, NetworkInterface};

use std::env;

fn main() {
    use pnet::datalink::Channel::Ethernet;

    let iface_name = env::args().nth(1).unwrap();
    let interface_names_match = |iface: &NetworkInterface| iface.name == iface_name;

    // Find the network interface with the provided name
    let interfaces = datalink::interfaces();
    let interface = interfaces.into_iter().filter(interface_names_match).next().unwrap();

    // Create a channel to receive on
    let mut rx = match datalink::channel(&interface, Default::default()) {
        Ok(Ethernet(_, rx)) => rx,
        Ok(_) => panic!("rs_sender: unhandled channel type"),
        Err(e) => panic!("rs_benchmark: unable to create channel: {}", e),
    };

    let mut i = 0usize;
    let mut timestamps = Vec::with_capacity(201);
    timestamps.push(time::precise_time_ns() / 1_000);

    loop {
        match rx.next() {
            Ok(_) => {
                i += 1;
                if i == 1_000_000 {
                    timestamps.push(time::precise_time_ns() / 1_000);
                    if timestamps.len() == 201 {
                        break;
                    }
                    i = 0;
                }
            }
            Err(e) => {
                println!("rs_benchmark: unable to receive packet: {}", e);
            }
        }
    }

    // We received 1_000_000 packets in ((b - a) * 1_000_000) seconds.
    for (a, b) in timestamps.iter().zip(timestamps.iter().skip(1)) {
        println!("{}", *b - *a);
    }
}
<commit_msg>Make rs_receiver.rs only use pnet_datalink<commit_after>// Copyright (c) 2014, 2015 Robert Clipsham <robert@octarineparrot.com>
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

extern crate pnet_datalink;
extern crate time;

use std::env;

fn main() {
    use pnet_datalink::Channel::Ethernet;

    let iface_name = env::args().nth(1).unwrap();

    // Find the network interface with the provided name
    let interfaces = pnet_datalink::interfaces();
    let interface = interfaces.into_iter()
        .filter(|iface| iface.name == iface_name)
        .next()
        .unwrap();

    // Create a channel to receive on
    let mut rx = match pnet_datalink::channel(&interface, Default::default()) {
        Ok(Ethernet(_, rx)) => rx,
        Ok(_) => panic!("rs_sender: unhandled channel type"),
        Err(e) => panic!("rs_sender: unable to create channel: {}", e),
    };

    let mut i = 0usize;
    let mut timestamps = Vec::with_capacity(201);
    timestamps.push(time::precise_time_ns() / 1_000);

    loop {
        match rx.next() {
            Ok(_) => {
                i += 1;
                if i == 1_000_000 {
                    timestamps.push(time::precise_time_ns() / 1_000);
                    if timestamps.len() == 201 {
                        break;
                    }
                    i = 0;
                }
            }
            Err(e) => {
                println!("rs_sender: unable to receive packet: {}", e);
            }
        }
    }

    // We received 1_000_000 packets in ((b - a) * 1_000_000) seconds.
    for (a, b) in timestamps.iter().zip(timestamps.iter().skip(1)) {
        println!("{}", *b - *a);
    }
}
<|endoftext|>