<commit_before>//! A DNS query.

use futures::{Async, Future, Poll};
use futures::task::TaskRc;
use rand::random;
use ::bits::{AsDName, ComposeMode, ComposeResult, MessageBuilder,
             MessageBuf, Question};
use ::iana::{Class, RRType};
use super::ResolverTask;
use super::conf::ResolvOptions;
use super::core::Core;
use super::error::Error;
use super::request::QueryRequest;

//------------ Query ---------------------------------------------------------

pub struct Query(State);

enum State {
    Real(RealQuery),
    Failed(Option<Error>),
}

impl Query {
    pub fn new<N: AsDName>(resolv: &ResolverTask, name: N, rtype: RRType,
                       class: Class) -> Self {
        let core = resolv.core.clone();
        let message = core.with(|core| {
            Self::build_message(name, rtype, class, core.options())
        });
        let message = match message {
            Ok(message) => message,
            Err(err) => return Query(State::Failed(Some(err.into())))
        };
        let dgram = core.with(|core| !core.conf().options.use_vc);
        let (index, request) = core.with(|c| RealQuery::start(c, dgram,
                                                              message));
        Query(State::Real(RealQuery {
            core: core, request: request,
            dgram: dgram, start_index: index, curr_index: index,
            attempt: 0
        }))
    }

    fn build_message<N: AsDName>(name: N, rtype: RRType, class: Class,
                                 opts: &ResolvOptions)
                                 -> ComposeResult<MessageBuilder> {
        let mut res = try!(MessageBuilder::new(ComposeMode::Stream, true));
        res.header_mut().set_rd(opts.recurse);
        try!(Question::push(&mut res, &name, rtype, class));
        Ok(res)
    }
}

//--- Future

impl Future for Query {
    type Item = MessageBuf;
    type Error = Error;

    fn poll(&mut self) -> Poll<Self::Item, Self::Error> {
        match self.0 {
            State::Real(ref mut query) => query.poll(),
            State::Failed(ref mut err) => {
                match err.take() {
                    Some(err) => Err(err),
                    None => panic!("polling a resolved Query")
                }
            }
        }
    }
}

//------------ RealQuery -----------------------------------------------------

/// The future of an ongoing query.
struct RealQuery {
    /// The resolver core we are working on.
    core: TaskRc<Core>,

    /// The request we are currently processing.
    request: QueryRequest,

    /// Are we still in datagram stage?
    ///
    /// Assuming the resolver config allows using datagram services at all,
    /// we’ll start with `true`. Only if a response is truncated do we have
    /// to switch to `false` and go through the stream services.
    dgram: bool,

    /// The index of the service we started at.
    ///
    /// The index references either `core.udp` or `core.tcp`, depending on
    /// the value of `dgram`. We need to keep the index we started at
    /// because of the rotate option which makes us start at a random
    /// service.
    start_index: usize,

    /// The index of the service we currently are using.
    curr_index: usize,

    /// The how-many-th attempt is this, starting at attempt 0.
    attempt: usize
}

impl RealQuery {
    fn start(core: &Core, dgram: bool, message: MessageBuilder)
             -> (usize, QueryRequest) {
        let track = if dgram { core.udp() }
                    else { core.tcp() };
        let index = if core.conf().options.rotate {
            random::<usize>() % track.len()
        }
        else { 0 };
        let request = QueryRequest::new(message, &track[index]);
        (index, request)
    }

    /// Processes a response.
    ///
    /// This will either resolve the future or switch to stream mode.
    fn response(&mut self, response: MessageBuf, message: MessageBuilder)
                -> Poll<MessageBuf, Error> {
        if response.header().tc() && self.dgram
                && !self.core.with(|core| core.options().ign_tc)
        {
            self.start_stream(message)
        }
        else { Ok(response.into()) }
    }

    /// Proceeds to the next request or errors out.
    fn next_request(&mut self, message: MessageBuilder)
                    -> Poll<MessageBuf, Error> {
        self.curr_index += 1;
        if (self.curr_index % self.track_len()) == self.start_index {
            self.attempt += 1;
            if self.attempt == self.core.with(|core| core.conf().attempts) {
                return Err(Error::Timeout)
            }
            let dgram = self.dgram;
            let (index, request) = self.core.with(|core| {
                Self::start(core, dgram, message)
            });
            self.start_index = index;
            self.curr_index = index;
            self.request = request;
        }
        else {
            let dgram = self.dgram;
            let curr_index = self.curr_index;
            self.request = self.core.with(|core| {
                let track = if dgram { core.udp() }
                            else { core.tcp() };
                QueryRequest::new(message, &track[curr_index])
            });
        }
        self.poll()
    }

    /// Switches to stream mode and starts the first request or errors out.
    fn start_stream(&mut self, message: MessageBuilder)
                    -> Poll<MessageBuf, Error> {
        self.dgram = false;
        let (index, request) = self.core.with(|core| {
            Self::start(core, false, message)
        });
        self.start_index = index;
        self.curr_index = index;
        self.request = request;
        self.poll()
    }

    fn track_len(&self) -> usize {
        if self.dgram { self.core.with(|core| core.udp().len()) }
        else { self.core.with(|core| core.tcp().len()) }
    }
}

//--- Future

impl Future for RealQuery {
    type Item = MessageBuf;
    type Error = Error;

    /// Polls for completion.
    ///
    /// Polls the current request. If that succeeded, returns the result. If
    /// it failed fatally, returns an error. If it failed non-fatally,
    /// proceeds to the next request.
    fn poll(&mut self) -> Poll<MessageBuf, Error> {
        match self.request.poll() {
            Ok(Async::NotReady) => Ok(Async::NotReady),
            Ok(Async::Ready((Ok(response), message))) => {
                self.response(response, message)
            }
            Ok(Async::Ready((Err(_), message))) => {
                self.next_request(message)
            }
            Err(err) => Err(err.into())
        }
    }
}

<commit_msg>Use Result instead of a State type for Query.<commit_after>//! A DNS query.

use futures::{Async, Future, Poll};
use futures::task::TaskRc;
use rand::random;
use ::bits::{AsDName, ComposeMode, ComposeResult, MessageBuilder,
             MessageBuf, Question};
use ::iana::{Class, RRType};
use super::ResolverTask;
use super::conf::ResolvOptions;
use super::core::Core;
use super::error::Error;
use super::request::QueryRequest;

//------------ Query ---------------------------------------------------------

/// A DNS query.
///
/// A query is a future that resolves a DNS question for all the resource
/// records for a given domain name, record type, and class into either a
/// response or a query error. It will follow the rules set down by the
/// resolver configuration in trying to find an upstream server that
/// replies to the question.
///
/// A query is actually a combinator of sorts. It needs a [`ResolverTask`]
/// to operate on which in itself is the result of a future, namely that
/// returned by [`Resolver::start()`]. As such, you cannot use a query by
/// itself. Please see the module level documentation for some examples on
/// how to start a query.
pub struct Query(Result<RealQuery, Option<Error>>);

impl Query {
    /// Creates a new query.
    pub fn new<N: AsDName>(resolv: &ResolverTask, name: N, rtype: RRType,
                       class: Class) -> Self {
        let core = resolv.core.clone();
        let message = core.with(|core| {
            Self::build_message(name, rtype, class, core.options())
        });
        let message = match message {
            Ok(message) => message,
            Err(err) => return Query(Err(Some(err.into())))
        };
        let dgram = core.with(|core| !core.conf().options.use_vc);
        let (index, request) = core.with(|c| RealQuery::start(c, dgram,
                                                              message));
        Query(Ok(RealQuery {
            core: core, request: request,
            dgram: dgram, start_index: index, curr_index: index,
            attempt: 0
        }))
    }

    /// Builds the message for this query.
    ///
    /// If this function fails, we are done.
    fn build_message<N: AsDName>(name: N, rtype: RRType, class: Class,
                                 opts: &ResolvOptions)
                                 -> ComposeResult<MessageBuilder> {
        let mut res = try!(MessageBuilder::new(ComposeMode::Stream, true));
        res.header_mut().set_rd(opts.recurse);
        try!(Question::push(&mut res, &name, rtype, class));
        Ok(res)
    }
}

//--- Future

impl Future for Query {
    type Item = MessageBuf;
    type Error = Error;

    fn poll(&mut self) -> Poll<Self::Item, Self::Error> {
        match self.0 {
            Ok(ref mut query) => query.poll(),
            Err(ref mut err) => {
                match err.take() {
                    Some(err) => Err(err),
                    None => panic!("polling a resolved Query")
                }
            }
        }
    }
}

//------------ RealQuery -----------------------------------------------------

/// An ongoing query.
struct RealQuery {
    /// The resolver core we are working on.
    core: TaskRc<Core>,

    /// The request we are currently processing.
    request: QueryRequest,

    /// Are we on datagram track?
    ///
    /// A query can be in one of two *tracks,* datagram track or stream
    /// track. Unless the config states that we should use the stream
    /// track only, we will start out in datagram track and only switch to
    /// the stream track if we receive a truncated response.
    dgram: bool,

    /// The index of the service we started at.
    ///
    /// The index references either `core.udp` or `core.tcp`, depending on
    /// the value of `dgram`. We need to keep the index we started at
    /// because of the rotate option which makes us start at a random
    /// service.
    start_index: usize,

    /// The index of the service we currently are using.
    curr_index: usize,

    /// The how-many-th attempt this is, starting at attempt 0.
    attempt: usize
}

impl RealQuery {
    /// Starts a new query.
    ///
    /// Takes a reference to the core, whether to use the datagram track,
    /// and the message and returns the index of the service used and a
    /// new request.
    fn start(core: &Core, dgram: bool, message: MessageBuilder)
             -> (usize, QueryRequest) {
        let track = if dgram { core.udp() }
                    else { core.tcp() };
        let index = if core.conf().options.rotate {
            random::<usize>() % track.len()
        }
        else { 0 };
        let request = QueryRequest::new(message, &track[index]);
        (index, request)
    }

    /// Processes a response.
    ///
    /// This will either resolve the future or switch to stream mode.
    fn response(&mut self, response: MessageBuf, message: MessageBuilder)
                -> Poll<MessageBuf, Error> {
        if response.header().tc() && self.dgram
                && !self.core.with(|core| core.options().ign_tc)
        {
            self.start_stream(message)
        }
        else { Ok(response.into()) }
    }

    /// Proceeds to the next request or errors out.
    fn next_request(&mut self, message: MessageBuilder)
                    -> Poll<MessageBuf, Error> {
        self.curr_index += 1;
        if (self.curr_index % self.track_len()) == self.start_index {
            self.attempt += 1;
            if self.attempt == self.core.with(|core| core.conf().attempts) {
                return Err(Error::Timeout)
            }
            let dgram = self.dgram;
            let (index, request) = self.core.with(|core| {
                Self::start(core, dgram, message)
            });
            self.start_index = index;
            self.curr_index = index;
            self.request = request;
        }
        else {
            let dgram = self.dgram;
            let curr_index = self.curr_index;
            self.request = self.core.with(|core| {
                let track = if dgram { core.udp() }
                            else { core.tcp() };
                QueryRequest::new(message, &track[curr_index])
            });
        }
        self.poll()
    }

