<commit_before>/* This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at http://mozilla.org/MPL/2.0/. */

use hyper::client::Pool;
use hyper::net::{HttpStream, HttpsConnector, SslClient};
use openssl::ssl::{SSL_OP_NO_SSLV2, SSL_OP_NO_SSLV3, SSL_VERIFY_PEER};
use openssl::ssl::{Ssl, SslContext, SslMethod, SslStream};
use std::sync::Arc;
use util::resource_files::resources_dir_path;

pub type Connector = HttpsConnector<ServoSslClient>;

// The basic logic here is to prefer ciphers with ECDSA certificates, Forward
// Secrecy, AES GCM ciphers, AES ciphers, and finally 3DES ciphers.
// A complete discussion of the issues involved in TLS configuration can be found here:
// https://wiki.mozilla.org/Security/Server_Side_TLS
const DEFAULT_CIPHERS: &'static str = concat!(
    "ECDHE-ECDSA-AES128-GCM-SHA256:ECDHE-RSA-AES128-GCM-SHA256:",
    "ECDHE-ECDSA-AES256-GCM-SHA384:ECDHE-RSA-AES256-GCM-SHA384:",
    "DHE-RSA-AES128-GCM-SHA256:ECDHE-ECDSA-AES128-SHA256:",
    "ECDHE-RSA-AES128-SHA256:ECDHE-ECDSA-AES256-SHA384:ECDHE-RSA-AES256-SHA384:",
    "ECDHE-ECDSA-AES128-SHA:ECDHE-RSA-AES128-SHA:ECDHE-ECDSA-AES256-SHA:",
    "ECDHE-RSA-AES256-SHA:DHE-RSA-AES128-SHA256:DHE-RSA-AES128-SHA:",
    "DHE-RSA-AES256-SHA256:DHE-RSA-AES256-SHA:ECDHE-RSA-DES-CBC3-SHA:",
    "ECDHE-ECDSA-DES-CBC3-SHA:AES128-GCM-SHA256:AES256-GCM-SHA384:",
    "AES128-SHA256:AES256-SHA256:AES128-SHA:AES256-SHA"
);

pub fn create_http_connector() -> Arc<Pool<Connector>> {
    let mut context = SslContext::new(SslMethod::Sslv23).unwrap();
    context.set_CA_file(&resources_dir_path().join("certs")).unwrap();
    context.set_cipher_list(DEFAULT_CIPHERS).unwrap();
    context.set_options(SSL_OP_NO_SSLV2 | SSL_OP_NO_SSLV3);
    let connector = HttpsConnector::new(ServoSslClient {
        context: Arc::new(context)
    });

    Arc::new(Pool::with_connector(Default::default(), connector))
}

pub struct ServoSslClient {
    context: Arc<SslContext>,
}

impl SslClient for ServoSslClient {
    type Stream = SslStream<HttpStream>;

    fn wrap_client(&self, stream: HttpStream, host: &str) -> Result<Self::Stream, ::hyper::Error> {
        let mut ssl = try!(Ssl::new(&self.context));
        try!(ssl.set_hostname(host));
        let host = host.to_owned();
        ssl.set_verify_callback(SSL_VERIFY_PEER, move |p, x| {
            ::openssl_verify::verify_callback(&host, p, x)
        });
        SslStream::connect(ssl, stream).map_err(From::from)
    }
}
<commit_msg>Auto merge of #12035 - alex:patch-1, r=frewsxcv<commit_after>/* This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at http://mozilla.org/MPL/2.0/. */

use hyper::client::Pool;
use hyper::net::{HttpStream, HttpsConnector, SslClient};
use openssl::ssl::{SSL_OP_NO_COMPRESSION, SSL_OP_NO_SSLV2, SSL_OP_NO_SSLV3, SSL_VERIFY_PEER};
use openssl::ssl::{Ssl, SslContext, SslMethod, SslStream};
use std::sync::Arc;
use util::resource_files::resources_dir_path;

pub type Connector = HttpsConnector<ServoSslClient>;

// The basic logic here is to prefer ciphers with ECDSA certificates, Forward
// Secrecy, AES GCM ciphers, AES ciphers, and finally 3DES ciphers.
// A complete discussion of the issues involved in TLS configuration can be found here:
// https://wiki.mozilla.org/Security/Server_Side_TLS
const DEFAULT_CIPHERS: &'static str = concat!(
    "ECDHE-ECDSA-AES128-GCM-SHA256:ECDHE-RSA-AES128-GCM-SHA256:",
    "ECDHE-ECDSA-AES256-GCM-SHA384:ECDHE-RSA-AES256-GCM-SHA384:",
    "DHE-RSA-AES128-GCM-SHA256:ECDHE-ECDSA-AES128-SHA256:",
    "ECDHE-RSA-AES128-SHA256:ECDHE-ECDSA-AES256-SHA384:ECDHE-RSA-AES256-SHA384:",
    "ECDHE-ECDSA-AES128-SHA:ECDHE-RSA-AES128-SHA:ECDHE-ECDSA-AES256-SHA:",
    "ECDHE-RSA-AES256-SHA:DHE-RSA-AES128-SHA256:DHE-RSA-AES128-SHA:",
    "DHE-RSA-AES256-SHA256:DHE-RSA-AES256-SHA:ECDHE-RSA-DES-CBC3-SHA:",
    "ECDHE-ECDSA-DES-CBC3-SHA:AES128-GCM-SHA256:AES256-GCM-SHA384:",
    "AES128-SHA256:AES256-SHA256:AES128-SHA:AES256-SHA"
);

pub fn create_http_connector() -> Arc<Pool<Connector>> {
    let mut context = SslContext::new(SslMethod::Sslv23).unwrap();
    context.set_CA_file(&resources_dir_path().join("certs")).unwrap();
    context.set_cipher_list(DEFAULT_CIPHERS).unwrap();
    context.set_options(SSL_OP_NO_SSLV2 | SSL_OP_NO_SSLV3 | SSL_OP_NO_COMPRESSION);
    let connector = HttpsConnector::new(ServoSslClient {
        context: Arc::new(context)
    });

    Arc::new(Pool::with_connector(Default::default(), connector))
}

pub struct ServoSslClient {
    context: Arc<SslContext>,
}

impl SslClient for ServoSslClient {
    type Stream = SslStream<HttpStream>;

    fn wrap_client(&self, stream: HttpStream, host: &str) -> Result<Self::Stream, ::hyper::Error> {
        let mut ssl = try!(Ssl::new(&self.context));
        try!(ssl.set_hostname(host));
        let host = host.to_owned();
        ssl.set_verify_callback(SSL_VERIFY_PEER, move |p, x| {
            ::openssl_verify::verify_callback(&host, p, x)
        });
        SslStream::connect(ssl, stream).map_err(From::from)
    }
}
<|endoftext|>


<commit_before>// Copyright (c) 2017-2019 Rene van der Meer
//
// Permission is hereby granted, free of charge, to any person obtaining a
// copy of this software and associated documentation files (the "Software"),
// to deal in the Software without restriction, including without limitation
// the rights to use, copy, modify, merge, publish, distribute, sublicense,
// and/or sell copies of the Software, and to permit persons to whom the
// Software is furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL
// THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
// FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
// DEALINGS IN THE SOFTWARE.

//! Implements `embedded_hal` traits.
//!
//! Include this module in your `use` statement to add the required traits
//! needed to use `I2c` with `embedded_hal` drivers.

use super::{Error, I2c, Result};

use embedded_hal::blocking::i2c::{Read, Write, WriteRead};

impl Write for I2c {
    type Error = Error;

    fn write(&mut self, address: u8, bytes: &[u8]) -> Result<()> {
        self.set_slave_address(u16::from(address))?;
        I2c::write(self, bytes)?;

        Ok(())
    }
}

impl Read for I2c {
    type Error = Error;

    fn read(&mut self, address: u8, buffer: &mut [u8]) -> Result<()> {
        self.set_slave_address(u16::from(address))?;
        I2c::read(self, buffer)?;

        Ok(())
    }
}

impl WriteRead for I2c {
    type Error = Error;

    fn write_read(&mut self, address: u8, bytes: &[u8], buffer: &mut [u8]) -> Result<()> {
        self.set_slave_address(u16::from(address))?;
        I2c::write_read(self, bytes, buffer)?;

        Ok(())
    }
}
<commit_msg>Change use statement order<commit_after>// Copyright (c) 2017-2019 Rene van der Meer
//
// Permission is hereby granted, free of charge, to any person obtaining a
// copy of this software and associated documentation files (the "Software"),
// to deal in the Software without restriction, including without limitation
// the rights to use, copy, modify, merge, publish, distribute, sublicense,
// and/or sell copies of the Software, and to permit persons to whom the
// Software is furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL
// THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
// FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
// DEALINGS IN THE SOFTWARE.

use embedded_hal::blocking::i2c::{Read, Write, WriteRead};

use super::{Error, I2c, Result};

impl Write for I2c {
    type Error = Error;

    fn write(&mut self, address: u8, bytes: &[u8]) -> Result<()> {
        self.set_slave_address(u16::from(address))?;
        I2c::write(self, bytes)?;

        Ok(())
    }
}

impl Read for I2c {
    type Error = Error;

    fn read(&mut self, address: u8, buffer: &mut [u8]) -> Result<()> {
        self.set_slave_address(u16::from(address))?;
        I2c::read(self, buffer)?;

        Ok(())
    }
}

impl WriteRead for I2c {
    type Error = Error;

