<commit_before>/*
Copyright (c) 2017, Robert Ou <rqou@robertou.com> and contributors
All rights reserved.

Redistribution and use in source and binary forms, with or without
modification, are permitted provided that the following conditions are met:

1. Redistributions of source code must retain the above copyright notice,
   this list of conditions and the following disclaimer.
2. Redistributions in binary form must reproduce the above copyright notice,
   this list of conditions and the following disclaimer in the documentation
   and/or other materials provided with the distribution.

THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
*/

use std::collections::HashMap;

#[macro_use]
extern crate serde_derive;
extern crate serde_json;

#[derive(Serialize, Deserialize, Debug)]
pub struct YosysNetlist {
    #[serde(default)]
    pub creator: String,
    #[serde(default)]
    pub modules: HashMap<String, YosysNetlistModule>,
}

#[derive(Serialize, Deserialize, Debug, Eq, PartialEq)]
pub enum YosysBitSpecial {
    #[serde(rename = "0")]
    _0,
    #[serde(rename = "1")]
    _1,
    #[serde(rename = "x")]
    X,
    #[serde(rename = "z")]
    Z,
}

#[derive(Serialize, Deserialize, Debug, Eq, PartialEq)]
pub enum YosysPortDirection {
    #[serde(rename = "input")]
    Input,
    #[serde(rename = "output")]
    Output,
    #[serde(rename = "inout")]
    InOut,
}

#[derive(Serialize, Deserialize, Debug, Eq, PartialEq)]
#[serde(untagged)]
pub enum YosysBit {
    N(usize),
    S(YosysBitSpecial)
}

#[derive(Serialize, Deserialize, Debug, Eq, PartialEq)]
#[serde(untagged)]
pub enum YosysAttribute {
    N(usize),
    S(String),
}

#[derive(Serialize, Deserialize, Debug)]
pub struct YosysNetlistModule {
    #[serde(default)]
    pub attributes: HashMap<String, YosysAttribute>,
    #[serde(default)]
    pub ports: HashMap<String, YosysNetlistPort>,
    #[serde(default)]
    pub cells: HashMap<String, YosysNetlistCell>,
    #[serde(default)]
    pub netnames: HashMap<String, YosysNetlistNetname>,
}

#[derive(Serialize, Deserialize, Debug)]
pub struct YosysNetlistPort {
    pub direction: YosysPortDirection,
    pub bits: Vec<YosysBit>,
}

#[derive(Serialize, Deserialize, Debug)]
pub struct YosysNetlistCell {
    #[serde(default)]
    pub hide_name: usize,
    #[serde(rename="type")]
    pub cell_type: String,
    #[serde(default)]
    pub parameters: HashMap<String, YosysAttribute>,
    #[serde(default)]
    pub attributes: HashMap<String, YosysAttribute>,
    #[serde(default)]
    pub port_directions: HashMap<String, String>,
    pub connections: HashMap<String, Vec<YosysBit>>,
}

#[derive(Serialize, Deserialize, Debug)]
pub struct YosysNetlistNetname {
    #[serde(default)]
    pub hide_name: usize,
    pub bits: Vec<YosysBit>,
    #[serde(default)]
    pub attributes: HashMap<String, YosysAttribute>,
}