    /// Switches to stream mode and starts the first request or errors out.
    fn start_stream(&mut self, message: MessageBuilder)
                    -> Poll<MessageBuf, Error> {
        self.dgram = false;
        let (index, request) = self.core.with(|core| {
            Self::start(core, false, message)
        });
        self.start_index = index;
        self.curr_index = index;
        self.request = request;
        self.poll()
    }

    /// Returns the number of services in the current track.
    fn track_len(&self) -> usize {
        if self.dgram { self.core.with(|core| core.udp().len()) }
        else { self.core.with(|core| core.tcp().len()) }
    }
}

//--- Future

impl Future for RealQuery {
    type Item = MessageBuf;
    type Error = Error;

    /// Polls for completion.
    ///
    /// Polls the current request. If that succeeded, returns the result. If
    /// it failed fatally, returns an error. If it failed non-fatally,
    /// proceeds to the next request.
    fn poll(&mut self) -> Poll<MessageBuf, Error> {
        match self.request.poll() {
            Ok(Async::NotReady) => Ok(Async::NotReady),
            Ok(Async::Ready((Ok(response), message))) => {
                self.response(response, message)
            }
            Ok(Async::Ready((Err(_), message))) => {
                self.next_request(message)
            }
            Err(err) => Err(err.into())
        }
    }
}

<|endoftext|>


<commit_before>//! Iron's HTTP Response representation and associated methods.

use std::io::{File, MemReader};
use std::io::util;
use std::path::BytesContainer;
use std::fmt::{Show, Formatter, FormatError};

use typemap::TypeMap;
use plugin::Extensible;

use http::headers::response::HeaderCollection;
use http::headers::content_type::MediaType;

use errors::FileError;
use status::{mod, Status};
use {Url};

pub use http::server::response::ResponseWriter as HttpResponse;

use content_type::get_content_type;

pub mod modifiers;

/// The response representation given to `Middleware`
pub struct Response {
    /// The body of the response.
    ///
    /// This is a Reader for generality, most data should
    /// be sent using either `serve` or `serve_file`.
    ///
    /// Arbitrary Readers can be sent by assigning to body.
    pub body: Option<Box<Reader + Send>>,

    /// The headers of the response.
    pub headers: HeaderCollection,

    /// The response status-code.
    pub status: Option<Status>,

    /// A TypeMap to be used as an extensible storage for data
    /// associated with this Response.
    pub extensions: TypeMap
}

impl Response {
    /// Construct a blank Response
    pub fn new() -> Response {
        Response {
            headers: HeaderCollection::new(),
            status: None, // Start with no response code.
            body: None, // Start with no body.
            extensions: TypeMap::new()
        }
    }

    /// Create a new response with the status.
    #[deprecated = "Use `Response::new().set(Status(status))` instead."]
    pub fn status(status: status::Status) -> Response {
        Response {
            body: None,
            headers: HeaderCollection::new(),
            status: Some(status),
            extensions: TypeMap::new()
        }
    }

    /// Create a new response with the specified body and status.
    #[deprecated = "Use `Response::new().set(Status(status)).set(Body(body))` instead."]
    pub fn with<B: BytesContainer>(status: status::Status, body: B) -> Response {
        Response {
            body: Some(box MemReader::new(body.container_as_bytes().to_vec()) as Box<Reader + Send>),
            headers: HeaderCollection::new(),
            status: Some(status),
            extensions: TypeMap::new()
        }
    }

    /// Create a new Response with the `location` header set to the specified url.
    #[deprecated = "Use `Response::new().set(Status(status)).set(Redirect(url))` instead."]
    pub fn redirect(status: status::Status, url: Url) -> Response {
        let mut headers = HeaderCollection::new();
        headers.location = Some(url.into_generic_url());
        Response {
            body: None,
            headers: headers,
            status: Some(status),
            extensions: TypeMap::new()
        }
    }

    /// Create a response from a file on disk.
    ///
    /// The status code is set to 200 OK and the content type is autodetected based on
    /// the file extension.
    #[deprecated = "Use `Response::new().set(Status(status)).set(Body(path))` instead"]
    pub fn from_file(path: &Path) -> Result<Response, FileError> {
        let file = try!(File::open(path).map_err(FileError::new));
        let mut response = Response::new();
        response.status = Some(status::Ok);
        response.body = Some(box file as Box<Reader + Send>);
        response.headers.content_type = path.extension_str().and_then(get_content_type);
        Ok(response)
    }

    // `write_back` is used to put all the data added to `self`
    // back onto an `HttpResponse` so that it is sent back to the
    // client.
    //
    // `write_back` consumes the `Response`.
    #[doc(hidden)]
    pub fn write_back(self, http_res: &mut HttpResponse) {
        http_res.headers = self.headers.clone();

        // Default to a 404 if no response code was set
        http_res.status = self.status.clone().unwrap_or(status::NotFound);

        match self.body {
            Some(body) => {
                http_res.headers.content_type =
                    Some(http_res.headers
                            .content_type
                            .unwrap_or_else(||
                                MediaType::new("text".into_string(), "plain".into_string(), vec![])
                            ));
                match util::copy(&mut body, http_res) {
                    Err(e) => {
                        error!("Error reading/writing body: {}", e);

                        // Can't do anything else here since all headers/status have
                        // already been sent.
                    },
                    _ => {}
                }
            },

            None => {
                http_res.headers.content_length = Some(0u);
            }
        }
    }
}

impl Show for Response {
    fn fmt(&self, f: &mut Formatter) -> Result<(), FormatError> {
        try!(writeln!(f, "Response {{"));

        try!(writeln!(f, "    status: {}", self.status));

        try!(write!(f, "}}"));
        Ok(())
    }
}

// Allow plugins to attach to responses.
impl Extensible for Response {
    fn extensions(&self) -> &TypeMap {
        &self.extensions
    }

    fn extensions_mut(&mut self) -> &mut TypeMap {
        &mut self.extensions
    }
}

#[test]
fn matches_content_type () {
    let path = &Path::new("test.txt");
    let content_type = path.extension_str().and_then(get_content_type).unwrap();

    assert_eq!(content_type.type_.as_slice(), "text");
    assert_eq!(content_type.subtype.as_slice(), "plain");
}
<commit_msg>(chore) Hack around lack of impl Reader for Box<Reader + Send>.<commit_after>//! Iron's HTTP Response representation and associated methods.

use std::io::{mod, File, MemReader};
use std::path::BytesContainer;
use std::fmt::{Show, Formatter, FormatError};

use typemap::TypeMap;
use plugin::Extensible;

use http::headers::response::HeaderCollection;
use http::headers::content_type::MediaType;

use errors::FileError;
use status::{mod, Status};
use {Url};

pub use http::server::response::ResponseWriter as HttpResponse;

use content_type::get_content_type;

pub mod modifiers;

/// The response representation given to `Middleware`
pub struct Response {
    /// The body of the response.
    ///
    /// This is a Reader for generality, most data should
    /// be sent using either `serve` or `serve_file`.
    ///
    /// Arbitrary Readers can be sent by assigning to body.
    pub body: Option<Box<Reader + Send>>,

    /// The headers of the response.
    pub headers: HeaderCollection,

    /// The response status-code.
    pub status: Option<Status>,

    /// A TypeMap to be used as an extensible storage for data
    /// associated with this Response.
    pub extensions: TypeMap
}

impl Response {
    /// Construct a blank Response
    pub fn new() -> Response {
        Response {
            headers: HeaderCollection::new(),
            status: None, // Start with no response code.
            body: None, // Start with no body.
            extensions: TypeMap::new()
        }
    }

    /// Create a new response with the status.
    #[deprecated = "Use `Response::new().set(Status(status))` instead."]
    pub fn status(status: status::Status) -> Response {
        Response {
            body: None,
            headers: HeaderCollection::new(),
            status: Some(status),
            extensions: TypeMap::new()
        }
    }

    /// Create a new response with the specified body and status.
    #[deprecated = "Use `Response::new().set(Status(status)).set(Body(body))` instead."]
    pub fn with<B: BytesContainer>(status: status::Status, body: B) -> Response {
        Response {
            body: Some(box MemReader::new(body.container_as_bytes().to_vec()) as Box<Reader + Send>),
            headers: HeaderCollection::new(),
            status: Some(status),
            extensions: TypeMap::new()
        }
    }

    /// Create a new Response with the `location` header set to the specified url.
    #[deprecated = "Use `Response::new().set(Status(status)).set(Redirect(url))` instead."]
    pub fn redirect(status: status::Status, url: Url) -> Response {
        let mut headers = HeaderCollection::new();
        headers.location = Some(url.into_generic_url());
        Response {
            body: None,
            headers: headers,
            status: Some(status),
            extensions: TypeMap::new()
        }
    }

    /// Create a response from a file on disk.
    ///
    /// The status code is set to 200 OK and the content type is autodetected based on
    /// the file extension.
    #[deprecated = "Use `Response::new().set(Status(status)).set(Body(path))` instead"]
    pub fn from_file(path: &Path) -> Result<Response, FileError> {
        let file = try!(File::open(path).map_err(FileError::new));
        let mut response = Response::new();
        response.status = Some(status::Ok);
        response.body = Some(box file as Box<Reader + Send>);
        response.headers.content_type = path.extension_str().and_then(get_content_type);
        Ok(response)
    }

    // `write_back` is used to put all the data added to `self`
    // back onto an `HttpResponse` so that it is sent back to the
    // client.
    //
    // `write_back` consumes the `Response`.
    #[doc(hidden)]
    pub fn write_back(self, http_res: &mut HttpResponse) {
        http_res.headers = self.headers.clone();

        // Default to a 404 if no response code was set
        http_res.status = self.status.clone().unwrap_or(status::NotFound);

        match self.body {
            Some(mut body) => {
                http_res.headers.content_type =
                    Some(http_res.headers
                            .content_type
                            .clone()
                            .unwrap_or_else(||
                                MediaType::new("text".into_string(), "plain".into_string(), vec![])
                            ));

                // FIXME: Manually inlined io::util::copy
                // because Box<Reader + Send> does not impl Reader.
                let mut buf = [0, ..1024 * 64];
                let mut out = Ok(());
                loop {
                    let len = match body.read(buf) {
                        Ok(len) => len,
                        Err(ref e) if e.kind == io::EndOfFile => break,
                        Err(e) => { out = Err(e); break; },
                    };

                    match http_res.write(buf[..len]) {
                        Err(e) => {
                            out = Err(e);
                            break;
                        },
                        _ => {}
                    };
                }

                match out {
                    Err(e) => {
                        error!("Error reading/writing body: {}", e);

                        // Can't do anything else here since all headers/status have
                        // already been sent.
                    },
                    _ => {}
                }
            },

            None => {
                http_res.headers.content_length = Some(0u);
            }
        }
    }
}