    fn write_read(&mut self, address: u8, bytes: &[u8], buffer: &mut [u8]) -> Result<()> {
        self.set_slave_address(u16::from(address))?;
        I2c::write_read(self, bytes, buffer)?;

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


<commit_before>//
//  SOS: the Stupid Operating System
//  by Eliza Weisman (hi@hawkweisman.me)
//
//  Copyright (c) 2015-2016 Eliza Weisman
//  Released under the terms of the MIT license. See `LICENSE` in the root
//  directory of this repository for more information.
//
//! `x86_64` paging.
use ::VAddr;
use ::paging::{Page, Mapper, FrameAllocator};

use super::{Frame, PAddr, PAGE_SIZE};
use self::table::*;

use core::ptr::Unique;

pub mod table;
pub mod entry;

/// Struct representing the currently active PML4 instance.
///
/// The `ActivePML4` is a `Unique` reference to a PML4-level page table. It's
/// unique because, well, there can only be one active PML4 at a given time.
///
///
pub struct ActivePML4(Unique<Table<PML4Level>>);

/// The active PML4 table is the single point of entry for page mapping.
impl Mapper for ActivePML4 {
    type Flags = EntryFlags;
    type Frame = Frame;

    fn translate(&self, vaddr: VAddr) -> Option<PAddr> {
        self.translate_page(Page::containing(vaddr))
            .map(|frame| {
                let offset = *vaddr % PAGE_SIZE as usize;
                PAddr::from(frame.number + offset as u64)
            })
    }

    fn translate_page(&self, page: Page) -> Option<Frame> {
        let pdpt = self.pml4().next_table(page.pml4_index());

        pdpt.and_then(|pdpt| pdpt.next_table(page.pdpt_index()))
            .and_then(|pd| pd.next_table(page.pd_index()))
            .and_then(|pt| pt[page.pt_index()].get_frame())
            .or_else( || {
                pdpt.and_then(|pdpt| {
                    let pdpt_entry = &pdpt[page.pdpt_index()];
                    pdpt_entry.do_huge(page.pd_index() + page.pt_index() )
                        .or_else(|| {
                            pdpt.next_table(page.pdpt_index())
                                .and_then(|pd| {
                                    let pd_entry = &pd[page.pd_index()];
                                    pd_entry.do_huge(page.pt_index())
                                })
                        })
                    })
                })
    }

    /// Modifies the page tables so that `page` maps to `frame`.
    ///
    /// # Arguments
    /// + `page`: the virtual `Page` to map
    /// + `frame`: the physical `Frame` that `Page` should map to.
    /// + `flags`: the page table entry flags.
    /// + `alloc`: a memory allocator
    fn map_to<A>( &mut self, page: Page, frame: Frame
                , flags: EntryFlags, alloc: &mut A)
    where A: FrameAllocator<Frame> {

       // get the page table index of the page to map
       let idx = page.pt_index();

        // access or create all the lower-level page tables.
        let mut page_table
            // get the PML4
            = self.pml4_mut()
                  // get or create the PDPT table at the page's PML4 index
                  .create_next(page.pml4_index(), alloc)
                  // get or create the PD table at the page's PDPT index
                  .create_next(page.pdpt_index(), alloc)
                  // get or create the page table at the  page's PD table index
                  .create_next(idx, alloc);

        // check if the page at that index is not currently in use, as we
        // cannot map a page which is currently in use.
        assert!(page_table[idx].is_unused()
               , "Could not map frame {:?}, page table entry {} is already \
                  in use!", frame, idx);
        // set the page table entry at that index
        page_table[idx].set(frame, flags | table::PRESENT);
    }

    fn identity_map<A>(&mut self, frame: Frame, flags: EntryFlags
                      , alloc: &mut A)
    where A: FrameAllocator<Frame> {
        self.map_to( Page::containing(VAddr::from(frame.base_addr().0 as usize))
                   , frame
                   , flags
                   , alloc )
    }

    fn map_to_any<A>(&mut self, page: Page, flags: EntryFlags, alloc: &mut A)
    where A: FrameAllocator<Frame> {
        self.map_to( page
                   , alloc.alloc_frame()
                            // TODO: would we rather rewrite this to return
                            // a `Result`? I think so.
                           .expect("Couldn't map page, out of frames!")
                   , flags
                   , alloc);
    }

}

impl ActivePML4 {

    pub unsafe fn new() -> Self {
        ActivePML4(Unique::new(PML4))
    }

    fn pml4(&self) -> &Table<PML4Level> {
        unsafe { self.0.get() }
    }

    fn pml4_mut(&mut self) -> &mut Table<PML4Level> {
        unsafe { self.0.get_mut() }
    }

}
<commit_msg>fix(memory::paging::ActivePML4): Unique::get/get_mut were renamed to as_ref/as_mut<commit_after>//
//  SOS: the Stupid Operating System
//  by Eliza Weisman (hi@hawkweisman.me)
//
//  Copyright (c) 2015-2016 Eliza Weisman
//  Released under the terms of the MIT license. See `LICENSE` in the root
//  directory of this repository for more information.
//
//! `x86_64` paging.
use ::VAddr;
use ::paging::{Page, Mapper, FrameAllocator};

use super::{Frame, PAddr, PAGE_SIZE};
use self::table::*;

use core::ptr::Unique;

pub mod table;
pub mod entry;

/// Struct representing the currently active PML4 instance.
///
/// The `ActivePML4` is a `Unique` reference to a PML4-level page table. It's
/// unique because, well, there can only be one active PML4 at a given time.
///
///
pub struct ActivePML4(Unique<Table<PML4Level>>);

/// The active PML4 table is the single point of entry for page mapping.
impl Mapper for ActivePML4 {
    type Flags = EntryFlags;
    type Frame = Frame;

    fn translate(&self, vaddr: VAddr) -> Option<PAddr> {
        self.translate_page(Page::containing(vaddr))
            .map(|frame| {
                let offset = *vaddr % PAGE_SIZE as usize;
                PAddr::from(frame.number + offset as u64)
            })
    }

    fn translate_page(&self, page: Page) -> Option<Frame> {
        let pdpt = self.pml4().next_table(page.pml4_index());

        pdpt.and_then(|pdpt| pdpt.next_table(page.pdpt_index()))
            .and_then(|pd| pd.next_table(page.pd_index()))
            .and_then(|pt| pt[page.pt_index()].get_frame())
            .or_else( || {
                pdpt.and_then(|pdpt| {
                    let pdpt_entry = &pdpt[page.pdpt_index()];
                    pdpt_entry.do_huge(page.pd_index() + page.pt_index() )
                        .or_else(|| {
                            pdpt.next_table(page.pdpt_index())
                                .and_then(|pd| {
                                    let pd_entry = &pd[page.pd_index()];
                                    pd_entry.do_huge(page.pt_index())
                                })
                        })
                    })
                })
    }

    /// Modifies the page tables so that `page` maps to `frame`.
    ///
    /// # Arguments
    /// + `page`: the virtual `Page` to map
    /// + `frame`: the physical `Frame` that `Page` should map to.
    /// + `flags`: the page table entry flags.
    /// + `alloc`: a memory allocator
    fn map_to<A>( &mut self, page: Page, frame: Frame
                , flags: EntryFlags, alloc: &mut A)
    where A: FrameAllocator<Frame> {

       // get the page table index of the page to map
       let idx = page.pt_index();

        // access or create all the lower-level page tables.
        let mut page_table
            // get the PML4
            = self.pml4_mut()
                  // get or create the PDPT table at the page's PML4 index
                  .create_next(page.pml4_index(), alloc)
                  // get or create the PD table at the page's PDPT index
                  .create_next(page.pdpt_index(), alloc)
                  // get or create the page table at the  page's PD table index
                  .create_next(idx, alloc);

        // check if the page at that index is not currently in use, as we
        // cannot map a page which is currently in use.
        assert!(page_table[idx].is_unused()
               , "Could not map frame {:?}, page table entry {} is already \
                  in use!", frame, idx);
        // set the page table entry at that index
        page_table[idx].set(frame, flags | table::PRESENT);
    }

    fn identity_map<A>(&mut self, frame: Frame, flags: EntryFlags
                      , alloc: &mut A)
    where A: FrameAllocator<Frame> {
        self.map_to( Page::containing(VAddr::from(frame.base_addr().0 as usize))
                   , frame
                   , flags
                   , alloc )
    }

    fn map_to_any<A>(&mut self, page: Page, flags: EntryFlags, alloc: &mut A)
    where A: FrameAllocator<Frame> {
        self.map_to( page
                   , alloc.alloc_frame()
                            // TODO: would we rather rewrite this to return
                            // a `Result`? I think so.
                           .expect("Couldn't map page, out of frames!")
                   , flags
                   , alloc);
    }

}

impl ActivePML4 {

    pub unsafe fn new() -> Self {
        ActivePML4(Unique::new(PML4))
    }
    // TODO: should this be replaced with an ops::Deref implementation?
    //          - eliza, 5/21/2017
    fn pml4(&self) -> &Table<PML4Level> {
        unsafe { self.0.as_ref() }
    }

    fn pml4_mut(&mut self) -> &mut Table<PML4Level> {
        unsafe { self.0.as_mut() }
    }

}
<|endoftext|>


<commit_before>use crate::{
    buffer::Buffer, connection::Connection, connection_status::ConnectionState, Error,
    PromiseResolver, Result,
};
use amq_protocol::frame::{gen_frame, parse_frame, AMQPFrame, GenError};
use log::{error, trace};
use mio::{event::Source, Events, Interest, Poll, Token, Waker};
use std::{
    collections::VecDeque,
    io::{self, Read, Write},
    sync::Arc,
    thread::Builder as ThreadBuilder,
    time::{Duration, Instant},
};

pub(crate) const SOCKET: Token = Token(1);
const WAKER: Token = Token(2);

const FRAMES_STORAGE: usize = 32;

#[derive(Debug, PartialEq)]
enum Status {
    Initial,
    Setup,
    Stop,
}

pub struct IoLoop<T> {
    connection: Connection,
    socket: T,
    status: Status,
    poll: Poll,
    waker: Arc<Waker>,
    frame_size: usize,
    receive_buffer: Buffer,
    send_buffer: Buffer,
    can_write: bool,
    can_read: bool,
    poll_timeout: Option<Duration>,
    serialized_frames: VecDeque<(u64, Option<PromiseResolver<()>>)>,
    last_write: Instant,
}

impl<T: Source + Read + Write + Send + 'static> IoLoop<T> {
    pub(crate) fn new(
        connection: Connection,
        socket: T,
        poll: Option<(Poll, Token)>,
    ) -> Result<Self> {
        let (poll, registered) = poll
            .map(|t| Ok((t.0, true)))
            .unwrap_or_else(|| Poll::new().map(|poll| (poll, false)))?;
        let frame_size = std::cmp::max(8192, connection.configuration().frame_max() as usize);
        let waker = Arc::new(Waker::new(poll.registry(), WAKER)?);
        let mut inner = Self {
            connection,
            socket,
            status: Status::Initial,
            poll,
            waker,
            frame_size,
            receive_buffer: Buffer::with_capacity(FRAMES_STORAGE * frame_size),
            send_buffer: Buffer::with_capacity(FRAMES_STORAGE * frame_size),
            can_write: false,
            can_read: false,
            poll_timeout: None,
            serialized_frames: VecDeque::default(),
            last_write: Instant::now(),
        };
        if registered {
            inner.poll.registry().reregister(
                &mut inner.socket,
                SOCKET,
                Interest::READABLE | Interest::WRITABLE,
            )?;
        } else {
            inner.poll.registry().register(
                &mut inner.socket,
                SOCKET,
                Interest::READABLE | Interest::WRITABLE,
            )?;
        }
        Ok(inner)
    }

    fn ensure_setup(&mut self) -> Result<()> {
        if self.status != Status::Setup && self.connection.status().connected() {
            let frame_max = self.connection.configuration().frame_max() as usize;
            self.frame_size = std::cmp::max(self.frame_size, frame_max);
            self.receive_buffer.grow(FRAMES_STORAGE * self.frame_size);
            self.send_buffer.grow(FRAMES_STORAGE * self.frame_size);
            let heartbeat = self.connection.configuration().heartbeat();
            if heartbeat != 0 {
                let heartbeat = Duration::from_millis(u64::from(heartbeat) * 500); // * 1000 (ms) / 2 (half the negociated timeout)
                self.poll_timeout = Some(heartbeat);
            }
            self.status = Status::Setup;
        }
        Ok(())
    }

    fn has_data(&self) -> bool {
        self.connection.has_pending_frames()
            || self.send_buffer.available_data() > 0
            || !self.serialized_frames.is_empty()
    }

    fn can_write(&self) -> bool {
        self.can_write && self.has_data() && !self.connection.status().blocked()
    }

    fn can_read(&self) -> bool {
        self.can_read
    }

    fn can_parse(&self) -> bool {
        self.receive_buffer.available_data() > 0
    }

    fn should_continue(&self) -> bool {
        let connection_status = self.connection.status();
        (self.status == Status::Initial
            || connection_status.connected()
            || connection_status.closing())
            && self.status != Status::Stop
            && !connection_status.errored()
    }

    pub fn start(mut self) -> Result<()> {
        let waker = self.waker.clone();
        self.connection.clone().set_io_loop(
            ThreadBuilder::new()
                .name("io_loop".to_owned())
                .spawn(move || {
                    let mut events = Events::with_capacity(1024);
                    while self.should_continue() {
                        if let Err(err) = self.run(&mut events) {
                            self.cancel_serilized_frames(err)?;
                        }
                    }
                    Ok(())
                })?,
            waker,
        )
    }

    fn cancel_serilized_frames(&self, error: Error) -> Result<()> {
        for (_, resolver) in self.serialized_frames.iter() {
            if let Some(resolver) = resolver {
                resolver.swear(Err(error.clone()));
            }
        }
        Err(error)
    }

    fn poll_timeout(&self) -> Option<Duration> {
        self.poll_timeout.map(|timeout| {
            timeout
                .checked_sub(self.last_write.elapsed())
                .unwrap_or_else(|| Duration::from_secs(0))
        })
    }

    fn poll(&mut self, events: &mut Events) -> Result<()> {
        trace!("io_loop poll");
        self.poll.poll(events, self.poll_timeout())?;
        trace!("io_loop poll done");
        for event in events.iter() {
            if event.token() == SOCKET {
                trace!("Got mio event for socket: {:?}", event);
                if event.is_read_closed() || event.is_write_closed() {
                    self.critical_error(io::Error::from(io::ErrorKind::ConnectionReset).into())?;
                }
                if event.is_error() {
                    self.critical_error(io::Error::from(io::ErrorKind::ConnectionAborted).into())?;
                }
                if event.is_readable() {
                    self.can_read = true;
                }
                if event.is_writable() {
                    self.can_write = true;
                }
            }
        }
        Ok(())
    }

    fn run(&mut self, events: &mut Events) -> Result<()> {
        trace!("io_loop run");
        self.ensure_setup()?;
        self.poll(events)?;
        self.do_run()
    }

    fn do_run(&mut self) -> Result<()> {
        trace!(
            "io_loop do_run; can_read={}, can_write={}, has_data={}",
            self.can_read,
            self.can_write,
            self.has_data()
        );
        loop {
            self.write()?;
            if self.connection.status().closed() {
                self.status = Status::Stop;
            }
            if self.should_continue() {
                self.read()?;
            }
            self.parse()?;
            self.connection.poll_internal_promises()?;
            if self.should_heartbeat() {
                self.connection.send_heartbeat()?;
                // Update last_write so that if we cannot write yet to the socket, we don't enqueue countless heartbeats
                self.last_write = Instant::now();
            }
            if self.stop_looping() {
                self.maybe_continue()?;
                break;
            }
        }
        trace!(
            "io_loop do_run done; can_read={}, can_write={}, has_data={}, status={:?}",
            self.can_read,
            self.can_write,
            self.has_data(),
            self.status
        );
        Ok(())
    }

    fn should_heartbeat(&self) -> bool {
        if let Some(heartbeat_timeout) = self.poll_timeout {
            self.last_write.elapsed() > heartbeat_timeout
        } else {
            false
        }
    }

    fn stop_looping(&self) -> bool {
        !self.can_read()
            || !self.can_write()
            || self.status == Status::Stop
            || self.connection.status().errored()
    }

    fn has_pending_operations(&self) -> bool {
        self.status != Status::Stop && (self.can_read() || self.can_parse() || self.can_write())
    }

    fn maybe_continue(&mut self) -> Result<()> {
        if self.has_pending_operations() {
            trace!(
                "io_loop send continue; can_read={}, can_write={}, has_data={}",
                self.can_read,
                self.can_write,
                self.has_data()
            );
            self.send_continue()?;
        }
        Ok(())
    }

    fn critical_error(&mut self, error: Error) -> Result<()> {
        if let ConnectionState::SentProtocolHeader(resolver, ..) = self.connection.status().state()
        {
            resolver.swear(Err(error.clone()));
        }
        self.status = Status::Stop;
        self.connection.set_error(error.clone())?;
        Err(error)
    }

    fn write(&mut self) -> Result<()> {
        if self.can_write() {
            if let Err(e) = self.write_to_stream() {
                if e.wouldblock() {
                    self.can_write = false
                } else {
                    error!("error writing: {:?}", e);
                    self.critical_error(e)?;
                }
            }
        }
        Ok(())
    }

    fn read(&mut self) -> Result<()> {
        if self.can_read() {
            if let Err(e) = self.read_from_stream() {
                if e.wouldblock() {
                    self.can_read = false
                } else {
                    error!("error reading: {:?}", e);
                    self.critical_error(e)?;
                }
            }
            self.receive_buffer.shift_unless_available(self.frame_size);
        }
        Ok(())
    }

    fn send_continue(&mut self) -> Result<()> {
        self.waker.wake()?;
        Ok(())
    }

    fn write_to_stream(&mut self) -> Result<()> {
        self.serialize()?;

        let sz = self.send_buffer.write_to(&mut self.socket)?;

        if sz > 0 {
            self.last_write = Instant::now();

            trace!("wrote {} bytes", sz);
            self.send_buffer.consume(sz, true);

            let mut written = sz as u64;
            while written > 0 {
                if let Some((to_write, resolver)) = self.serialized_frames.pop_front() {
                    if written < to_write {
                        self.serialized_frames
                            .push_front((to_write - written, resolver));
                        trace!("{} to write to complete this frame", to_write - written);
                        written = 0;
                    } else {
                        if let Some(resolver) = resolver {
                            resolver.swear(Ok(()));
                        }
                        written -= to_write;
                    }
                } else {
                    error!("We've written {} but didn't expect to write anything", written);
                    break;
                }
            }

            if self.send_buffer.available_data() > 0 {
                // We didn't write all the data yet
                trace!("Still {} to write", self.send_buffer.available_data());
                self.send_continue()?;
            }
        } else {
            error!("Socket was writable but we wrote 0, marking as wouldblock");
            self.can_write = false;
        }
        Ok(())
    }

    fn read_from_stream(&mut self) -> Result<()> {
        match self.connection.status().state() {
            ConnectionState::Closed => Ok(()),
            ConnectionState::Error => Err(Error::InvalidConnectionState(ConnectionState::Error)),
            _ => {
                self.receive_buffer
                    .read_from(&mut self.socket, false)
                    .map(|sz| {
                        trace!("read {} bytes", sz);
                        self.receive_buffer.fill(sz, false);
                    })?;
                Ok(())
            }
        }
    }

    fn serialize(&mut self) -> Result<()> {
        if let Some((next_msg, resolver)) = self.connection.next_frame() {
            // FIXME: having the body here is too verbose, impl Display on frame instead: trace!("will write to buffer: {:?}", next_msg);
            trace!("Will write to buffer");
            let checkpoint = self.send_buffer.checkpoint();
            let res = gen_frame(&next_msg)((&mut self.send_buffer).into());
            match res.map(|w| w.into_inner().1) {
                Ok(sz) => {
                    self.serialized_frames.push_back((sz, resolver));
                    Ok(())
                }
                Err(e) => {
                    self.send_buffer.rollback(checkpoint);
                    match e {
                        GenError::BufferTooSmall(_) => {
                            // Requeue msg
                            self.connection.requeue_frame((next_msg, resolver))?;
                            Ok(())
                        }
                        e => {
                            error!("error generating frame: {:?}", e);
                            let error = Error::SerialisationError(Arc::new(e));
                            self.connection.set_error(error.clone())?;
                            Err(error)
                        }
                    }
                }
            }
        } else {
            Ok(())
        }
    }

    fn parse(&mut self) -> Result<()> {
        if self.can_parse() {
            if let Some(frame) = self.do_parse()? {
                self.connection.handle_frame(frame)?;
            }
        }
        Ok(())
    }

    fn do_parse(&mut self) -> Result<Option<AMQPFrame>> {
        match parse_frame(self.receive_buffer.data()) {
            Ok((i, f)) => {
                let consumed = self.receive_buffer.offset(i);
                self.receive_buffer.consume(consumed, false);
                Ok(Some(f))
            }
            Err(e) => {
                if e.is_incomplete() {
                    self.receive_buffer.shift();
                    Ok(None)
                } else {
                    error!("parse error: {:?}", e);
                    let error = Error::ParsingError(e);
                    self.connection.set_error(error.clone())?;
                    Err(error)
                }
            }
        }
    }
}
<commit_msg>rustfmt<commit_after>use crate::{
    buffer::Buffer, connection::Connection, connection_status::ConnectionState, Error,
    PromiseResolver, Result,
};
use amq_protocol::frame::{gen_frame, parse_frame, AMQPFrame, GenError};
use log::{error, trace};
use mio::{event::Source, Events, Interest, Poll, Token, Waker};
use std::{
    collections::VecDeque,
    io::{self, Read, Write},
    sync::Arc,
    thread::Builder as ThreadBuilder,
    time::{Duration, Instant},
};

pub(crate) const SOCKET: Token = Token(1);
const WAKER: Token = Token(2);

const FRAMES_STORAGE: usize = 32;

#[derive(Debug, PartialEq)]
enum Status {
    Initial,
    Setup,
    Stop,
}

pub struct IoLoop<T> {
    connection: Connection,
    socket: T,
    status: Status,
    poll: Poll,
    waker: Arc<Waker>,
    frame_size: usize,
    receive_buffer: Buffer,
    send_buffer: Buffer,
    can_write: bool,
    can_read: bool,
    poll_timeout: Option<Duration>,
    serialized_frames: VecDeque<(u64, Option<PromiseResolver<()>>)>,
    last_write: Instant,
}

impl<T: Source + Read + Write + Send + 'static> IoLoop<T> {
    pub(crate) fn new(
        connection: Connection,
        socket: T,
        poll: Option<(Poll, Token)>,
    ) -> Result<Self> {
        let (poll, registered) = poll
            .map(|t| Ok((t.0, true)))
            .unwrap_or_else(|| Poll::new().map(|poll| (poll, false)))?;
        let frame_size = std::cmp::max(8192, connection.configuration().frame_max() as usize);
        let waker = Arc::new(Waker::new(poll.registry(), WAKER)?);
        let mut inner = Self {
            connection,
            socket,
            status: Status::Initial,
            poll,
            waker,
            frame_size,
            receive_buffer: Buffer::with_capacity(FRAMES_STORAGE * frame_size),
            send_buffer: Buffer::with_capacity(FRAMES_STORAGE * frame_size),
            can_write: false,
            can_read: false,
            poll_timeout: None,
            serialized_frames: VecDeque::default(),
            last_write: Instant::now(),
        };
        if registered {
            inner.poll.registry().reregister(
                &mut inner.socket,
                SOCKET,
                Interest::READABLE | Interest::WRITABLE,
            )?;
        } else {
            inner.poll.registry().register(
                &mut inner.socket,
                SOCKET,
                Interest::READABLE | Interest::WRITABLE,
            )?;
        }
        Ok(inner)
    }