pub fn read_yosys_netlist(input: &[u8]) -> Result<YosysNetlist, serde_json::Error> {
    serde_json::from_slice(input)
}

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

    #[test]
    fn super_empty_json() {
        let result = read_yosys_netlist(br#"
            {}"#).unwrap();
        assert_eq!(result.creator, "");
        assert_eq!(result.modules.len(), 0);
    }

    #[test]
    fn empty_json() {
        let result = read_yosys_netlist(br#"
            {
              "creator": "this is a test",
              "modules": {
              }
            }"#).unwrap();
        assert_eq!(result.creator, "this is a test");
        assert_eq!(result.modules.len(), 0);
    }

    #[test]
    fn empty_json_2() {
        let result = read_yosys_netlist(br#"
            {
              "modules": {
              }
            }"#).unwrap();
        assert_eq!(result.creator, "");
        assert_eq!(result.modules.len(), 0);
    }

    #[test]
    fn bit_values_test() {
        let result = read_yosys_netlist(br#"
            {
              "modules": {
                "mymodule": {
                  "cells": {
                    "mycell": {
                      "type": "celltype",
                      "connections": {
                        "IN": [ "x", 0, "z", 234, "1", "0" ]
                      }
                    }
                  }
                }
              }
            }"#).unwrap();
        assert_eq!(result.modules.get("mymodule").unwrap().cells.get("mycell").unwrap().connections.get("IN").unwrap(),
            &vec![YosysBit::S(YosysBitSpecial::X), YosysBit::N(0), YosysBit::S(YosysBitSpecial::Z), YosysBit::N(234),
                YosysBit::S(YosysBitSpecial::_1), YosysBit::S(YosysBitSpecial::_0)]);
    }
}
<commit_msg>yosys-netlist-json: Do some renaming and commenting<commit_after>/*
Copyright (c) 2017, Robert Ou <rqou@robertou.com> and contributors
All rights reserved.

Redistribution and use in source and binary forms, with or without
modification, are permitted provided that the following conditions are met:

1. Redistributions of source code must retain the above copyright notice,
   this list of conditions and the following disclaimer.
2. Redistributions in binary form must reproduce the above copyright notice,
   this list of conditions and the following disclaimer in the documentation
   and/or other materials provided with the distribution.

THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
*/

use std::collections::HashMap;

#[macro_use]
extern crate serde_derive;
extern crate serde_json;

/// Legal values for the direction of a port on a module
#[derive(Serialize, Deserialize, Debug, Eq, PartialEq)]
pub enum PortDirection {
    #[serde(rename = "input")]
    Input,
    #[serde(rename = "output")]
    Output,
    #[serde(rename = "inout")]
    InOut,
}

/// Special constant bit values
#[derive(Serialize, Deserialize, Debug, Eq, PartialEq)]
pub enum SpecialBit {
    /// Constant 0
    #[serde(rename = "0")]
    _0,
    /// Constant 1
    #[serde(rename = "1")]
    _1,
    /// Constant X (invalid)
    #[serde(rename = "x")]
    X,
    /// Constant Z (tri-state)
    #[serde(rename = "z")]
    Z,
}

/// A number representing a single bit of a wire
#[derive(Serialize, Deserialize, Debug, Eq, PartialEq)]
#[serde(untagged)]
pub enum BitVal {
    /// An actual signal number
    N(usize),
    /// A special constant value
    S(SpecialBit)
}

/// The value of an attribute/parameter
#[derive(Serialize, Deserialize, Debug, Eq, PartialEq)]
#[serde(untagged)]
pub enum AttributeVal {
    /// Numeric attribute value
    N(usize),
    /// String attribute value
    S(String),
}

/// Represents an entire .json file used by Yosys
#[derive(Serialize, Deserialize, Debug)]
pub struct Netlist {
    /// The program that created this file.
    #[serde(default)]
    pub creator: String,
    /// A map from module names to module objects contained in this .json file
    #[serde(default)]
    pub modules: HashMap<String, Module>,
}

/// Represents one module in the Yosys hierarchy
#[derive(Serialize, Deserialize, Debug)]
pub struct Module {
    /// Module attributes
    #[serde(default)]
    pub attributes: HashMap<String, AttributeVal>,
    /// Module ports (interfaces to other modules)
    #[serde(default)]
    pub ports: HashMap<String, Port>,
    /// Module cells (objects inside this module)
    #[serde(default)]
    pub cells: HashMap<String, Cell>,
    /// Module netnames (names of wires in this module)
    #[serde(default)]
    pub netnames: HashMap<String, Netname>,
}

/// Represents a port on a module
#[derive(Serialize, Deserialize, Debug)]
pub struct Port {
    /// Port direction
    pub direction: PortDirection,
    /// Bit value(s) representing the wire(s) connected to this port
    pub bits: Vec<BitVal>,
}

/// Represents a cell in a module
#[derive(Serialize, Deserialize, Debug)]
pub struct Cell {
    /// Indicates an internal/auto-generated name that starts with `$`
    #[serde(default)]
    pub hide_name: usize,
    /// Name of the type of this cell
    #[serde(rename="type")]
    pub cell_type: String,
    /// Parameters specified on this cell
    #[serde(default)]
    pub parameters: HashMap<String, AttributeVal>,
    /// Attributes specified on this cell
    #[serde(default)]
    pub attributes: HashMap<String, AttributeVal>,
    /// The direction of the ports on this cell
    #[serde(default)]
    pub port_directions: HashMap<String, String>,
    /// Bit value(s) representing the wire(s) connected to the inputs/outputs of this cell
    pub connections: HashMap<String, Vec<BitVal>>,
}

/// Represents the name of a net in a module
#[derive(Serialize, Deserialize, Debug)]
pub struct Netname {
    /// Indicates an internal/auto-generated name that starts with `$`
    #[serde(default)]
    pub hide_name: usize,
    /// Bit value(s) that should be given this name
    pub bits: Vec<BitVal>,
    /// Attributes for this netname
    #[serde(default)]
    pub attributes: HashMap<String, AttributeVal>,
}

pub fn read_yosys_netlist(input: &[u8]) -> Result<Netlist, serde_json::Error> {
    serde_json::from_slice(input)
}

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

    #[test]
    fn super_empty_json() {
        let result = read_yosys_netlist(br#"
            {}"#).unwrap();
        assert_eq!(result.creator, "");
        assert_eq!(result.modules.len(), 0);
    }

    #[test]
    fn empty_json() {
        let result = read_yosys_netlist(br#"
            {
              "creator": "this is a test",
              "modules": {
              }
            }"#).unwrap();
        assert_eq!(result.creator, "this is a test");
        assert_eq!(result.modules.len(), 0);
    }

    #[test]
    fn empty_json_2() {
        let result = read_yosys_netlist(br#"
            {
              "modules": {
              }
            }"#).unwrap();
        assert_eq!(result.creator, "");
        assert_eq!(result.modules.len(), 0);
    }

    #[test]
    fn bit_values_test() {
        let result = read_yosys_netlist(br#"
            {
              "modules": {
                "mymodule": {
                  "cells": {
                    "mycell": {
                      "type": "celltype",
                      "connections": {
                        "IN": [ "x", 0, "z", 234, "1", "0" ]
                      }
                    }
                  }
                }
              }
            }"#).unwrap();
        assert_eq!(result.modules.get("mymodule").unwrap().cells.get("mycell").unwrap().connections.get("IN").unwrap(),
            &vec![BitVal::S(SpecialBit::X), BitVal::N(0), BitVal::S(SpecialBit::Z), BitVal::N(234),
                BitVal::S(SpecialBit::_1), BitVal::S(SpecialBit::_0)]);
    }
}
<|endoftext|>


<commit_before>//! Low level and internal http and https implementation.

use std::collections::HashMap;
use std::io::{stderr, Read, Result as IoResult, Write};
use std::net::{SocketAddr, TcpStream};
use std::sync::Arc;
use std::time::Duration;

use log::Level::Info;
use rustls::{ClientConfig, ClientSession, ProtocolVersion, Session};
use webpki::DNSNameRef;
use webpki_roots;

use super::constants;
use super::dns::Resolver;
use super::request::Request;
use super::results::{CabotError, CabotResult};

const BUFFER_PAGE_SIZE: usize = 2048;
const RESPONSE_BUFFER_SIZE: usize = 1024;

#[derive(Debug, PartialEq)]
enum TransferEncoding {
    ContentLength(usize),
    Chunked,
    Unkown,
    None,
}

impl From<&[u8]> for TransferEncoding {
    fn from(hdr: &[u8]) -> Self {
        let hdr = String::from_utf8_lossy(hdr);
        let hdrup = hdr.to_ascii_uppercase();
        match hdrup.as_str() {
            "CHUNKED" => TransferEncoding::Chunked,
            _ => TransferEncoding::Unkown,
        }
    }
}

struct HttpDecoder<'a, T> {
    writer: &'a mut Write,
    reader: &'a mut T,
    buffer: Vec<u8>,
    transfer_encoding: TransferEncoding,
}

impl<'a, T> HttpDecoder<'a, T> {
    fn process_headers(&mut self) -> Option<usize> {
        let ret = {
            let resp_header: Vec<&[u8]> = constants::SPLIT_HEADERS_RE
                .splitn(self.buffer.as_slice(), 2)
                .collect();
            if resp_header.len() == 2 {
                // We have headers
                let headers = resp_header.get(0).unwrap();
                if let Some(header) = constants::TRANSFER_ENCODING.captures(headers) {
                    if let Some(tenc) = header.get(1) {
                        self.transfer_encoding = TransferEncoding::from(tenc.as_bytes());
                    }
                } else if let Some(header) = constants::CONTENT_LENGTH.captures(headers) {
                    if let Some(clength) = header.get(1) {
                        let clength = String::from_utf8_lossy(clength.as_bytes()).into_owned();
                        let clength = usize::from_str_radix(clength.as_str(), 10).unwrap();
                        self.transfer_encoding = TransferEncoding::ContentLength(clength);
                    }
                }
                let resp = headers.len();
                self.writer.write(headers).unwrap();
                self.writer.write(b"\r\n\r\n").unwrap();
                Some(resp + 4) // + CRLF CRLF
            } else {
                None
            }
        };
        if let Some(to_drain) = ret {
            let buffer = self.buffer.drain(to_drain..).collect();
            self.buffer = buffer;
            info!("Headers read");
        }
        ret
    }
}

impl<'a, T> HttpDecoder<'a, T>
where
    for<'b> T: Read + Sized,
{
    fn new(writer: &'a mut Write, reader: &'a mut T) -> Self {
        HttpDecoder {
            writer,
            reader,
            buffer: Vec::with_capacity(BUFFER_PAGE_SIZE),
            transfer_encoding: TransferEncoding::None,
        }
    }
    fn chunk_read(&mut self) -> IoResult<usize> {
        let mut buf = [0; BUFFER_PAGE_SIZE];
        let ret = self.reader.read(&mut buf[..]);
        if let Ok(count) = ret {
            if count > 0 {
                self.buffer.extend_from_slice(&buf[0..count]);
            }
        }
        ret
    }
    fn read_write_headers(&mut self) -> IoResult<()> {
        info!("Reading headers");
        loop {
            let count = self.chunk_read()?;
            info!("Count: {}", count);
            if count == 0 {
                break;
            }
            if let Some(_) = self.process_headers() {
                break;
            }
        }
        Ok(())
    }

    fn read_write_no_transfer_encoding(&mut self) -> IoResult<()> {
        loop {
            self.writer.write(self.buffer.as_slice()).unwrap();
            self.buffer.clear();

            let cnt = self.chunk_read()?;
            if cnt == 0 {
                break;
            }
        }

        self.writer.flush()?;
        Ok(())
    }

    fn read_content_length(&mut self, size: usize) -> IoResult<()> {
        let mut read_count = self.buffer.len();
        loop {
            self.writer.write(self.buffer.as_slice()).unwrap();
            self.buffer.clear();

            if read_count >= size {
                break;
            }
            read_count = read_count + self.chunk_read()?;
        }
        Ok(())
    }

    fn read_write_chunk(&mut self) -> IoResult<()> {
        // vars for chunk encoding only
        let mut read_chunk = true;
        let mut read_size = 0;
        let mut header_len: usize;

        'outer: loop {
            while read_chunk {
                header_len = 0;
                {
                    // we read the chunk size to drain
                    let header: Vec<&[u8]> = constants::SPLIT_HEADER_BRE
                        .splitn(self.buffer.as_slice(), 2)
                        .collect();
                    if header.len() == 2 {
                        if let Some(header) = constants::GET_CHUNK_SIZE.captures(header[0]) {
                            if let Some(size) = header.get(1) {
                                let size = size.as_bytes();
                                header_len = size.len() + 2;
                                let size = String::from_utf8_lossy(size).into_owned();
                                read_size = usize::from_str_radix(size.as_str(), 16).unwrap();
                                read_chunk = false;
                            }
                        } else {
                            // else return Error
                            break;
                        }
                    }
                }

                if read_size == 0 {
                    self.buffer.clear(); // should we check that it is '0\r\n' ?
                    break 'outer;
                }

                let buffer: Vec<u8> = self.buffer.drain(header_len..).collect();
                self.buffer = buffer;

                if !read_chunk && self.buffer.len() >= read_size + 2 {
                    let mut buffer: Vec<u8> = self.buffer.drain(read_size..).collect();
                    self.writer.write(self.buffer.as_slice()).unwrap();
                    self.writer.flush()?;

                    let buffer2 = buffer.drain(2..).collect(); // CRLF
                    self.buffer = buffer2;
                    read_chunk = true;
                    read_size = 0;
                }
            }

            let cnt = self.chunk_read()?;
            if cnt == 0 {
                break;
            }
        }

        Ok(())
    }

    fn read_write_body(&mut self) -> IoResult<()> {
        info!("Reading body");
        match self.transfer_encoding {
            TransferEncoding::ContentLength(size) => {
                self.read_content_length(size)?;
            }
            TransferEncoding::Chunked => {
                self.read_write_chunk()?;
            }
            _ => {
                warn!("Neither Content-Length, not chunk is response header");
                self.read_write_no_transfer_encoding()?;
            }
        }

        if self.buffer.len() > 0 {
            let b = String::from_utf8_lossy(self.buffer.as_slice());
            error!("Buffer not clear: {}", b);
        }

        Ok(())
    }

    fn read_write(&mut self) -> IoResult<()> {
        self.read_write_headers()?;
        self.read_write_body()
    }
}

fn log_request(request: &[u8], verbose: bool) {
    if !log_enabled!(Info) && !verbose {
        return;
    }
    let request: Vec<&[u8]> = constants::SPLIT_HEADERS_RE.splitn(request, 2).collect();
    let headers = String::from_utf8_lossy(&request[0]);
    let headers: Vec<&str> = constants::SPLIT_HEADER_RE.split(&headers).collect();
    let bodylen = if request.len() == 2 {
        let body = &request[1];
        body.len()
    } else {
        0
    };
    if log_enabled!(Info) {
        for header in headers {
            info!("> {}", header);
        }
        if bodylen > 0 {
            info!("> [{} bytes]", bodylen);
        }
        info!(">");
    } else if verbose {
        for header in headers {
            writeln!(&mut stderr(), "> {}", header).unwrap();
        }
        if bodylen > 0 {
            writeln!(&mut stderr(), "> [{} bytes]", bodylen).unwrap();
        }
        writeln!(&mut stderr(), ">").unwrap();
    }
}

fn read_buf<T>(client: &mut T, buf: &mut [u8]) -> Vec<u8>
where
    T: Read + Sized,
{
    let mut response: Vec<u8> = Vec::with_capacity(RESPONSE_BUFFER_SIZE);
    while let Ok(count) = client.read(&mut buf[..]) {
        if count > 0 {
            response.extend_from_slice(&buf[0..count]);
        } else {
            break;
        }
    }
    response
}

fn from_http(
    request: &Request,
    client: &mut TcpStream,
    out: &mut Write,
    verbose: bool,
) -> CabotResult<()> {
    let request_bytes = request.to_bytes();
    let raw_request = request_bytes.as_slice();
    log_request(&raw_request, verbose);

    debug!("Sending request...");
    client.write_all(&raw_request).unwrap();

    debug!("Reading response headers...");

    let mut http_decoder = HttpDecoder::new(out, client);
    http_decoder.read_write()?;
    Ok(())
}

fn from_https(
    request: &Request,
    mut client: &mut TcpStream,
    out: &mut Write,
    verbose: bool,
) -> CabotResult<()> {
    let request_bytes = request.to_bytes();
    let raw_request = request_bytes.as_slice();
    let mut response: Vec<u8> = Vec::with_capacity(RESPONSE_BUFFER_SIZE);
    let mut buf = [0; BUFFER_PAGE_SIZE];

    let mut config = ClientConfig::new();
    config
        .root_store
        .add_server_trust_anchors(&webpki_roots::TLS_SERVER_ROOTS);
    let rc_config = Arc::new(config);
    let host = DNSNameRef::try_from_ascii_str(request.host())
        .map_err(|_| CabotError::HostnameParseError(request.host().to_owned()))?;
    let mut tlsclient = ClientSession::new(&rc_config, host);
    let mut is_handshaking = true;

    loop {
        while tlsclient.wants_write() {
            let count = tlsclient.write_tls(&mut client).unwrap();
            debug!("Write {} TLS bytes", count);
        }
        if is_handshaking && !tlsclient.is_handshaking() {
            info!("Handshake complete");
            is_handshaking = false;
            let protocol = tlsclient.get_protocol_version();
            match protocol {
                Some(ProtocolVersion::SSLv2) => {
                    info!("Protocol SSL v2 negociated");
                }
                Some(ProtocolVersion::SSLv3) => {
                    info!("Protocol SSL v3 negociated");
                }
                Some(ProtocolVersion::TLSv1_0) => {
                    info!("Protocol TLS v1.0 negociated");
                }
                Some(ProtocolVersion::TLSv1_1) => {
                    info!("Protocol TLS v1.1 negociated");
                }
                Some(ProtocolVersion::TLSv1_2) => {
                    info!("Protocol TLS v1.2 negociated");
                }
                Some(ProtocolVersion::TLSv1_3) => {
                    info!("Protocol TLS v1.3 negociated");
                }
                Some(ProtocolVersion::Unknown(num)) => {
                    info!("Unknown TLS Protocol negociated: {}", num);
                }
                None => {
                    info!("No TLS Protocol negociated");
                }
            }
            log_request(&raw_request, verbose);
            tlsclient.write_all(&raw_request).unwrap();
            let count = tlsclient.write_tls(&mut client).unwrap();
            debug!("Write {} TLS bytes", count);
        }

        if tlsclient.wants_read() {
            let count = tlsclient.read_tls(&mut client)?;
            debug!("Read {} TLS bytes", count);
            if count == 0 {
                break;
            }

            tlsclient.process_new_packets()?;

            let mut part: Vec<u8> = read_buf(&mut tlsclient, &mut buf);
            response.append(&mut part);
        } else {
            break;
        }
    }

    let mut http_decoder = HttpDecoder::new(out, client);
    http_decoder.buffer = response;
    http_decoder.process_headers().unwrap();
    http_decoder.read_write_body()?;
    Ok(())
}

pub fn http_query(
    request: &Request,
    mut out: &mut Write,
    authorities: &HashMap<String, SocketAddr>,
    verbose: bool,
    ipv4: bool,
    ipv6: bool,
) -> CabotResult<()> {
    debug!(
        "HTTP Query {} {}",
        request.http_method(),
        request.request_uri()
    );

    let authority = request.authority();

    let addr = match authorities.get(authority) {
        Some(val) => {
            info!("Fetch authority {} using autorities map", authority);
            *val
        }
        None => {
            info!("Fetch authority {} using resolver", authority);
            let resolver = Resolver::new(verbose);
            resolver.get_addr(authority, ipv4, ipv6)?
        }
    };

    info!("Connecting to {}", addr);
    let mut client = TcpStream::connect(addr)?;

    client.set_read_timeout(Some(Duration::new(5, 0))).unwrap();

    match request.scheme() {
        "http" => from_http(request, &mut client, &mut out, verbose)?,
        "https" => from_https(request, &mut client, &mut out, verbose)?,
        _ => {
            return Err(CabotError::SchemeError(format!(
                "Unrecognized scheme {}",
                request.scheme()
            )))
        }
    };

    out.flush().unwrap();

    Ok(())
}
<commit_msg>Chunk request in https<commit_after>//! Low level and internal http and https implementation.

use std::collections::HashMap;
use std::io::{stderr, Read, Result as IoResult, Write};
use std::net::{SocketAddr, TcpStream};
use std::sync::Arc;
use std::time::Duration;

use log::Level::Info;
use rustls::{ClientConfig, ClientSession, ProtocolVersion, Session};
use webpki::DNSNameRef;
use webpki_roots;

use super::constants;
use super::dns::Resolver;
use super::request::Request;
use super::results::{CabotError, CabotResult};

const BUFFER_PAGE_SIZE: usize = 2048;
const RESPONSE_BUFFER_SIZE: usize = 1024;

#[derive(Debug, PartialEq)]
enum TransferEncoding {
    ContentLength(usize),
    Chunked,
    Unkown,
    None,
}

impl From<&[u8]> for TransferEncoding {
    fn from(hdr: &[u8]) -> Self {
        let hdr = String::from_utf8_lossy(hdr);
        let hdrup = hdr.to_ascii_uppercase();
        match hdrup.as_str() {
            "CHUNKED" => TransferEncoding::Chunked,
            _ => TransferEncoding::Unkown,
        }
    }
}

struct HttpDecoder<'a, T> {
    writer: &'a mut Write,
    reader: &'a mut T,
    tls_session: Option<&'a mut ClientSession>,
    buffer: Vec<u8>,
    transfer_encoding: TransferEncoding,
}

impl<'a, T> HttpDecoder<'a, T> {
    fn process_headers(&mut self) -> Option<usize> {
        let ret = {
            let resp_header: Vec<&[u8]> = constants::SPLIT_HEADERS_RE
                .splitn(self.buffer.as_slice(), 2)
                .collect();
            if resp_header.len() == 2 {
                // We have headers
                let headers = resp_header.get(0).unwrap();
                if let Some(header) = constants::TRANSFER_ENCODING.captures(headers) {
                    if let Some(tenc) = header.get(1) {
                        self.transfer_encoding = TransferEncoding::from(tenc.as_bytes());
                    }
                } else if let Some(header) = constants::CONTENT_LENGTH.captures(headers) {
                    if let Some(clength) = header.get(1) {
                        let clength = String::from_utf8_lossy(clength.as_bytes()).into_owned();
                        let clength = usize::from_str_radix(clength.as_str(), 10).unwrap();
                        self.transfer_encoding = TransferEncoding::ContentLength(clength);
                    }
                }
                let resp = headers.len();
                self.writer.write(headers).unwrap();
                self.writer.write(b"\r\n\r\n").unwrap();
                Some(resp + 4) // + CRLF CRLF
            } else {
                None
            }
        };
        if let Some(to_drain) = ret {
            let buffer = self.buffer.drain(to_drain..).collect();
            self.buffer = buffer;
            info!("Headers read");
        }
        ret
    }
    fn process_chunk(&mut self) -> IoResult<bool> {
        let mut read_chunk = true;
        let mut read_size = 0;
        let mut header_len: usize;
        while read_chunk {
            header_len = 0;
            {
                // we read the chunk size to drain
                let header: Vec<&[u8]> = constants::SPLIT_HEADER_BRE
                    .splitn(self.buffer.as_slice(), 2)
                    .collect();
                if header.len() == 2 {
                    if let Some(header) = constants::GET_CHUNK_SIZE.captures(header[0]) {
                        if let Some(size) = header.get(1) {
                            let size = size.as_bytes();
                            header_len = size.len() + 2;
                            let size = String::from_utf8_lossy(size).into_owned();
                            read_size = usize::from_str_radix(size.as_str(), 16).unwrap();
                            read_chunk = false;
                        }
                    } else {
                        // else return Error
                        break;
                    }
                }
            }

            if read_size == 0 {
                self.buffer.clear(); // should we check that it is '0\r\n' ?
                return Ok(true);
            }

            let buffer: Vec<u8> = self.buffer.drain(header_len..).collect();
            self.buffer = buffer;

            if !read_chunk && self.buffer.len() >= read_size + 2 {
                let mut buffer: Vec<u8> = self.buffer.drain(read_size..).collect();
                self.writer.write(self.buffer.as_slice()).unwrap();
                self.writer.flush()?;

                let buffer2 = buffer.drain(2..).collect(); // CRLF
                self.buffer = buffer2;
                read_chunk = true;
                read_size = 0;
            }
        }
        return Ok(false);
    }
}

impl<'a, T> HttpDecoder<'a, T>
where
    for<'b> T: Read + Write + Sized,
{
    fn new(writer: &'a mut Write, reader: &'a mut T) -> Self {
        HttpDecoder {
            writer,
            reader,
            tls_session: None,
            buffer: Vec::with_capacity(BUFFER_PAGE_SIZE),
            transfer_encoding: TransferEncoding::None,
        }
    }

    fn chunk_read(&mut self) -> IoResult<usize> {
        let mut buf = [0; BUFFER_PAGE_SIZE];
        let ret = self.reader.read(&mut buf[..]);
        if let Ok(count) = ret {
            if count > 0 {
                self.buffer.extend_from_slice(&buf[0..count]);
            }
        }
        ret
    }
    fn read_write_headers(&mut self) -> IoResult<()> {
        info!("Reading headers");
        loop {
            let count = self.chunk_read()?;
            info!("Count: {}", count);
            if count == 0 {
                break;
            }
            if let Some(_) = self.process_headers() {
                break;
            }
        }
        Ok(())
    }

    fn read_write_no_transfer_encoding(&mut self) -> IoResult<()> {
        loop {
            self.writer.write(self.buffer.as_slice()).unwrap();
            self.buffer.clear();

            let cnt = self.chunk_read()?;
            if cnt == 0 {
                break;
            }
        }

        self.writer.flush()?;
        Ok(())
    }

    fn read_content_length(&mut self, size: usize) -> IoResult<()> {
        let mut read_count = self.buffer.len();
        loop {
            self.writer.write(self.buffer.as_slice()).unwrap();
            self.buffer.clear();

            if read_count >= size {
                break;
            }
            read_count = read_count + self.chunk_read()?;
        }
        Ok(())
    }

    fn read_write_chunk(&mut self) -> IoResult<()> {
        loop {
            let done = self.process_chunk()?;
            if done {
                break;
            }
            let cnt = self.chunk_read()?;
            if cnt == 0 {
                break;
            }
        }

        Ok(())
    }

    fn read_write_body(&mut self) -> IoResult<()> {
        info!("Reading body");
        match self.transfer_encoding {
            TransferEncoding::ContentLength(size) => {
                self.read_content_length(size)?;
            }
            TransferEncoding::Chunked => {
                self.read_write_chunk()?;
            }
            _ => {
                warn!("Neither Content-Length, not chunk is response header");
                self.read_write_no_transfer_encoding()?;
            }
        }

        if self.buffer.len() > 0 {
            let b = String::from_utf8_lossy(self.buffer.as_slice());
            error!("Buffer not clear: {}", b);
        }

        Ok(())
    }

    fn read_write(&mut self) -> IoResult<()> {
        self.read_write_headers()?;
        self.read_write_body()
    }

    /// TLS

    fn new_with_tls(
        writer: &'a mut Write,
        reader: &'a mut T,
        tls_session: &'a mut ClientSession,
    ) -> Self {
        HttpDecoder {
            writer,
            reader,
            tls_session: Some(tls_session),
            buffer: Vec::with_capacity(BUFFER_PAGE_SIZE),
            transfer_encoding: TransferEncoding::None,
        }
    }

    fn chunk_read_tls(&mut self) -> IoResult<usize> {
        let mut read_len = 0;
        let mut read_len_tls = 0;
        let reader = &mut self.reader;
        let tls_session = self.tls_session.as_mut().unwrap();
        while tls_session.wants_write() {
            let count = tls_session.write_tls(reader).unwrap();
            debug!("Write {} TLS bytes", count);
        }
        while tls_session.wants_read() {
            let count = tls_session.read_tls(reader)?;
            read_len_tls = read_len_tls + count;
            debug!("Read {} TLS bytes", count);
            if count == 0 {
                break;
            }
        }
        tls_session.process_new_packets().unwrap();

        let mut buf = [0; BUFFER_PAGE_SIZE];
        let ret = tls_session.read(&mut buf[..]);
        if let Ok(count) = ret {
            read_len = read_len + count;
            if count > 0 {
                self.buffer.extend_from_slice(&buf[0..count]);
            }
        }
        Ok(read_len)
    }

    fn read_write_headers_tls(&mut self) -> IoResult<()> {
        info!("Reading headers");
        loop {
            let count = self.chunk_read_tls()?;
            info!("Count: {}", count);
            if count == 0 {
                break;
            }
            if let Some(_) = self.process_headers() {
                break;
            }
        }
        Ok(())
    }

    fn read_write_no_transfer_encoding_tls(&mut self) -> IoResult<()> {
        loop {
            self.writer.write(self.buffer.as_slice()).unwrap();
            self.buffer.clear();

            let cnt = self.chunk_read_tls()?;
            if cnt == 0 {
                break;
            }
        }

        self.writer.flush()?;
        Ok(())
    }

    fn read_content_length_tls(&mut self, size: usize) -> IoResult<()> {
        let mut read_count = self.buffer.len();
        loop {
            self.writer.write(self.buffer.as_slice()).unwrap();
            self.buffer.clear();

            if read_count >= size {
                break;
            }
            read_count = read_count + self.chunk_read_tls()?;
        }
        Ok(())
    }

    fn read_write_chunk_tls(&mut self) -> IoResult<()> {
        loop {
            let done = self.process_chunk()?;
            if done {
                break;
            }
            let cnt = self.chunk_read_tls()?;
            if cnt == 0 {
                break;
            }
        }

        Ok(())
    }

    fn read_write_body_tls(&mut self) -> IoResult<()> {
        info!("Reading body");
        match self.transfer_encoding {
            TransferEncoding::ContentLength(size) => {
                warn!("Using Content-Length to determinate the end of the query");
                self.read_content_length_tls(size)?;
            }
            TransferEncoding::Chunked => {
                self.read_write_chunk_tls()?;
            }
            _ => {
                warn!("Neither Content-Length, not chunk is response header");
                self.read_write_no_transfer_encoding_tls()?;
            }
        }

        if self.buffer.len() > 0 {
            let b = String::from_utf8_lossy(self.buffer.as_slice());
            error!("Buffer not clear: {}", b);
        }

        Ok(())
    }

    fn read_write_tls(&mut self) -> IoResult<()> {
        self.read_write_headers_tls()?;
        self.read_write_body_tls()
    }
}

fn log_request(request: &[u8], verbose: bool) {
    if !log_enabled!(Info) && !verbose {
        return;
    }
    let request: Vec<&[u8]> = constants::SPLIT_HEADERS_RE.splitn(request, 2).collect();
    let headers = String::from_utf8_lossy(&request[0]);
    let headers: Vec<&str> = constants::SPLIT_HEADER_RE.split(&headers).collect();
    let bodylen = if request.len() == 2 {
        let body = &request[1];
        body.len()
    } else {
        0
    };
    if log_enabled!(Info) {
        for header in headers {
            info!("> {}", header);
        }
        if bodylen > 0 {
            info!("> [{} bytes]", bodylen);
        }
        info!(">");
    } else if verbose {
        for header in headers {
            writeln!(&mut stderr(), "> {}", header).unwrap();
        }
        if bodylen > 0 {
            writeln!(&mut stderr(), "> [{} bytes]", bodylen).unwrap();
        }
        writeln!(&mut stderr(), ">").unwrap();
    }
}

fn read_buf<T>(client: &mut T, buf: &mut [u8]) -> Vec<u8>
where
    T: Read + Sized,
{
    let mut response: Vec<u8> = Vec::with_capacity(RESPONSE_BUFFER_SIZE);
    while let Ok(count) = client.read(&mut buf[..]) {
        if count > 0 {
            response.extend_from_slice(&buf[0..count]);
        } else {
            break;
        }
    }
    response
}

fn from_http(
    request: &Request,
    client: &mut TcpStream,
    out: &mut Write,
    verbose: bool,
) -> CabotResult<()> {
    let request_bytes = request.to_bytes();
    let raw_request = request_bytes.as_slice();
    log_request(&raw_request, verbose);

    debug!("Sending request...");
    client.write_all(&raw_request).unwrap();

    debug!("Reading response headers...");

    let mut http_decoder = HttpDecoder::new(out, client);
    http_decoder.read_write()?;
    Ok(())
}

fn from_https(
    request: &Request,
    mut client: &mut TcpStream,
    out: &mut Write,
    verbose: bool,
) -> CabotResult<()> {
    let request_bytes = request.to_bytes();
    let raw_request = request_bytes.as_slice();
    let mut response: Vec<u8> = Vec::with_capacity(RESPONSE_BUFFER_SIZE);
    let mut buf = [0; BUFFER_PAGE_SIZE];

    let mut config = ClientConfig::new();
    config
        .root_store
        .add_server_trust_anchors(&webpki_roots::TLS_SERVER_ROOTS);
    let rc_config = Arc::new(config);
    let host = DNSNameRef::try_from_ascii_str(request.host())
        .map_err(|_| CabotError::HostnameParseError(request.host().to_owned()))?;
    let mut tlsclient = ClientSession::new(&rc_config, host);
    let mut is_handshaking = true;
    while is_handshaking {
        while tlsclient.wants_write() {
            let count = tlsclient.write_tls(&mut client).unwrap();
            debug!("Write {} TLS bytes", count);
        }
        if tlsclient.wants_read() {
            let count = tlsclient.read_tls(&mut client)?;
            debug!("Read {} TLS bytes", count);
            tlsclient.process_new_packets()?;

            let mut part: Vec<u8> = read_buf(&mut tlsclient, &mut buf);
            response.append(&mut part);
        }
        if is_handshaking && !tlsclient.is_handshaking() {
            info!("Handshake complete");
            is_handshaking = false;
            let protocol = tlsclient.get_protocol_version();
            match protocol {
                Some(ProtocolVersion::SSLv2) => {
                    info!("Protocol SSL v2 negociated");
                }
                Some(ProtocolVersion::SSLv3) => {
                    info!("Protocol SSL v3 negociated");
                }
                Some(ProtocolVersion::TLSv1_0) => {
                    info!("Protocol TLS v1.0 negociated");
                }
                Some(ProtocolVersion::TLSv1_1) => {
                    info!("Protocol TLS v1.1 negociated");
                }
                Some(ProtocolVersion::TLSv1_2) => {
                    info!("Protocol TLS v1.2 negociated");
                }
                Some(ProtocolVersion::TLSv1_3) => {
                    info!("Protocol TLS v1.3 negociated");
                }
                Some(ProtocolVersion::Unknown(num)) => {
                    info!("Unknown TLS Protocol negociated: {}", num);
                }
                None => {
                    info!("No TLS Protocol negociated");
                }
            }
        }
    }
    log_request(&raw_request, verbose);
    tlsclient.write_all(&raw_request).unwrap();

    let mut http_decoder = HttpDecoder::new_with_tls(out, client, &mut tlsclient);
    http_decoder.read_write_tls()?;
    Ok(())
}

pub fn http_query(
    request: &Request,
    mut out: &mut Write,
    authorities: &HashMap<String, SocketAddr>,
    verbose: bool,
    ipv4: bool,
    ipv6: bool,
) -> CabotResult<()> {
    debug!(
        "HTTP Query {} {}",
        request.http_method(),
        request.request_uri()
    );

    let authority = request.authority();

    let addr = match authorities.get(authority) {
        Some(val) => {
            info!("Fetch authority {} using autorities map", authority);
            *val
        }
        None => {
            info!("Fetch authority {} using resolver", authority);
            let resolver = Resolver::new(verbose);
            resolver.get_addr(authority, ipv4, ipv6)?
        }
    };

    info!("Connecting to {}", addr);
    let mut client = TcpStream::connect(addr)?;

    client.set_read_timeout(Some(Duration::new(5, 0))).unwrap();

    match request.scheme() {
        "http" => from_http(request, &mut client, &mut out, verbose)?,
        "https" => from_https(request, &mut client, &mut out, verbose)?,
        _ => {
            return Err(CabotError::SchemeError(format!(
                "Unrecognized scheme {}",
                request.scheme()
            )))
        }
    };

    out.flush().unwrap();

    Ok(())
}
<|endoftext|>


<commit_before>//! Exposes the `Iron` type, the main entrance point of the
//! `Iron` library.

use std::net::{ToSocketAddrs, SocketAddr};
#[cfg(feature = "ssl")]
use std::path::PathBuf;

pub use hyper::server::Listening;
use hyper::server::Server;
use hyper::net::Fresh;

use request::HttpRequest;
use response::HttpResponse;

use error::HttpResult;

use {Request, Handler};
use status;

/// The primary entrance point to `Iron`, a `struct` to instantiate a new server.
///
/// `Iron` contains the `Handler` which takes a `Request` and produces a
/// `Response`.
pub struct Iron<H> {
    /// Iron contains a `Handler`, which it uses to create responses for client
    /// requests.
    pub handler: H,

    /// Once listening, the local address that this server is bound to.
    addr: Option<SocketAddr>,

    /// Once listening, the protocol used to serve content.
    protocol: Option<Protocol>
}

/// Protocol used to serve content. Future versions of Iron may add new protocols
/// to this enum. Thus you should not exhaustively match on its variants.
#[derive(Clone)]
pub enum Protocol {
    /// Plaintext HTTP/1
    Http,
    /// HTTP/1 over SSL/TLS
    #[cfg(feature = "ssl")]
    Https {
        /// Path to SSL certificate file
        certificate: PathBuf,
        /// Path to SSL private key file
        key: PathBuf
    }
}

impl Protocol {
    /// Return the name used for this protocol in a URI's scheme part.
    pub fn name(&self) -> &'static str {
        match *self {
            Protocol::Http => "http",
            #[cfg(feature = "ssl")]
            Protocol::Https { .. } => "https"
        }
    }
}

impl<H: Handler> Iron<H> {
    /// Kick off the server process using the HTTP protocol.
    ///
    /// Call this once to begin listening for requests on the server.
    /// This consumes the Iron instance, but does the listening on
    /// another task, so is not blocking.
    ///
    /// The thread returns a guard that will automatically join with the parent
    /// once it is dropped, blocking until this happens.
    ///
    /// Defaults to a threadpool of size `2 * num_cpus`.
    ///
    /// ## Panics
    ///
    /// Panics if the provided address does not parse. To avoid this
    /// call `to_socket_addrs` yourself and pass a parsed `SocketAddr`.
    pub fn http<A: ToSocketAddrs>(self, addr: A) -> HttpResult<Listening> {
        self.listen_with(addr, 2 * ::num_cpus::get(), Protocol::Http)
    }

    /// Kick off the server process using the HTTPS protocol.
    ///
    /// Call this once to begin listening for requests on the server.
    /// This consumes the Iron instance, but does the listening on
    /// another task, so is not blocking.
    ///
    /// The thread returns a guard that will automatically join with the parent
    /// once it is dropped, blocking until this happens.
    ///
    /// Defaults to a threadpool of size `2 * num_cpus`.
    ///
    /// ## Panics
    ///
    /// Panics if the provided address does not parse. To avoid this
    /// call `to_socket_addrs` yourself and pass a parsed `SocketAddr`.
    #[cfg(feature = "ssl")]
    pub fn https<A: ToSocketAddrs>(self, addr: A, certificate: PathBuf, key: PathBuf)
                                   -> HttpResult<Listening> {
        self.listen_with(addr, 2 * ::num_cpus::get(),
                         Protocol::Https { certificate: certificate, key: key })
    }

    /// Kick off the server process with X threads.
    ///
    /// ## Panics
    ///
    /// Panics if the provided address does not parse. To avoid this
    /// call `to_socket_addrs` yourself and pass a parsed `SocketAddr`.
    pub fn listen_with<A: ToSocketAddrs>(mut self, addr: A, threads: usize,
                                         protocol: Protocol) -> HttpResult<Listening> {
        let sock_addr = addr.to_socket_addrs()
            .ok().and_then(|mut addrs| addrs.next()).expect("Could not parse socket address.");

        self.addr = Some(sock_addr);

        self.protocol = Some(protocol.clone());

        match protocol {
            Protocol::Http => try!(Server::http(sock_addr)).handle_threads(self, threads),

            #[cfg(feature = "ssl")]
            Protocol::Https { ref certificate, ref key } => {
                use hyper::net::Openssl;

                try!(Server::https(sock_addr,
                                   try!(Openssl::with_cert_and_key(certificate, key))))
                    .handle_threads(self, threads)
            }
        }
    }

    /// Instantiate a new instance of `Iron`.
    ///
    /// This will create a new `Iron`, the base unit of the server, using the
    /// passed in `Handler`.
    pub fn new(handler: H) -> Iron<H> {
        Iron { handler: handler, addr: None, protocol: None }
    }

    fn bad_request(&self, mut http_res: HttpResponse<Fresh>) {
        *http_res.status_mut() = status::BadRequest;

        // Consume and flush the response.
        // We would like this to work, but can't do anything if it doesn't.
        if let Ok(res) = http_res.start()
        {
            let _ = res.end();
        }
    }
}

impl<H: Handler> ::hyper::server::Handler for Iron<H> {
    fn handle(&self, http_req: HttpRequest, mut http_res: HttpResponse<Fresh>) {
        // Set some defaults in case request handler panics.
        // This should not be necessary anymore once stdlib's catch_panic becomes stable.
        *http_res.status_mut() = status::InternalServerError;

        // Create `Request` wrapper.
        match Request::from_http(http_req, self.addr.clone().unwrap(),
                                 self.protocol.as_ref().unwrap()) {
            Ok(mut req) => {
                // Dispatch the request, write the response back to http_res
                self.handler.handle(&mut req).unwrap_or_else(|e| {
                    error!("Error handling:\n{:?}\nError was: {:?}", req, e.error);
                    e.response
                }).write_back(http_res)
            },
            Err(e) => {
                error!("Error creating request:\n    {}", e);
                self.bad_request(http_res)
            }
        }
    }
}
<commit_msg>Re-enable keep-alive by default, allow setting read, write, and keep-alive timeouts.<commit_after>//! Exposes the `Iron` type, the main entrance point of the
//! `Iron` library.

use std::net::{ToSocketAddrs, SocketAddr};
use std::time::Duration;
#[cfg(feature = "ssl")]
use std::path::PathBuf;

pub use hyper::server::Listening;
use hyper::server::Server;
use hyper::net::Fresh;

use request::HttpRequest;
use response::HttpResponse;

use error::HttpResult;

use {Request, Handler};
use status;

/// The primary entrance point to `Iron`, a `struct` to instantiate a new server.
///
/// `Iron` contains the `Handler` which takes a `Request` and produces a
/// `Response`.
pub struct Iron<H> {
    /// Iron contains a `Handler`, which it uses to create responses for client
    /// requests.
    pub handler: H,

    /// Once listening, the local address that this server is bound to.
    addr: Option<SocketAddr>,

    /// Once listening, the protocol used to serve content.
    protocol: Option<Protocol>
}

/// A settings struct containing a set of timeouts which can be applied to a server.
#[derive(Debug, PartialEq, Clone, Copy)]
pub struct Timeouts {
    /// Controls the timeout for keep alive connections.
    ///
    /// The default is `Some(Duration::from_secs(5))`.
    ///
    /// NOTE: Setting this to None will have the effect of turning off keep alive.
    pub keep_alive: Option<Duration>,

    /// Controls the timeout for reads on existing connections.
    ///
    /// The default is `Some(Duration::from_secs(30))`
    pub read: Option<Duration>,

    /// Controls the timeout for writes on existing conncetions.
    ///
    /// The default is `Some(Duration::from_secs(1))`
    pub write: Option<Duration>
}

impl Default for Timeouts {
    fn default() -> Self {
        Timeouts {
            keep_alive: Some(Duration::from_secs(5)),
            read: Some(Duration::from_secs(30)),
            write: Some(Duration::from_secs(1))
        }
    }
}

/// Protocol used to serve content. Future versions of Iron may add new protocols
/// to this enum. Thus you should not exhaustively match on its variants.
#[derive(Clone)]
pub enum Protocol {
    /// Plaintext HTTP/1
    Http,
    /// HTTP/1 over SSL/TLS
    #[cfg(feature = "ssl")]
    Https {
        /// Path to SSL certificate file
        certificate: PathBuf,
        /// Path to SSL private key file
        key: PathBuf
    }
}

impl Protocol {
    /// Return the name used for this protocol in a URI's scheme part.
    pub fn name(&self) -> &'static str {
        match *self {
            Protocol::Http => "http",
            #[cfg(feature = "ssl")]
            Protocol::Https { .. } => "https"
        }
    }
}

impl<H: Handler> Iron<H> {
    /// Kick off the server process using the HTTP protocol.
    ///
    /// Call this once to begin listening for requests on the server.
    /// This consumes the Iron instance, but does the listening on
    /// another task, so is not blocking.
    ///
    /// The thread returns a guard that will automatically join with the parent
    /// once it is dropped, blocking until this happens.
    ///
    /// Defaults to a threadpool of size `2 * num_cpus`.
    ///
    /// ## Panics
    ///
    /// Panics if the provided address does not parse. To avoid this
    /// call `to_socket_addrs` yourself and pass a parsed `SocketAddr`.
    pub fn http<A: ToSocketAddrs>(self, addr: A) -> HttpResult<Listening> {
        self.listen_with(addr, 8 * ::num_cpus::get(), Protocol::Http, None)
    }

    /// Kick off the server process using the HTTPS protocol.
    ///
    /// Call this once to begin listening for requests on the server.
    /// This consumes the Iron instance, but does the listening on
    /// another task, so is not blocking.
    ///
    /// The thread returns a guard that will automatically join with the parent
    /// once it is dropped, blocking until this happens.
    ///
    /// Defaults to a threadpool of size `2 * num_cpus`.
    ///
    /// ## Panics
    ///
    /// Panics if the provided address does not parse. To avoid this
    /// call `to_socket_addrs` yourself and pass a parsed `SocketAddr`.
    #[cfg(feature = "ssl")]
    pub fn https<A: ToSocketAddrs>(self, addr: A, certificate: PathBuf, key: PathBuf)
                                   -> HttpResult<Listening> {
        self.listen_with(addr, 8 * ::num_cpus::get(),
                         Protocol::Https { certificate: certificate, key: key }, None)
    }

    /// Kick off the server process with X threads.
    ///
    /// ## Panics
    ///
    /// Panics if the provided address does not parse. To avoid this
    /// call `to_socket_addrs` yourself and pass a parsed `SocketAddr`.
    pub fn listen_with<A: ToSocketAddrs>(mut self, addr: A, threads: usize,
                                         protocol: Protocol,
                                         timeouts: Option<Timeouts>) -> HttpResult<Listening> {
        let sock_addr = addr.to_socket_addrs()
            .ok().and_then(|mut addrs| addrs.next()).expect("Could not parse socket address.");

        self.addr = Some(sock_addr);
        self.protocol = Some(protocol.clone());

        match protocol {
            Protocol::Http => {
                let mut server = try!(Server::http(sock_addr));
                let timeouts = timeouts.unwrap_or_default();
                server.keep_alive(timeouts.keep_alive);
                server.set_read_timeout(timeouts.read);
                server.set_write_timeout(timeouts.write);
                server.handle_threads(self, threads)
            },

            #[cfg(feature = "ssl")]
            Protocol::Https { ref certificate, ref key } => {
                use hyper::net::Openssl;

                let ssl = try!(Openssl::with_cert_and_key(certificate, key));
                let mut server = try!(Server::https(sock_addr, ssl));
                let timeouts = timeouts.unwrap_or_default();
                server.keep_alive(timeouts.keep_alive);
                server.set_read_timeout(timeouts.read);
                server.set_write_timeout(timeouts.write);
                server.handle_threads(self, threads)
            }
        }
    }

    /// Instantiate a new instance of `Iron`.
    ///
    /// This will create a new `Iron`, the base unit of the server, using the
    /// passed in `Handler`.
    pub fn new(handler: H) -> Iron<H> {
        Iron { handler: handler, addr: None, protocol: None }
    }
}

impl<H: Handler> ::hyper::server::Handler for Iron<H> {
    fn handle(&self, http_req: HttpRequest, mut http_res: HttpResponse<Fresh>) {
        // Set some defaults in case request handler panics.
        // This should not be necessary anymore once stdlib's catch_panic becomes stable.
        *http_res.status_mut() = status::InternalServerError;

        // Create `Request` wrapper.
        match Request::from_http(http_req, self.addr.clone().unwrap(),
                                 self.protocol.as_ref().unwrap()) {
            Ok(mut req) => {
                // Dispatch the request, write the response back to http_res
                self.handler.handle(&mut req).unwrap_or_else(|e| {
                    error!("Error handling:\n{:?}\nError was: {:?}", req, e.error);
                    e.response
                }).write_back(http_res)
            },
            Err(e) => {
                error!("Error creating request:\n    {}", e);
                bad_request(http_res)
            }
        }
    }
}

fn bad_request(mut http_res: HttpResponse<Fresh>) {
    *http_res.status_mut() = status::BadRequest;

    // Consume and flush the response.
    // We would like this to work, but can't do anything if it doesn't.
    if let Ok(res) = http_res.start()
    {
        let _ = res.end();
    }
}

<|endoftext|>


<commit_before>//! Exposes the `Iron` type, the main entrance point of the
//! `Iron` library.

use std::io::net::ip::{SocketAddr, IpAddr};
use std::sync::Arc;

use http::server as http;
use super::{Request, Handler};
use super::status;

use super::response::HttpResponse;
use super::request::HttpRequest;

/// The primary entrance point to `Iron`, a `struct` to instantiate a new server.
///
/// `Iron` contains the `Handler` which takes a `Request` and produces a
/// `Response`.
pub struct Iron<H> {
    /// Iron contains a Handler, which it uses to create responses for client
    /// requests.
    pub handler: H,
}

// The struct which actually listens and serves requests.
struct IronListener<H> {
    handler: Arc<H>,
    ip: IpAddr,
    port: u16
}

impl<H: Send + Sync> Clone for IronListener<H> {
    fn clone(&self) -> IronListener<H> {
        IronListener {
            // Just increment a reference count.
            handler: self.handler.clone(),
            ip: self.ip.clone(),
            port: self.port.clone()
        }
    }
}

impl<H: Handler> Iron<H> {
    /// Kick off the server process.
    ///
    /// Call this once to begin listening for requests on the server.
    /// This consumes the Iron instance, but does the listening on
    /// another task, so is not blocking.
    pub fn listen(self, ip: IpAddr, port: u16) {
        use http::server::Server;

        spawn(proc() {
            IronListener {
                handler: Arc::new(self.handler),
                ip: ip,
                port: port
            }.serve_forever();
        });
    }

    /// Instantiate a new instance of `Iron`.
    ///
    /// This will create a new `Iron`, the base unit of the server, using the
    /// passed in Handler.
    pub fn new(handler: H) -> Iron<H> {
        Iron { handler: handler }
    }
}

impl<H: Handler> http::Server for IronListener<H> {
    fn get_config(&self) -> http::Config {
        http::Config {
            bind_address: SocketAddr {
                ip: self.ip,
                port: self.port
            }
        }
    }

    fn handle_request(&self, http_req: HttpRequest, http_res: &mut HttpResponse) {
        // Create wrapper Request
        let mut req = match Request::from_http(http_req) {
            Ok(req) => req,
            Err(e) => {
                error!("Error getting request: {}", e);
                http_res.status = status::InternalServerError;
                let _ = http_res.write(b"Internal Server Error");
                return;
            }
        };

        // Dispatch the request
        let res = self.handler.call(&mut req);

        match res {
            // Write the response back to http_res
            Ok(res) => res.write_back(http_res),
            Err(e) => {
                error!("Error handling {}: {}", req, e);
                http_res.status = status::InternalServerError;
                let _ = http_res.write(b"Internal Server Error");
            }
        }
    }
}
<commit_msg>(doc) Small doc fixes in iron.<commit_after>//! Exposes the `Iron` type, the main entrance point of the
//! `Iron` library.

use std::io::net::ip::{SocketAddr, IpAddr};
use std::sync::Arc;

use http::server as http;
use super::{Request, Handler};
use super::status;

use super::response::HttpResponse;
use super::request::HttpRequest;

/// The primary entrance point to `Iron`, a `struct` to instantiate a new server.
///
/// `Iron` contains the `Handler` which takes a `Request` and produces a
/// `Response`.
pub struct Iron<H> {
    /// Iron contains a `Handler`, which it uses to create responses for client
    /// requests.
    pub handler: H,
}

// The struct which actually listens and serves requests.
struct IronListener<H> {
    handler: Arc<H>,
    ip: IpAddr,
    port: u16
}

impl<H: Send + Sync> Clone for IronListener<H> {
    fn clone(&self) -> IronListener<H> {
        IronListener {
            // Just increment a reference count.
            handler: self.handler.clone(),
            ip: self.ip.clone(),
            port: self.port.clone()
        }
    }
}

impl<H: Handler> Iron<H> {
    /// Kick off the server process.
    ///
    /// Call this once to begin listening for requests on the server.
    /// This consumes the Iron instance, but does the listening on
    /// another task, so is not blocking.
    pub fn listen(self, ip: IpAddr, port: u16) {
        use http::server::Server;

        spawn(proc() {
            IronListener {
                handler: Arc::new(self.handler),
                ip: ip,
                port: port
            }.serve_forever();
        });
    }

    /// Instantiate a new instance of `Iron`.
    ///
    /// This will create a new `Iron`, the base unit of the server, using the
    /// passed in `Handler`.
    pub fn new(handler: H) -> Iron<H> {
        Iron { handler: handler }
    }
}

impl<H: Handler> http::Server for IronListener<H> {
    fn get_config(&self) -> http::Config {
        http::Config {
            bind_address: SocketAddr {
                ip: self.ip,
                port: self.port
            }
        }
    }

    fn handle_request(&self, http_req: HttpRequest, http_res: &mut HttpResponse) {
        // Create `Request` wrapper.
        let mut req = match Request::from_http(http_req) {
            Ok(req) => req,
            Err(e) => {
                error!("Error getting request: {}", e);
                http_res.status = status::InternalServerError;
                let _ = http_res.write(b"Internal Server Error");
                return;
            }
        };

        // Dispatch the request
        let res = self.handler.call(&mut req);

        match res {
            // Write the response back to http_res
            Ok(res) => res.write_back(http_res),
            Err(e) => {
                error!("Error handling {}: {}", req, e);
                http_res.status = status::InternalServerError;
                let _ = http_res.write(b"Internal Server Error");
            }
        }
    }
}
<|endoftext|>


<commit_before>#![feature(non_ascii_idents)]

extern crate num;

extern crate special_fun;
extern crate roots;

use special_fun::FloatSpecial;
use std::f64::consts::PI;
use roots::SimpleConvergency;
use roots::find_root_brent;

fn jinc(r: f64) -> f64 {
    let π2r = 2.0f64*PI*r;
    2.0f64*π2r.besselj(1.0) / π2r
}

fn find_roots(a: f64, b: f64, f: &Fn(f64)->f64) -> Vec<f64> {
    Vec::new()
}

fn main() {
    let v = 0.5f64.besselj(1.0);
    println!("besselj(1.0, 0.5) : {}", v);
    println!("jinc(0.61) : {}", jinc(0.609835));

}
<commit_msg>compute `jinc` roots.<commit_after>#![feature(non_ascii_idents)]

extern crate num;

extern crate special_fun;
extern crate roots;

use special_fun::FloatSpecial;
use std::f64::consts::PI;
use roots::SimpleConvergency;
use roots::find_root_brent;

fn jinc(r: f64) -> f64 {
    let π2r = 2.0f64*PI*r;
    2.0f64*π2r.besselj(1.0) / π2r
}

fn find_roots(a: f64, b: f64, f: &Fn(f64)->f64) -> Vec<f64> {
    let mut roots = Vec::new();
    const RANGES : i32 = 100;
    let dx = (b - a) / RANGES as f64;
    let conv = SimpleConvergency{eps:1e-12f64, max_iter:80};
    for r in 0..RANGES {
        let beg = a + (r as f64)*dx;
        let maybe_root = find_root_brent(beg, beg+dx, &f, &conv).ok();
        match maybe_root {
            Some(root) => roots.push(root),
            None => {}
        }
    }

    roots
}

fn main() {
    let v = 0.5f64.besselj(1.0);
    println!("besselj(1.0, 0.5) : {}", v);
    println!("jinc(0.61) : {}", jinc(0.609835));

    let jinc_roots = find_roots(0f64, 5f64, &jinc);
    println!("jinc roots : {:?}", jinc_roots);
}
<|endoftext|>


<commit_before>use todo::Todo;

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct TodoList {
    pub name: String,
    pub list: Vec<Todo>,
}

impl TodoList {
    pub fn new(name: &str) -> TodoList {
        TodoList {
            name: name.to_string(),
            list: Vec::new(),
        }
    }
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Tdo {
    pub lists: Vec<TodoList>,
}

impl Tdo {
    pub fn new() -> Tdo {
        Tdo {
            lists: Vec::new(),
        }
    }
}
<commit_msg>add save and load function to Tdo<commit_after>use todo::Todo;
use std::fs::File;
use storage::StorageError;

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct TodoList {
    pub name: String,
    pub list: Vec<Todo>,
}

impl TodoList {
    pub fn new(name: &str) -> TodoList {
        TodoList {
            name: name.to_string(),
            list: Vec::new(),
        }
    }
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Tdo {
    pub lists: Vec<TodoList>,
}

impl Tdo {
    pub fn new() -> Tdo {
        Tdo { lists: Vec::new() }
    }

    pub fn load(path: &str) -> Tdo {
        let file = File::open(path).unwrap();
        super::serde_json::from_reader(&file).unwrap()
    }

    pub fn save(path: &str, tdo: Tdo) -> Result<(), StorageError> {
        match File::create(path) {
            Ok(mut f) => {
                let _ = super::serde_json::to_writer_pretty(&mut f, &tdo);
                Ok(())
            },
            Err(_) => Err(StorageError::FileCorrupted)
        }

    }
}
<|endoftext|>


<commit_before>extern crate lancelot;
extern crate log;

use std::env;
use std::process;

use log::{error, info, trace};
use better_panic;
use failure::{Error, Fail};

use lancelot::workspace::Workspace;

#[derive(Debug, Fail)]
pub enum MainError {
    #[fail(display = "foo")]
    Foo,
}

pub struct Config {
    pub filename: String,
}

impl Config {
    pub fn from_args(args: env::Args) -> Result<Config, &'static str> {
        let args: Vec<String> = args.collect();

        if args.len() < 2 {
            return Err("not enough arguments: provide `lancelot.exe /path/to/input`");
        }

        let filename = args[1].clone();
        trace!("config: parsed filename: {:?}", filename);

        Ok(Config { filename })
    }
}

pub fn setup_logging(_args: &Config) {
    simplelog::TermLogger::init(simplelog::LevelFilter::Info, simplelog::Config::default())
        .expect("failed to setup logging");
}

pub fn run(args: &Config) -> Result<(), Error> {
    info!("filename: {:?}", args.filename);

    let ws = Workspace::from_file(&args.filename)?.load()?;

    let mut functions = ws.get_functions().collect::<Vec<_>>();
    functions.sort();

    info!("found {} functions", functions.len());
    for rva in functions.iter() {
        println!("{:#x}", ws.va(**rva).unwrap());
    }

    Ok(())
}

fn main() {
    better_panic::install();

    let args = Config::from_args(env::args()).unwrap_or_else(|err| {
        eprintln!("error parsing arguments: {}", err);
        process::exit(1);
    });

    setup_logging(&args);

    if let Err(e) = run(&args) {
        error!("{:?}", e)
    }
}
<commit_msg>main: emit bb count for each found insn<commit_after>extern crate lancelot;
extern crate log;

use std::env;
use std::process;

use log::{error, info, trace};
use better_panic;
use failure::{Error, Fail};

use lancelot::workspace::Workspace;

#[derive(Debug, Fail)]
pub enum MainError {
    #[fail(display = "foo")]
    Foo,
}

pub struct Config {
    pub filename: String,
}

impl Config {
    pub fn from_args(args: env::Args) -> Result<Config, &'static str> {
        let args: Vec<String> = args.collect();

        if args.len() < 2 {
            return Err("not enough arguments: provide `lancelot.exe /path/to/input`");
        }

        let filename = args[1].clone();
        trace!("config: parsed filename: {:?}", filename);

        Ok(Config { filename })
    }
}

pub fn setup_logging(_args: &Config) {
    simplelog::TermLogger::init(simplelog::LevelFilter::Info, simplelog::Config::default())
        .expect("failed to setup logging");
}

pub fn run(args: &Config) -> Result<(), Error> {
    info!("filename: {:?}", args.filename);

    let ws = Workspace::from_file(&args.filename)?.load()?;

    let mut functions = ws.get_functions().collect::<Vec<_>>();
    functions.sort();

    info!("found {} functions", functions.len());
    for rva in functions.iter() {
        let basic_blocks = ws.get_basic_blocks(**rva)?;
        println!("{:#x} found {} basic blocks", **rva, basic_blocks.len());
    }

    Ok(())
}

fn main() {
    better_panic::install();

    let args = Config::from_args(env::args()).unwrap_or_else(|err| {
        eprintln!("error parsing arguments: {}", err);
        process::exit(1);
    });

    setup_logging(&args);

    if let Err(e) = run(&args) {
        error!("{:?}", e)
    }
}
<|endoftext|>


<commit_before>use std::env;
use std::fs::File;

fn _normalize_col(col: &str) -> String {
    lazy_static::lazy_static! {
        static ref RE: regex::Regex = regex::Regex::new(r"\(.*?\)").unwrap();
    }
    RE.replace_all(col, "")
        .to_lowercase()
        .trim()
        .replace(" ", "_")
        .replace(".", "_")
        .replace("-", "_")
        .replace("/", "_")
        .replace("?", "")
        .replace(",", "_")
        .replace("&", "_")
}

fn _create_table(db: &mut rusqlite::Connection, table_name: &str, cols: &[String]) {
    let create_columns = cols
        .iter()
        .map(|c| format!("{} varchar", c))
        .collect::<Vec<String>>()
        .join(", ");
    db.execute(
        &format!("CREATE TABLE {} ({})", table_name, create_columns),
        &[] as &[&dyn rusqlite::types::ToSql],
    )
    .unwrap();
}

fn _load_table_from_path(
    db: &mut rusqlite::Connection,
    table_name: &str,
    path: String,
) -> Result<Vec<String>, Box<dyn std::error::Error>> {
    let mut num_rows = 0;
    let f = File::open(path)?;
    let file_size = f.metadata().unwrap().len();
    let mut reader = csv::Reader::from_reader(f);

    let normalized_cols =
        reader
            .headers()?
            .iter()
            .map(_normalize_col)
            .fold(vec![], |mut v, orig_col| {
                let mut col = orig_col.clone();
                let mut i = 1;
                while v.contains(&col) {
                    col = format!("{}_{}", orig_col, i);
                    i += 1
                }
                v.push(col);
                v
            });
    _create_table(db, table_name, &normalized_cols);

    let insert_query = format!(
        "INSERT INTO {} VALUES ({})",
        table_name,
        normalized_cols
            .iter()
            .map(|_| "?")
            .collect::<Vec<_>>()
            .join(", ")
    );
    let pb = indicatif::ProgressBar::new(file_size);
    pb.set_style(
        indicatif::ProgressStyle::default_bar()
            .template("[{elapsed_precise}] [{wide_bar:.cyan/blue}] {bytes}/{total_bytes} ({eta})")
            .progress_chars("#>-"),
    );
    let mut records = reader.records();
    let tx = db.transaction().unwrap();
    {
        let mut stmt = tx.prepare(&insert_query).expect("tx.prepare() failed");
        while let Some(row) = records.next() {
            stmt.execute(&row?).unwrap();

            num_rows += 1;
            if num_rows % 10000 == 0 {
                pb.set_position(records.reader().position().byte())
            }
        }
    }
    tx.commit().unwrap();
    pb.finish();

    println!(
        "Loaded {} rows into {}({})",
        num_rows,
        table_name,
        normalized_cols.join(", "),
    );
    Ok(normalized_cols)
}

struct FromAnySqlType {
    value: String,
}

impl rusqlite::types::FromSql for FromAnySqlType {
    fn column_result(
        value: rusqlite::types::ValueRef<'_>,
    ) -> Result<FromAnySqlType, rusqlite::types::FromSqlError> {
        let result = match value {
            rusqlite::types::ValueRef::Null => "null".to_string(),
            rusqlite::types::ValueRef::Integer(v) => v.to_string(),
            rusqlite::types::ValueRef::Real(v) => v.to_string(),
            rusqlite::types::ValueRef::Blob(v) | rusqlite::types::ValueRef::Text(v) => {
                String::from_utf8(v.to_vec()).unwrap()
            }
        };
        Ok(FromAnySqlType { value: result })
    }
}

fn _prepare_query<'a>(
    conn: &'a mut rusqlite::Connection,
    query: &str,
) -> Result<rusqlite::Statement<'a>, String> {
    conn.prepare(&query).map_err(|e| e.to_string())
}