impl Show for Response {
    fn fmt(&self, f: &mut Formatter) -> Result<(), FormatError> {
        try!(writeln!(f, "Response {{"));

        try!(writeln!(f, "    status: {}", self.status));

        try!(write!(f, "}}"));
        Ok(())
    }
}

// Allow plugins to attach to responses.
impl Extensible for Response {
    fn extensions(&self) -> &TypeMap {
        &self.extensions
    }

    fn extensions_mut(&mut self) -> &mut TypeMap {
        &mut self.extensions
    }
}

#[test]
fn matches_content_type () {
    let path = &Path::new("test.txt");
    let content_type = path.extension_str().and_then(get_content_type).unwrap();

    assert_eq!(content_type.type_.as_slice(), "text");
    assert_eq!(content_type.subtype.as_slice(), "plain");
}
<|endoftext|>


<commit_before>// Copyright 2018 Developers of the Rand project.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// https://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or https://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

//! Entropy generator, or wrapper around external generators

use rand_core::{RngCore, CryptoRng, Error, ErrorKind, impls};
#[allow(unused)]
use rngs;

/// An interface returning random data from external source(s), provided
/// specifically for securely seeding algorithmic generators (PRNGs).
///
/// Where possible, `EntropyRng` retrieves random data from the operating
/// system's interface for random numbers ([`OsRng`]); if that fails it will
/// fall back to the [`JitterRng`] entropy collector. In the latter case it will
/// still try to use [`OsRng`] on the next usage.
///
/// If no secure source of entropy is available `EntropyRng` will panic on use;
/// i.e. it should never output predictable data.
///
/// This is either a little slow ([`OsRng`] requires a system call) or extremely
/// slow ([`JitterRng`] must use significant CPU time to generate sufficient
/// jitter); for better performance it is common to seed a local PRNG from
/// external entropy then primarily use the local PRNG ([`thread_rng`] is
/// provided as a convenient, local, automatically-seeded CSPRNG).
///
/// # Panics
///
/// On most systems, like Windows, Linux, macOS and *BSD on common hardware, it
/// is highly unlikely for both [`OsRng`] and [`JitterRng`] to fail. But on
/// combinations like webassembly without Emscripten or stdweb both sources are
/// unavailable. If both sources fail, only [`try_fill_bytes`] is able to
/// report the error, and only the one from `OsRng`. The other [`RngCore`]
/// methods will panic in case of an error.
///
/// [`OsRng`]: struct.OsRng.html
/// [`JitterRng`]: jitter/struct.JitterRng.html
/// [`thread_rng`]: ../fn.thread_rng.html
/// [`RngCore`]: ../trait.RngCore.html
/// [`try_fill_bytes`]: ../trait.RngCore.html#method.tymethod.try_fill_bytes
#[derive(Debug)]
pub struct EntropyRng {
    source: Source,
}

#[derive(Debug)]
enum Source {
    Os(Os),
    Custom(Custom),
    Jitter(Jitter),
    None,
}

impl EntropyRng {
    /// Create a new `EntropyRng`.
    ///
    /// This method will do no system calls or other initialization routines,
    /// those are done on first use. This is done to make `new` infallible,
    /// and `try_fill_bytes` the only place to report errors.
    pub fn new() -> Self {
        EntropyRng { source: Source::None }
    }
}

impl Default for EntropyRng {
    fn default() -> Self {
        EntropyRng::new()
    }
}

impl RngCore for EntropyRng {
    fn next_u32(&mut self) -> u32 {
        impls::next_u32_via_fill(self)
    }

    fn next_u64(&mut self) -> u64 {
        impls::next_u64_via_fill(self)
    }

    fn fill_bytes(&mut self, dest: &mut [u8]) {
        self.try_fill_bytes(dest).unwrap_or_else(|err|
                panic!("all entropy sources failed; first error: {}", err))
    }

    fn try_fill_bytes(&mut self, dest: &mut [u8]) -> Result<(), Error> {
        let mut reported_error = None;

        if let Source::Os(ref mut os_rng) = self.source {
            match os_rng.fill(dest) {
                Ok(()) => return Ok(()),
                Err(err) => {
                    warn!("EntropyRng: OsRng failed \
                          [trying other entropy sources]: {}", err);
                    reported_error = Some(err);
                },
            }
        } else if Os::is_supported() {
            match Os::new_and_fill(dest) {
                Ok(os_rng) => {
                    debug!("EntropyRng: using OsRng");
                    self.source = Source::Os(os_rng);
                    return Ok(());
                },
                Err(err) => { reported_error = reported_error.or(Some(err)) },
            }
        }

        if let Source::Custom(ref mut rng) = self.source {
            match rng.fill(dest) {
                Ok(()) => return Ok(()),
                Err(err) => {
                    warn!("EntropyRng: custom entropy source failed \
                          [trying other entropy sources]: {}", err);
                    reported_error = Some(err);
                },
            }
        } else if Custom::is_supported() {
            match Custom::new_and_fill(dest) {
                Ok(custom) => {
                    debug!("EntropyRng: using custom entropy source");
                    self.source = Source::Custom(custom);
                    return Ok(());
                },
                Err(err) => { reported_error = reported_error.or(Some(err)) },
            }
        }

        if let Source::Jitter(ref mut jitter_rng) = self.source {
            match jitter_rng.fill(dest) {
                Ok(()) => return Ok(()),
                Err(err) => {
                    warn!("EntropyRng: JitterRng failed: {}", err);
                    reported_error = Some(err);
                },
            }
        } else if Jitter::is_supported() {
            match Jitter::new_and_fill(dest) {
                Ok(jitter_rng) => {
                    debug!("EntropyRng: using JitterRng");
                    self.source = Source::Jitter(jitter_rng);
                    return Ok(());
                },
                Err(err) => { reported_error = reported_error.or(Some(err)) },
            }
        }

        if let Some(err) = reported_error {
            Err(Error::with_cause(ErrorKind::Unavailable,
                                  "All entropy sources failed",
                                  err))
        } else {
            Err(Error::new(ErrorKind::Unavailable,
                           "No entropy sources available"))
        }
    }
}

impl CryptoRng for EntropyRng {}

trait EntropySource {
    fn new_and_fill(dest: &mut [u8]) -> Result<Self, Error>
        where Self: Sized;

    fn fill(&mut self, dest: &mut [u8]) -> Result<(), Error>;

    fn is_supported() -> bool { true }
}

#[allow(unused)]
#[derive(Clone, Debug)]
struct NoSource;

#[allow(unused)]
impl EntropySource for NoSource {
    fn new_and_fill(dest: &mut [u8]) -> Result<Self, Error> {
        Err(Error::new(ErrorKind::Unavailable, "Source not supported"))
    }

    fn fill(&mut self, dest: &mut [u8]) -> Result<(), Error> {
        unreachable!()
    }

    fn is_supported() -> bool { false }
}

#[cfg(all(feature="std",
          any(target_os = "linux", target_os = "android",
              target_os = "netbsd",
              target_os = "dragonfly",
              target_os = "haiku",
              target_os = "emscripten",
              target_os = "solaris",
              target_os = "cloudabi",
              target_os = "macos", target_os = "ios",
              target_os = "freebsd",
              target_os = "openbsd", target_os = "bitrig",
              target_os = "redox",
              target_os = "fuchsia",
              windows,
              all(target_arch = "wasm32", feature = "stdweb"),
              all(target_arch = "wasm32", feature = "wasm-bindgen"),
)))]
#[derive(Clone, Debug)]
pub struct Os(rngs::OsRng);

#[cfg(all(feature="std",
          any(target_os = "linux", target_os = "android",
              target_os = "netbsd",
              target_os = "dragonfly",
              target_os = "haiku",
              target_os = "emscripten",
              target_os = "solaris",
              target_os = "cloudabi",
              target_os = "macos", target_os = "ios",
              target_os = "freebsd",
              target_os = "openbsd", target_os = "bitrig",
              target_os = "redox",
              target_os = "fuchsia",
              windows,
              all(target_arch = "wasm32", feature = "stdweb"),
              all(target_arch = "wasm32", feature = "wasm-bindgen"),
)))]
impl EntropySource for Os {
    fn new_and_fill(dest: &mut [u8]) -> Result<Self, Error> {
        let mut rng = rngs::OsRng::new()?;
        rng.try_fill_bytes(dest)?;
        Ok(Os(rng))
    }

    fn fill(&mut self, dest: &mut [u8]) -> Result<(), Error> {
        self.0.try_fill_bytes(dest)
    }
}

#[cfg(not(all(feature="std",
              any(target_os = "linux", target_os = "android",
                  target_os = "netbsd",
                  target_os = "dragonfly",
                  target_os = "haiku",
                  target_os = "emscripten",
                  target_os = "solaris",
                  target_os = "cloudabi",
                  target_os = "macos", target_os = "ios",
                  target_os = "freebsd",
                  target_os = "openbsd", target_os = "bitrig",
                  target_os = "redox",
                  target_os = "fuchsia",
                  windows,
                  all(target_arch = "wasm32", feature = "stdweb"),
                  all(target_arch = "wasm32", feature = "wasm-bindgen"),
))))]
type Os = NoSource;

type Custom = NoSource;

#[cfg(not(target_arch = "wasm32"))]
#[derive(Clone, Debug)]
pub struct Jitter(rngs::JitterRng);

#[cfg(not(target_arch = "wasm32"))]
impl EntropySource for Jitter {
    fn new_and_fill(dest: &mut [u8]) -> Result<Self, Error> {
        let mut rng = rngs::JitterRng::new()?;
        rng.try_fill_bytes(dest)?;
        Ok(Jitter(rng))
    }

    fn fill(&mut self, dest: &mut [u8]) -> Result<(), Error> {
        self.0.try_fill_bytes(dest)
    }
}

#[cfg(target_arch = "wasm32")]
type Jitter = NoSource;

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

    #[test]
    fn test_entropy() {
        let mut rng = EntropyRng::new();
        let n = (rng.next_u32() ^ rng.next_u32()).count_ones();
        assert!(n >= 2);    // p(failure) approx 1e-7
    }
}
<commit_msg>fix doclink<commit_after>// Copyright 2018 Developers of the Rand project.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// https://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or https://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

//! Entropy generator, or wrapper around external generators

use rand_core::{RngCore, CryptoRng, Error, ErrorKind, impls};
#[allow(unused)]
use rngs;