    fn ensure_setup(&mut self) -> Result<()> {
        if self.status != Status::Setup && self.connection.status().connected() {
            let frame_max = self.connection.configuration().frame_max() as usize;
            self.frame_size = std::cmp::max(self.frame_size, frame_max);
            self.receive_buffer.grow(FRAMES_STORAGE * self.frame_size);
            self.send_buffer.grow(FRAMES_STORAGE * self.frame_size);
            let heartbeat = self.connection.configuration().heartbeat();
            if heartbeat != 0 {
                let heartbeat = Duration::from_millis(u64::from(heartbeat) * 500); // * 1000 (ms) / 2 (half the negociated timeout)
                self.poll_timeout = Some(heartbeat);
            }
            self.status = Status::Setup;
        }
        Ok(())
    }

    fn has_data(&self) -> bool {
        self.connection.has_pending_frames()
            || self.send_buffer.available_data() > 0
            || !self.serialized_frames.is_empty()
    }

    fn can_write(&self) -> bool {
        self.can_write && self.has_data() && !self.connection.status().blocked()
    }

    fn can_read(&self) -> bool {
        self.can_read
    }

    fn can_parse(&self) -> bool {
        self.receive_buffer.available_data() > 0
    }

    fn should_continue(&self) -> bool {
        let connection_status = self.connection.status();
        (self.status == Status::Initial
            || connection_status.connected()
            || connection_status.closing())
            && self.status != Status::Stop
            && !connection_status.errored()
    }