fn _handle_query(conn: &mut rusqlite::Connection, line: &str) -> Result<(), String> {
    let mut stmt = _prepare_query(conn, line)?;

    let mut table = comfy_table::Table::new();
    table.load_preset("││──╞═╪╡┆    ┬┴┌┐└┘");
    table.set_content_arrangement(comfy_table::ContentArrangement::Dynamic);
    let mut title_row = comfy_table::Row::new();
    for col in stmt.column_names() {
        title_row.add_cell(comfy_table::Cell::new(col));
    }
    table.set_header(title_row);

    let mut results = stmt.query(&[] as &[&dyn rusqlite::types::ToSql]).unwrap();
    while let Ok(Some(r)) = results.next() {
        let mut row = comfy_table::Row::new();
        for i in 0..r.column_count() {
            let cell: FromAnySqlType = r.get(i).unwrap();
            row.add_cell(comfy_table::Cell::new(&cell.value));
        }
        table.add_row(row);
    }
    println!("{}", table);
    Ok(())
}

fn _handle_export(conn: &mut rusqlite::Connection, line: &str) -> Result<(), String> {
    lazy_static::lazy_static! {
        static ref RE: regex::Regex = regex::Regex::new(r"^\.export\(([\w_\-\./]+)\) (.*)").unwrap();
    }
    let caps = RE
        .captures(line)
        .ok_or_else(|| "Must match `.export(file-name) SQL`".to_owned())?;
    let destination_path = &caps[1];
    let query = &caps[2];

    let mut stmt = _prepare_query(conn, query)?;

    let mut writer = csv::Writer::from_path(destination_path).unwrap();
    writer.write_record(stmt.column_names()).unwrap();

    let mut results = stmt.query(&[] as &[&dyn rusqlite::types::ToSql]).unwrap();
    while let Ok(Some(r)) = results.next() {
        writer
            .write_record((0..r.column_count()).map(|i| {
                let cell: FromAnySqlType = r.get(i).unwrap();
                cell.value
            }))
            .unwrap();
    }

    Ok(())
}