/// An interface returning random data from external source(s), provided
/// specifically for securely seeding algorithmic generators (PRNGs).
///
/// Where possible, `EntropyRng` retrieves random data from the operating
/// system's interface for random numbers ([`OsRng`]); if that fails it will
/// fall back to the [`JitterRng`] entropy collector. In the latter case it will
/// still try to use [`OsRng`] on the next usage.
///
/// If no secure source of entropy is available `EntropyRng` will panic on use;
/// i.e. it should never output predictable data.
///
/// This is either a little slow ([`OsRng`] requires a system call) or extremely
/// slow ([`JitterRng`] must use significant CPU time to generate sufficient
/// jitter); for better performance it is common to seed a local PRNG from
/// external entropy then primarily use the local PRNG ([`thread_rng`] is
/// provided as a convenient, local, automatically-seeded CSPRNG).
///
/// # Panics
///
/// On most systems, like Windows, Linux, macOS and *BSD on common hardware, it
/// is highly unlikely for both [`OsRng`] and [`JitterRng`] to fail. But on
/// combinations like webassembly without Emscripten or stdweb both sources are
/// unavailable. If both sources fail, only [`try_fill_bytes`] is able to
/// report the error, and only the one from `OsRng`. The other [`RngCore`]
/// methods will panic in case of an error.
///
/// [`OsRng`]: struct.OsRng.html
/// [`JitterRng`]: struct.JitterRng.html
/// [`thread_rng`]: ../fn.thread_rng.html
/// [`RngCore`]: ../trait.RngCore.html
/// [`try_fill_bytes`]: ../trait.RngCore.html#method.tymethod.try_fill_bytes
#[derive(Debug)]
pub struct EntropyRng {
    source: Source,
}

#[derive(Debug)]
enum Source {
    Os(Os),
    Custom(Custom),
    Jitter(Jitter),
    None,
}

impl EntropyRng {
    /// Create a new `EntropyRng`.
    ///
    /// This method will do no system calls or other initialization routines,
    /// those are done on first use. This is done to make `new` infallible,
    /// and `try_fill_bytes` the only place to report errors.
    pub fn new() -> Self {
        EntropyRng { source: Source::None }
    }
}

impl Default for EntropyRng {
    fn default() -> Self {
        EntropyRng::new()
    }
}

impl RngCore for EntropyRng {
    fn next_u32(&mut self) -> u32 {
        impls::next_u32_via_fill(self)
    }

    fn next_u64(&mut self) -> u64 {
        impls::next_u64_via_fill(self)
    }

    fn fill_bytes(&mut self, dest: &mut [u8]) {
        self.try_fill_bytes(dest).unwrap_or_else(|err|
                panic!("all entropy sources failed; first error: {}", err))
    }

    fn try_fill_bytes(&mut self, dest: &mut [u8]) -> Result<(), Error> {
        let mut reported_error = None;

        if let Source::Os(ref mut os_rng) = self.source {
            match os_rng.fill(dest) {
                Ok(()) => return Ok(()),
                Err(err) => {
                    warn!("EntropyRng: OsRng failed \
                          [trying other entropy sources]: {}", err);
                    reported_error = Some(err);
                },
            }
        } else if Os::is_supported() {
            match Os::new_and_fill(dest) {
                Ok(os_rng) => {
                    debug!("EntropyRng: using OsRng");
                    self.source = Source::Os(os_rng);
                    return Ok(());
                },
                Err(err) => { reported_error = reported_error.or(Some(err)) },
            }
        }

        if let Source::Custom(ref mut rng) = self.source {
            match rng.fill(dest) {
                Ok(()) => return Ok(()),
                Err(err) => {
                    warn!("EntropyRng: custom entropy source failed \
                          [trying other entropy sources]: {}", err);
                    reported_error = Some(err);
                },
            }
        } else if Custom::is_supported() {
            match Custom::new_and_fill(dest) {
                Ok(custom) => {
                    debug!("EntropyRng: using custom entropy source");
                    self.source = Source::Custom(custom);
                    return Ok(());
                },
                Err(err) => { reported_error = reported_error.or(Some(err)) },
            }
        }

        if let Source::Jitter(ref mut jitter_rng) = self.source {
            match jitter_rng.fill(dest) {
                Ok(()) => return Ok(()),
                Err(err) => {
                    warn!("EntropyRng: JitterRng failed: {}", err);
                    reported_error = Some(err);
                },
            }
        } else if Jitter::is_supported() {
            match Jitter::new_and_fill(dest) {
                Ok(jitter_rng) => {
                    debug!("EntropyRng: using JitterRng");
                    self.source = Source::Jitter(jitter_rng);
                    return Ok(());
                },
                Err(err) => { reported_error = reported_error.or(Some(err)) },
            }
        }

        if let Some(err) = reported_error {
            Err(Error::with_cause(ErrorKind::Unavailable,
                                  "All entropy sources failed",
                                  err))
        } else {
            Err(Error::new(ErrorKind::Unavailable,
                           "No entropy sources available"))
        }
    }
}

impl CryptoRng for EntropyRng {}

trait EntropySource {
    fn new_and_fill(dest: &mut [u8]) -> Result<Self, Error>
        where Self: Sized;

    fn fill(&mut self, dest: &mut [u8]) -> Result<(), Error>;

    fn is_supported() -> bool { true }
}

#[allow(unused)]
#[derive(Clone, Debug)]
struct NoSource;

#[allow(unused)]
impl EntropySource for NoSource {
    fn new_and_fill(dest: &mut [u8]) -> Result<Self, Error> {
        Err(Error::new(ErrorKind::Unavailable, "Source not supported"))
    }

    fn fill(&mut self, dest: &mut [u8]) -> Result<(), Error> {
        unreachable!()
    }

    fn is_supported() -> bool { false }
}

#[cfg(all(feature="std",
          any(target_os = "linux", target_os = "android",
              target_os = "netbsd",
              target_os = "dragonfly",
              target_os = "haiku",
              target_os = "emscripten",
              target_os = "solaris",
              target_os = "cloudabi",
              target_os = "macos", target_os = "ios",
              target_os = "freebsd",
              target_os = "openbsd", target_os = "bitrig",
              target_os = "redox",
              target_os = "fuchsia",
              windows,
              all(target_arch = "wasm32", feature = "stdweb"),
              all(target_arch = "wasm32", feature = "wasm-bindgen"),
)))]
#[derive(Clone, Debug)]
pub struct Os(rngs::OsRng);

#[cfg(all(feature="std",
          any(target_os = "linux", target_os = "android",
              target_os = "netbsd",
              target_os = "dragonfly",
              target_os = "haiku",
              target_os = "emscripten",
              target_os = "solaris",
              target_os = "cloudabi",
              target_os = "macos", target_os = "ios",
              target_os = "freebsd",
              target_os = "openbsd", target_os = "bitrig",
              target_os = "redox",
              target_os = "fuchsia",
              windows,
              all(target_arch = "wasm32", feature = "stdweb"),
              all(target_arch = "wasm32", feature = "wasm-bindgen"),
)))]
impl EntropySource for Os {
    fn new_and_fill(dest: &mut [u8]) -> Result<Self, Error> {
        let mut rng = rngs::OsRng::new()?;
        rng.try_fill_bytes(dest)?;
        Ok(Os(rng))
    }

    fn fill(&mut self, dest: &mut [u8]) -> Result<(), Error> {
        self.0.try_fill_bytes(dest)
    }
}

#[cfg(not(all(feature="std",
              any(target_os = "linux", target_os = "android",
                  target_os = "netbsd",
                  target_os = "dragonfly",
                  target_os = "haiku",
                  target_os = "emscripten",
                  target_os = "solaris",
                  target_os = "cloudabi",
                  target_os = "macos", target_os = "ios",
                  target_os = "freebsd",
                  target_os = "openbsd", target_os = "bitrig",
                  target_os = "redox",
                  target_os = "fuchsia",
                  windows,
                  all(target_arch = "wasm32", feature = "stdweb"),
                  all(target_arch = "wasm32", feature = "wasm-bindgen"),
))))]
type Os = NoSource;

type Custom = NoSource;

#[cfg(not(target_arch = "wasm32"))]
#[derive(Clone, Debug)]
pub struct Jitter(rngs::JitterRng);

#[cfg(not(target_arch = "wasm32"))]
impl EntropySource for Jitter {
    fn new_and_fill(dest: &mut [u8]) -> Result<Self, Error> {
        let mut rng = rngs::JitterRng::new()?;
        rng.try_fill_bytes(dest)?;
        Ok(Jitter(rng))
    }

    fn fill(&mut self, dest: &mut [u8]) -> Result<(), Error> {
        self.0.try_fill_bytes(dest)
    }
}

#[cfg(target_arch = "wasm32")]
type Jitter = NoSource;

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

    #[test]
    fn test_entropy() {
        let mut rng = EntropyRng::new();
        let n = (rng.next_u32() ^ rng.next_u32()).count_ones();
        assert!(n >= 2);    // p(failure) approx 1e-7
    }
}
<|endoftext|>


<commit_before>extern crate hyper;
extern crate serde_json;

use request;

pub use types::{Account, CreateAccountResponse};

impl Account {
    pub fn new(short_name: &str, author_name: &str, author_url: &str) -> Account {
        Account {
            short_name: short_name.to_string(),
            author_name: author_name.to_string(),
            author_url: author_url.to_string(),
            access_token: "".to_string(),
            auth_url: "".to_string(),
            page_count: 0,
        }
    }

    pub fn set_short_name(&mut self, short_name: &str) -> &mut Account {
        self.short_name = short_name.to_string();
        self
    }

    pub fn set_author_name(&mut self, author_name: &str) -> &mut Account {
        self.author_name = author_name.to_string();
        self
    }

    pub fn set_author_url(&mut self, author_url: &str) -> &mut Account {
        self.author_url = author_url.to_string();
        self
    }

    pub fn set_auth_url(&mut self, auth_url: &str) -> &mut Account {
        self.auth_url = auth_url.to_string();
        self
    }

    pub fn set_page_count(&mut self, page_count: &i32) -> &mut Account {
        self.page_count = page_count.clone();
        self
    }

    pub fn set_access_token(&mut self, access_token: &str) -> &mut Account {
        self.access_token = access_token.to_string();
        self
    }

    pub fn create_account(&mut self) -> String {
        let url = format!(
            "https://api.telegra.ph/createAccount?short_name={}&author_name={}",
            self.short_name,
            self.author_name
        );
        let res_json = request::get(&url);
        let decoded: CreateAccountResponse = serde_json::from_str(&res_json).unwrap();

        self.set_access_token(&decoded.result.access_token)
            .set_auth_url(&decoded.result.auth_url);

        res_json
    }

    pub fn edit_accout_info(&self, acc: &Account) -> String {
        let url = String::from("https://api.telegra.ph/editAccountInfo?");
        let mut params: Vec<String> = vec![];

        let access_token = "access_token=".to_string() + &if acc.access_token.len() == 0 {
                self.access_token.clone()
            } else {
                acc.access_token.clone()
            };
        params.push(access_token);

        if acc.short_name.len() > 0 {
            let short_name = "short_name=".to_string() + &acc.short_name;
            params.push(short_name);
        }
        if acc.author_url.len() > 0 {
            let author_url = "author_url=".to_string() + &acc.author_url;
            params.push(author_url);
        }
        if acc.author_name.len() > 0 {
            let author_name = "author_name=".to_string() + &acc.author_name;
            params.push(author_name);
        }
        if acc.auth_url.len() > 0 {
            let auth_url = "auth_url=".to_string() + &acc.auth_url;
            params.push(auth_url);
        }

        let url = url + &params.join("&");

        request::get(&url)
    }