    pub fn start(mut self) -> Result<()> {
        let waker = self.waker.clone();
        self.connection.clone().set_io_loop(
            ThreadBuilder::new()
                .name("io_loop".to_owned())
                .spawn(move || {
                    let mut events = Events::with_capacity(1024);
                    while self.should_continue() {
                        if let Err(err) = self.run(&mut events) {
                            self.cancel_serilized_frames(err)?;
                        }
                    }
                    Ok(())
                })?,
            waker,
        )
    }

    fn cancel_serilized_frames(&self, error: Error) -> Result<()> {
        for (_, resolver) in self.serialized_frames.iter() {
            if let Some(resolver) = resolver {
                resolver.swear(Err(error.clone()));
            }
        }
        Err(error)
    }

    fn poll_timeout(&self) -> Option<Duration> {
        self.poll_timeout.map(|timeout| {
            timeout
                .checked_sub(self.last_write.elapsed())
                .unwrap_or_else(|| Duration::from_secs(0))
        })
    }

    fn poll(&mut self, events: &mut Events) -> Result<()> {
        trace!("io_loop poll");
        self.poll.poll(events, self.poll_timeout())?;
        trace!("io_loop poll done");
        for event in events.iter() {
            if event.token() == SOCKET {
                trace!("Got mio event for socket: {:?}", event);
                if event.is_read_closed() || event.is_write_closed() {
                    self.critical_error(io::Error::from(io::ErrorKind::ConnectionReset).into())?;
                }
                if event.is_error() {
                    self.critical_error(io::Error::from(io::ErrorKind::ConnectionAborted).into())?;
                }
                if event.is_readable() {
                    self.can_read = true;
                }
                if event.is_writable() {
                    self.can_write = true;
                }
            }
        }
        Ok(())
    }

    fn run(&mut self, events: &mut Events) -> Result<()> {
        trace!("io_loop run");
        self.ensure_setup()?;
        self.poll(events)?;
        self.do_run()
    }

    fn do_run(&mut self) -> Result<()> {
        trace!(
            "io_loop do_run; can_read={}, can_write={}, has_data={}",
            self.can_read,
            self.can_write,
            self.has_data()
        );
        loop {
            self.write()?;
            if self.connection.status().closed() {
                self.status = Status::Stop;
            }
            if self.should_continue() {
                self.read()?;
            }
            self.parse()?;
            self.connection.poll_internal_promises()?;
            if self.should_heartbeat() {
                self.connection.send_heartbeat()?;
                // Update last_write so that if we cannot write yet to the socket, we don't enqueue countless heartbeats
                self.last_write = Instant::now();
            }
            if self.stop_looping() {
                self.maybe_continue()?;
                break;
            }
        }
        trace!(
            "io_loop do_run done; can_read={}, can_write={}, has_data={}, status={:?}",
            self.can_read,
            self.can_write,
            self.has_data(),
            self.status
        );
        Ok(())
    }

    fn should_heartbeat(&self) -> bool {
        if let Some(heartbeat_timeout) = self.poll_timeout {
            self.last_write.elapsed() > heartbeat_timeout
        } else {
            false
        }
    }

    fn stop_looping(&self) -> bool {
        !self.can_read()
            || !self.can_write()
            || self.status == Status::Stop
            || self.connection.status().errored()
    }

    fn has_pending_operations(&self) -> bool {
        self.status != Status::Stop && (self.can_read() || self.can_parse() || self.can_write())
    }

    fn maybe_continue(&mut self) -> Result<()> {
        if self.has_pending_operations() {
            trace!(
                "io_loop send continue; can_read={}, can_write={}, has_data={}",
                self.can_read,
                self.can_write,
                self.has_data()
            );
            self.send_continue()?;
        }
        Ok(())
    }

    fn critical_error(&mut self, error: Error) -> Result<()> {
        if let ConnectionState::SentProtocolHeader(resolver, ..) = self.connection.status().state()
        {
            resolver.swear(Err(error.clone()));
        }
        self.status = Status::Stop;
        self.connection.set_error(error.clone())?;
        Err(error)
    }

    fn write(&mut self) -> Result<()> {
        if self.can_write() {
            if let Err(e) = self.write_to_stream() {
                if e.wouldblock() {
                    self.can_write = false
                } else {
                    error!("error writing: {:?}", e);
                    self.critical_error(e)?;
                }
            }
        }
        Ok(())
    }

    fn read(&mut self) -> Result<()> {
        if self.can_read() {
            if let Err(e) = self.read_from_stream() {
                if e.wouldblock() {
                    self.can_read = false
                } else {
                    error!("error reading: {:?}", e);
                    self.critical_error(e)?;
                }
            }
            self.receive_buffer.shift_unless_available(self.frame_size);
        }
        Ok(())
    }

    fn send_continue(&mut self) -> Result<()> {
        self.waker.wake()?;
        Ok(())
    }

    fn write_to_stream(&mut self) -> Result<()> {
        self.serialize()?;

        let sz = self.send_buffer.write_to(&mut self.socket)?;

        if sz > 0 {
            self.last_write = Instant::now();

            trace!("wrote {} bytes", sz);
            self.send_buffer.consume(sz, true);

            let mut written = sz as u64;
            while written > 0 {
                if let Some((to_write, resolver)) = self.serialized_frames.pop_front() {
                    if written < to_write {
                        self.serialized_frames
                            .push_front((to_write - written, resolver));
                        trace!("{} to write to complete this frame", to_write - written);
                        written = 0;
                    } else {
                        if let Some(resolver) = resolver {
                            resolver.swear(Ok(()));
                        }
                        written -= to_write;
                    }
                } else {
                    error!(
                        "We've written {} but didn't expect to write anything",
                        written
                    );
                    break;
                }
            }

            if self.send_buffer.available_data() > 0 {
                // We didn't write all the data yet
                trace!("Still {} to write", self.send_buffer.available_data());
                self.send_continue()?;
            }
        } else {
            error!("Socket was writable but we wrote 0, marking as wouldblock");
            self.can_write = false;
        }
        Ok(())
    }

    fn read_from_stream(&mut self) -> Result<()> {
        match self.connection.status().state() {
            ConnectionState::Closed => Ok(()),
            ConnectionState::Error => Err(Error::InvalidConnectionState(ConnectionState::Error)),
            _ => {
                self.receive_buffer
                    .read_from(&mut self.socket, false)
                    .map(|sz| {
                        trace!("read {} bytes", sz);
                        self.receive_buffer.fill(sz, false);
                    })?;
                Ok(())
            }
        }
    }

    fn serialize(&mut self) -> Result<()> {
        if let Some((next_msg, resolver)) = self.connection.next_frame() {
            // FIXME: having the body here is too verbose, impl Display on frame instead: trace!("will write to buffer: {:?}", next_msg);
            trace!("Will write to buffer");
            let checkpoint = self.send_buffer.checkpoint();
            let res = gen_frame(&next_msg)((&mut self.send_buffer).into());
            match res.map(|w| w.into_inner().1) {
                Ok(sz) => {
                    self.serialized_frames.push_back((sz, resolver));
                    Ok(())
                }
                Err(e) => {
                    self.send_buffer.rollback(checkpoint);
                    match e {
                        GenError::BufferTooSmall(_) => {
                            // Requeue msg
                            self.connection.requeue_frame((next_msg, resolver))?;
                            Ok(())
                        }
                        e => {
                            error!("error generating frame: {:?}", e);
                            let error = Error::SerialisationError(Arc::new(e));
                            self.connection.set_error(error.clone())?;
                            Err(error)
                        }
                    }
                }
            }
        } else {
            Ok(())
        }
    }

    fn parse(&mut self) -> Result<()> {
        if self.can_parse() {
            if let Some(frame) = self.do_parse()? {
                self.connection.handle_frame(frame)?;
            }
        }
        Ok(())
    }

    fn do_parse(&mut self) -> Result<Option<AMQPFrame>> {
        match parse_frame(self.receive_buffer.data()) {
            Ok((i, f)) => {
                let consumed = self.receive_buffer.offset(i);
                self.receive_buffer.consume(consumed, false);
                Ok(Some(f))
            }
            Err(e) => {
                if e.is_incomplete() {
                    self.receive_buffer.shift();
                    Ok(None)
                } else {
                    error!("parse error: {:?}", e);
                    let error = Error::ParsingError(e);
                    self.connection.set_error(error.clone())?;
                    Err(error)
                }
            }
        }
    }
}
<|endoftext|>


<commit_before>/*!
 * Compiled representation of a grammar. Simplified, normalized
 * version of `parse_tree`. The normalization passes produce this
 * representation incrementally.
 */

use intern::{InternedString};
use grammar::pattern::{Pattern, PatternKind};
use std::iter::once;
use std::fmt::{Debug, Display, Formatter, Error};
use util::{map, Map, Sep};

// These concepts we re-use wholesale
pub use grammar::parse_tree::{NonterminalString,
                              Span,
                              TerminalString, TypeParameter};

#[derive(Clone, Debug)]
pub struct Grammar {
    // a unique prefix that can be appended to identifiers to ensure
    // that they do not conflict with any action strings
    pub prefix: String,