fn _process_query(conn: &mut rusqlite::Connection, line: &str) {
    let result = if line.starts_with(".export") {
        _handle_export(conn, line)
    } else {
        _handle_query(conn, line)
    };
    if let Err(e) = result {
        println!("{}", e);
    }
}

struct SimpleWordCompleter {
    words: Vec<String>,
}

static BREAK_CHARS: [u8; 4] = [b' ', b'(', b')', b','];
impl SimpleWordCompleter {
    fn new(words: Vec<String>) -> SimpleWordCompleter {
        SimpleWordCompleter { words }
    }
}

impl rustyline::Helper for SimpleWordCompleter {}

impl rustyline::hint::Hinter for SimpleWordCompleter {
    fn hint(&self, _line: &str, _pos: usize, _ctx: &rustyline::Context<'_>) -> Option<String> {
        None
    }
}

impl rustyline::highlight::Highlighter for SimpleWordCompleter {}

impl rustyline::validate::Validator for SimpleWordCompleter {}

impl rustyline::completion::Completer for SimpleWordCompleter {
    type Candidate = String;

    fn complete(
        &self,
        line: &str,
        pos: usize,
        _ctx: &rustyline::Context<'_>,
    ) -> rustyline::Result<(usize, Vec<String>)> {
        let (start, word) = rustyline::completion::extract_word(line, pos, None, &BREAK_CHARS);

        let matches = self
            .words
            .iter()
            .filter(|w| w.starts_with(word))
            .cloned()
            .collect();
        Ok((start, matches))
    }
}