    pub fn get_account_info(access_token: &String, fields_list: &Vec<&str>) -> String {
        let fields = fields_list.iter()
            .map(|&field| "\"".to_string() + field + "\"")
            .collect::<Vec<String>>()
            .join(",");
        let fields = "[".to_string() + &fields + "]";

        let url = format!(
            "https://api.telegra.ph/getAccountInfo?access_token={}&fields={}",
            access_token,
            fields
        );

        request::get(&url)
    }

}
<commit_msg>added update method<commit_after>extern crate hyper;
extern crate serde_json;

use request;

pub use types::{Account, CreateAccountResponse};

impl Account {
    pub fn new(short_name: &str, author_name: &str, author_url: &str) -> Account {
        Account {
            short_name: short_name.to_string(),
            author_name: author_name.to_string(),
            author_url: author_url.to_string(),
            access_token: "".to_string(),
            auth_url: "".to_string(),
            page_count: 0,
        }
    }

    pub fn set_short_name(&mut self, short_name: &str) -> &mut Account {
        self.short_name = short_name.to_string();
        self
    }

    pub fn set_author_name(&mut self, author_name: &str) -> &mut Account {
        self.author_name = author_name.to_string();
        self
    }

    pub fn set_author_url(&mut self, author_url: &str) -> &mut Account {
        self.author_url = author_url.to_string();
        self
    }

    pub fn set_auth_url(&mut self, auth_url: &str) -> &mut Account {
        self.auth_url = auth_url.to_string();
        self
    }

    pub fn set_page_count(&mut self, page_count: &i32) -> &mut Account {
        self.page_count = page_count.clone();
        self
    }

    pub fn set_access_token(&mut self, access_token: &str) -> &mut Account {
        self.access_token = access_token.to_string();
        self
    }

    pub fn create_account(&mut self) -> String {
    fn update(&mut self, new_acc: &Account) -> &mut Account {
        self.set_short_name(&new_acc.short_name);
        self.set_author_name(&new_acc.author_name);
        self.set_author_url(&new_acc.author_url);

        self
    }

        let url = format!(
            "https://api.telegra.ph/createAccount?short_name={}&author_name={}",
            self.short_name,
            self.author_name
        );
        let res_json = request::get(&url);
        let decoded: CreateAccountResponse = serde_json::from_str(&res_json).unwrap();

        self.set_access_token(&decoded.result.access_token)
            .set_auth_url(&decoded.result.auth_url);

        res_json
    }

    pub fn edit_accout_info(&self, acc: &Account) -> String {
        let url = String::from("https://api.telegra.ph/editAccountInfo?");
        let mut params: Vec<String> = vec![];

        let access_token = "access_token=".to_string() + &if acc.access_token.len() == 0 {
                self.access_token.clone()
            } else {
                acc.access_token.clone()
            };
        params.push(access_token);

        if acc.short_name.len() > 0 {
            let short_name = "short_name=".to_string() + &acc.short_name;
            params.push(short_name);
        }
        if acc.author_url.len() > 0 {
            let author_url = "author_url=".to_string() + &acc.author_url;
            params.push(author_url);
        }
        if acc.author_name.len() > 0 {
            let author_name = "author_name=".to_string() + &acc.author_name;
            params.push(author_name);
        }
        if acc.auth_url.len() > 0 {
            let auth_url = "auth_url=".to_string() + &acc.auth_url;
            params.push(auth_url);
        }

        let url = url + &params.join("&");

        request::get(&url)
    }

    pub fn get_account_info(access_token: &String, fields_list: &Vec<&str>) -> String {
        let fields = fields_list.iter()
            .map(|&field| "\"".to_string() + field + "\"")
            .collect::<Vec<String>>()
            .join(",");
        let fields = "[".to_string() + &fields + "]";

        let url = format!(
            "https://api.telegra.ph/getAccountInfo?access_token={}&fields={}",
            access_token,
            fields
        );

        request::get(&url)
    }

}
<|endoftext|>


<commit_before>// Copyright 2017 `multipart-async` Crate Developers
//
// Licensed under the Apache License, Version 2.0, <LICENSE-APACHE or
// http://apache.org/licenses/LICENSE-2.0> or the MIT license <LICENSE-MIT or
// http://opensource.org/licenses/MIT>, at your option. This file may not be
// copied, modified, or distributed except according to those terms.
use futures::{Future, Stream, Async, Poll};
use futures::Async::*;

use mime::Mime;

use std::{io, mem};

use server::boundary::BoundaryFinder;
use server::{BodyChunk, StreamError, Multipart, httparse, twoway};

use self::httparse::EMPTY_HEADER;

use helpers::*;

const MAX_BUF_LEN: usize = 1024;
const MAX_HEADERS: usize = 4;

pub struct NextField<S: Stream> {
    stream: SomeCell<BoundaryFinder<S>>,
    accumulator: Vec<u8>,
    headers: Option<Headers>,
}

pub fn next_field<S: Stream>(stream: BoundaryFinder<S>) -> NextField<S> {
    NextField {
        stream: SomeCell::new(stream, "NextField value already taken"),
        accumulator: Vec::new(),
        headers: None,
    }
}

impl<S: Stream> NextField<S> where S::Item: BodyChunk, S::Error: StreamError {
    fn read_headers(&mut self) -> Poll<Headers, S::Error> {
        if let Some(headers) = self.headers.take() {
            return ready(headers);
        }

        loop {
            let chunk = try_ready!(self.stream.poll())
                .ok_or_else(|| io_error("unexpected end of stream"))?;

            // End of the headers section is signalled by a double-CRLF
            if let Some(header_end) = twoway::find_bytes(chunk.as_slice(), b"\r\n\r\n") {
                // Split after the double-CRLF because we don't want to yield it and httparse expects it
                let (headers, rem) = chunk.split_at(header_end + 4);
                self.stream.try_as_mut()?.push_chunk(rem);

                if !self.accumulator.is_empty() {
                    self.accumulator.extend_from_slice(headers.as_slice());
                    let headers = parse_headers(&self.accumulator)?;
                    self.accumulator.clear();

                    break ready(headers);
                } else {
                    break ready(parse_headers(headers.as_slice())?);
                }
            }

            if self.accumulator.len().saturating_add(chunk.len()) > MAX_BUF_LEN {
                return error("headers section too long or trailing double-CRLF missing");
            }

            self.accumulator.extend_from_slice(chunk.as_slice());
        }
    }

    fn read_to_string(&mut self) -> Poll<String, S::Error> {
        loop {
            if self.accumulator.is_empty() {
                self.accumulator = try_ready!(self.stream.poll())
                    .map_or(Vec::new(), BodyChunk::into_vec);
            }

            match try_ready!(self.stream.poll()) {
                Some(chunk) => self.accumulator.extend_from_slice(chunk.as_slice()),
                None => return ready(String::from_utf8(self.take_accumulator())?),
            }
         }
    }

    fn take_accumulator(&mut self) -> Vec<u8> {
        mem::replace(&mut self.accumulator, Vec::new())
    }
}
impl<S: Stream> Future for NextField<S> where S::Item: BodyChunk, S::Error: StreamError {
    type Item = Field<S>;
    type Error = S::Error;

    fn poll(&mut self) -> Poll<Self::Item, Self::Error> {
        let headers = try_ready!(self.read_headers());

        let data = if let Some(cont_type) = headers.cont_type {
            FieldData::File(FileData {
                content_type: cont_type,
                filename: headers.filename,
                multipart: Multipart { stream: self.stream.try_take()? },
                _priv: (),
            })
        } else if let Ready(text) = self.read_to_string()? {
            FieldData::Text(
                TextData {
                    text: text,
                    multipart: Multipart { stream: self.stream.try_take()? },
                    _priv: (),
                }
            )
        } else {
            self.headers = Some(headers);
            return not_ready();
        };

        ready(Field {
            name: headers.name,
            data: data,
            _priv: (),
        })
    }
}

fn parse_headers(bytes: &[u8]) -> io::Result<Headers> {
        let mut header_buf = httparse::parse_headers(bytes).map_err(io_error)?;

}

struct Headers {
    name: String,
    cont_type: Option<Mime>,
    filename: Option<String>,
}

pub struct Field<S: Stream> {
    pub name: String,
    pub data: FieldData<S>,
    _priv: (),
}

pub enum FieldData<S: Stream> {
    Text(TextData<S>),
    File(FileData<S>),
}

pub struct TextData<S: Stream> {
    pub text: String,
    pub multipart: Multipart<S>,
    _priv: (),
}

pub struct FileData<S: Stream> {
    pub filename: Option<String>,
    pub content_type: Mime,
    pub multipart: Multipart<S>,
    _priv: (),
}

const CONT_DISP: &str = "Content-Disposition";
<commit_msg>WIP: add edge case TODO<commit_after>// Copyright 2017 `multipart-async` Crate Developers
//
// Licensed under the Apache License, Version 2.0, <LICENSE-APACHE or
// http://apache.org/licenses/LICENSE-2.0> or the MIT license <LICENSE-MIT or
// http://opensource.org/licenses/MIT>, at your option. This file may not be
// copied, modified, or distributed except according to those terms.
use futures::{Future, Stream, Async, Poll};
use futures::Async::*;

use mime::Mime;

use std::{io, mem};

use server::boundary::BoundaryFinder;
use server::{BodyChunk, StreamError, Multipart, httparse, twoway};

use self::httparse::EMPTY_HEADER;

use helpers::*;

const MAX_BUF_LEN: usize = 1024;
const MAX_HEADERS: usize = 4;

pub struct NextField<S: Stream> {
    stream: SomeCell<BoundaryFinder<S>>,
    accumulator: Vec<u8>,
    headers: Option<Headers>,
}

pub fn next_field<S: Stream>(stream: BoundaryFinder<S>) -> NextField<S> {
    NextField {
        stream: SomeCell::new(stream, "NextField value already taken"),
        accumulator: Vec::new(),
        headers: None,
    }
}

impl<S: Stream> NextField<S> where S::Item: BodyChunk, S::Error: StreamError {
    fn read_headers(&mut self) -> Poll<Headers, S::Error> {
        if let Some(headers) = self.headers.take() {
            return ready(headers);
        }

        loop {
            let chunk = try_ready!(self.stream.poll())
                .ok_or_else(|| io_error("unexpected end of stream"))?;

            // End of the headers section is signalled by a double-CRLF
            if let Some(header_end) = twoway::find_bytes(chunk.as_slice(), b"\r\n\r\n") {
                // Split after the double-CRLF because we don't want to yield it and httparse expects it
                let (headers, rem) = chunk.split_at(header_end + 4);
                self.stream.try_as_mut()?.push_chunk(rem);

                if !self.accumulator.is_empty() {
                    self.accumulator.extend_from_slice(headers.as_slice());
                    let headers = parse_headers(&self.accumulator)?;
                    self.accumulator.clear();

                    break ready(headers);
                } else {
                    break ready(parse_headers(headers.as_slice())?);
                }
            }

            // TODO: edge case where double-CRLF falls on chunk boundaries

            if self.accumulator.len().saturating_add(chunk.len()) > MAX_BUF_LEN {
                return error("headers section too long or trailing double-CRLF missing");
            }

            self.accumulator.extend_from_slice(chunk.as_slice());
        }
    }