    // these are the nonterminals that were declared to be public; the
    // key is the user's name for the symbol, the value is the
    // artificial symbol we introduce, which will always have a single
    // production like `Foo' = Foo`.
    pub start_nonterminals: Map<NonterminalString, NonterminalString>,

    // the "use foo;" statements that the user declared
    pub uses: Vec<String>,
    pub type_parameters: Vec<TypeParameter>,
    pub parameters: Vec<Parameter>,
    pub where_clauses: Vec<String>,

    // the grammar proper:

    pub action_fn_defns: Vec<ActionFnDefn>,
    pub productions: Map<NonterminalString, Vec<Production>>,
    pub token_span: Span,
    pub conversions: Map<TerminalString, Pattern<TypeRepr>>,
    pub types: Types,
}

#[derive(Clone, Debug, PartialEq, Eq)]
pub struct Parameter {
    pub name: InternedString,
    pub ty: TypeRepr,
}

#[derive(Clone, Hash, PartialEq, Eq, PartialOrd, Ord)]
pub struct Production {
    // this overlaps with the key in the hashmap, obviously, but it's
    // handy to have it
    pub nonterminal: NonterminalString,
    pub symbols: Vec<Symbol>,
    pub action_fn: ActionFn,
    pub span: Span,
}

#[derive(Copy, Clone, Hash, PartialEq, Eq, PartialOrd, Ord)]
pub enum Symbol {
    Nonterminal(NonterminalString),
    Terminal(TerminalString),
}

#[derive(Clone, PartialEq, Eq)]
pub struct ActionFnDefn {
    pub arg_patterns: Vec<InternedString>,
    pub arg_types: Vec<TypeRepr>,
    pub ret_type: TypeRepr,
    pub code: String,
}

#[derive(Clone, PartialEq, Eq)]
pub enum TypeRepr {
    Tuple(Vec<TypeRepr>),
    Nominal(NominalTypeRepr),
    Lifetime(InternedString),
}

#[derive(Clone, PartialEq, Eq)]
pub struct NominalTypeRepr {
    pub path: Vec<InternedString>,
    pub types: Vec<TypeRepr>
}

#[derive(Clone, Debug)]
pub struct Types {
    terminal_enum_type: NominalTypeRepr,
    default_terminal_type: TypeRepr,
    terminal_types: Map<TerminalString, TypeRepr>,
    nonterminal_types: Map<NonterminalString, TypeRepr>
}

impl Types {
    pub fn new(terminal_enum_type: NominalTypeRepr) -> Types {
        Types { terminal_enum_type: terminal_enum_type.clone(),
                terminal_types: map(),
                default_terminal_type: TypeRepr::Nominal(terminal_enum_type),
                nonterminal_types: map() }
    }

    pub fn add_type(&mut self, nt_id: NonterminalString, ty: TypeRepr) {
        assert!(self.nonterminal_types.insert(nt_id, ty).is_none());
    }

    pub fn add_term_type(&mut self, term: TerminalString, ty: TypeRepr) {
        assert!(self.terminal_types.insert(term, ty).is_none());
    }

    pub fn terminal_enum_type(&self) -> &NominalTypeRepr {
        &self.terminal_enum_type
    }

    pub fn terminal_type(&self, id: TerminalString) -> &TypeRepr {
        self.terminal_types.get(&id).unwrap_or(&self.default_terminal_type)
    }

    pub fn lookup_nonterminal_type(&self, id: NonterminalString) -> Option<&TypeRepr> {
        self.nonterminal_types.get(&id)
    }

    pub fn nonterminal_type(&self, id: NonterminalString) -> &TypeRepr {
        &self.nonterminal_types[&id]
    }
}

impl Display for Parameter {
    fn fmt(&self, fmt: &mut Formatter) -> Result<(), Error> {
        write!(fmt, "{}: {}", self.name, self.ty)
    }
}

impl Display for TypeRepr {
    fn fmt(&self, fmt: &mut Formatter) -> Result<(), Error> {
        match *self {
            TypeRepr::Tuple(ref types) =>
                write!(fmt, "({})", Sep(", ", types)),
            TypeRepr::Nominal(ref data) =>
                write!(fmt, "{}", data),
            TypeRepr::Lifetime(id) =>
                write!(fmt, "{}", id),
        }
    }
}

impl Debug for TypeRepr {
    fn fmt(&self, fmt: &mut Formatter) -> Result<(), Error> {
        Display::fmt(self, fmt)
    }
}

impl Display for NominalTypeRepr {
    fn fmt(&self, fmt: &mut Formatter) -> Result<(), Error> {
        if self.types.len() == 0 {
            write!(fmt, "{}", Sep("::", &self.path))
        } else {
            write!(fmt, "{}<{}>", Sep("::", &self.path), Sep(", ", &self.types))
        }
    }
}

impl Debug for NominalTypeRepr {
    fn fmt(&self, fmt: &mut Formatter) -> Result<(), Error> {
        Display::fmt(self, fmt)
    }
}

#[derive(Copy, Clone, Debug, Hash, PartialOrd, Ord, PartialEq, Eq)]
pub struct ActionFn(u32);

impl ActionFn {
    pub fn new(x: usize) -> ActionFn {
        ActionFn(x as u32)
    }

    pub fn index(&self) -> usize {
        self.0 as usize
    }
}

impl Symbol {
    pub fn ty<'ty>(&self, t: &'ty Types) -> &'ty TypeRepr {
        match *self {
            Symbol::Terminal(id) => t.terminal_type(id),
            Symbol::Nonterminal(id) => t.nonterminal_type(id),
        }
    }
}

impl Display for Symbol {
    fn fmt(&self, fmt: &mut Formatter) -> Result<(), Error> {
        match *self {
            Symbol::Nonterminal(id) => write!(fmt, "{}", id),
            Symbol::Terminal(id) => write!(fmt, "{}", id),
        }
    }
}

impl Debug for Symbol {
    fn fmt(&self, fmt: &mut Formatter) -> Result<(), Error> {
        Display::fmt(self, fmt)
    }
}

impl Debug for Production {
    fn fmt(&self, fmt: &mut Formatter) -> Result<(), Error> {
        write!(fmt,
               "{} = {} => {:?};",
               self.nonterminal, Sep(", ", &self.symbols), self.action_fn)
    }
}

impl Debug for ActionFnDefn {
    fn fmt(&self, fmt: &mut Formatter) -> Result<(), Error> {
        write!(fmt, "{}", self.to_fn_string("_"))
    }
}

impl ActionFnDefn {
    fn to_fn_string(&self, name: &str) -> String {
        let arg_strings: Vec<String> =
               self.arg_patterns
                   .iter()
                   .zip(self.arg_types.iter())
                   .map(|(p, t)| format!("{}: {}", p, t))
                   .collect();

        format!("fn {}({}) -> {} {{ {} }}",
                name, Sep(", ", &arg_strings), self.ret_type, self.code)
    }
}

impl Grammar {
    pub fn default_pattern(&self, id: InternedString) -> Pattern<TypeRepr> {
        let path: Vec<InternedString> =
            self.types.terminal_enum_type().path.iter()
                                                .cloned()
                                                .chain(once(id))
                                                .collect();
        Pattern {
            span: self.token_span,
            kind: PatternKind::Enum(path, vec![
                Pattern {
                    span: self.token_span,
                    kind: PatternKind::DotDot
                }
            ])
        }
    }

    pub fn pattern(&self, t: TerminalString) -> Pattern<TypeRepr> {
        match self.conversions.get(&t).cloned() {
            Some(p) => p,
            None => self.default_pattern(t.0),
        }
    }

    pub fn productions_for(&self, nonterminal: NonterminalString) -> &[Production] {
        match self.productions.get(&nonterminal) {
            Some(v) => &v[..],
            None => &[], // this...probably shouldn't happen actually?
        }
    }
}

<commit_msg>add some comments<commit_after>/*!
 * Compiled representation of a grammar. Simplified, normalized
 * version of `parse_tree`. The normalization passes produce this
 * representation incrementally.
 */

use intern::{InternedString};
use grammar::pattern::{Pattern, PatternKind};
use std::iter::once;
use std::fmt::{Debug, Display, Formatter, Error};
use util::{map, Map, Sep};

// These concepts we re-use wholesale
pub use grammar::parse_tree::{NonterminalString,
                              Span,
                              TerminalString, TypeParameter};

#[derive(Clone, Debug)]
pub struct Grammar {
    // a unique prefix that can be appended to identifiers to ensure
    // that they do not conflict with any action strings
    pub prefix: String,

    // these are the nonterminals that were declared to be public; the
    // key is the user's name for the symbol, the value is the
    // artificial symbol we introduce, which will always have a single
    // production like `Foo' = Foo`.
    pub start_nonterminals: Map<NonterminalString, NonterminalString>,

    // the "use foo;" statements that the user declared
    pub uses: Vec<String>,

    // type parameters declared on the grammar, like `grammar<T>;`
    pub type_parameters: Vec<TypeParameter>,

    // actual parameters declared on the grammar, like the `x: u32` in `grammar(x: u32);`
    pub parameters: Vec<Parameter>,

    // where clauses declared on the grammar, like `grammar<T> where T: Sized`
    pub where_clauses: Vec<String>,

    // the grammar proper:

    pub action_fn_defns: Vec<ActionFnDefn>,
    pub productions: Map<NonterminalString, Vec<Production>>,
    pub token_span: Span,
    pub conversions: Map<TerminalString, Pattern<TypeRepr>>,
    pub types: Types,
}

#[derive(Clone, Debug, PartialEq, Eq)]
pub struct Parameter {
    pub name: InternedString,
    pub ty: TypeRepr,
}

#[derive(Clone, Hash, PartialEq, Eq, PartialOrd, Ord)]
pub struct Production {
    // this overlaps with the key in the hashmap, obviously, but it's
    // handy to have it
    pub nonterminal: NonterminalString,
    pub symbols: Vec<Symbol>,
    pub action_fn: ActionFn,
    pub span: Span,
}

#[derive(Copy, Clone, Hash, PartialEq, Eq, PartialOrd, Ord)]
pub enum Symbol {
    Nonterminal(NonterminalString),
    Terminal(TerminalString),
}

#[derive(Clone, PartialEq, Eq)]
pub struct ActionFnDefn {
    pub arg_patterns: Vec<InternedString>,
    pub arg_types: Vec<TypeRepr>,
    pub ret_type: TypeRepr,
    pub code: String,
}

#[derive(Clone, PartialEq, Eq)]
pub enum TypeRepr {
    Tuple(Vec<TypeRepr>),
    Nominal(NominalTypeRepr),
    Lifetime(InternedString),
}

#[derive(Clone, PartialEq, Eq)]
pub struct NominalTypeRepr {
    pub path: Vec<InternedString>,
    pub types: Vec<TypeRepr>
}

#[derive(Clone, Debug)]
pub struct Types {
    terminal_enum_type: NominalTypeRepr,
    default_terminal_type: TypeRepr,
    terminal_types: Map<TerminalString, TypeRepr>,
    nonterminal_types: Map<NonterminalString, TypeRepr>
}