fn main() -> Result<(), Box<dyn std::error::Error>> {
    let mut paths = env::args().skip(1);

    let mut conn = rusqlite::Connection::open_in_memory().unwrap();

    let mut base_words = [
        "distinct", "select", "from", "group", "by", "order", "where", "count", "limit", "offset",
        ".export",
    ]
    .iter()
    .map(|&s| s.to_string())
    .collect::<Vec<String>>();

    if paths.len() == 1 {
        let mut col_names = _load_table_from_path(&mut conn, "t", paths.next().unwrap())?;
        base_words.append(&mut col_names);
    } else {
        for (idx, path) in paths.enumerate() {
            let mut col_names = _load_table_from_path(&mut conn, &format!("t{}", idx + 1), path)?;
            base_words.append(&mut col_names);
        }
    }

    let completer = SimpleWordCompleter::new(base_words);
    let mut rl = rustyline::Editor::new();
    rl.set_helper(Some(completer));
    loop {
        match rl.readline("> ") {
            Ok(line) => {
                if line.trim().is_empty() {
                    continue;
                }
                _process_query(&mut conn, &line);
                rl.add_history_entry(line);
            }
            Err(rustyline::error::ReadlineError::Interrupted) => {
                println!("Interrupted");
                continue;
            }
            Err(rustyline::error::ReadlineError::Eof) => {
                break;
            }
            Err(err) => {
                println!("Error: {}", err);
                break;
            }
        }
    }

    Ok(())
}
<commit_msg>Fixed running with all numeric columns<commit_after>use std::env;
use std::fs::File;