    fn read_to_string(&mut self) -> Poll<String, S::Error> {
        loop {
            if self.accumulator.is_empty() {
                self.accumulator = try_ready!(self.stream.poll())
                    .map_or(Vec::new(), BodyChunk::into_vec);
            }

            match try_ready!(self.stream.poll()) {
                Some(chunk) => self.accumulator.extend_from_slice(chunk.as_slice()),
                None => return ready(String::from_utf8(self.take_accumulator())?),
            }
         }
    }

    fn take_accumulator(&mut self) -> Vec<u8> {
        mem::replace(&mut self.accumulator, Vec::new())
    }
}
impl<S: Stream> Future for NextField<S> where S::Item: BodyChunk, S::Error: StreamError {
    type Item = Field<S>;
    type Error = S::Error;

    fn poll(&mut self) -> Poll<Self::Item, Self::Error> {
        let headers = try_ready!(self.read_headers());

        let data = if let Some(cont_type) = headers.cont_type {
            FieldData::File(FileData {
                content_type: cont_type,
                filename: headers.filename,
                multipart: Multipart { stream: self.stream.try_take()? },
                _priv: (),
            })
        } else if let Ready(text) = self.read_to_string()? {
            FieldData::Text(
                TextData {
                    text: text,
                    multipart: Multipart { stream: self.stream.try_take()? },
                    _priv: (),
                }
            )
        } else {
            self.headers = Some(headers);
            return not_ready();
        };

        ready(Field {
            name: headers.name,
            data: data,
            _priv: (),
        })
    }
}

fn parse_headers(bytes: &[u8]) -> io::Result<Headers> {
        let mut header_buf = httparse::parse_headers(bytes).map_err(io_error)?;

}

struct Headers {
    name: String,
    cont_type: Option<Mime>,
    filename: Option<String>,
}

pub struct Field<S: Stream> {
    pub name: String,
    pub data: FieldData<S>,
    _priv: (),
}

pub enum FieldData<S: Stream> {
    Text(TextData<S>),
    File(FileData<S>),
}

pub struct TextData<S: Stream> {
    pub text: String,
    pub multipart: Multipart<S>,
    _priv: (),
}

pub struct FileData<S: Stream> {
    pub filename: Option<String>,
    pub content_type: Mime,
    pub multipart: Multipart<S>,
    _priv: (),
}

const CONT_DISP: &str = "Content-Disposition";
<|endoftext|>


<commit_before>use std;
use std::num::One;

/// Calculate ceil(a / b).
pub fn div_ceil<T>(a: T, b: T) -> T where
    T: std::ops::Add<T, Output = T> + std::ops::Sub<T, Output = T> +
       std::ops::Div<T, Output = T> + std::num::One + Copy
{
    (a + b - T::one()) / b
}

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

    #[test]
    fn test_div_ceil() {
        assert_eq!(div_ceil(13, 12), 2);
        assert_eq!(div_ceil(1, 2), 1);
        assert_eq!(div_ceil(0, 3), 0);
    }
}
<commit_msg>Add CollectSlice<commit_after>use std;
use std::num::One;

/// Calculate ceil(a / b).
pub fn div_ceil<T>(a: T, b: T) -> T where
    T: std::ops::Add<T, Output = T> + std::ops::Sub<T, Output = T> +
       std::ops::Div<T, Output = T> + std::num::One + Copy
{
    (a + b - T::one()) / b
}

pub trait CollectSlice: Iterator {
    /// Loop through the iterator, writing items into the given slice until either the
    /// iterator runs out or the slice fills up.
    fn collect_slice(&mut self, slice: &mut [Self::Item]);
}

impl<T, I: Iterator<Item = T>> CollectSlice for I {
    fn collect_slice(&mut self, slice: &mut [Self::Item]) {
        for (s, i) in slice.iter_mut().zip(self) {
            *s = i;
        }
    }
}

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

    #[test]
    fn test_div_ceil() {
        assert_eq!(div_ceil(13, 12), 2);
        assert_eq!(div_ceil(1, 2), 1);
        assert_eq!(div_ceil(0, 3), 0);
    }

    #[test]
    fn test_collect_slice() {
        {
            let mut buf = [0; 5];

            (0..5).map(|i| {
                i + 1
            }).collect_slice(&mut buf[..]);

            assert_eq!(buf[0], 1);
            assert_eq!(buf[1], 2);
            assert_eq!(buf[2], 3);
            assert_eq!(buf[3], 4);
            assert_eq!(buf[4], 5);
        }

        {
            let mut buf = [0; 5];

            (0..3).map(|i| {
                i + 1
            }).collect_slice(&mut buf[..]);

            assert_eq!(buf[0], 1);
            assert_eq!(buf[1], 2);
            assert_eq!(buf[2], 3);
            assert_eq!(buf[3], 0);
            assert_eq!(buf[4], 0);
        }

        {
            let mut buf = [0; 5];

            let mut iter = (0..3).map(|i| {
                i + 1
            });

            iter.collect_slice(&mut buf[2..]);

            assert_eq!(buf[0], 0);
            assert_eq!(buf[1], 0);
            assert_eq!(buf[2], 1);
            assert_eq!(buf[3], 2);
            assert_eq!(buf[4], 3);
        }

        {
            let mut buf = [0; 3];

            let mut iter = (0..5).map(|i| {
                i + 1
            });

            iter.collect_slice(&mut buf[..]);

            assert_eq!(buf[0], 1);
            assert_eq!(buf[1], 2);
            assert_eq!(buf[2], 3);

            assert_eq!(iter.next().unwrap(), 4);
            assert_eq!(iter.next().unwrap(), 5);
        }
    }
}
<|endoftext|>


<commit_before>use std::str;
use std::collections::{HashMap, BTreeSet};

use rocksdb::{self, Writable, IteratorMode, Direction, WriteBatch};
use byteorder::{ByteOrder, BigEndian, WriteBytesExt};

use RocksDBIndexStore;
use document_index::DocRef;
use key_builder::KeyBuilder;
use doc_id_set::DocIdSet;

#[derive(Debug)]
pub enum SegmentMergeError {
    TooManyDocs,
    RocksDBError(rocksdb::Error),
}

impl From<rocksdb::Error> for SegmentMergeError {
    fn from(e: rocksdb::Error) -> SegmentMergeError {
        SegmentMergeError::RocksDBError(e)
    }
}

impl RocksDBIndexStore {
    fn merge_segment_data(&self, source_segments: &Vec<u32>, dest_segment: u32, doc_ref_mapping: &HashMap<DocRef, u16>) -> Result<(), SegmentMergeError> {
        // Put source_segments in a BTreeSet as this is much faster for performing contains queries against
        let source_segments_btree = source_segments.iter().collect::<BTreeSet<_>>();

        // Merge the term directories
        // The term directory keys are ordered to be most convenient for retrieving all the segments
        // of for a term/field combination in one go (field/term/segment). So we don't end up pulling
        // in a lot of unwanted data, we firstly iterate the keys, it they one of the source segments
        // looking for then we load them and append them to our new segment.

        /// Converts term directory key strings "d1/2/3" into tuples of 3 i32s (1, 2, 3)
        fn parse_term_directory_key(key: &[u8]) -> (u32, u32, u32) {
            let mut nums_iter = key[1..].split(|b| *b == b'/').map(|s| str::from_utf8(s).unwrap().parse::<u32>().unwrap());
            (nums_iter.next().unwrap(), nums_iter.next().unwrap(), nums_iter.next().unwrap())
        }

        let mut current_td_key: Option<(u32, u32)> = None;
        let mut current_td = Vec::new();

        for (k, v) in self.db.iterator(IteratorMode::From(b"d", Direction::Forward)) {
            if k[0] != b'd' {
                // No more term directories to merge
                break;
            }

            let (field, term, segment) = parse_term_directory_key(&k);

            if !source_segments_btree.contains(&segment) {
                continue;
            }

            if current_td_key != Some((field, term)) {
                // Finished current term directory. Write it to the DB and start the next one
                if let Some((field, term)) = current_td_key {
                    let kb = KeyBuilder::segment_dir_list(dest_segment, field, term);
                    try!(self.db.put(&kb.key(), &current_td));
                    current_td.clear();
                }

                current_td_key = Some((field, term));
            }

            // Merge term directory into the new one (and remap the doc ids)
            for doc_id in DocIdSet::Owned(v.to_vec()).iter() {
                let doc_ref = DocRef::from_segment_ord(segment, doc_id);
                let new_doc_id = doc_ref_mapping.get(&doc_ref).unwrap();
                current_td.write_u16::<BigEndian>(*new_doc_id).unwrap();
            }
        }

        // All done, write the last term directory
        if let Some((field, term)) = current_td_key {
            let kb = KeyBuilder::segment_dir_list(dest_segment, field, term);
            try!(self.db.put(&kb.key(), &current_td));
            current_td.clear();
        }

        // Merge the stored values
        // All stored value keys start with the segment id. So we need to:
        // - Iterate all stored value keys that are prefixed by one of the stored segment ids
        // - Remap their doc ids to the one in the new segment
        // - Write the value back with the new segment/doc ids in the key

        /// Converts stored value key strings "v1/2/3/v" into tuples of 3 i32s and a Vec<u8> (1, 2, 3, vec![b'v', b'a', b'l'])
        fn parse_stored_value_key(key: &[u8]) -> (u32, u32, u32, Vec<u8>) {
            let mut parts_iter = key[1..].split(|b| *b == b'/');
            let segment = str::from_utf8(parts_iter.next().unwrap()).unwrap().parse::<u32>().unwrap();
            let doc_id = str::from_utf8(parts_iter.next().unwrap()).unwrap().parse::<u32>().unwrap();
            let field_ord = str::from_utf8(parts_iter.next().unwrap()).unwrap().parse::<u32>().unwrap();
            let value_type = parts_iter.next().unwrap().to_vec();

            (segment, doc_id, field_ord, value_type)
        }

        for source_segment in source_segments.iter() {
            let kb = KeyBuilder::segment_stored_values_prefix(*source_segment);
            for (k, v) in self.db.iterator(IteratorMode::From(&kb.key(), Direction::Forward)) {
                if k[0] != b'v' {
                    // No more stored values to move
                    break;
                }

                let (segment, doc_id, field, value_type) = parse_stored_value_key(&k);

                if segment != *source_segment {
                    // Segment finished
                    break;
                }

                // Remap doc id
                let doc_ref = DocRef::from_segment_ord(segment, doc_id as u16);
                let new_doc_id = doc_ref_mapping.get(&doc_ref).unwrap();

                // Write value into new segment
                let kb = KeyBuilder::stored_field_value(dest_segment, *new_doc_id, field, &value_type);
                try!(self.db.put(&kb.key(), &v));
            }
        }

        // Merge the statistics
        // Like stored values, these start with segment ids. But instead of just rewriting the
        // key, we need to sum up all the statistics across the segments being merged.

        let mut statistics = HashMap::new();

        /// Converts statistic key strings "s1/total_docs" into tuples of 1 i32 and a Vec<u8> (1, ['t', 'o', 't', ...])
        fn parse_statistic_key(key: &[u8]) -> (u32, Vec<u8>) {
            let mut parts_iter = key[1..].split(|b| *b == b'/');
            let segment = str::from_utf8(parts_iter.next().unwrap()).unwrap().parse::<u32>().unwrap();
            let statistic_name = parts_iter.next().unwrap().to_vec();

            (segment, statistic_name)
        }

        // Fetch and merge statistics
        for source_segment in source_segments.iter() {
            let kb = KeyBuilder::segment_stat_prefix(*source_segment);
            for (k, v) in self.db.iterator(IteratorMode::From(&kb.key(), Direction::Forward)) {
                if k[0] != b's' {
                    // No more statistics to merge
                    break;
                }

                let (segment, statistic_name) = parse_statistic_key(&k);

                if segment != *source_segment {
                    // Segment finished
                    break;
                }

                let mut stat = statistics.entry(statistic_name).or_insert(0);
                *stat += BigEndian::read_i64(&v);
            }
        }

        // Write merged statistics to new segment
        for (stat_name, stat_value) in statistics {
            let kb = KeyBuilder::segment_stat(dest_segment, &stat_name);
            let mut val_bytes = [0; 8];
            BigEndian::write_i64(&mut val_bytes, stat_value);
            try!(self.db.put(&kb.key(), &val_bytes));
        }

        // Note: Don't merge the deletion lists
        // Deletion lists can change at any time so we must lock the "document index"
        // before merging them so they can't be altered during merge. we cannot lock
        // this until the commit phase though.

        Ok(())
    }