impl Types {
    pub fn new(terminal_enum_type: NominalTypeRepr) -> Types {
        Types { terminal_enum_type: terminal_enum_type.clone(),
                terminal_types: map(),
                default_terminal_type: TypeRepr::Nominal(terminal_enum_type),
                nonterminal_types: map() }
    }

    pub fn add_type(&mut self, nt_id: NonterminalString, ty: TypeRepr) {
        assert!(self.nonterminal_types.insert(nt_id, ty).is_none());
    }

    pub fn add_term_type(&mut self, term: TerminalString, ty: TypeRepr) {
        assert!(self.terminal_types.insert(term, ty).is_none());
    }

    pub fn terminal_enum_type(&self) -> &NominalTypeRepr {
        &self.terminal_enum_type
    }

    pub fn terminal_type(&self, id: TerminalString) -> &TypeRepr {
        self.terminal_types.get(&id).unwrap_or(&self.default_terminal_type)
    }

    pub fn lookup_nonterminal_type(&self, id: NonterminalString) -> Option<&TypeRepr> {
        self.nonterminal_types.get(&id)
    }

    pub fn nonterminal_type(&self, id: NonterminalString) -> &TypeRepr {
        &self.nonterminal_types[&id]
    }
}

impl Display for Parameter {
    fn fmt(&self, fmt: &mut Formatter) -> Result<(), Error> {
        write!(fmt, "{}: {}", self.name, self.ty)
    }
}

impl Display for TypeRepr {
    fn fmt(&self, fmt: &mut Formatter) -> Result<(), Error> {
        match *self {
            TypeRepr::Tuple(ref types) =>
                write!(fmt, "({})", Sep(", ", types)),
            TypeRepr::Nominal(ref data) =>
                write!(fmt, "{}", data),
            TypeRepr::Lifetime(id) =>
                write!(fmt, "{}", id),
        }
    }
}

impl Debug for TypeRepr {
    fn fmt(&self, fmt: &mut Formatter) -> Result<(), Error> {
        Display::fmt(self, fmt)
    }
}

impl Display for NominalTypeRepr {
    fn fmt(&self, fmt: &mut Formatter) -> Result<(), Error> {
        if self.types.len() == 0 {
            write!(fmt, "{}", Sep("::", &self.path))
        } else {
            write!(fmt, "{}<{}>", Sep("::", &self.path), Sep(", ", &self.types))
        }
    }
}

impl Debug for NominalTypeRepr {
    fn fmt(&self, fmt: &mut Formatter) -> Result<(), Error> {
        Display::fmt(self, fmt)
    }
}

#[derive(Copy, Clone, Debug, Hash, PartialOrd, Ord, PartialEq, Eq)]
pub struct ActionFn(u32);

impl ActionFn {
    pub fn new(x: usize) -> ActionFn {
        ActionFn(x as u32)
    }

    pub fn index(&self) -> usize {
        self.0 as usize
    }
}

impl Symbol {
    pub fn ty<'ty>(&self, t: &'ty Types) -> &'ty TypeRepr {
        match *self {
            Symbol::Terminal(id) => t.terminal_type(id),
            Symbol::Nonterminal(id) => t.nonterminal_type(id),
        }
    }
}

impl Display for Symbol {
    fn fmt(&self, fmt: &mut Formatter) -> Result<(), Error> {
        match *self {
            Symbol::Nonterminal(id) => write!(fmt, "{}", id),
            Symbol::Terminal(id) => write!(fmt, "{}", id),
        }
    }
}

impl Debug for Symbol {
    fn fmt(&self, fmt: &mut Formatter) -> Result<(), Error> {
        Display::fmt(self, fmt)
    }
}

impl Debug for Production {
    fn fmt(&self, fmt: &mut Formatter) -> Result<(), Error> {
        write!(fmt,
               "{} = {} => {:?};",
               self.nonterminal, Sep(", ", &self.symbols), self.action_fn)
    }
}

impl Debug for ActionFnDefn {
    fn fmt(&self, fmt: &mut Formatter) -> Result<(), Error> {
        write!(fmt, "{}", self.to_fn_string("_"))
    }
}

impl ActionFnDefn {
    fn to_fn_string(&self, name: &str) -> String {
        let arg_strings: Vec<String> =
               self.arg_patterns
                   .iter()
                   .zip(self.arg_types.iter())
                   .map(|(p, t)| format!("{}: {}", p, t))
                   .collect();

        format!("fn {}({}) -> {} {{ {} }}",
                name, Sep(", ", &arg_strings), self.ret_type, self.code)
    }
}

impl Grammar {
    pub fn default_pattern(&self, id: InternedString) -> Pattern<TypeRepr> {
        let path: Vec<InternedString> =
            self.types.terminal_enum_type().path.iter()
                                                .cloned()
                                                .chain(once(id))
                                                .collect();
        Pattern {
            span: self.token_span,
            kind: PatternKind::Enum(path, vec![
                Pattern {
                    span: self.token_span,
                    kind: PatternKind::DotDot
                }
            ])
        }
    }

    pub fn pattern(&self, t: TerminalString) -> Pattern<TypeRepr> {
        match self.conversions.get(&t).cloned() {
            Some(p) => p,
            None => self.default_pattern(t.0),
        }
    }

    pub fn productions_for(&self, nonterminal: NonterminalString) -> &[Production] {
        match self.productions.get(&nonterminal) {
            Some(v) => &v[..],
            None => &[], // this...probably shouldn't happen actually?
        }
    }
}

<|endoftext|>


<commit_before>//! HTTP Request handling.
//!
//! # Example
//! ```
//! use cabot::RequestBuilder;
//!
//! let request = RequestBuilder::new("http://localhost/")
//!     .set_http_method("POST")
//!     .add_header("Content-Type: application/json")
//!     .set_body_as_str("{}")
//!     .build()
//!     .unwrap();
//!     let attempt = "POST / HTTP/1.1\r\nContent-Type: \
//!                    application/json\r\nUser-Agent: cabot/0.2.1\r\n\
//!                    Host: localhost\r\nConnection: \
//!                    close\r\nContent-Length: 2\r\n\r\n{}";
//! assert_eq!(request.to_string(), attempt.to_string());
//! ```

use url::{self, Url};

use super::constants;
use super::results::{CabotError, CabotResult};

/// An HTTP Request representation.
///
/// Request is build using [RequestBuilder](../request/struct.RequestBuilder.html)
/// and them consume by the [Client](../client/struct.Client.html)
/// to perform the query.
#[derive(Default)]
pub struct Request {
    host: String,
    port: u16,
    authority: String,
    is_domain: bool,
    scheme: String,
    http_method: String,
    request_uri: String,
    http_version: String,
    headers: Vec<String>,
    body: Option<Vec<u8>>,
}

impl Request {
    fn new(
        host: String,
        port: u16,
        authority: String,
        is_domain: bool,
        scheme: String,
        http_method: String,
        request_uri: String,
        http_version: String,
        headers: Vec<String>,
        body: Option<Vec<u8>>,
    ) -> Request {
        Request {
            host,
            port,
            authority,
            is_domain,
            scheme,
            http_method,
            request_uri,
            http_version,
            headers,
            body,
        }
    }

    /// The HTTP verb used to perform the request,
    /// such as GET, POST, ...
    pub fn http_method(&self) -> &str {
        self.http_method.as_str()
    }

    /// The HTTP Body of the request.
    pub fn body(&self) -> Option<&[u8]> {
        match self.body {
            None => None,
            Some(ref body) => Some(body.as_slice()),
        }
    }

    /// Clone the body and retrieve it in a String object.
    ///
    /// Important: Currently assume the body is encoded in utf-8.
    ///
    /// Errors:
    ///
    ///  - CabotError::EncodingError in case the body is not an utf-8 string
    ///
    pub fn body_as_string(&self) -> CabotResult<Option<String>> {
        let body = match self.body {
            None => return Ok(None),
            Some(ref body) => {
                let mut body_vec: Vec<u8> = Vec::new();
                body_vec.extend_from_slice(body);
                String::from_utf8(body_vec)?
            }
        };
        Ok(Some(body))
    }

    /// The Version of the HTTP to perform the request.
    pub fn http_version(&self) -> &str {
        self.http_version.as_str()
    }

    /// The server name to connect. can be a name to resolve or an IP address.
    pub fn host(&self) -> &str {
        self.host.as_str()
    }

    /// The TCP server port to connect.
    pub fn port(&self) -> u16 {
        self.port
    }

    /// The authority part of the url (`host`:`port`).
    pub fn authority(&self) -> &str {
        self.authority.as_str()
    }

    /// The protocol scheme, can be http or https.
    pub fn scheme(&self) -> &str {
        self.scheme.as_str()
    }

    /// The URI to send, something like a PATH_INFO and a querystring.
    pub fn request_uri(&self) -> &str {
        self.request_uri.as_str()
    }
    /// The Bytes representation of the query to send to the server.
    pub fn to_bytes(&self) -> Vec<u8> {
        let mut resp = Vec::with_capacity(
            1024 + match self.body() {
                Some(bytes) => bytes.len(),
                None => 0,
            },
        );
        resp.extend_from_slice(
            format!(
                "{} {} {}\r\n",
                self.http_method(),
                self.request_uri(),
                self.http_version()
            )
            .as_bytes(),
        );

        for header in self.headers.as_slice() {
            resp.extend_from_slice(format!("{}\r\n", header).as_bytes());
        }
        if self.is_domain {
            resp.extend_from_slice(format!("Host: {}\r\n", self.host()).as_bytes());
        }
        resp.extend_from_slice(b"Connection: close\r\n");
        if let Some(payload) = self.body() {
            resp.extend_from_slice(format!("Content-Length: {}\r\n\r\n", payload.len()).as_bytes());
            resp.extend_from_slice(payload);
        } else {
            resp.extend_from_slice(b"\r\n");
        }
        resp
    }

    /// The String representation of the query.
    /// This method could panic in case the request is not utf-8 bytes.
    pub fn to_string(&self) -> String {
        let req = self.to_bytes();
        String::from_utf8(req).unwrap()
    }
}

/// Construct [Request](../request/struct.Request.html)
pub struct RequestBuilder {
    http_method: String,
    user_agent: String,
    url: Result<Url, url::ParseError>,
    http_version: String,
    headers: Vec<String>,
    body: Option<Vec<u8>>,
}

impl RequestBuilder {
    /// Create a new RequestBuilder with the given url.
    ///
    /// `url` will be parsed to get the host to contact and the uri to send.
    /// When building a request object after creating the builder,
    /// the http method will be `GET` with neither header nor body and
    /// the http version will be `HTTP/1.1`
    pub fn new(url: &str) -> Self {
        let url = url.parse::<Url>();
        RequestBuilder {
            http_method: "GET".to_owned(),
            url,
            user_agent: constants::user_agent(),
            http_version: "HTTP/1.1".to_owned(),
            headers: Vec::new(),
            body: None,
        }
    }