fn normalize_col(col: &str) -> String {
    lazy_static::lazy_static! {
        static ref RE: regex::Regex = regex::Regex::new(r"\(.*?\)").unwrap();
    }
    let mut col = RE
        .replace_all(col, "")
        .to_lowercase()
        .trim()
        .replace(" ", "_")
        .replace(".", "_")
        .replace("-", "_")
        .replace("/", "_")
        .replace("?", "")
        .replace(",", "_")
        .replace("&", "_");
    if !col.chars().next().unwrap().is_alphabetic() {
        col = format!("c_{}", col)
    }
    col
}

fn _create_table(db: &mut rusqlite::Connection, table_name: &str, cols: &[String]) {
    let create_columns = cols
        .iter()
        .map(|c| format!("{} varchar", c))
        .collect::<Vec<String>>()
        .join(", ");
    db.execute(
        &format!("CREATE TABLE {} ({})", table_name, create_columns),
        &[] as &[&dyn rusqlite::types::ToSql],
    )
    .unwrap();
}

fn _load_table_from_path(
    db: &mut rusqlite::Connection,
    table_name: &str,
    path: String,
) -> Result<Vec<String>, Box<dyn std::error::Error>> {
    let mut num_rows = 0;
    let f = File::open(path)?;
    let file_size = f.metadata().unwrap().len();
    let mut reader = csv::Reader::from_reader(f);

    let normalized_cols =
        reader
            .headers()?
            .iter()
            .map(normalize_col)
            .fold(vec![], |mut v, orig_col| {
                let mut col = orig_col.clone();
                let mut i = 1;
                while v.contains(&col) {
                    col = format!("{}_{}", orig_col, i);
                    i += 1
                }
                v.push(col);
                v
            });
    _create_table(db, table_name, &normalized_cols);

    let insert_query = format!(
        "INSERT INTO {} VALUES ({})",
        table_name,
        normalized_cols
            .iter()
            .map(|_| "?")
            .collect::<Vec<_>>()
            .join(", ")
    );
    let pb = indicatif::ProgressBar::new(file_size);
    pb.set_style(
        indicatif::ProgressStyle::default_bar()
            .template("[{elapsed_precise}] [{wide_bar:.cyan/blue}] {bytes}/{total_bytes} ({eta})")
            .progress_chars("#>-"),
    );
    let mut records = reader.records();
    let tx = db.transaction().unwrap();
    {
        let mut stmt = tx.prepare(&insert_query).expect("tx.prepare() failed");
        while let Some(row) = records.next() {
            stmt.execute(&row?).unwrap();

            num_rows += 1;
            if num_rows % 10000 == 0 {
                pb.set_position(records.reader().position().byte())
            }
        }
    }
    tx.commit().unwrap();
    pb.finish();

    println!(
        "Loaded {} rows into {}({})",
        num_rows,
        table_name,
        normalized_cols.join(", "),
    );
    Ok(normalized_cols)
}