    fn commit_segment_merge(&self, source_segments: &Vec<u32>, dest_segment: u32, doc_ref_mapping: &HashMap<DocRef, u16>) -> Result<(), SegmentMergeError> {
        let write_batch = WriteBatch::default();

        // Activate new segment
        let kb = KeyBuilder::segment_active(dest_segment);
        try!(write_batch.put(&kb.key(), b""));

        // Deactivate old segments
        for source_segment in source_segments.iter() {
            // Activate new segment
            let kb = KeyBuilder::segment_active(*source_segment);
            try!(write_batch.delete(&kb.key()));
        }

        // Update document index and commit
        // This will write the write batch
        try!(self.document_index.commit_segment_merge(&self.db, write_batch, source_segments, dest_segment, doc_ref_mapping));

        Ok(())
    }

    pub fn merge_segments(&self, source_segments: &Vec<u32>) -> Result<u32, SegmentMergeError> {
        let dest_segment = try!(self.segments.new_segment(&self.db));

        // Generate a mapping between the ids of the documents in the old segments to the new one
        // This packs the id spaces of the old segments together:
        // For example, say we have to merge 3 segments with 100 documents each:
        //  - The first segment's ids will be the same as before
        //  - The second segment's ids will be remapped to 100 - 199
        //  - The third segment's ids will be remapped to 200 - 299

        let mut doc_ref_mapping: HashMap<DocRef, u16> = HashMap::new();
        let mut current_ord: u32 = 0;

        for source_segment in source_segments.iter() {
            let kb = KeyBuilder::segment_stat(*source_segment, b"total_docs");
            let total_docs = match try!(self.db.get(&kb.key())) {
                Some(total_docs_bytes) => {
                    BigEndian::read_i64(&total_docs_bytes)
                }
                None => continue,
            };

            for source_ord in 0..total_docs {
                if current_ord >= 65536 {
                    return Err(SegmentMergeError::TooManyDocs);
                }

                let from = DocRef::from_segment_ord(*source_segment, source_ord as u16);
                doc_ref_mapping.insert(from, current_ord as u16);
                current_ord += 1;
            }
        }

        // Merge segment data
        // Most of the heavy lifting happens here. This merges all the immutable parts of
        // the segment (which is everything but the deletion list). It does not activate the
        // segment.
        // This means that nothing bad will happen if it crashes half way through -- the
        // worst that could happen is we're left with a partially-written segment that we
        // have to clean up.
        try!(self.merge_segment_data(&source_segments, dest_segment, &doc_ref_mapping));

        // Commit the merge
        // This activates the new segment and updates the document index. Effectively committing
        // the merge.
        // Throughout this stage we need an exclusive lock to the document index. This is to
        // prevent documents in the source segments being deleted/updated so we don't accidentally
        // undelete them (this will block until the merge is complete so they delete/update from
        // the new segment).
        try!(self.commit_segment_merge(&source_segments, dest_segment, &doc_ref_mapping));

        Ok(dest_segment)
    }

    pub fn purge_segments(&self, segments: &Vec<u32>) -> Result<(), rocksdb::Error> {
        // Put segments in a BTreeSet as this is much faster for performing contains queries against
        let segments_btree = segments.iter().collect::<BTreeSet<_>>();

        // Purge term directories

        /// Converts term directory key strings "d1/2/3" into tuples of 3 i32s (1, 2, 3)
        fn parse_term_directory_key(key: &[u8]) -> (u32, u32, u32) {
            let mut nums_iter = key[1..].split(|b| *b == b'/').map(|s| str::from_utf8(s).unwrap().parse::<u32>().unwrap());
            (nums_iter.next().unwrap(), nums_iter.next().unwrap(), nums_iter.next().unwrap())
        }

        for (k, _) in self.db.iterator(IteratorMode::From(b"d", Direction::Forward)) {
            if k[0] != b'd' {
                // No more term directories to delete
                break;
            }

            let (_, _, segment) = parse_term_directory_key(&k);

            if segments_btree.contains(&segment) {
                try!(self.db.delete(&k));
            }
        }

        // Purge the stored values

        /// Converts stored value key strings "v1/2/3/v" into tuples of 3 i32s and a Vec<u8> (1, 2, 3, vec![b'v', b'a', b'l'])
        fn parse_stored_value_key(key: &[u8]) -> (u32, u32, u32, Vec<u8>) {
            let mut parts_iter = key[1..].split(|b| *b == b'/');
            let segment = str::from_utf8(parts_iter.next().unwrap()).unwrap().parse::<u32>().unwrap();
            let doc_id = str::from_utf8(parts_iter.next().unwrap()).unwrap().parse::<u32>().unwrap();
            let field_ord = str::from_utf8(parts_iter.next().unwrap()).unwrap().parse::<u32>().unwrap();
            let value_type = parts_iter.next().unwrap().to_vec();

            (segment, doc_id, field_ord, value_type)
        }

        for source_segment in segments.iter() {
            let kb = KeyBuilder::segment_stored_values_prefix(*source_segment);
            for (k, _) in self.db.iterator(IteratorMode::From(&kb.key(), Direction::Forward)) {
                if k[0] != b'v' {
                    // No more stored values to delete
                    break;
                }

                let (segment, _, _, _) = parse_stored_value_key(&k);

                if segment != *source_segment {
                    // Segment finished
                    break;
                }

                try!(self.db.delete(&k));
            }
        }

        // Purge the statistics

        /// Converts statistic key strings "s1/total_docs" into tuples of 1 i32 and a Vec<u8> (1, ['t', 'o', 't', ...])
        fn parse_statistic_key(key: &[u8]) -> (u32, Vec<u8>) {
            let mut parts_iter = key[1..].split(|b| *b == b'/');
            let segment = str::from_utf8(parts_iter.next().unwrap()).unwrap().parse::<u32>().unwrap();
            let statistic_name = parts_iter.next().unwrap().to_vec();

            (segment, statistic_name)
        }

        for source_segment in segments.iter() {
            let kb = KeyBuilder::segment_stat_prefix(*source_segment);
            for (k, _) in self.db.iterator(IteratorMode::From(&kb.key(), Direction::Forward)) {
                if k[0] != b's' {
                    // No more statistics to purge
                    break;
                }

                let (segment, _) = parse_statistic_key(&k);

                if segment != *source_segment {
                    // Segment finished
                    break;
                }

                try!(self.db.delete(&k));
            }
        }

        // Purge the deletion lists
        for source_segment in segments.iter() {
            let kb = KeyBuilder::segment_del_list(*source_segment);
            try!(self.db.delete(&kb.key()));
        }

        Ok(())
    }
}
<commit_msg>Make SegmentMergeError convertable to String<commit_after>use std::str;
use std::collections::{HashMap, BTreeSet};

use rocksdb::{self, Writable, IteratorMode, Direction, WriteBatch};
use byteorder::{ByteOrder, BigEndian, WriteBytesExt};

use RocksDBIndexStore;
use document_index::DocRef;
use key_builder::KeyBuilder;
use doc_id_set::DocIdSet;

#[derive(Debug)]
pub enum SegmentMergeError {
    TooManyDocs,
    RocksDBError(rocksdb::Error),
}

impl From<rocksdb::Error> for SegmentMergeError {
    fn from(e: rocksdb::Error) -> SegmentMergeError {
        SegmentMergeError::RocksDBError(e)
    }
}

impl From<SegmentMergeError> for String {
    fn from(e: SegmentMergeError) -> String {
        match e {
            SegmentMergeError::TooManyDocs => "Too many docs".to_string(),
            SegmentMergeError::RocksDBError(e) => e.into(),
        }
    }
}