    /// Replace the url in case the RequestBuilder is used many time
    /// for multiple query.
    pub fn set_url(mut self, url: &str) -> Self {
        self.url = url.parse::<Url>();
        self
    }

    /// Set the http method such as `GET` `POST`. Default value is `GET`.
    pub fn set_http_method(mut self, http_method: &str) -> Self {
        self.http_method = http_method.to_owned();
        self
    }

    /// Set the protocol version to use.. Default value is `HTTP/1.1`.
    pub fn set_http_version(mut self, http_version: &str) -> Self {
        self.http_version = http_version.to_owned();
        self
    }

    /// Add a HTTP header.
    pub fn add_header(mut self, header: &str) -> Self {
        self.headers.push(header.to_owned());
        self
    }

    /// Add many headers.
    pub fn add_headers(mut self, headers: &[&str]) -> Self {
        for header in headers {
            self.headers.push(header.to_string());
        }
        self
    }

    /// Override the [default user-agent](../constants/index.html)
    ///
    /// Important: don't add a user_agent usering the `add_header` function
    ///            to avoid a duplicate header `User-Agent`
    pub fn set_user_agent(mut self, user_agent: &str) -> Self {
        self.user_agent = user_agent.to_owned();
        self
    }

    /// Set a response body.
    ///
    /// If a body is set, the `Content-Length` headers is added by cabot.
    pub fn set_body(mut self, buf: &[u8]) -> Self {
        let mut body = Vec::with_capacity(buf.len());
        body.extend_from_slice(buf);
        self.body = Some(body);
        self
    }

    /// Set a body to send in the query. By default a query has no body.
    pub fn set_body_as_str(self, body: &str) -> Self {
        self.set_body(body.as_bytes())
    }

    /// Construct the [Request](../request/struct.Request.html).
    /// To perform the query, a [Client](../client/struct.Client.html)
    /// has to be created.
    ///
    /// Errors:
    ///
    ///   - CabotError::ParseUrlError in case the `url` is not parsable
    ///   - CabotError::OpaqueUrlError in case the `url` is parsed but miss informations such as hostname.
    ///
    pub fn build(&self) -> CabotResult<Request> {
        let url = self.url.as_ref().map_err(|err| *err)?;

        let host = url.host_str().ok_or(CabotError::OpaqueUrlError(
            "Unable to find host".to_string(),
        ))?;

        let port = url
            .port_or_known_default()
            .ok_or(CabotError::OpaqueUrlError(
                "Unable to determine a port".to_string(),
            ))?;

        let query = url.query();
        let mut request_uri = url.path().to_owned();
        if let Some(querystring) = query {
            request_uri.push_str("?");
            request_uri.push_str(querystring);
        }
        let is_domain = url.domain().is_some();

        let mut headers = self.headers.clone();
        headers.push(format!("User-Agent: {}", self.user_agent));

        Ok(Request::new(
            host.to_owned(),
            port,
            format!("{}:{}", host, port),
            is_domain,
            url.scheme().to_owned(),
            self.http_method.clone(),
            request_uri,
            self.http_version.clone(),
            headers,
            match self.body {
                Some(ref body) => Some(body.clone()),
                None => None,
            },
        ))
    }
}

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

    #[test]
    fn test_get_request_to_string() {
        let request = Request::new(
            "127.0.0.1".to_owned(),
            80,
            "127.0.0.1:80".to_owned(),
            false,
            "http".to_owned(),
            "GET".to_owned(),
            "/path?query".to_owned(),
            "HTTP/1.1".to_owned(),
            Vec::new(),
            None,
        );
        let attempt = "GET /path?query HTTP/1.1\r\nConnection: close\r\n\r\n";
        assert_eq!(request.to_string(), attempt);
    }

    #[test]
    fn test_get_request_with_host_to_string() {
        let request = Request::new(
            "localhost".to_owned(),
            80,
            "localhost:80".to_owned(),
            true,
            "http".to_owned(),
            "GET".to_owned(),
            "/path?query".to_owned(),
            "HTTP/1.1".to_owned(),
            Vec::new(),
            None,
        );
        let attempt = "GET /path?query HTTP/1.1\r\nHost: localhost\r\nConnection: close\r\n\r\n";
        assert_eq!(request.to_string(), attempt);
    }

    #[test]
    fn test_get_request_with_headers_to_string() {
        let request = Request::new(
            "localhost".to_owned(),
            80,
            "localhost:80".to_owned(),
            true,
            "http".to_owned(),
            "GET".to_owned(),
            "/path?query".to_owned(),
            "HTTP/1.1".to_owned(),
            vec![
                "Accept-Language: fr".to_owned(),
                "Accept-Encoding: gzip".to_owned(),
            ],
            None,
        );
        let attempt = "GET /path?query HTTP/1.1\r\nAccept-Language: fr\r\nAccept-Encoding: \
                       gzip\r\nHost: localhost\r\nConnection: close\r\n\r\n";
        assert_eq!(request.to_string(), attempt);
    }

    #[test]
    fn test_post_request_with_headers_to_string() {
        let body: Vec<u8> = vec![123, 125];
        let request = Request::new(
            "localhost".to_owned(),
            80,
            "localhost:80".to_owned(),
            true,
            "http".to_owned(),
            "POST".to_owned(),
            "/".to_owned(),
            "HTTP/1.1".to_owned(),
            vec![
                "Accept-Language: fr".to_owned(),
                "Content-Type: application/json".to_owned(),
            ],
            Some(body),
        );
        let attempt = "POST / HTTP/1.1\r\nAccept-Language: fr\r\nContent-Type: \
                       application/json\r\nHost: localhost\r\nConnection: \
                       close\r\nContent-Length: 2\r\n\r\n{}";
        assert_eq!(request.to_string(), attempt);
    }

    #[test]
    fn test_request_builder_simple() {
        let request = RequestBuilder::new("http://localhost/").build().unwrap();
        assert_eq!(request.host(), "localhost".to_string());
        assert_eq!(request.scheme(), "http".to_string());
        assert_eq!(request.body, None);
        assert_eq!(request.http_method(), "GET".to_string());
        assert_eq!(request.http_version(), "HTTP/1.1".to_string());
        let headers: Vec<String> = vec![format!("User-Agent: {}", constants::user_agent())];
        assert_eq!(request.headers, headers);
    }

    #[test]
    fn test_request_builder_complete() {
        let builder = RequestBuilder::new("http://localhost/")
            .set_http_method("POST")
            .set_http_version("HTTP/1.0")
            .set_user_agent("anonymized")
            .add_header("Content-Type: application/json")
            .add_headers(&["Accept-Encoding: deflate", "Accept-Language: fr"])
            .set_body_as_str("{}");
        let body: &[u8] = &[123, 125];
        let request = builder.build().unwrap();
        assert_eq!(request.host(), "localhost".to_string());
        assert_eq!(request.body(), Some(body));
        assert_eq!(request.body_as_string().unwrap().unwrap(), "{}".to_string());
        assert_eq!(request.scheme(), "http".to_string());
        assert_eq!(request.http_method(), "POST".to_string());
        assert_eq!(request.request_uri(), "/");
        assert_eq!(request.http_version(), "HTTP/1.0".to_string());
        assert_eq!(
            request.headers,
            vec![
                "Content-Type: application/json".to_string(),
                "Accept-Encoding: deflate".to_string(),
                "Accept-Language: fr".to_string(),
                "User-Agent: anonymized".to_string(),
            ]
        );

        let builder = builder.set_url("http://[::1]/path");
        let request = builder.build().unwrap();
        assert_eq!(request.host(), "[::1]".to_string());
        assert_eq!(request.request_uri(), "/path");
        assert_eq!(request.body(), Some(body));
        assert_eq!(request.body_as_string().unwrap().unwrap(), "{}".to_string());
        assert_eq!(request.scheme(), "http".to_string());
        assert_eq!(request.http_method(), "POST".to_string());
        assert_eq!(request.http_version(), "HTTP/1.0".to_string());
        assert_eq!(
            request.headers,
            vec![
                "Content-Type: application/json".to_string(),
                "Accept-Encoding: deflate".to_string(),
                "Accept-Language: fr".to_string(),
                "User-Agent: anonymized".to_string(),
            ]
        );

        let builder = builder.set_url("not_an_url");
        let err = builder.build();
        assert!(err.is_err());
    }

}
<commit_msg>Fix UserAgent version<commit_after>//! HTTP Request handling.
//!
//! # Example
//! ```
//! use cabot::RequestBuilder;
//!
//! let request = RequestBuilder::new("http://localhost/")
//!     .set_http_method("POST")
//!     .add_header("Content-Type: application/json")
//!     .set_body_as_str("{}")
//!     .build()
//!     .unwrap();
//!     let attempt = "POST / HTTP/1.1\r\nContent-Type: \
//!                    application/json\r\nUser-Agent: cabot/0.3.0\r\n\
//!                    Host: localhost\r\nConnection: \
//!                    close\r\nContent-Length: 2\r\n\r\n{}";
//! assert_eq!(request.to_string(), attempt.to_string());
//! ```

use url::{self, Url};

use super::constants;
use super::results::{CabotError, CabotResult};

/// An HTTP Request representation.
///
/// Request is build using [RequestBuilder](../request/struct.RequestBuilder.html)
/// and them consume by the [Client](../client/struct.Client.html)
/// to perform the query.
#[derive(Default)]
pub struct Request {
    host: String,
    port: u16,
    authority: String,
    is_domain: bool,
    scheme: String,
    http_method: String,
    request_uri: String,
    http_version: String,
    headers: Vec<String>,
    body: Option<Vec<u8>>,
}

impl Request {
    fn new(
        host: String,
        port: u16,
        authority: String,
        is_domain: bool,
        scheme: String,
        http_method: String,
        request_uri: String,
        http_version: String,
        headers: Vec<String>,
        body: Option<Vec<u8>>,
    ) -> Request {
        Request {
            host,
            port,
            authority,
            is_domain,
            scheme,
            http_method,
            request_uri,
            http_version,
            headers,
            body,
        }
    }

    /// The HTTP verb used to perform the request,
    /// such as GET, POST, ...
    pub fn http_method(&self) -> &str {
        self.http_method.as_str()
    }

    /// The HTTP Body of the request.
    pub fn body(&self) -> Option<&[u8]> {
        match self.body {
            None => None,
            Some(ref body) => Some(body.as_slice()),
        }
    }

    /// Clone the body and retrieve it in a String object.
    ///
    /// Important: Currently assume the body is encoded in utf-8.
    ///
    /// Errors:
    ///
    ///  - CabotError::EncodingError in case the body is not an utf-8 string
    ///
    pub fn body_as_string(&self) -> CabotResult<Option<String>> {
        let body = match self.body {
            None => return Ok(None),
            Some(ref body) => {
                let mut body_vec: Vec<u8> = Vec::new();
                body_vec.extend_from_slice(body);
                String::from_utf8(body_vec)?
            }
        };
        Ok(Some(body))
    }

    /// The Version of the HTTP to perform the request.
    pub fn http_version(&self) -> &str {
        self.http_version.as_str()
    }