struct FromAnySqlType {
    value: String,
}

impl rusqlite::types::FromSql for FromAnySqlType {
    fn column_result(
        value: rusqlite::types::ValueRef<'_>,
    ) -> Result<FromAnySqlType, rusqlite::types::FromSqlError> {
        let result = match value {
            rusqlite::types::ValueRef::Null => "null".to_string(),
            rusqlite::types::ValueRef::Integer(v) => v.to_string(),
            rusqlite::types::ValueRef::Real(v) => v.to_string(),
            rusqlite::types::ValueRef::Blob(v) | rusqlite::types::ValueRef::Text(v) => {
                String::from_utf8(v.to_vec()).unwrap()
            }
        };
        Ok(FromAnySqlType { value: result })
    }
}

fn _prepare_query<'a>(
    conn: &'a mut rusqlite::Connection,
    query: &str,
) -> Result<rusqlite::Statement<'a>, String> {
    conn.prepare(&query).map_err(|e| e.to_string())
}

fn _handle_query(conn: &mut rusqlite::Connection, line: &str) -> Result<(), String> {
    let mut stmt = _prepare_query(conn, line)?;

    let mut table = comfy_table::Table::new();
    table.load_preset("││──╞═╪╡┆    ┬┴┌┐└┘");
    table.set_content_arrangement(comfy_table::ContentArrangement::Dynamic);
    let mut title_row = comfy_table::Row::new();
    for col in stmt.column_names() {
        title_row.add_cell(comfy_table::Cell::new(col));
    }
    table.set_header(title_row);

    let mut results = stmt.query(&[] as &[&dyn rusqlite::types::ToSql]).unwrap();
    while let Ok(Some(r)) = results.next() {
        let mut row = comfy_table::Row::new();
        for i in 0..r.column_count() {
            let cell: FromAnySqlType = r.get(i).unwrap();
            row.add_cell(comfy_table::Cell::new(&cell.value));
        }
        table.add_row(row);
    }
    println!("{}", table);
    Ok(())
}