impl RocksDBIndexStore {
    fn merge_segment_data(&self, source_segments: &Vec<u32>, dest_segment: u32, doc_ref_mapping: &HashMap<DocRef, u16>) -> Result<(), SegmentMergeError> {
        // Put source_segments in a BTreeSet as this is much faster for performing contains queries against
        let source_segments_btree = source_segments.iter().collect::<BTreeSet<_>>();

        // Merge the term directories
        // The term directory keys are ordered to be most convenient for retrieving all the segments
        // of for a term/field combination in one go (field/term/segment). So we don't end up pulling
        // in a lot of unwanted data, we firstly iterate the keys, it they one of the source segments
        // looking for then we load them and append them to our new segment.

        /// Converts term directory key strings "d1/2/3" into tuples of 3 i32s (1, 2, 3)
        fn parse_term_directory_key(key: &[u8]) -> (u32, u32, u32) {
            let mut nums_iter = key[1..].split(|b| *b == b'/').map(|s| str::from_utf8(s).unwrap().parse::<u32>().unwrap());
            (nums_iter.next().unwrap(), nums_iter.next().unwrap(), nums_iter.next().unwrap())
        }

        let mut current_td_key: Option<(u32, u32)> = None;
        let mut current_td = Vec::new();

        for (k, v) in self.db.iterator(IteratorMode::From(b"d", Direction::Forward)) {
            if k[0] != b'd' {
                // No more term directories to merge
                break;
            }

            let (field, term, segment) = parse_term_directory_key(&k);

            if !source_segments_btree.contains(&segment) {
                continue;
            }

            if current_td_key != Some((field, term)) {
                // Finished current term directory. Write it to the DB and start the next one
                if let Some((field, term)) = current_td_key {
                    let kb = KeyBuilder::segment_dir_list(dest_segment, field, term);
                    try!(self.db.put(&kb.key(), &current_td));
                    current_td.clear();
                }

                current_td_key = Some((field, term));
            }

            // Merge term directory into the new one (and remap the doc ids)
            for doc_id in DocIdSet::Owned(v.to_vec()).iter() {
                let doc_ref = DocRef::from_segment_ord(segment, doc_id);
                let new_doc_id = doc_ref_mapping.get(&doc_ref).unwrap();
                current_td.write_u16::<BigEndian>(*new_doc_id).unwrap();
            }
        }

        // All done, write the last term directory
        if let Some((field, term)) = current_td_key {
            let kb = KeyBuilder::segment_dir_list(dest_segment, field, term);
            try!(self.db.put(&kb.key(), &current_td));
            current_td.clear();
        }

        // Merge the stored values
        // All stored value keys start with the segment id. So we need to:
        // - Iterate all stored value keys that are prefixed by one of the stored segment ids
        // - Remap their doc ids to the one in the new segment
        // - Write the value back with the new segment/doc ids in the key

        /// Converts stored value key strings "v1/2/3/v" into tuples of 3 i32s and a Vec<u8> (1, 2, 3, vec![b'v', b'a', b'l'])
        fn parse_stored_value_key(key: &[u8]) -> (u32, u32, u32, Vec<u8>) {
            let mut parts_iter = key[1..].split(|b| *b == b'/');
            let segment = str::from_utf8(parts_iter.next().unwrap()).unwrap().parse::<u32>().unwrap();
            let doc_id = str::from_utf8(parts_iter.next().unwrap()).unwrap().parse::<u32>().unwrap();
            let field_ord = str::from_utf8(parts_iter.next().unwrap()).unwrap().parse::<u32>().unwrap();
            let value_type = parts_iter.next().unwrap().to_vec();

            (segment, doc_id, field_ord, value_type)
        }

        for source_segment in source_segments.iter() {
            let kb = KeyBuilder::segment_stored_values_prefix(*source_segment);
            for (k, v) in self.db.iterator(IteratorMode::From(&kb.key(), Direction::Forward)) {
                if k[0] != b'v' {
                    // No more stored values to move
                    break;
                }

                let (segment, doc_id, field, value_type) = parse_stored_value_key(&k);

                if segment != *source_segment {
                    // Segment finished
                    break;
                }

                // Remap doc id
                let doc_ref = DocRef::from_segment_ord(segment, doc_id as u16);
                let new_doc_id = doc_ref_mapping.get(&doc_ref).unwrap();

                // Write value into new segment
                let kb = KeyBuilder::stored_field_value(dest_segment, *new_doc_id, field, &value_type);
                try!(self.db.put(&kb.key(), &v));
            }
        }

        // Merge the statistics
        // Like stored values, these start with segment ids. But instead of just rewriting the
        // key, we need to sum up all the statistics across the segments being merged.

        let mut statistics = HashMap::new();

        /// Converts statistic key strings "s1/total_docs" into tuples of 1 i32 and a Vec<u8> (1, ['t', 'o', 't', ...])
        fn parse_statistic_key(key: &[u8]) -> (u32, Vec<u8>) {
            let mut parts_iter = key[1..].split(|b| *b == b'/');
            let segment = str::from_utf8(parts_iter.next().unwrap()).unwrap().parse::<u32>().unwrap();
            let statistic_name = parts_iter.next().unwrap().to_vec();

            (segment, statistic_name)
        }

        // Fetch and merge statistics
        for source_segment in source_segments.iter() {
            let kb = KeyBuilder::segment_stat_prefix(*source_segment);
            for (k, v) in self.db.iterator(IteratorMode::From(&kb.key(), Direction::Forward)) {
                if k[0] != b's' {
                    // No more statistics to merge
                    break;
                }

                let (segment, statistic_name) = parse_statistic_key(&k);

                if segment != *source_segment {
                    // Segment finished
                    break;
                }

                let mut stat = statistics.entry(statistic_name).or_insert(0);
                *stat += BigEndian::read_i64(&v);
            }
        }

        // Write merged statistics to new segment
        for (stat_name, stat_value) in statistics {
            let kb = KeyBuilder::segment_stat(dest_segment, &stat_name);
            let mut val_bytes = [0; 8];
            BigEndian::write_i64(&mut val_bytes, stat_value);
            try!(self.db.put(&kb.key(), &val_bytes));
        }

        // Note: Don't merge the deletion lists
        // Deletion lists can change at any time so we must lock the "document index"
        // before merging them so they can't be altered during merge. we cannot lock
        // this until the commit phase though.

        Ok(())
    }

    fn commit_segment_merge(&self, source_segments: &Vec<u32>, dest_segment: u32, doc_ref_mapping: &HashMap<DocRef, u16>) -> Result<(), SegmentMergeError> {
        let write_batch = WriteBatch::default();

        // Activate new segment
        let kb = KeyBuilder::segment_active(dest_segment);
        try!(write_batch.put(&kb.key(), b""));

        // Deactivate old segments
        for source_segment in source_segments.iter() {
            // Activate new segment
            let kb = KeyBuilder::segment_active(*source_segment);
            try!(write_batch.delete(&kb.key()));
        }

        // Update document index and commit
        // This will write the write batch
        try!(self.document_index.commit_segment_merge(&self.db, write_batch, source_segments, dest_segment, doc_ref_mapping));

        Ok(())
    }

    pub fn merge_segments(&self, source_segments: &Vec<u32>) -> Result<u32, SegmentMergeError> {
        let dest_segment = try!(self.segments.new_segment(&self.db));

        // Generate a mapping between the ids of the documents in the old segments to the new one
        // This packs the id spaces of the old segments together:
        // For example, say we have to merge 3 segments with 100 documents each:
        //  - The first segment's ids will be the same as before
        //  - The second segment's ids will be remapped to 100 - 199
        //  - The third segment's ids will be remapped to 200 - 299

        let mut doc_ref_mapping: HashMap<DocRef, u16> = HashMap::new();
        let mut current_ord: u32 = 0;

        for source_segment in source_segments.iter() {
            let kb = KeyBuilder::segment_stat(*source_segment, b"total_docs");
            let total_docs = match try!(self.db.get(&kb.key())) {
                Some(total_docs_bytes) => {
                    BigEndian::read_i64(&total_docs_bytes)
                }
                None => continue,
            };

            for source_ord in 0..total_docs {
                if current_ord >= 65536 {
                    return Err(SegmentMergeError::TooManyDocs);
                }

                let from = DocRef::from_segment_ord(*source_segment, source_ord as u16);
                doc_ref_mapping.insert(from, current_ord as u16);
                current_ord += 1;
            }
        }

        // Merge segment data
        // Most of the heavy lifting happens here. This merges all the immutable parts of
        // the segment (which is everything but the deletion list). It does not activate the
        // segment.
        // This means that nothing bad will happen if it crashes half way through -- the
        // worst that could happen is we're left with a partially-written segment that we
        // have to clean up.
        try!(self.merge_segment_data(&source_segments, dest_segment, &doc_ref_mapping));

        // Commit the merge
        // This activates the new segment and updates the document index. Effectively committing
        // the merge.
        // Throughout this stage we need an exclusive lock to the document index. This is to
        // prevent documents in the source segments being deleted/updated so we don't accidentally
        // undelete them (this will block until the merge is complete so they delete/update from
        // the new segment).
        try!(self.commit_segment_merge(&source_segments, dest_segment, &doc_ref_mapping));

        Ok(dest_segment)
    }

    pub fn purge_segments(&self, segments: &Vec<u32>) -> Result<(), rocksdb::Error> {
        // Put segments in a BTreeSet as this is much faster for performing contains queries against
        let segments_btree = segments.iter().collect::<BTreeSet<_>>();

        // Purge term directories

        /// Converts term directory key strings "d1/2/3" into tuples of 3 i32s (1, 2, 3)
        fn parse_term_directory_key(key: &[u8]) -> (u32, u32, u32) {
            let mut nums_iter = key[1..].split(|b| *b == b'/').map(|s| str::from_utf8(s).unwrap().parse::<u32>().unwrap());
            (nums_iter.next().unwrap(), nums_iter.next().unwrap(), nums_iter.next().unwrap())
        }

        for (k, _) in self.db.iterator(IteratorMode::From(b"d", Direction::Forward)) {
            if k[0] != b'd' {
                // No more term directories to delete
                break;
            }

            let (_, _, segment) = parse_term_directory_key(&k);

            if segments_btree.contains(&segment) {
                try!(self.db.delete(&k));
            }
        }

        // Purge the stored values

        /// Converts stored value key strings "v1/2/3/v" into tuples of 3 i32s and a Vec<u8> (1, 2, 3, vec![b'v', b'a', b'l'])
        fn parse_stored_value_key(key: &[u8]) -> (u32, u32, u32, Vec<u8>) {
            let mut parts_iter = key[1..].split(|b| *b == b'/');
            let segment = str::from_utf8(parts_iter.next().unwrap()).unwrap().parse::<u32>().unwrap();
            let doc_id = str::from_utf8(parts_iter.next().unwrap()).unwrap().parse::<u32>().unwrap();
            let field_ord = str::from_utf8(parts_iter.next().unwrap()).unwrap().parse::<u32>().unwrap();
            let value_type = parts_iter.next().unwrap().to_vec();

            (segment, doc_id, field_ord, value_type)
        }

        for source_segment in segments.iter() {
            let kb = KeyBuilder::segment_stored_values_prefix(*source_segment);
            for (k, _) in self.db.iterator(IteratorMode::From(&kb.key(), Direction::Forward)) {
                if k[0] != b'v' {
                    // No more stored values to delete
                    break;
                }

                let (segment, _, _, _) = parse_stored_value_key(&k);

                if segment != *source_segment {
                    // Segment finished
                    break;
                }

                try!(self.db.delete(&k));
            }
        }

        // Purge the statistics

        /// Converts statistic key strings "s1/total_docs" into tuples of 1 i32 and a Vec<u8> (1, ['t', 'o', 't', ...])
        fn parse_statistic_key(key: &[u8]) -> (u32, Vec<u8>) {
            let mut parts_iter = key[1..].split(|b| *b == b'/');
            let segment = str::from_utf8(parts_iter.next().unwrap()).unwrap().parse::<u32>().unwrap();
            let statistic_name = parts_iter.next().unwrap().to_vec();

            (segment, statistic_name)
        }

        for source_segment in segments.iter() {
            let kb = KeyBuilder::segment_stat_prefix(*source_segment);
            for (k, _) in self.db.iterator(IteratorMode::From(&kb.key(), Direction::Forward)) {
                if k[0] != b's' {
                    // No more statistics to purge
                    break;
                }

                let (segment, _) = parse_statistic_key(&k);

                if segment != *source_segment {
                    // Segment finished
                    break;
                }

                try!(self.db.delete(&k));
            }
        }

        // Purge the deletion lists
        for source_segment in segments.iter() {
            let kb = KeyBuilder::segment_del_list(*source_segment);
            try!(self.db.delete(&kb.key()));
        }

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