    /// The server name to connect. can be a name to resolve or an IP address.
    pub fn host(&self) -> &str {
        self.host.as_str()
    }

    /// The TCP server port to connect.
    pub fn port(&self) -> u16 {
        self.port
    }

    /// The authority part of the url (`host`:`port`).
    pub fn authority(&self) -> &str {
        self.authority.as_str()
    }

    /// The protocol scheme, can be http or https.
    pub fn scheme(&self) -> &str {
        self.scheme.as_str()
    }

    /// The URI to send, something like a PATH_INFO and a querystring.
    pub fn request_uri(&self) -> &str {
        self.request_uri.as_str()
    }
    /// The Bytes representation of the query to send to the server.
    pub fn to_bytes(&self) -> Vec<u8> {
        let mut resp = Vec::with_capacity(
            1024 + match self.body() {
                Some(bytes) => bytes.len(),
                None => 0,
            },
        );
        resp.extend_from_slice(
            format!(
                "{} {} {}\r\n",
                self.http_method(),
                self.request_uri(),
                self.http_version()
            )
            .as_bytes(),
        );

        for header in self.headers.as_slice() {
            resp.extend_from_slice(format!("{}\r\n", header).as_bytes());
        }
        if self.is_domain {
            resp.extend_from_slice(format!("Host: {}\r\n", self.host()).as_bytes());
        }
        resp.extend_from_slice(b"Connection: close\r\n");
        if let Some(payload) = self.body() {
            resp.extend_from_slice(format!("Content-Length: {}\r\n\r\n", payload.len()).as_bytes());
            resp.extend_from_slice(payload);
        } else {
            resp.extend_from_slice(b"\r\n");
        }
        resp
    }

    /// The String representation of the query.
    /// This method could panic in case the request is not utf-8 bytes.
    pub fn to_string(&self) -> String {
        let req = self.to_bytes();
        String::from_utf8(req).unwrap()
    }
}

/// Construct [Request](../request/struct.Request.html)
pub struct RequestBuilder {
    http_method: String,
    user_agent: String,
    url: Result<Url, url::ParseError>,
    http_version: String,
    headers: Vec<String>,
    body: Option<Vec<u8>>,
}

impl RequestBuilder {
    /// Create a new RequestBuilder with the given url.
    ///
    /// `url` will be parsed to get the host to contact and the uri to send.
    /// When building a request object after creating the builder,
    /// the http method will be `GET` with neither header nor body and
    /// the http version will be `HTTP/1.1`
    pub fn new(url: &str) -> Self {
        let url = url.parse::<Url>();
        RequestBuilder {
            http_method: "GET".to_owned(),
            url,
            user_agent: constants::user_agent(),
            http_version: "HTTP/1.1".to_owned(),
            headers: Vec::new(),
            body: None,
        }
    }

    /// Replace the url in case the RequestBuilder is used many time
    /// for multiple query.
    pub fn set_url(mut self, url: &str) -> Self {
        self.url = url.parse::<Url>();
        self
    }

    /// Set the http method such as `GET` `POST`. Default value is `GET`.
    pub fn set_http_method(mut self, http_method: &str) -> Self {
        self.http_method = http_method.to_owned();
        self
    }

    /// Set the protocol version to use.. Default value is `HTTP/1.1`.
    pub fn set_http_version(mut self, http_version: &str) -> Self {
        self.http_version = http_version.to_owned();
        self
    }

    /// Add a HTTP header.
    pub fn add_header(mut self, header: &str) -> Self {
        self.headers.push(header.to_owned());
        self
    }

    /// Add many headers.
    pub fn add_headers(mut self, headers: &[&str]) -> Self {
        for header in headers {
            self.headers.push(header.to_string());
        }
        self
    }

    /// Override the [default user-agent](../constants/index.html)
    ///
    /// Important: don't add a user_agent usering the `add_header` function
    ///            to avoid a duplicate header `User-Agent`
    pub fn set_user_agent(mut self, user_agent: &str) -> Self {
        self.user_agent = user_agent.to_owned();
        self
    }

    /// Set a response body.
    ///
    /// If a body is set, the `Content-Length` headers is added by cabot.
    pub fn set_body(mut self, buf: &[u8]) -> Self {
        let mut body = Vec::with_capacity(buf.len());
        body.extend_from_slice(buf);
        self.body = Some(body);
        self
    }

    /// Set a body to send in the query. By default a query has no body.
    pub fn set_body_as_str(self, body: &str) -> Self {
        self.set_body(body.as_bytes())
    }

    /// Construct the [Request](../request/struct.Request.html).
    /// To perform the query, a [Client](../client/struct.Client.html)
    /// has to be created.
    ///
    /// Errors:
    ///
    ///   - CabotError::ParseUrlError in case the `url` is not parsable
    ///   - CabotError::OpaqueUrlError in case the `url` is parsed but miss informations such as hostname.
    ///
    pub fn build(&self) -> CabotResult<Request> {
        let url = self.url.as_ref().map_err(|err| *err)?;

        let host = url.host_str().ok_or(CabotError::OpaqueUrlError(
            "Unable to find host".to_string(),
        ))?;

        let port = url
            .port_or_known_default()
            .ok_or(CabotError::OpaqueUrlError(
                "Unable to determine a port".to_string(),
            ))?;

        let query = url.query();
        let mut request_uri = url.path().to_owned();
        if let Some(querystring) = query {
            request_uri.push_str("?");
            request_uri.push_str(querystring);
        }
        let is_domain = url.domain().is_some();

        let mut headers = self.headers.clone();
        headers.push(format!("User-Agent: {}", self.user_agent));

        Ok(Request::new(
            host.to_owned(),
            port,
            format!("{}:{}", host, port),
            is_domain,
            url.scheme().to_owned(),
            self.http_method.clone(),
            request_uri,
            self.http_version.clone(),
            headers,
            match self.body {
                Some(ref body) => Some(body.clone()),
                None => None,
            },
        ))
    }
}

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

    #[test]
    fn test_get_request_to_string() {
        let request = Request::new(
            "127.0.0.1".to_owned(),
            80,
            "127.0.0.1:80".to_owned(),
            false,
            "http".to_owned(),
            "GET".to_owned(),
            "/path?query".to_owned(),
            "HTTP/1.1".to_owned(),
            Vec::new(),
            None,
        );
        let attempt = "GET /path?query HTTP/1.1\r\nConnection: close\r\n\r\n";
        assert_eq!(request.to_string(), attempt);
    }

    #[test]
    fn test_get_request_with_host_to_string() {
        let request = Request::new(
            "localhost".to_owned(),
            80,
            "localhost:80".to_owned(),
            true,
            "http".to_owned(),
            "GET".to_owned(),
            "/path?query".to_owned(),
            "HTTP/1.1".to_owned(),
            Vec::new(),
            None,
        );
        let attempt = "GET /path?query HTTP/1.1\r\nHost: localhost\r\nConnection: close\r\n\r\n";
        assert_eq!(request.to_string(), attempt);
    }

    #[test]
    fn test_get_request_with_headers_to_string() {
        let request = Request::new(
            "localhost".to_owned(),
            80,
            "localhost:80".to_owned(),
            true,
            "http".to_owned(),
            "GET".to_owned(),
            "/path?query".to_owned(),
            "HTTP/1.1".to_owned(),
            vec![
                "Accept-Language: fr".to_owned(),
                "Accept-Encoding: gzip".to_owned(),
            ],
            None,
        );
        let attempt = "GET /path?query HTTP/1.1\r\nAccept-Language: fr\r\nAccept-Encoding: \
                       gzip\r\nHost: localhost\r\nConnection: close\r\n\r\n";
        assert_eq!(request.to_string(), attempt);
    }

    #[test]
    fn test_post_request_with_headers_to_string() {
        let body: Vec<u8> = vec![123, 125];
        let request = Request::new(
            "localhost".to_owned(),
            80,
            "localhost:80".to_owned(),
            true,
            "http".to_owned(),
            "POST".to_owned(),
            "/".to_owned(),
            "HTTP/1.1".to_owned(),
            vec![
                "Accept-Language: fr".to_owned(),
                "Content-Type: application/json".to_owned(),
            ],
            Some(body),
        );
        let attempt = "POST / HTTP/1.1\r\nAccept-Language: fr\r\nContent-Type: \
                       application/json\r\nHost: localhost\r\nConnection: \
                       close\r\nContent-Length: 2\r\n\r\n{}";
        assert_eq!(request.to_string(), attempt);
    }

    #[test]
    fn test_request_builder_simple() {
        let request = RequestBuilder::new("http://localhost/").build().unwrap();
        assert_eq!(request.host(), "localhost".to_string());
        assert_eq!(request.scheme(), "http".to_string());
        assert_eq!(request.body, None);
        assert_eq!(request.http_method(), "GET".to_string());
        assert_eq!(request.http_version(), "HTTP/1.1".to_string());
        let headers: Vec<String> = vec![format!("User-Agent: {}", constants::user_agent())];
        assert_eq!(request.headers, headers);
    }

    #[test]
    fn test_request_builder_complete() {
        let builder = RequestBuilder::new("http://localhost/")
            .set_http_method("POST")
            .set_http_version("HTTP/1.0")
            .set_user_agent("anonymized")
            .add_header("Content-Type: application/json")
            .add_headers(&["Accept-Encoding: deflate", "Accept-Language: fr"])
            .set_body_as_str("{}");
        let body: &[u8] = &[123, 125];
        let request = builder.build().unwrap();
        assert_eq!(request.host(), "localhost".to_string());
        assert_eq!(request.body(), Some(body));
        assert_eq!(request.body_as_string().unwrap().unwrap(), "{}".to_string());
        assert_eq!(request.scheme(), "http".to_string());
        assert_eq!(request.http_method(), "POST".to_string());
        assert_eq!(request.request_uri(), "/");
        assert_eq!(request.http_version(), "HTTP/1.0".to_string());
        assert_eq!(
            request.headers,
            vec![
                "Content-Type: application/json".to_string(),
                "Accept-Encoding: deflate".to_string(),
                "Accept-Language: fr".to_string(),
                "User-Agent: anonymized".to_string(),
            ]
        );

        let builder = builder.set_url("http://[::1]/path");
        let request = builder.build().unwrap();
        assert_eq!(request.host(), "[::1]".to_string());
        assert_eq!(request.request_uri(), "/path");
        assert_eq!(request.body(), Some(body));
        assert_eq!(request.body_as_string().unwrap().unwrap(), "{}".to_string());
        assert_eq!(request.scheme(), "http".to_string());
        assert_eq!(request.http_method(), "POST".to_string());
        assert_eq!(request.http_version(), "HTTP/1.0".to_string());
        assert_eq!(
            request.headers,
            vec![
                "Content-Type: application/json".to_string(),
                "Accept-Encoding: deflate".to_string(),
                "Accept-Language: fr".to_string(),
                "User-Agent: anonymized".to_string(),
            ]
        );

        let builder = builder.set_url("not_an_url");
        let err = builder.build();
        assert!(err.is_err());
    }

}
<|endoftext|>