fn _handle_export(conn: &mut rusqlite::Connection, line: &str) -> Result<(), String> {
    lazy_static::lazy_static! {
        static ref RE: regex::Regex = regex::Regex::new(r"^\.export\(([\w_\-\./]+)\) (.*)").unwrap();
    }
    let caps = RE
        .captures(line)
        .ok_or_else(|| "Must match `.export(file-name) SQL`".to_owned())?;
    let destination_path = &caps[1];
    let query = &caps[2];

    let mut stmt = _prepare_query(conn, query)?;

    let mut writer = csv::Writer::from_path(destination_path).unwrap();
    writer.write_record(stmt.column_names()).unwrap();

    let mut results = stmt.query(&[] as &[&dyn rusqlite::types::ToSql]).unwrap();
    while let Ok(Some(r)) = results.next() {
        writer
            .write_record((0..r.column_count()).map(|i| {
                let cell: FromAnySqlType = r.get(i).unwrap();
                cell.value
            }))
            .unwrap();
    }

    Ok(())
}

fn _process_query(conn: &mut rusqlite::Connection, line: &str) {
    let result = if line.starts_with(".export") {
        _handle_export(conn, line)
    } else {
        _handle_query(conn, line)
    };
    if let Err(e) = result {
        println!("{}", e);
    }
}

struct SimpleWordCompleter {
    words: Vec<String>,
}

static BREAK_CHARS: [u8; 4] = [b' ', b'(', b')', b','];
impl SimpleWordCompleter {
    fn new(words: Vec<String>) -> SimpleWordCompleter {
        SimpleWordCompleter { words }
    }
}

impl rustyline::Helper for SimpleWordCompleter {}

impl rustyline::hint::Hinter for SimpleWordCompleter {
    fn hint(&self, _line: &str, _pos: usize, _ctx: &rustyline::Context<'_>) -> Option<String> {
        None
    }
}

impl rustyline::highlight::Highlighter for SimpleWordCompleter {}

impl rustyline::validate::Validator for SimpleWordCompleter {}

impl rustyline::completion::Completer for SimpleWordCompleter {
    type Candidate = String;

    fn complete(
        &self,
        line: &str,
        pos: usize,
        _ctx: &rustyline::Context<'_>,
    ) -> rustyline::Result<(usize, Vec<String>)> {
        let (start, word) = rustyline::completion::extract_word(line, pos, None, &BREAK_CHARS);

        let matches = self
            .words
            .iter()
            .filter(|w| w.starts_with(word))
            .cloned()
            .collect();
        Ok((start, matches))
    }
}

fn main() -> Result<(), Box<dyn std::error::Error>> {
    let mut paths = env::args().skip(1);

    let mut conn = rusqlite::Connection::open_in_memory().unwrap();

    let mut base_words = [
        "distinct", "select", "from", "group", "by", "order", "where", "count", "limit", "offset",
        ".export",
    ]
    .iter()
    .map(|&s| s.to_string())
    .collect::<Vec<String>>();

    if paths.len() == 1 {
        let mut col_names = _load_table_from_path(&mut conn, "t", paths.next().unwrap())?;
        base_words.append(&mut col_names);
    } else {
        for (idx, path) in paths.enumerate() {
            let mut col_names = _load_table_from_path(&mut conn, &format!("t{}", idx + 1), path)?;
            base_words.append(&mut col_names);
        }
    }

    let completer = SimpleWordCompleter::new(base_words);
    let mut rl = rustyline::Editor::new();
    rl.set_helper(Some(completer));
    loop {
        match rl.readline("> ") {
            Ok(line) => {
                if line.trim().is_empty() {
                    continue;
                }
                _process_query(&mut conn, &line);
                rl.add_history_entry(line);
            }
            Err(rustyline::error::ReadlineError::Interrupted) => {
                println!("Interrupted");
                continue;
            }
            Err(rustyline::error::ReadlineError::Eof) => {
                break;
            }
            Err(err) => {
                println!("Error: {}", err);
                break;
            }
        }
    }

    Ok(())
}

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

    #[test]
    fn test_normalize_col() {
        for (value, expected) in &[("abc", "abc"), ("2/6/2000", "c_2_6_2000")] {
            assert_eq!(&&normalize_col(value), expected);
        }
    }
}
<|endoftext|>


<commit_before>#![warn(rust_2018_idioms)]

use tokio::io::{AsyncReadExt, AsyncWriteExt};
use tokio::net::{TcpListener, TcpStream};

use futures::FutureExt;
use std::env;
use std::error::Error;

use nix::sys::socket::{getsockopt, sockopt, InetAddr};
use std::net::SocketAddr;
use std::os::unix::io::AsRawFd;

use bytes::BytesMut;

const BUF_SIZE: usize = 1024 * 1024;

#[tokio::main]
async fn main() -> Result<(), Box<dyn Error>> {
    let listen_addr = env::args()
        .nth(1)
        .unwrap_or_else(|| "127.0.0.1:1080".to_string());

    println!("Listening on: {}", listen_addr);

    let listener = TcpListener::bind(listen_addr).await?;

    while let Ok((inbound, client_addr)) = listener.accept().await {
        let server_addr = getsockopt(inbound.as_raw_fd(), sockopt::OriginalDst).unwrap();
        let server_addr = InetAddr::V4(server_addr).to_std();
        eprintln!(
            "[INFO]{} -> {}: connection incoming",
            client_addr, server_addr
        );
        if inbound.local_addr().unwrap() == server_addr {
            eprintln!(
                "[WARM]{} -> {}: connection from local network",
                client_addr, server_addr
            );
            continue;
        }
        let transfer = transfer(inbound, server_addr).map(move |r| {
            if let Err(e) = r {
                println!(
                    "[INFO]{} -> {}: connection closed with error; {}",
                    client_addr, server_addr, e
                );
            }
        });
        tokio::spawn(transfer);
    }

    Ok(())
}

async fn transfer(mut inbound: TcpStream, proxy_addr: SocketAddr) -> Result<(), Box<dyn Error>> {
    let mut outbound = TcpStream::connect(proxy_addr).await?;

    inbound.set_nodelay(true)?;
    outbound.set_nodelay(true)?;

    let client_addr = inbound.peer_addr().unwrap();
    let server_addr = outbound.peer_addr().unwrap();

    let (ri, wi) = inbound.split();
    let (ro, wo) = outbound.split();

    tokio::try_join!(copy_one_direction(ro, wi), copy_one_direction(ri, wo))?;
    eprintln!(
        "[INFO]{} -> {}: connection closed",
        client_addr, server_addr
    );

    Ok(())
}

async fn copy_one_direction<T, U>(mut from: T, mut to: U) -> std::io::Result<()>
where
    T: AsyncReadExt + Unpin,
    U: AsyncWriteExt + Unpin,
{
    let mut buf = BytesMut::with_capacity(BUF_SIZE);
    loop {
        let n = from.read_buf(&mut buf).await?;
        if n == 0 {
            break;
        }
        let n = to.write_buf(&mut buf).await?;
        if n == 0 {
            break;
        }
    }
    to.shutdown().await
}
<commit_msg>Lazy allocate buffer until readable event<commit_after>#![warn(rust_2018_idioms)]

use tokio::io::AsyncWriteExt;
use tokio::net::{TcpListener, TcpStream};

use futures::FutureExt;
use std::error::Error;
use std::{env, io};

use nix::sys::socket::{getsockopt, sockopt, InetAddr};
use std::net::SocketAddr;
use std::os::unix::io::AsRawFd;

use bytes::BytesMut;
use tokio::net::tcp::{ReadHalf, WriteHalf};

const BUF_SIZE: usize = 1024 * 1024;

#[tokio::main]
async fn main() -> Result<(), Box<dyn Error>> {
    let listen_addr = env::args()
        .nth(1)
        .unwrap_or_else(|| "127.0.0.1:1080".to_string());

    println!("Listening on: {}", listen_addr);

    let listener = TcpListener::bind(listen_addr).await?;

    while let Ok((inbound, client_addr)) = listener.accept().await {
        let server_addr = getsockopt(inbound.as_raw_fd(), sockopt::OriginalDst).unwrap();
        let server_addr = InetAddr::V4(server_addr).to_std();
        eprintln!(
            "[INFO]{} -> {}: connection incoming",
            client_addr, server_addr
        );
        if inbound.local_addr().unwrap() == server_addr {
            eprintln!(
                "[WARM]{} -> {}: connection from local network",
                client_addr, server_addr
            );
            continue;
        }
        let transfer = transfer(inbound, server_addr).map(move |r| {
            if let Err(e) = r {
                println!(
                    "[INFO]{} -> {}: connection closed with error; {}",
                    client_addr, server_addr, e
                );
            }
        });
        tokio::spawn(transfer);
    }

    Ok(())
}

async fn transfer(mut inbound: TcpStream, proxy_addr: SocketAddr) -> Result<(), Box<dyn Error>> {
    let mut outbound = TcpStream::connect(proxy_addr).await?;

    inbound.set_nodelay(true)?;
    outbound.set_nodelay(true)?;

    let client_addr = inbound.peer_addr().unwrap();
    let server_addr = outbound.peer_addr().unwrap();

    let (mut ri, mut wi) = inbound.split();
    let (mut ro, mut wo) = outbound.split();

    tokio::try_join!(
        copy_one_direction(&mut ro, &mut wi),
        copy_one_direction(&mut ri, &mut wo)
    )?;
    eprintln!(
        "[INFO]{} -> {}: connection closed",
        client_addr, server_addr
    );

    Ok(())
}

async fn copy_one_direction(
    from: &mut ReadHalf<'_>,
    to: &mut WriteHalf<'_>,
) -> std::io::Result<()> {
    loop {
        from.readable().await?;
        {
            let mut buf = BytesMut::with_capacity(BUF_SIZE);
            match from.try_read_buf(&mut buf) {
                Ok(0) => {
                    break;
                }
                Ok(_n) => {
                    to.write_all_buf(&mut buf).await?;
                }
                Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => {
                    continue;
                }
                Err(e) => {
                    return Err(e.into());
                }
            }
        }
    }
    to.shutdown().await
}
<|endoftext|>


<commit_before>mod ast;
mod parser;

fn main() {
    println!("Hello, world!");
}
<commit_msg>Make `main()` actually parse and print an AST<commit_after>use std::process::exit;
use std::{env, fs};

mod ast;
mod parser;

fn main() {
    let args: Vec<String> = env::args().collect();
    if args.len() < 2 || args.len() > 2 {
        println!("Usage: hummingbird [file]");
        exit(-1);
    }
    let filename = &args[1];
    let source = fs::read_to_string(filename).expect("Unable to read source file");

    let program = parser::parse(source);
    println!("Program:\n{:?}", program);
}
<|endoftext|>