<commit_before>use url::percent_encoding::percent_encode_byte;

#[derive(Debug)]
pub struct RequestBuilder<'a, T, Q> {
    method: &'a str,
    path: &'a str,
    query: Q,
    headers: T,
}

#[derive(Debug)]
pub struct HttpHeader<'a, T> {
    name: &'a str,
    value: &'a str,
    headers: T,
}

#[derive(Debug)]
pub struct QueryPair<'a, T> {
    name: &'a str,
    value: &'a [u8],
    pairs: T,
}

#[derive(Debug)]
pub struct EmptyHeader {}

#[derive(Debug)]
pub struct BaseQuery<'a>(&'a str);

pub trait EncodeToBuf {
    fn encode(&self, buf: &mut Vec<u8>);
}

impl<'a> RequestBuilder<'a, EmptyHeader, BaseQuery<'a>> {
    pub fn new(
        method: &'a str,
        path: &'a str,
        query: Option<&'a str>,
    ) -> RequestBuilder<'a, EmptyHeader, BaseQuery<'a>> {
        RequestBuilder {
            method,
            path,
            query: BaseQuery(query.unwrap_or("")),
            headers: EmptyHeader {},
        }
    }
}

impl<'a, 'b, 'c, T, Q> RequestBuilder<'a, T, Q> {
    pub fn query(self, name: &'c str, value: &'c [u8]) -> RequestBuilder<'a, T, QueryPair<'c, Q>> {
        RequestBuilder {
            method: self.method,
            path: self.path,
            query: QueryPair {
                name,
                value,
                pairs: self.query,
            },
            headers: self.headers,
        }
    }

    pub fn query_opt(
        self,
        name: &'c str,
        value: Option<&'c [u8]>,
    ) -> RequestBuilder<'a, T, QueryPair<'c, Q>> {
        RequestBuilder {
            method: self.method,
            path: self.path,
            query: QueryPair {
                name: if value.is_some() { name } else { "" },
                value: value.unwrap_or(b""),
                pairs: self.query,
            },
            headers: self.headers,
        }
    }

    pub fn header(self, name: &'b str, value: &'b str) -> RequestBuilder<'a, HttpHeader<'b, T>, Q> {
        RequestBuilder {
            method: self.method,
            path: self.path,
            query: self.query,
            headers: HttpHeader {
                name,
                value,
                headers: self.headers,
            },
        }
    }
}

impl<'a, T: EncodeToBuf, Q: EncodeToBuf> RequestBuilder<'a, T, Q> {
    pub fn encode(&self, buf: &mut Vec<u8>) {
        buf.extend_from_slice(self.method.as_bytes());
        buf.extend_from_slice(b" ");
        buf.extend_from_slice(self.path.as_bytes());
        // The query either encodes an extra ? or an extra &, pop either off
        self.query.encode(buf);
        buf.pop();
        buf.extend_from_slice(b" HTTP/1.1");
        buf.extend_from_slice(b"\r\n");
        self.headers.encode(buf);
        buf.extend_from_slice(b"\r\n");
    }
}

impl<'a, T: EncodeToBuf> EncodeToBuf for HttpHeader<'a, T> {
    fn encode(&self, buf: &mut Vec<u8>) {
        self.headers.encode(buf);
        buf.extend_from_slice(self.name.as_bytes());
        buf.extend_from_slice(b": ");
        buf.extend_from_slice(self.value.as_bytes());
        buf.extend_from_slice(b"\r\n");
    }
}

impl EncodeToBuf for EmptyHeader {
    fn encode(&self, _: &mut Vec<u8>) {}
}

impl<'a, T: EncodeToBuf> EncodeToBuf for QueryPair<'a, T> {
    fn encode(&self, buf: &mut Vec<u8>) {
        self.pairs.encode(buf);
        if self.name != "" {
            buf.extend_from_slice(self.name.as_bytes());
            buf.extend_from_slice(b"=");
            encode_param(self.value, buf);
            buf.extend_from_slice(b"&");
        }
    }
}

impl<'a> EncodeToBuf for BaseQuery<'a> {
    fn encode(&self, buf: &mut Vec<u8>) {
        buf.extend_from_slice(b"?");
        buf.extend_from_slice(self.0.as_bytes());
        // Add an & if it the base query didn't contain it.
        if !self.0.is_empty() && !self.0.ends_with("&") {
            buf.extend_from_slice(b"&");
        }
    }
}

fn encode_param(param: &[u8], buf: &mut Vec<u8>) {
    for byte in param {
        let c = char::from(*byte);
        let mut char_buf = [0u8; 4];
        if (*byte > 0x20 && *byte < 0x7E) && (c.is_numeric() || c.is_alphabetic() || c == '-') {
            c.encode_utf8(&mut char_buf);
            buf.extend_from_slice(&char_buf[0..c.len_utf8()]);
        } else {
            buf.extend_from_slice(percent_encode_byte(*byte).as_bytes());
        }
    }
}

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

    #[test]
    fn test_encode_http_request() {
        let req = RequestBuilder::new("GET", "/foobar/baz", Some("a=b"))
            .query("b", "c".as_bytes())
            .header("header1", "value1")
            .header("header2", "value2");
        let mut encoded = Vec::new();
        req.encode(&mut encoded);
        assert_eq!(
            String::from_utf8(encoded).unwrap(),
            vec![
                "GET /foobar/baz?a=b&b=c HTTP/1.1",
                "header1: value1",
                "header2: value2",
                "\r\n",
            ]
            .join("\r\n")
        );
    }

    #[test]
    fn test_encode_opt_param() {
        let req = RequestBuilder::new("GET", "/foobar/baz", Some("a=b"))
            .query_opt("a", None)
            .query_opt("b", Some(b"c"));
        let mut encoded = Vec::new();
        req.encode(&mut encoded);
        assert_eq!(
            String::from_utf8(encoded).unwrap(),
            vec!["GET /foobar/baz?a=b&b=c HTTP/1.1", "\r\n",].join("\r\n")
        );
    }

    #[test]
    fn test_percent_encode_query() {
        let req = RequestBuilder::new("GET", "/foobar/baz", None).query("a", "&".as_bytes());
        let mut encoded = Vec::new();
        req.encode(&mut encoded);
        assert_eq!(
            String::from_utf8(encoded).unwrap(),
            vec!["GET /foobar/baz?a=%26 HTTP/1.1", "\r\n",].join("\r\n")
        );
    }
}
<commit_msg>Use Vec based HTTP Builder for simplicity<commit_after>use url::percent_encoding::percent_encode_byte;

#[derive(Debug)]
pub struct RequestBuilder<'a> {
    method: &'a str,
    path: &'a str,
    base_query: Option<&'a str>,
    query_pairs: Vec<QueryPair<'a>>,
    headers: Vec<HttpHeader<'a>>,
}

#[derive(Debug)]
pub struct HttpHeader<'a> {
    name: &'a str,
    value: &'a str,
}

#[derive(Debug)]
pub struct QueryPair<'a> {
    name: &'a str,
    value: &'a [u8],
}

pub trait EncodeToBuf {
    fn encode(&self, buf: &mut Vec<u8>);
}

impl<'a> RequestBuilder<'a> {
    pub fn new(method: &'a str, path: &'a str, query: Option<&'a str>) -> RequestBuilder<'a> {
        RequestBuilder {
            method,
            path,
            base_query: query,
            query_pairs: Vec::with_capacity(10),
            headers: Vec::with_capacity(10),
        }
    }
}

impl<'a> RequestBuilder<'a> {
    pub fn query(&mut self, name: &'a str, value: &'a [u8]) -> &mut RequestBuilder<'a> {
        self.query_pairs.push(QueryPair { name, value });
        self
    }

    pub fn query_opt(&mut self, name: &'a str, value: Option<&'a [u8]>) -> &mut RequestBuilder<'a> {
        if let Some(v) = value {
            self.query_pairs.push(QueryPair { name, value: v });
        }
        self
    }

    pub fn header(&mut self, name: &'a str, value: &'a str) -> &mut RequestBuilder<'a> {
        self.headers.push(HttpHeader { name, value });
        self
    }
}

impl<'a> RequestBuilder<'a> {
    pub fn encode(&self, buf: &mut Vec<u8>) {
        buf.extend_from_slice(self.method.as_bytes());
        buf.extend_from_slice(b" ");
        buf.extend_from_slice(self.path.as_bytes());

        buf.extend_from_slice(b"?");
        if let Some(q) = self.base_query {
            buf.extend_from_slice(q.as_bytes());
            buf.extend_from_slice(b"&");
        }
        for pair in &self.query_pairs {
            buf.extend_from_slice(pair.name.as_bytes());
            buf.extend_from_slice(b"=");
            encode_param(pair.value, buf);
            buf.extend_from_slice(b"&");
        }
        // The query either encodes an extra ? or an extra &, pop either off
        buf.pop();
        buf.extend_from_slice(b" HTTP/1.1");
        buf.extend_from_slice(b"\r\n");
        for header in &self.headers {
            buf.extend_from_slice(header.name.as_bytes());
            buf.extend_from_slice(b": ");
            buf.extend_from_slice(header.value.as_bytes());
            buf.extend_from_slice(b"\r\n");
        }
        buf.extend_from_slice(b"\r\n");
    }
}

fn encode_param(param: &[u8], buf: &mut Vec<u8>) {
    for byte in param {
        let c = char::from(*byte);
        let mut char_buf = [0u8; 4];
        if (*byte > 0x20 && *byte < 0x7E) && (c.is_numeric() || c.is_alphabetic() || c == '-') {
            c.encode_utf8(&mut char_buf);
            buf.extend_from_slice(&char_buf[0..c.len_utf8()]);
        } else {
            buf.extend_from_slice(percent_encode_byte(*byte).as_bytes());
        }
    }
}

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

    #[test]
    fn test_encode_http_request() {
        let mut encoded = Vec::new();
        RequestBuilder::new("GET", "/foobar/baz", Some("a=b"))
            .query("b", "c".as_bytes())
            .header("header1", "value1")
            .header("header2", "value2")
            .encode(&mut encoded);
        assert_eq!(
            String::from_utf8(encoded).unwrap(),
            vec![
                "GET /foobar/baz?a=b&b=c HTTP/1.1",
                "header1: value1",
                "header2: value2",
                "\r\n",
            ]
            .join("\r\n")
        );
    }

    #[test]
    fn test_encode_opt_param() {
        let mut encoded = Vec::new();
        RequestBuilder::new("GET", "/foobar/baz", Some("a=b"))
            .query_opt("a", None)
            .query_opt("b", Some(b"c"))
            .encode(&mut encoded);
        assert_eq!(
            String::from_utf8(encoded).unwrap(),
            vec!["GET /foobar/baz?a=b&b=c HTTP/1.1", "\r\n",].join("\r\n")
        );
    }

    #[test]
    fn test_percent_encode_query() {
        let mut encoded = Vec::new();
        let req = RequestBuilder::new("GET", "/foobar/baz", None)
            .query("a", "&".as_bytes())
            .encode(&mut encoded);
        assert_eq!(
            String::from_utf8(encoded).unwrap(),
            vec!["GET /foobar/baz?a=%26 HTTP/1.1", "\r\n",].join("\r\n")
        );
    }
}
<|endoftext|>


<commit_before>#![cfg_attr(debug_assertions, allow(unused_imports, dead_code))]

use std::collections::HashSet;
use std::time::{UNIX_EPOCH, Duration};

use trust_dns::op::{Message, MessageType, ResponseCode, Query, OpCode, Edns};
use trust_dns::rr::{DNSClass, RecordType, Record, RData, Name};
use trust_dns::rr::rdata::sig::SIG;
use trust_dns::rr::dnssec::{Signer, TrustAnchor};
use trust_dns::serialize::binary::{BinEncoder, BinSerializable};
use openssl::crypto::pkey::Role;
use itertools::Itertools;

use resolver::RcResolver;
use utils::Result;
use recursive::RecursiveResolver;

lazy_static! {
    static ref TRUST_ANCHOR: TrustAnchor = TrustAnchor::default();
}

#[derive(Eq, PartialEq, Debug)]
enum ValidationResult {
    NonAuthenticated,
    Authenticated,
    Bogus,
}
pub trait ResponseValidator {
    fn is_valid(&mut self, msg: &Message) -> bool;
}
pub struct DummyValidator;
impl ResponseValidator for DummyValidator {
    fn is_valid(&mut self, _: &Message) -> bool {
        true
    }
}
pub trait SubqueryResolver {
    fn resolve_sub(&self, q: Query) -> Result<Message>;
}
impl SubqueryResolver for RcResolver {
    fn resolve_sub(&self, q: Query) -> Result<Message> {
        RecursiveResolver::resolve(&q, self.clone())
    }
}
pub struct DnssecValidator<T: SubqueryResolver> {
    resolver: T,
}
impl<T: SubqueryResolver> DnssecValidator<T> {
    pub fn new(resolver: T) -> Self {
        DnssecValidator {
            resolver: resolver,
        }
    }
    fn verify_sig_time(sig: &SIG) -> bool {
        let mut cur_time = UNIX_EPOCH.elapsed()
        .unwrap_or_else(|_| Duration::new(0, 0)).as_secs() as u64;
        cur_time &= 0xffff_ffff;
        let inception = sig.get_sig_inception() as u64;
        if cur_time < inception {
            cur_time += 0x1_0000_0000;
        }
        let expiration = {
            let t = sig.get_sig_expiration() as u64;
            if t < inception {
                t + 0x1_0000_0000
            } else {
                t
            }
        };
        cur_time < expiration
    }
    fn verify_ds(&self, dnskey: &Record) -> Result<bool> {
        let (public_key, algorithm) = if let RData::DNSKEY(ref rdata) = *dnskey.get_rdata() {
            (rdata.get_public_key(), rdata.get_algorithm())
        } else {
            panic!("Invalid record type")
        };
        if TRUST_ANCHOR.contains(public_key) {
            return Ok(true);
        }
        let name = dnskey.get_name();
        if name.is_root() {
            // Avoid infinite loop
            return Ok(false);
        }
        let mut q = Query::new();
        q.name(dnskey.get_name().clone());
        q.query_type(RecordType::DS);
        q.query_class(dnskey.get_dns_class());
        let ds_resp = try!(self.resolver.resolve_sub(q));
        let mut buf = Vec::<u8>::new();
        {
            let mut encoder = BinEncoder::new(&mut buf);
            encoder.set_canonical_names(true);
            try!(dnskey.get_name().emit(&mut encoder));
            try!(dnskey.get_rdata().emit(&mut encoder));
        }
        Ok(ds_resp.get_answers().iter().any(|x| match *x.get_rdata() {
            RData::DS(ref ds) if ds.get_algorithm() == algorithm => {
                let hash = ds.get_digest_type().hash(&buf);
                hash == ds.get_digest()
            },
            _ => false,
        }))
    }
    fn verify_rrsig(&self, rrset: &[Record], rrsigs: &[SIG]) -> Result<ValidationResult> {
        assert!(!rrset.is_empty());
        let rr_name = rrset[0].get_name();
        let rr_class = rrset[0].get_dns_class();
        let rr_type = rrset[0].get_rr_type();
        if rrsigs.is_empty() {
            debug!("No RRSIG for {} {:?} {:?}", rr_name, rr_class, rr_type);
            return Ok(ValidationResult::Bogus);
        }
        debug_assert!(rrset.iter().all(|x| x.get_rr_type() == rr_type));
        debug_assert!(rrset.iter().all(|x| x.get_name() == rr_name));
        debug_assert!(rrsigs.iter().all(|x| x.get_type_covered() == rr_type));
        for sig in rrsigs {
            if !Self::verify_sig_time(sig) {
                debug!("Not in validity time ({} -> {}).",
                       sig.get_sig_inception(), sig.get_sig_expiration());
                continue;
            }
            let signer_name = sig.get_signer_name();
            // The number of labels in the RRset owner name MUST be greater than
            // or equal to the value in the RRSIG RR's Labels field.
            if rr_name.num_labels() < sig.get_num_labels() {
                debug!("Number of labels: name = {} ({}), sig = {}",
                       rr_name.num_labels(), rr_name, sig.get_num_labels());
                continue;
            }
            if !signer_name.zone_of(rr_name) {
                debug!("Signer's name ({}) is not zone of RRset owner name ({})",
                       signer_name, rr_name);
                continue;
            }
            let dnskeys = if rr_type == RecordType::DNSKEY && signer_name == rr_name {
                rrset.iter()
                .filter(|rec| self.verify_ds(rec).unwrap_or(false))
                .map(|rec| if let RData::DNSKEY(ref dnskey) = *rec.get_rdata() {
                    dnskey.clone()
                } else {
                    panic!("Invalid record type")
                }).collect_vec()
            } else {
                let mut q = Query::new();
                q.name(signer_name.clone())
                .query_class(rr_class)
                .query_type(RecordType::DNSKEY);
                let resp = try!(self.resolver.resolve_sub(q));
                resp.get_answers().iter().filter_map(|rec| match *rec.get_rdata() {
                    RData::DNSKEY(ref dnskey) => Some(dnskey.clone()),
                    _ => None,
                }).collect_vec()
            };
            for key in dnskeys {
                if key.is_revoke() {
                    debug!("DNSKEY is revoked: {:?}", key);
                    continue;
                }
                if !key.is_zone_key() {
                    continue;
                }
                if *key.get_algorithm() != sig.get_algorithm() {
                    continue;
                }
                let pkey = key.get_algorithm().public_key_from_vec(key.get_public_key());
                let pkey = match pkey {
                    Ok(k) => k,
                    Err(e) => {
                        warn!("Failed to get public key from DNSKEY: {:?}", e);
                        continue
                    },
                };
                if !pkey.can(Role::Verify) {
                    debug!("PKey can't be used to verify: {:?}", key);
                }
                let signer = Signer::new_verifier(
                    *key.get_algorithm(), pkey, signer_name.clone(),
                );
                let rrset_hash: Vec<u8> = signer.hash_rrset(
                    rr_name, rr_class,
                    sig.get_num_labels(), sig.get_type_covered(), sig.get_algorithm(),
                    sig.get_original_ttl(), sig.get_sig_expiration(), sig.get_sig_inception(),
                    sig.get_key_tag(), signer_name, rrset,
                );
                if signer.verify(&rrset_hash, sig.get_sig()) {
                    return Ok(ValidationResult::Authenticated);
                } else {
                    trace!("Failed to verify {} {:?} with {:?}", rr_name, rr_type, key);
                }
            }
        }
        debug!("No valid RRSIG for {} {:?} {:?}", rr_name, rr_class, rr_type);
        Ok(ValidationResult::Bogus)
    }
    // Largely stolen from trust-dns
    fn verify_rrsigs(&self, msg: &Message) -> Result<ValidationResult> {
        let rrsigs = msg.get_answers().iter()
        .chain(msg.get_name_servers())
        .filter(|rr| rr.get_rr_type() == RecordType::RRSIG)
        .collect_vec();

        if rrsigs.is_empty() {
            // TODO: Check whether this zone is really not signed
            debug!("No Sig");
            return Ok(ValidationResult::NonAuthenticated);
        }
        // Group the record sets by name and type
        let mut rrset_types: HashSet<(Name, RecordType)> = HashSet::new();
        msg.get_answers().iter()
        .chain(msg.get_name_servers())
        .filter(|rr| rr.get_rr_type() != RecordType::RRSIG)
        .map(|rr| (rr.get_name().clone(), rr.get_rr_type()))
        .foreach(|rr| { rrset_types.insert(rr); });

        for (name, rr_type) in rrset_types {
            let rrset = msg.get_answers().iter()
            .chain(msg.get_name_servers())
            .filter(|rr| rr.get_rr_type() == rr_type && *rr.get_name() == name)
            .cloned()
            .collect_vec();
            let rrsig = rrsigs.iter()
            .filter(|rr| *rr.get_name() == name)
            .filter_map(|rr| if let RData::SIG(ref sig) = *rr.get_rdata() {
                if sig.get_type_covered() == rr_type {
                    Some(sig.clone())
                } else {
                    None
                }
            } else {
                panic!("Unexpected RData: {:?}", rr)
            })
            .collect_vec();
            if rr_type == RecordType::NS && rrsig.is_empty() && msg.get_answers().is_empty() {
                // NS records are not signed in referral
                continue;
            }
            if try!(self.verify_rrsig(&rrset, &rrsig)) != ValidationResult::Authenticated {
                return Ok(ValidationResult::Bogus);
            }
        }

        // TODO: Verify NSEC/NSEC3
        Ok(ValidationResult::Authenticated)
    }
}
impl<T: SubqueryResolver> ResponseValidator for DnssecValidator<T> {
    fn is_valid(&mut self, msg: &Message) -> bool {
        self.verify_rrsigs(msg).unwrap_or(ValidationResult::Bogus) != ValidationResult::Bogus
    }
}

#[cfg(test)]
mod tests {
    extern crate env_logger;
    use std::time::Duration;
    use trust_dns::op::*;
    use trust_dns::rr::*;
    use super::*;
    use ::query::query;
    use ::resolver::*;
    use ::mioco_config_start;

    fn test_query(domain: &str) -> Message {
        let mut q = Query::new();
        q.name(Name::parse(domain, Some(&Name::root())).unwrap())
        .query_class(DNSClass::IN)
        .query_type(RecordType::A);
        let msg = query(q, "8.8.8.8".parse().unwrap(), Duration::from_secs(5)).unwrap();
        assert!(msg.get_answers().iter().any(|r| r.get_rr_type() == RecordType::RRSIG));
        msg
    }

    #[test]
    fn test_dnssec_valid() {
        env_logger::init().is_ok();
        mioco_config_start(move || {
            let resolver = RcResolver::default();
            let msg = test_query("sigok.verteiltesysteme.net");
            let mut validator = DnssecValidator::new(resolver.clone());
            assert!(validator.is_valid(&msg));
        }).unwrap();
    }
    #[test]
    fn test_dnssec_invalid() {
        env_logger::init().is_ok();
        mioco_config_start(move || {
            let resolver = RcResolver::default();
            let msg = test_query("sigfail.verteiltesysteme.net");
            let mut validator = DnssecValidator::new(resolver.clone());
            assert!(!validator.is_valid(&msg));
        }).unwrap();
    }
}
<commit_msg>Fix clippy warning (code complexity)<commit_after>#![cfg_attr(debug_assertions, allow(unused_imports, dead_code))]

use std::collections::HashSet;
use std::time::{UNIX_EPOCH, Duration};

use trust_dns::op::{Message, MessageType, ResponseCode, Query, OpCode, Edns};
use trust_dns::rr::{DNSClass, RecordType, Record, RData, Name};
use trust_dns::rr::rdata::{SIG, DNSKEY};
use trust_dns::rr::dnssec::{Signer, TrustAnchor};
use trust_dns::serialize::binary::{BinEncoder, BinSerializable};
use openssl::crypto::pkey::Role;
use itertools::Itertools;

use resolver::RcResolver;
use utils::Result;
use recursive::RecursiveResolver;

lazy_static! {
    static ref TRUST_ANCHOR: TrustAnchor = TrustAnchor::default();
}

#[derive(Eq, PartialEq, Debug)]
enum ValidationResult {
    NonAuthenticated,
    Authenticated,
    Bogus,
}
pub trait ResponseValidator {
    fn is_valid(&mut self, msg: &Message) -> bool;
}
pub struct DummyValidator;
impl ResponseValidator for DummyValidator {
    fn is_valid(&mut self, _: &Message) -> bool {
        true
    }
}
pub trait SubqueryResolver {
    fn resolve_sub(&self, q: Query) -> Result<Message>;
}
impl SubqueryResolver for RcResolver {
    fn resolve_sub(&self, q: Query) -> Result<Message> {
        RecursiveResolver::resolve(&q, self.clone())
    }
}
pub struct DnssecValidator<T: SubqueryResolver> {
    resolver: T,
}
impl<T: SubqueryResolver> DnssecValidator<T> {
    pub fn new(resolver: T) -> Self {
        DnssecValidator {
            resolver: resolver,
        }
    }
    fn verify_sig_time(sig: &SIG) -> bool {
        let mut cur_time = UNIX_EPOCH.elapsed()
        .unwrap_or_else(|_| Duration::new(0, 0)).as_secs() as u64;
        cur_time &= 0xffff_ffff;
        let inception = sig.get_sig_inception() as u64;
        if cur_time < inception {
            cur_time += 0x1_0000_0000;
        }
        let expiration = {
            let t = sig.get_sig_expiration() as u64;
            if t < inception {
                t + 0x1_0000_0000
            } else {
                t
            }
        };
        cur_time < expiration
    }
    fn verify_ds(&self, dnskey: &Record) -> Result<bool> {
        let (public_key, algorithm) = if let RData::DNSKEY(ref rdata) = *dnskey.get_rdata() {
            (rdata.get_public_key(), rdata.get_algorithm())
        } else {
            panic!("Invalid record type")
        };
        if TRUST_ANCHOR.contains(public_key) {
            return Ok(true);
        }
        let name = dnskey.get_name();
        if name.is_root() {
            // Avoid infinite loop
            return Ok(false);
        }
        let mut q = Query::new();
        q.name(dnskey.get_name().clone());
        q.query_type(RecordType::DS);
        q.query_class(dnskey.get_dns_class());
        let ds_resp = try!(self.resolver.resolve_sub(q));
        let mut buf = Vec::<u8>::new();
        {
            let mut encoder = BinEncoder::new(&mut buf);
            encoder.set_canonical_names(true);
            try!(dnskey.get_name().emit(&mut encoder));
            try!(dnskey.get_rdata().emit(&mut encoder));
        }
        Ok(ds_resp.get_answers().iter().any(|x| match *x.get_rdata() {
            RData::DS(ref ds) if ds.get_algorithm() == algorithm => {
                let hash = ds.get_digest_type().hash(&buf);
                hash == ds.get_digest()
            },
            _ => false,
        }))
    }
    fn create_verifier(key: &DNSKEY, signer_name: &Name) -> Option<Signer> {
        if key.is_revoke() {
            debug!("DNSKEY is revoked: {:?}", key);
            return None;
        }
        if !key.is_zone_key() {
            return None;
        }
        let pkey = key.get_algorithm().public_key_from_vec(key.get_public_key());
        let pkey = match pkey {
            Ok(k) => k,
            Err(e) => {
                warn!("Failed to get public key from DNSKEY: {:?}", e);
                return None
            },
        };
        if !pkey.can(Role::Verify) {
            debug!("PKey can't be used to verify: {:?}", key);
            return None;
        }
        Some(Signer::new_verifier(
            *key.get_algorithm(), pkey, signer_name.clone(),
        ))
    }
    fn verify_rrsig(&self, rrset: &[Record], rrsigs: &[SIG]) -> Result<ValidationResult> {
        assert!(!rrset.is_empty());
        let rr_name = rrset[0].get_name();
        let rr_class = rrset[0].get_dns_class();
        let rr_type = rrset[0].get_rr_type();
        if rrsigs.is_empty() {
            debug!("No RRSIG for {} {:?} {:?}", rr_name, rr_class, rr_type);
            return Ok(ValidationResult::Bogus);
        }
        debug_assert!(rrset.iter().all(|x| x.get_rr_type() == rr_type));
        debug_assert!(rrset.iter().all(|x| x.get_name() == rr_name));
        debug_assert!(rrsigs.iter().all(|x| x.get_type_covered() == rr_type));
        for sig in rrsigs {
            if !Self::verify_sig_time(sig) {
                debug!("Not in validity time ({} -> {}).",
                       sig.get_sig_inception(), sig.get_sig_expiration());
                continue;
            }
            let signer_name = sig.get_signer_name();
            // The number of labels in the RRset owner name MUST be greater than
            // or equal to the value in the RRSIG RR's Labels field.
            if rr_name.num_labels() < sig.get_num_labels() {
                debug!("Number of labels: name = {} ({}), sig = {}",
                       rr_name.num_labels(), rr_name, sig.get_num_labels());
                continue;
            }
            if !signer_name.zone_of(rr_name) {
                debug!("Signer's name ({}) is not zone of RRset owner name ({})",
                       signer_name, rr_name);
                continue;
            }
            let dnskeys = if rr_type == RecordType::DNSKEY && signer_name == rr_name {
                rrset.iter()
                .filter(|rec| self.verify_ds(rec).unwrap_or(false))
                .map(|rec| if let RData::DNSKEY(ref dnskey) = *rec.get_rdata() {
                    dnskey.clone()
                } else {
                    panic!("Invalid record type")
                }).collect_vec()
            } else {
                let mut q = Query::new();
                q.name(signer_name.clone())
                .query_class(rr_class)
                .query_type(RecordType::DNSKEY);
                let resp = try!(self.resolver.resolve_sub(q));
                resp.get_answers().iter().filter_map(|rec| match *rec.get_rdata() {
                    RData::DNSKEY(ref dnskey) => Some(dnskey.clone()),
                    _ => None,
                }).collect_vec()
            };
            for key in dnskeys {
                if *key.get_algorithm() != sig.get_algorithm() {
                    continue;
                }
                let signer = match Self::create_verifier(&key, signer_name) {
                    Some(x) => x,
                    None => continue,
                };
                let rrset_hash: Vec<u8> = signer.hash_rrset(
                    rr_name, rr_class,
                    sig.get_num_labels(), sig.get_type_covered(), sig.get_algorithm(),
                    sig.get_original_ttl(), sig.get_sig_expiration(), sig.get_sig_inception(),
                    sig.get_key_tag(), signer_name, rrset,
                );
                if signer.verify(&rrset_hash, sig.get_sig()) {
                    return Ok(ValidationResult::Authenticated);
                } else {
                    trace!("Failed to verify {} {:?} with {:?}", rr_name, rr_type, key);
                }
            }
        }
        debug!("No valid RRSIG for {} {:?} {:?}", rr_name, rr_class, rr_type);
        Ok(ValidationResult::Bogus)
    }
    // Largely stolen from trust-dns
    fn verify_rrsigs(&self, msg: &Message) -> Result<ValidationResult> {
        let rrsigs = msg.get_answers().iter()
        .chain(msg.get_name_servers())
        .filter(|rr| rr.get_rr_type() == RecordType::RRSIG)
        .collect_vec();

        if rrsigs.is_empty() {
            // TODO: Check whether this zone is really not signed
            debug!("No Sig");
            return Ok(ValidationResult::NonAuthenticated);
        }
        // Group the record sets by name and type
        let mut rrset_types: HashSet<(Name, RecordType)> = HashSet::new();
        msg.get_answers().iter()
        .chain(msg.get_name_servers())
        .filter(|rr| rr.get_rr_type() != RecordType::RRSIG)
        .map(|rr| (rr.get_name().clone(), rr.get_rr_type()))
        .foreach(|rr| { rrset_types.insert(rr); });

        for (name, rr_type) in rrset_types {
            let rrset = msg.get_answers().iter()
            .chain(msg.get_name_servers())
            .filter(|rr| rr.get_rr_type() == rr_type && *rr.get_name() == name)
            .cloned()
            .collect_vec();
            let rrsig = rrsigs.iter()
            .filter(|rr| *rr.get_name() == name)
            .filter_map(|rr| if let RData::SIG(ref sig) = *rr.get_rdata() {
                if sig.get_type_covered() == rr_type {
                    Some(sig.clone())
                } else {
                    None
                }
            } else {
                panic!("Unexpected RData: {:?}", rr)
            })
            .collect_vec();
            if rr_type == RecordType::NS && rrsig.is_empty() && msg.get_answers().is_empty() {
                // NS records are not signed in referral
                continue;
            }
            if try!(self.verify_rrsig(&rrset, &rrsig)) != ValidationResult::Authenticated {
                return Ok(ValidationResult::Bogus);
            }
        }

        // TODO: Verify NSEC/NSEC3
        Ok(ValidationResult::Authenticated)
    }
}
impl<T: SubqueryResolver> ResponseValidator for DnssecValidator<T> {
    fn is_valid(&mut self, msg: &Message) -> bool {
        self.verify_rrsigs(msg).unwrap_or(ValidationResult::Bogus) != ValidationResult::Bogus
    }
}

#[cfg(test)]
mod tests {
    extern crate env_logger;
    use std::time::Duration;
    use trust_dns::op::*;
    use trust_dns::rr::*;
    use super::*;
    use ::query::query;
    use ::resolver::*;
    use ::mioco_config_start;

    fn test_query(domain: &str) -> Message {
        let mut q = Query::new();
        q.name(Name::parse(domain, Some(&Name::root())).unwrap())
        .query_class(DNSClass::IN)
        .query_type(RecordType::A);
        let msg = query(q, "8.8.8.8".parse().unwrap(), Duration::from_secs(5)).unwrap();
        assert!(msg.get_answers().iter().any(|r| r.get_rr_type() == RecordType::RRSIG));
        msg
    }

    #[test]
    fn test_dnssec_valid() {
        env_logger::init().is_ok();
        mioco_config_start(move || {
            let resolver = RcResolver::default();
            let msg = test_query("sigok.verteiltesysteme.net");
            let mut validator = DnssecValidator::new(resolver.clone());
            assert!(validator.is_valid(&msg));
        }).unwrap();
    }
    #[test]
    fn test_dnssec_invalid() {
        env_logger::init().is_ok();
        mioco_config_start(move || {
            let resolver = RcResolver::default();
            let msg = test_query("sigfail.verteiltesysteme.net");
            let mut validator = DnssecValidator::new(resolver.clone());
            assert!(!validator.is_valid(&msg));
        }).unwrap();
    }
}
<|endoftext|>


<commit_before>#[cfg(test)]
#[macro_export]
macro_rules! ssz_tests {
    ($type: ident) => {
        #[test]
        pub fn test_ssz_round_trip() {
            use crate::test_utils::{SeedableRng, TestRandom, XorShiftRng};
            use ssz::{ssz_encode, Decodable};

            let mut rng = XorShiftRng::from_seed([42; 16]);
            let original = $type::random_for_test(&mut rng);

            let bytes = ssz_encode(&original);
            let (decoded, _): ($type, usize) = <_>::ssz_decode(&bytes, 0).unwrap();

            assert_eq!(original, decoded);
        }

        #[test]
        pub fn test_tree_hash_root() {
            use crate::test_utils::{SeedableRng, TestRandom, XorShiftRng};
            use tree_hash::TreeHash;

            let mut rng = XorShiftRng::from_seed([42; 16]);
            let original = $type::random_for_test(&mut rng);

            let result = original.tree_hash_root();

            assert_eq!(result.len(), 32);
            // TODO: Add further tests
            // https://github.com/sigp/lighthouse/issues/170
        }
    };
}

#[cfg(test)]
#[macro_export]
macro_rules! cached_tree_hash_tests {
    ($type: ident) => {
        #[test]
        pub fn test_cached_tree_hash() {
            use crate::test_utils::{SeedableRng, TestRandom, XorShiftRng};
            use tree_hash::TreeHash;

            let mut rng = XorShiftRng::from_seed([42; 16]);

            // Test the original hash
            let original = $type::random_for_test(&mut rng);
            let mut hasher = cached_tree_hash::CachedTreeHasher::new(&original).unwrap();
            assert_eq!(
                hasher.tree_hash_root().unwrap(),
                original.tree_hash_root(),
                "Original hash failed."
            );

            // Test the updated hash
            let modified = $type::random_for_test(&mut rng);
            hasher.update(&modified).unwrap();
            assert_eq!(
                hasher.tree_hash_root().unwrap(),
                modified.tree_hash_root(),
                "Modification hash failed"
            );
        }
    };
}
<commit_msg>Improve cached hash testing in `types`<commit_after>#[cfg(test)]
#[macro_export]
macro_rules! ssz_tests {
    ($type: ident) => {
        #[test]
        pub fn test_ssz_round_trip() {
            use crate::test_utils::{SeedableRng, TestRandom, XorShiftRng};
            use ssz::{ssz_encode, Decodable};

            let mut rng = XorShiftRng::from_seed([42; 16]);
            let original = $type::random_for_test(&mut rng);

            let bytes = ssz_encode(&original);
            let (decoded, _): ($type, usize) = <_>::ssz_decode(&bytes, 0).unwrap();

            assert_eq!(original, decoded);
        }

        #[test]
        pub fn test_tree_hash_root() {
            use crate::test_utils::{SeedableRng, TestRandom, XorShiftRng};
            use tree_hash::TreeHash;

            let mut rng = XorShiftRng::from_seed([42; 16]);
            let original = $type::random_for_test(&mut rng);

            let result = original.tree_hash_root();

            assert_eq!(result.len(), 32);
            // TODO: Add further tests
            // https://github.com/sigp/lighthouse/issues/170
        }
    };
}

#[cfg(test)]
#[macro_export]
macro_rules! cached_tree_hash_tests {
    ($type: ident) => {
        #[test]
        pub fn test_cached_tree_hash() {
            use crate::test_utils::{SeedableRng, TestRandom, XorShiftRng};
            use tree_hash::TreeHash;

            let mut rng = XorShiftRng::from_seed([42; 16]);

            // Test the original hash
            let original = $type::random_for_test(&mut rng);
            let mut cache = cached_tree_hash::TreeHashCache::new(&original, 0).unwrap();

            assert_eq!(
                cache.tree_hash_root().unwrap().to_vec(),
                original.tree_hash_root(),
                "Original hash failed."
            );

            // Test the updated hash
            let modified = $type::random_for_test(&mut rng);
            cache.update(&modified).unwrap();
            assert_eq!(
                cache.tree_hash_root().unwrap().to_vec(),
                modified.tree_hash_root(),
                "Modification hash failed"
            );

            // Produce a new cache for the modified object and compare it to the updated cache.
            let mut modified_cache = cached_tree_hash::TreeHashCache::new(&modified, 0).unwrap();

            // Reset the caches.
            cache.reset_modifications();
            modified_cache.reset_modifications();

            // Ensure the modified cache is the same as a newly created cache. This is a sanity
            // check to make sure there are no artifacts of the original cache remaining after an
            // update.
            assert_eq!(
                modified_cache, cache,
                "The modified cache does not match a new cache."
            )
        }
    };
}
<|endoftext|>


<commit_before>use cudart::ffi;

#[allow(dead_code)]
pub struct Memory<T> {
    pub data: *mut ::libc::c_void,
    size: usize,
    dummy: [T; 0]
}

impl<T> Memory<T> {
    pub fn new(size: usize) -> Result<Memory<T>, &'static str> {
        let mut data: *mut ::libc::c_void = ::std::ptr::null_mut();
        match unsafe { ffi::cudaMalloc(&mut data, (size * ::std::mem::size_of::<T>()) as u64) } {
            ffi::Error::Success => Ok(Memory::<T> {
                data: data,
                size: size,
                dummy: []
            }),
            e => Err(e.to_str())
        }
    }

    pub fn write(&self, data: &Vec<T>) -> Result<(), &'static str> {
        match unsafe { ffi::cudaMemcpy(self.data,
                                       data.as_ptr() as *mut ::libc::c_void,
                                       (data.len() * ::std::mem::size_of::<T>()) as u64,
                                       ffi::MemcpyKind::HostToDevice) } {
            ffi::Error::Success => Ok(()),
            e => Err(e.to_str())
        }
    }

    pub fn read(&self, data: &Vec<T>) -> Result<(), &'static str> {
        match unsafe { ffi::cudaMemcpy(data.as_ptr() as *mut ::libc::c_void,
                                       self.data,
                                       (data.len() * ::std::mem::size_of::<T>()) as u64,
                                       ffi::MemcpyKind::DeviceToHost) } {
            ffi::Error::Success => Ok(()),
            e => Err(e.to_str())
        }
    }
}
<commit_msg>cleanup memory codes<commit_after>use cudart::ffi;
use libc::c_void;
use std::ptr::null_mut;
use std::mem::size_of;

#[allow(dead_code)]
pub struct Memory<T> {
    pub data: *mut c_void,
    size: usize,
    dummy: [T; 0]
}

impl<T> Memory<T> {
    pub fn new(size: usize) -> Result<Memory<T>, &'static str> {
        let mut data: *mut c_void = null_mut();
        match unsafe { ffi::cudaMalloc(&mut data, (size * size_of::<T>()) as u64) } {
            ffi::Error::Success => Ok(Memory::<T> {
                data: data,
                size: size,
                dummy: []
            }),
            e => Err(e.to_str())
        }
    }

    pub fn write(&self, data: &Vec<T>) -> Result<(), &'static str> {
        match unsafe { ffi::cudaMemcpy(self.data,
                                       data.as_ptr() as *mut c_void,
                                       (data.len() * size_of::<T>()) as u64,
                                       ffi::MemcpyKind::HostToDevice) } {
            ffi::Error::Success => Ok(()),
            e => Err(e.to_str())
        }
    }

    pub fn read(&self, data: &Vec<T>) -> Result<(), &'static str> {
        match unsafe { ffi::cudaMemcpy(data.as_ptr() as *mut c_void,
                                       self.data,
                                       (data.len() * size_of::<T>()) as u64,
                                       ffi::MemcpyKind::DeviceToHost) } {
            ffi::Error::Success => Ok(()),
            e => Err(e.to_str())
        }
    }
}
<|endoftext|>


<commit_before>//! Support for sending and receiving data link layer packets using libpcap.
//! Also has support for reading pcap files.
extern crate pcap;

use std::marker::{Send, Sync};
use std::io;
use std::iter::repeat;
use self::pcap::{Active, Activated};
use datalink::{self, NetworkInterface};
use datalink::{EthernetDataLinkChannelIterator, EthernetDataLinkReceiver, EthernetDataLinkSender};
use datalink::Channel::Ethernet;
use packet::Packet;
use packet::ethernet::{EthernetPacket, MutableEthernetPacket};
use std::sync::{Arc, Mutex};
use std::time::Duration;
use std::path::Path;

/// Configuration for the pcap datalink backend
#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)]
pub struct Config {
    /// The size of buffer to use when reading packets. Defaults to 4096
    pub read_buffer_size: usize,

    /// The read timeout. Defaults to None.
    pub read_timeout: Option<Duration>,
}

impl<'a> From<&'a datalink::Config> for Config {
    fn from(config: &datalink::Config) -> Config {
        Config{
            read_buffer_size: config.read_buffer_size,
            read_timeout: config.read_timeout,
        }
    }
}

impl Default for Config {
    fn default() -> Config {
        Config{
            read_buffer_size: 4096,
            read_timeout: None,
        }
    }
}

/// Create a datalink channel from the provided pcap device
#[inline]
pub fn channel(network_interface: &NetworkInterface,
               config: Config) -> io::Result<datalink::Channel> {
    let cap = match pcap::Capture::from_device(&*network_interface.name) {
        Ok(cap) => cap,
        Err(e) => return Err(io::Error::new(io::ErrorKind::Other, e)),
    }.buffer_size(config.read_buffer_size as i32);
    let cap = match config.read_timeout {
        Some(to) => cap.timeout(
            (to.as_secs() * 1000 + (to.subsec_nanos() / 1000) as u64) as i32
        ),
        None => cap
    };
    let cap = match cap.open() {
        Ok(cap) => cap,
        Err(e) => return Err(io::Error::new(io::ErrorKind::Other, e)),
    };
    let cap = Arc::new(Mutex::new(cap));
    Ok(Ethernet(
        Box::new(DataLinkSenderImpl {
            capture: cap.clone(),
        }),
        Box::new(DataLinkReceiverImpl {
            capture: cap.clone(),
            read_buffer: repeat(0u8).take(config.read_buffer_size).collect(),
        })
    ))
}

/// Create a datalink channel from a pcap file
#[inline]
pub fn from_file<P: AsRef<Path>>(path: P, config: Config) -> io::Result<datalink::Channel> {
    let cap = match pcap::Capture::from_file(path) {
        Ok(cap) => cap,
        Err(e) => return Err(io::Error::new(io::ErrorKind::Other, e)),
    };
    let cap = Arc::new(Mutex::new(cap));
    Ok(Ethernet(
        Box::new(InvalidDataLinkSenderImpl {}),
        Box::new(DataLinkReceiverImpl {
            capture: cap.clone(),
            read_buffer: repeat(0u8).take(config.read_buffer_size).collect(),
        })
    ))
}

struct DataLinkSenderImpl {
    capture: Arc<Mutex<pcap::Capture<Active>>>,
}

impl EthernetDataLinkSender for DataLinkSenderImpl {
    #[inline]
    fn build_and_send(&mut self,
                      num_packets: usize,
                      packet_size: usize,
                      func: &mut FnMut(MutableEthernetPacket))
                      -> Option<io::Result<()>> {
        for _ in 0..num_packets {
            let mut data = vec![0; packet_size];
            {
                let eh = MutableEthernetPacket::new(&mut data).unwrap();
                func(eh);
            }
            let mut cap = self.capture.lock().unwrap();
            if let Err(e) = cap.sendpacket(&data) {
                return Some(Err(io::Error::new(io::ErrorKind::Other, e)))
            }
        }
        Some(Ok(()))
    }

    #[inline]
    fn send_to(&mut self,
               packet: &EthernetPacket,
               _dst: Option<NetworkInterface>) -> Option<io::Result<()>> {
        let mut cap = self.capture.lock().unwrap();
        Some(match cap.sendpacket(packet.packet()) {
            Ok(()) => Ok(()),
            Err(e) => Err(io::Error::new(io::ErrorKind::Other, e)),
        })
    }
}

struct InvalidDataLinkSenderImpl {}

impl EthernetDataLinkSender for InvalidDataLinkSenderImpl {
    #[inline]
    fn build_and_send(&mut self,
                      _num_packets: usize,
                      _packet_size: usize,
                      _func: &mut FnMut(MutableEthernetPacket))
                      -> Option<io::Result<()>> {
        None
    }

    #[inline]
    fn send_to(&mut self,
               _packet: &EthernetPacket,
               _dst: Option<NetworkInterface>) -> Option<io::Result<()>> {
        None
    }
}

struct DataLinkReceiverImpl<T: Activated + Send + Sync> {
    capture: Arc<Mutex<pcap::Capture<T>>>,
    read_buffer: Vec<u8>,
}

impl <T: Activated + Send + Sync> EthernetDataLinkReceiver for DataLinkReceiverImpl<T> {
    fn iter<'a>(&'a mut self) -> Box<EthernetDataLinkChannelIterator + 'a> {
        Box::new(DataLinkChannelIteratorImpl { pc: self })
    }
}

struct DataLinkChannelIteratorImpl<'a, T: Activated + Send + Sync + 'a> {
    pc: &'a mut DataLinkReceiverImpl<T>,
}

impl<'a, T: Activated + Send + Sync> EthernetDataLinkChannelIterator<'a> for DataLinkChannelIteratorImpl<'a, T> {
    fn next(&mut self) -> io::Result<EthernetPacket> {
        let mut cap = self.pc.capture.lock().unwrap();
        match cap.next() {
            Ok(pkt) => {
                self.pc.read_buffer.truncate(0);
                self.pc.read_buffer.extend(pkt.data);
            },
            Err(e) => return Err(io::Error::new(io::ErrorKind::Other, e)),
        };
        Ok(EthernetPacket::new(&self.pc.read_buffer).unwrap())
    }
}

/// Get a list of available network interfaces for the current machine.
pub fn interfaces() -> Vec<NetworkInterface> {
    if let Ok(devices) = pcap::Device::list() {
        devices.iter().enumerate().map(|(i, dev)| {
            NetworkInterface {
                name: dev.name.clone(),
                index: i as u32,
                mac: None,
                ips: Vec::new(),
                flags: 0,
            }
        }).collect()
    } else {
        vec![]
    }
}
<commit_msg>Update pcap datalink to pcap API changes introduced in 0.6.0.<commit_after>//! Support for sending and receiving data link layer packets using libpcap.
//! Also has support for reading pcap files.
extern crate pcap;

use std::marker::{Send, Sync};
use std::io;
use std::iter::repeat;
use self::pcap::{Active, Activated};
use datalink::{self, NetworkInterface};
use datalink::{EthernetDataLinkChannelIterator, EthernetDataLinkReceiver, EthernetDataLinkSender};
use datalink::Channel::Ethernet;
use packet::Packet;
use packet::ethernet::{EthernetPacket, MutableEthernetPacket};
use std::sync::{Arc, Mutex};
use std::time::Duration;
use std::path::Path;

/// Configuration for the pcap datalink backend
#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)]
pub struct Config {
    /// The size of buffer to use when reading packets. Defaults to 4096
    pub read_buffer_size: usize,

    /// The read timeout. Defaults to None.
    pub read_timeout: Option<Duration>,
}

impl<'a> From<&'a datalink::Config> for Config {
    fn from(config: &datalink::Config) -> Config {
        Config{
            read_buffer_size: config.read_buffer_size,
            read_timeout: config.read_timeout,
        }
    }
}

impl Default for Config {
    fn default() -> Config {
        Config{
            read_buffer_size: 4096,
            read_timeout: None,
        }
    }
}

/// Create a datalink channel from the provided pcap device
#[inline]
pub fn channel(network_interface: &NetworkInterface,
               config: Config) -> io::Result<datalink::Channel> {
    let cap = match pcap::Capture::from_device(&*network_interface.name) {
        Ok(cap) => cap,
        Err(e) => return Err(io::Error::new(io::ErrorKind::Other, e)),
    }.buffer_size(config.read_buffer_size as i32);
    let cap = match config.read_timeout {
        Some(to) => cap.timeout(
            (to.as_secs() * 1000 + (to.subsec_nanos() / 1000) as u64) as i32
        ),
        None => cap
    };
    let cap = match cap.open() {
        Ok(cap) => cap,
        Err(e) => return Err(io::Error::new(io::ErrorKind::Other, e)),
    };
    let cap = Arc::new(Mutex::new(cap));
    Ok(Ethernet(
        Box::new(DataLinkSenderImpl {
            capture: cap.clone(),
        }),
        Box::new(DataLinkReceiverImpl {
            capture: cap.clone(),
            read_buffer: repeat(0u8).take(config.read_buffer_size).collect(),
        })
    ))
}

/// Create a datalink channel from a pcap file
#[inline]
pub fn from_file<P: AsRef<Path>>(path: P, config: Config) -> io::Result<datalink::Channel> {
    let cap = match pcap::Capture::from_file(path) {
        Ok(cap) => cap,
        Err(e) => return Err(io::Error::new(io::ErrorKind::Other, e)),
    };
    let cap = Arc::new(Mutex::new(cap));
    Ok(Ethernet(
        Box::new(InvalidDataLinkSenderImpl {}),
        Box::new(DataLinkReceiverImpl {
            capture: cap.clone(),
            read_buffer: repeat(0u8).take(config.read_buffer_size).collect(),
        })
    ))
}

struct DataLinkSenderImpl {
    capture: Arc<Mutex<pcap::Capture<Active>>>,
}

impl EthernetDataLinkSender for DataLinkSenderImpl {
    #[inline]
    fn build_and_send(&mut self,
                      num_packets: usize,
                      packet_size: usize,
                      func: &mut FnMut(MutableEthernetPacket))
                      -> Option<io::Result<()>> {
        for _ in 0..num_packets {
            let mut data = vec![0; packet_size];
            {
                let eh = MutableEthernetPacket::new(&mut data).unwrap();
                func(eh);
            }
            let mut cap = self.capture.lock().unwrap();
            if let Err(e) = cap.sendpacket(data) {
                return Some(Err(io::Error::new(io::ErrorKind::Other, e)))
            }
        }
        Some(Ok(()))
    }

    #[inline]
    fn send_to(&mut self,
               packet: &EthernetPacket,
               _dst: Option<NetworkInterface>) -> Option<io::Result<()>> {
        let mut cap = self.capture.lock().unwrap();
        Some(match cap.sendpacket(packet.packet()) {
            Ok(()) => Ok(()),
            Err(e) => Err(io::Error::new(io::ErrorKind::Other, e)),
        })
    }
}

struct InvalidDataLinkSenderImpl {}

impl EthernetDataLinkSender for InvalidDataLinkSenderImpl {
    #[inline]
    fn build_and_send(&mut self,
                      _num_packets: usize,
                      _packet_size: usize,
                      _func: &mut FnMut(MutableEthernetPacket))
                      -> Option<io::Result<()>> {
        None
    }

    #[inline]
    fn send_to(&mut self,
               _packet: &EthernetPacket,
               _dst: Option<NetworkInterface>) -> Option<io::Result<()>> {
        None
    }
}

struct DataLinkReceiverImpl<T: Activated + Send + Sync> {
    capture: Arc<Mutex<pcap::Capture<T>>>,
    read_buffer: Vec<u8>,
}

impl <T: Activated + Send + Sync> EthernetDataLinkReceiver for DataLinkReceiverImpl<T> {
    fn iter<'a>(&'a mut self) -> Box<EthernetDataLinkChannelIterator + 'a> {
        Box::new(DataLinkChannelIteratorImpl { pc: self })
    }
}

struct DataLinkChannelIteratorImpl<'a, T: Activated + Send + Sync + 'a> {
    pc: &'a mut DataLinkReceiverImpl<T>,
}

impl<'a, T: Activated + Send + Sync> EthernetDataLinkChannelIterator<'a> for DataLinkChannelIteratorImpl<'a, T> {
    fn next(&mut self) -> io::Result<EthernetPacket> {
        let mut cap = self.pc.capture.lock().unwrap();
        match cap.next() {
            Ok(pkt) => {
                self.pc.read_buffer.truncate(0);
                self.pc.read_buffer.extend(pkt.data);
            },
            Err(e) => return Err(io::Error::new(io::ErrorKind::Other, e)),
        };
        Ok(EthernetPacket::new(&self.pc.read_buffer).unwrap())
    }
}

/// Get a list of available network interfaces for the current machine.
pub fn interfaces() -> Vec<NetworkInterface> {
    if let Ok(devices) = pcap::Device::list() {
        devices.iter().enumerate().map(|(i, dev)| {
            NetworkInterface {
                name: dev.name.clone(),
                index: i as u32,
                mac: None,
                ips: Vec::new(),
                flags: 0,
            }
        }).collect()
    } else {
        vec![]
    }
}
<|endoftext|>


<commit_before>use super::merge_policy::{MergeCandidate, MergePolicy};
use crate::core::SegmentMeta;
use std::cmp;
use std::f64;

const DEFAULT_LEVEL_LOG_SIZE: f64 = 0.75;
const DEFAULT_MIN_LAYER_SIZE: u32 = 10_000;
const DEFAULT_MIN_MERGE_SIZE: usize = 8;
const DEFAULT_MAX_MERGE_SIZE: usize = 10_000_000;

/// `LogMergePolicy` tries tries to merge segments that have a similar number of
/// documents.
#[derive(Debug, Clone)]
pub struct LogMergePolicy {
    min_merge_size: usize,
    max_merge_size: usize,
    min_layer_size: u32,
    level_log_size: f64,
}

impl LogMergePolicy {
    fn clip_min_size(&self, size: u32) -> u32 {
        cmp::max(self.min_layer_size, size)
    }

    /// Set the minimum number of segment that may be merge together.
    pub fn set_min_merge_size(&mut self, min_merge_size: usize) {
        self.min_merge_size = min_merge_size;
    }

    /// Set the maximum number docs in a segment for it to be considered for
    /// merging.
    pub fn set_max_merge_size(&mut self, max_merge_size: usize) {
        self.max_merge_size = max_merge_size;
    }

    /// Set the minimum segment size under which all segment belong
    /// to the same level.
    pub fn set_min_layer_size(&mut self, min_layer_size: u32) {
        self.min_layer_size = min_layer_size;
    }

    /// Set the ratio between two consecutive levels.
    ///
    /// Segment are group in levels according to their sizes.
    /// These levels are defined as intervals of exponentially growing sizes.
    /// level_log_size define the factor by which one should multiply the limit
    /// to reach a level, in order to get the limit to reach the following
    /// level.
    pub fn set_level_log_size(&mut self, level_log_size: f64) {
        self.level_log_size = level_log_size;
    }
}

impl MergePolicy for LogMergePolicy {
    fn compute_merge_candidates(&self, segments: &[SegmentMeta]) -> Vec<MergeCandidate> {
        if segments.is_empty() {
            return Vec::new();
        }

        let mut size_sorted_tuples = segments
            .iter()
            .map(SegmentMeta::num_docs)
            .filter(|s| s <= &(self.max_merge_size as u32))
            .enumerate()
            .collect::<Vec<(usize, u32)>>();

        size_sorted_tuples.sort_by(|x, y| y.1.cmp(&(x.1)));

        let size_sorted_log_tuples: Vec<_> = size_sorted_tuples
            .into_iter()
            .map(|(ind, num_docs)| (ind, f64::from(self.clip_min_size(num_docs)).log2()))
            .collect();

        let (first_ind, first_score) = size_sorted_log_tuples[0];
        let mut current_max_log_size = first_score;
        let mut levels = vec![vec![first_ind]];
        for &(ind, score) in (&size_sorted_log_tuples).iter().skip(1) {
            if score < (current_max_log_size - self.level_log_size) {
                current_max_log_size = score;
                levels.push(Vec::new());
            }
            levels.last_mut().unwrap().push(ind);
        }

        levels
            .iter()
            .filter(|level| level.len() >= self.min_merge_size)
            .map(|ind_vec| MergeCandidate(ind_vec.iter().map(|&ind| segments[ind].id()).collect()))
            .collect()
    }
}

impl Default for LogMergePolicy {
    fn default() -> LogMergePolicy {
        LogMergePolicy {
            min_merge_size: DEFAULT_MIN_MERGE_SIZE,
            max_merge_size: DEFAULT_MAX_MERGE_SIZE,
            min_layer_size: DEFAULT_MIN_LAYER_SIZE,
            level_log_size: DEFAULT_LEVEL_LOG_SIZE,
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::core::{SegmentId, SegmentMeta, SegmentMetaInventory};
    use crate::indexer::merge_policy::MergePolicy;
    use once_cell::sync::Lazy;

    static INVENTORY: Lazy<SegmentMetaInventory> = Lazy::new(SegmentMetaInventory::default);

    fn test_merge_policy() -> LogMergePolicy {
        let mut log_merge_policy = LogMergePolicy::default();
        log_merge_policy.set_min_merge_size(3);
        log_merge_policy.set_max_merge_size(100_000);
        log_merge_policy.set_min_layer_size(2);
        log_merge_policy
    }

    #[test]
    fn test_log_merge_policy_empty() {
        let y = Vec::new();
        let result_list = test_merge_policy().compute_merge_candidates(&y);
        assert!(result_list.is_empty());
    }

    fn create_random_segment_meta(num_docs: u32) -> SegmentMeta {
        INVENTORY.new_segment_meta(SegmentId::generate_random(), num_docs)
    }

    #[test]
    fn test_log_merge_policy_pair() {
        let test_input = vec![
            create_random_segment_meta(10),
            create_random_segment_meta(10),
            create_random_segment_meta(10),
        ];
        let result_list = test_merge_policy().compute_merge_candidates(&test_input);
        assert_eq!(result_list.len(), 1);
    }

    #[test]
    fn test_log_merge_policy_levels() {
        // multiple levels all get merged correctly
        // 2 MergeCandidates expected:
        // * one with the 6 * 10-docs segments
        // * one with the 3 * 1000-docs segments
        // no MergeCandidate expected for the 2 * 10_000-docs segments as min_merge_size=3
        let test_input = vec![
            create_random_segment_meta(10),
            create_random_segment_meta(10),
            create_random_segment_meta(10),
            create_random_segment_meta(1_000),
            create_random_segment_meta(1_000),
            create_random_segment_meta(1_000),
            create_random_segment_meta(10_000),
            create_random_segment_meta(10_000),
            create_random_segment_meta(10),
            create_random_segment_meta(10),
            create_random_segment_meta(10),
        ];
        let result_list = test_merge_policy().compute_merge_candidates(&test_input);
        assert_eq!(result_list.len(), 2);
    }

    #[test]
    fn test_log_merge_policy_within_levels() {
        // multiple levels all get merged correctly
        let test_input = vec![
            create_random_segment_meta(10),   // log2(10) = ~3.32 (> 3.58 - 0.75)
            create_random_segment_meta(11),   // log2(11) = ~3.46
            create_random_segment_meta(12),   // log2(12) = ~3.58
            create_random_segment_meta(800),  // log2(800) = ~9.64 (> 9.97 - 0.75)
            create_random_segment_meta(1000), // log2(1000) = ~9.97
            create_random_segment_meta(1000),
        ]; // log2(1000) = ~9.97
        let result_list = test_merge_policy().compute_merge_candidates(&test_input);
        assert_eq!(result_list.len(), 2);
    }
    #[test]
    fn test_log_merge_policy_small_segments() {
        // segments under min_layer_size are merged together
        let test_input = vec![
            create_random_segment_meta(1),
            create_random_segment_meta(1),
            create_random_segment_meta(1),
            create_random_segment_meta(2),
            create_random_segment_meta(2),
            create_random_segment_meta(2),
        ];
        let result_list = test_merge_policy().compute_merge_candidates(&test_input);
        assert_eq!(result_list.len(), 1);
    }

    #[test]
    fn test_large_merge_segments() {
        let test_input = vec![
            create_random_segment_meta(1_000_000),
            create_random_segment_meta(100_001),
            create_random_segment_meta(100_000),
            create_random_segment_meta(100_000),
            create_random_segment_meta(100_000),
        ];
        let result_list = test_merge_policy().compute_merge_candidates(&test_input);
        // Do not include large segments
        assert_eq!(result_list.len(), 1);
        assert_eq!(result_list[0].0.len(), 3)
    }
}
<commit_msg>Closes #828. (#829)<commit_after>use super::merge_policy::{MergeCandidate, MergePolicy};
use crate::core::SegmentMeta;
use std::cmp;
use std::f64;

const DEFAULT_LEVEL_LOG_SIZE: f64 = 0.75;
const DEFAULT_MIN_LAYER_SIZE: u32 = 10_000;
const DEFAULT_MIN_MERGE_SIZE: usize = 8;
const DEFAULT_MAX_MERGE_SIZE: usize = 10_000_000;

/// `LogMergePolicy` tries tries to merge segments that have a similar number of
/// documents.
#[derive(Debug, Clone)]
pub struct LogMergePolicy {
    min_merge_size: usize,
    max_merge_size: usize,
    min_layer_size: u32,
    level_log_size: f64,
}

impl LogMergePolicy {
    fn clip_min_size(&self, size: u32) -> u32 {
        cmp::max(self.min_layer_size, size)
    }

    /// Set the minimum number of segment that may be merge together.
    pub fn set_min_merge_size(&mut self, min_merge_size: usize) {
        self.min_merge_size = min_merge_size;
    }

    /// Set the maximum number docs in a segment for it to be considered for
    /// merging.
    pub fn set_max_merge_size(&mut self, max_merge_size: usize) {
        self.max_merge_size = max_merge_size;
    }

    /// Set the minimum segment size under which all segment belong
    /// to the same level.
    pub fn set_min_layer_size(&mut self, min_layer_size: u32) {
        self.min_layer_size = min_layer_size;
    }

    /// Set the ratio between two consecutive levels.
    ///
    /// Segment are group in levels according to their sizes.
    /// These levels are defined as intervals of exponentially growing sizes.
    /// level_log_size define the factor by which one should multiply the limit
    /// to reach a level, in order to get the limit to reach the following
    /// level.
    pub fn set_level_log_size(&mut self, level_log_size: f64) {
        self.level_log_size = level_log_size;
    }
}

impl MergePolicy for LogMergePolicy {
    fn compute_merge_candidates(&self, segments: &[SegmentMeta]) -> Vec<MergeCandidate> {
        let mut size_sorted_tuples = segments
            .iter()
            .map(SegmentMeta::num_docs)
            .filter(|s| s <= &(self.max_merge_size as u32))
            .enumerate()
            .collect::<Vec<(usize, u32)>>();

        size_sorted_tuples.sort_by(|x, y| y.1.cmp(&(x.1)));

        if size_sorted_tuples.len() <= 1 {
            return Vec::new();
        }

        let size_sorted_log_tuples: Vec<_> = size_sorted_tuples
            .into_iter()
            .map(|(ind, num_docs)| (ind, f64::from(self.clip_min_size(num_docs)).log2()))
            .collect();

        if let Some(&(first_ind, first_score)) = size_sorted_log_tuples.first() {
            let mut current_max_log_size = first_score;
            let mut levels = vec![vec![first_ind]];
            for &(ind, score) in (&size_sorted_log_tuples).iter().skip(1) {
                if score < (current_max_log_size - self.level_log_size) {
                    current_max_log_size = score;
                    levels.push(Vec::new());
                }
                levels.last_mut().unwrap().push(ind);
            }
            levels
                .iter()
                .filter(|level| level.len() >= self.min_merge_size)
                .map(|ind_vec| {
                    MergeCandidate(ind_vec.iter().map(|&ind| segments[ind].id()).collect())
                })
                .collect()
        } else {
            return vec![];
        }
    }
}

impl Default for LogMergePolicy {
    fn default() -> LogMergePolicy {
        LogMergePolicy {
            min_merge_size: DEFAULT_MIN_MERGE_SIZE,
            max_merge_size: DEFAULT_MAX_MERGE_SIZE,
            min_layer_size: DEFAULT_MIN_LAYER_SIZE,
            level_log_size: DEFAULT_LEVEL_LOG_SIZE,
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::core::{SegmentId, SegmentMeta, SegmentMetaInventory};
    use crate::indexer::merge_policy::MergePolicy;
    use once_cell::sync::Lazy;

    static INVENTORY: Lazy<SegmentMetaInventory> = Lazy::new(SegmentMetaInventory::default);

    fn test_merge_policy() -> LogMergePolicy {
        let mut log_merge_policy = LogMergePolicy::default();
        log_merge_policy.set_min_merge_size(3);
        log_merge_policy.set_max_merge_size(100_000);
        log_merge_policy.set_min_layer_size(2);
        log_merge_policy
    }

    #[test]
    fn test_log_merge_policy_empty() {
        let y = Vec::new();
        let result_list = test_merge_policy().compute_merge_candidates(&y);
        assert!(result_list.is_empty());
    }

    fn create_random_segment_meta(num_docs: u32) -> SegmentMeta {
        INVENTORY.new_segment_meta(SegmentId::generate_random(), num_docs)
    }

    #[test]
    fn test_log_merge_policy_pair() {
        let test_input = vec![
            create_random_segment_meta(10),
            create_random_segment_meta(10),
            create_random_segment_meta(10),
        ];
        let result_list = test_merge_policy().compute_merge_candidates(&test_input);
        assert_eq!(result_list.len(), 1);
    }

    #[test]
    fn test_log_merge_policy_levels() {
        // multiple levels all get merged correctly
        // 2 MergeCandidates expected:
        // * one with the 6 * 10-docs segments
        // * one with the 3 * 1000-docs segments
        // no MergeCandidate expected for the 2 * 10_000-docs segments as min_merge_size=3
        let test_input = vec![
            create_random_segment_meta(10),
            create_random_segment_meta(10),
            create_random_segment_meta(10),
            create_random_segment_meta(1_000),
            create_random_segment_meta(1_000),
            create_random_segment_meta(1_000),
            create_random_segment_meta(10_000),
            create_random_segment_meta(10_000),
            create_random_segment_meta(10),
            create_random_segment_meta(10),
            create_random_segment_meta(10),
        ];
        let result_list = test_merge_policy().compute_merge_candidates(&test_input);
        assert_eq!(result_list.len(), 2);
    }

    #[test]
    fn test_log_merge_policy_within_levels() {
        // multiple levels all get merged correctly
        let test_input = vec![
            create_random_segment_meta(10),   // log2(10) = ~3.32 (> 3.58 - 0.75)
            create_random_segment_meta(11),   // log2(11) = ~3.46
            create_random_segment_meta(12),   // log2(12) = ~3.58
            create_random_segment_meta(800),  // log2(800) = ~9.64 (> 9.97 - 0.75)
            create_random_segment_meta(1000), // log2(1000) = ~9.97
            create_random_segment_meta(1000),
        ]; // log2(1000) = ~9.97
        let result_list = test_merge_policy().compute_merge_candidates(&test_input);
        assert_eq!(result_list.len(), 2);
    }

    #[test]
    fn test_log_merge_policy_small_segments() {
        // segments under min_layer_size are merged together
        let test_input = vec![
            create_random_segment_meta(1),
            create_random_segment_meta(1),
            create_random_segment_meta(1),
            create_random_segment_meta(2),
            create_random_segment_meta(2),
            create_random_segment_meta(2),
        ];
        let result_list = test_merge_policy().compute_merge_candidates(&test_input);
        assert_eq!(result_list.len(), 1);
    }

    #[test]
    fn test_log_merge_policy_all_segments_too_large_to_merge() {
        let eight_large_segments: Vec<SegmentMeta> =
            std::iter::repeat_with(|| create_random_segment_meta(100_001))
                .take(8)
                .collect();
        assert!(test_merge_policy()
            .compute_merge_candidates(&eight_large_segments)
            .is_empty());
    }

    #[test]
    fn test_large_merge_segments() {
        let test_input = vec![
            create_random_segment_meta(1_000_000),
            create_random_segment_meta(100_001),
            create_random_segment_meta(100_000),
            create_random_segment_meta(100_000),
            create_random_segment_meta(100_000),
        ];
        let result_list = test_merge_policy().compute_merge_candidates(&test_input);
        // Do not include large segments
        assert_eq!(result_list.len(), 1);
        assert_eq!(result_list[0].0.len(), 3)
    }
}
<|endoftext|>


<commit_before>// Copyright 2015 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

//! This module contains the code to convert from the wacky tcx data
//! structures into the hair. The `builder` is generally ignorant of
//! the tcx etc, and instead goes through the `Cx` for most of its
//! work.
//!

use hair::*;

use rustc_data_structures::indexed_vec::Idx;
use rustc::hir::def_id::{DefId, LOCAL_CRATE};
use rustc::hir::map::blocks::FnLikeNode;
use rustc::middle::region;
use rustc::infer::InferCtxt;
use rustc::ty::subst::Subst;
use rustc::ty::{self, Ty, TyCtxt};
use rustc::ty::subst::{Kind, Substs};
use syntax::ast::{self, LitKind};
use syntax::attr;
use syntax::symbol::Symbol;
use rustc::hir;
use rustc_data_structures::sync::Lrc;
use hair::pattern::parse_float;

#[derive(Clone)]
pub struct Cx<'a, 'gcx: 'a + 'tcx, 'tcx: 'a> {
    tcx: TyCtxt<'a, 'gcx, 'tcx>,
    infcx: &'a InferCtxt<'a, 'gcx, 'tcx>,

    pub root_lint_level: ast::NodeId,
    pub param_env: ty::ParamEnv<'gcx>,

    /// Identity `Substs` for use with const-evaluation.
    pub identity_substs: &'gcx Substs<'gcx>,

    pub region_scope_tree: Lrc<region::ScopeTree>,
    pub tables: &'a ty::TypeckTables<'gcx>,

    /// This is `Constness::Const` if we are compiling a `static`,
    /// `const`, or the body of a `const fn`.
    constness: hir::Constness,

    /// What kind of body is being compiled.
    pub body_owner_kind: hir::BodyOwnerKind,

    /// True if this constant/function needs overflow checks.
    check_overflow: bool,
}

impl<'a, 'gcx, 'tcx> Cx<'a, 'gcx, 'tcx> {
    pub fn new(infcx: &'a InferCtxt<'a, 'gcx, 'tcx>,
               src_id: ast::NodeId) -> Cx<'a, 'gcx, 'tcx> {
        let tcx = infcx.tcx;
        let src_def_id = tcx.hir.local_def_id(src_id);
        let body_owner_kind = tcx.hir.body_owner_kind(src_id);

        let constness = match body_owner_kind {
            hir::BodyOwnerKind::Const |
            hir::BodyOwnerKind::Static(_) => hir::Constness::Const,
            hir::BodyOwnerKind::Fn => {
                let fn_like = FnLikeNode::from_node(infcx.tcx.hir.get(src_id));
                fn_like.map_or(hir::Constness::NotConst, |f| f.constness())
            }
        };

        let attrs = tcx.hir.attrs(src_id);

        // Some functions always have overflow checks enabled,
        // however, they may not get codegen'd, depending on
        // the settings for the crate they are codegened in.
        let mut check_overflow = attr::contains_name(attrs, "rustc_inherit_overflow_checks");

        // Respect -C overflow-checks.
        check_overflow |= tcx.sess.overflow_checks();

        // Constants and const fn's always need overflow checks.
        check_overflow |= constness == hir::Constness::Const;

        let lint_level = lint_level_for_hir_id(tcx, src_id);
        Cx {
            tcx,
            infcx,
            root_lint_level: lint_level,
            param_env: tcx.param_env(src_def_id),
            identity_substs: Substs::identity_for_item(tcx.global_tcx(), src_def_id),
            region_scope_tree: tcx.region_scope_tree(src_def_id),
            tables: tcx.typeck_tables_of(src_def_id),
            constness,
            body_owner_kind,
            check_overflow,
        }
    }

}

impl<'a, 'gcx, 'tcx> Cx<'a, 'gcx, 'tcx> {
    /// Normalizes `ast` into the appropriate `mirror` type.
    pub fn mirror<M: Mirror<'tcx>>(&mut self, ast: M) -> M::Output {
        ast.make_mirror(self)
    }

    pub fn usize_ty(&mut self) -> Ty<'tcx> {
        self.tcx.types.usize
    }

    pub fn usize_literal(&mut self, value: u64) -> Literal<'tcx> {
        Literal::Value {
            value: ty::Const::from_usize(self.tcx, value),
        }
    }

    pub fn bool_ty(&mut self) -> Ty<'tcx> {
        self.tcx.types.bool
    }

    pub fn unit_ty(&mut self) -> Ty<'tcx> {
        self.tcx.mk_nil()
    }

    pub fn true_literal(&mut self) -> Literal<'tcx> {
        Literal::Value {
            value: ty::Const::from_bool(self.tcx, true),
        }
    }

    pub fn false_literal(&mut self) -> Literal<'tcx> {
        Literal::Value {
            value: ty::Const::from_bool(self.tcx, false),
        }
    }

    // FIXME: Combine with rustc_mir::hair::pattern::lit_to_const
    pub fn const_eval_literal(
        &mut self,
        lit: &'tcx ast::LitKind,
        ty: Ty<'tcx>,
        sp: Span,
        neg: bool,
    ) -> Literal<'tcx> {
        trace!("const_eval_literal: {:#?}, {:?}, {:?}, {:?}", lit, ty, sp, neg);

        let parse_float = |num, fty| -> ConstValue<'tcx> {
            parse_float(num, fty, neg).unwrap_or_else(|_| {
                // FIXME(#31407) this is only necessary because float parsing is buggy
                self.tcx.sess.span_fatal(sp, "could not evaluate float literal (see issue #31407)");
            })
        };

        let trunc = |n| {
            let param_ty = self.param_env.and(self.tcx.lift_to_global(&ty).unwrap());
            let bit_width = self.tcx.layout_of(param_ty).unwrap().size.bits();
            trace!("trunc {} with size {} and amt {}", n, bit_width, 128 - bit_width);
            let amt = 128 - bit_width;
            let result = (n << amt) >> amt;
            trace!("trunc result: {}", result);
            ConstValue::Scalar(Scalar::Bits {
                bits: result,
                defined: bit_width as u8,
            })
        };

        use rustc::mir::interpret::*;
        let lit = match *lit {
            LitKind::Str(ref s, _) => {
                let s = s.as_str();
                let id = self.tcx.allocate_bytes(s.as_bytes());
                ConstValue::ScalarPair(
                    Scalar::Ptr(id.into()),
                    Scalar::Bits {
                        bits: s.len() as u128,
                        defined: self.tcx.data_layout.pointer_size.bits() as u8,
                    }
                )
            },
            LitKind::ByteStr(ref data) => {
                let id = self.tcx.allocate_bytes(data);
                ConstValue::Scalar(Scalar::Ptr(id.into()))
            },
            LitKind::Byte(n) => ConstValue::Scalar(Scalar::Bits {
                bits: n as u128,
                defined: 8,
            }),
            LitKind::Int(n, _) if neg => {
                let n = n as i128;
                let n = n.overflowing_neg().0;
                trunc(n as u128)
            },
            LitKind::Int(n, _) => trunc(n),
            LitKind::Float(n, fty) => {
                parse_float(n, fty)
            }
            LitKind::FloatUnsuffixed(n) => {
                let fty = match ty.sty {
                    ty::TyFloat(fty) => fty,
                    _ => bug!()
                };
                parse_float(n, fty)
            }
            LitKind::Bool(b) => ConstValue::Scalar(Scalar::Bits {
                bits: b as u128,
                defined: 8,
            }),
            LitKind::Char(c) => ConstValue::Scalar(Scalar::Bits {
                bits: c as u128,
                defined: 32,
            }),
        };
        Literal::Value {
            value: ty::Const::from_const_value(self.tcx, lit, ty)
        }
    }

    pub fn pattern_from_hir(&mut self, p: &hir::Pat) -> Pattern<'tcx> {
        let tcx = self.tcx.global_tcx();
        let p = match tcx.hir.get(p.id) {
            hir::map::NodePat(p) | hir::map::NodeBinding(p) => p,
            node => bug!("pattern became {:?}", node)
        };
        Pattern::from_hir(tcx,
                          self.param_env.and(self.identity_substs),
                          self.tables(),
                          p)
    }

    pub fn trait_method(&mut self,
                        trait_def_id: DefId,
                        method_name: &str,
                        self_ty: Ty<'tcx>,
                        params: &[Kind<'tcx>])
                        -> (Ty<'tcx>, Literal<'tcx>) {
        let method_name = Symbol::intern(method_name);
        let substs = self.tcx.mk_substs_trait(self_ty, params);
        for item in self.tcx.associated_items(trait_def_id) {
            if item.kind == ty::AssociatedKind::Method && item.name == method_name {
                let method_ty = self.tcx.type_of(item.def_id);
                let method_ty = method_ty.subst(self.tcx, substs);
                return (method_ty,
                        Literal::Value {
                            value: ty::Const::zero_sized(self.tcx, method_ty)
                        });
            }
        }

        bug!("found no method `{}` in `{:?}`", method_name, trait_def_id);
    }

    pub fn all_fields(&mut self, adt_def: &ty::AdtDef, variant_index: usize) -> Vec<Field> {
        (0..adt_def.variants[variant_index].fields.len())
            .map(Field::new)
            .collect()
    }

    pub fn needs_drop(&mut self, ty: Ty<'tcx>) -> bool {
        let (ty, param_env) = self.tcx.lift_to_global(&(ty, self.param_env)).unwrap_or_else(|| {
            bug!("MIR: Cx::needs_drop({:?}, {:?}) got \
                  type with inference types/regions",
                 ty, self.param_env);
        });
        ty.needs_drop(self.tcx.global_tcx(), param_env)
    }

    fn lint_level_of(&self, node_id: ast::NodeId) -> LintLevel {
        let hir_id = self.tcx.hir.definitions().node_to_hir_id(node_id);
        let has_lint_level = self.tcx.dep_graph.with_ignore(|| {
            self.tcx.lint_levels(LOCAL_CRATE).lint_level_set(hir_id).is_some()
        });

        if has_lint_level {
            LintLevel::Explicit(node_id)
        } else {
            LintLevel::Inherited
        }
    }

    pub fn tcx(&self) -> TyCtxt<'a, 'gcx, 'tcx> {
        self.tcx
    }

    pub fn tables(&self) -> &'a ty::TypeckTables<'gcx> {
        self.tables
    }

    pub fn check_overflow(&self) -> bool {
        self.check_overflow
    }

    pub fn type_moves_by_default(&self, ty: Ty<'tcx>, span: Span) -> bool {
        self.infcx.type_moves_by_default(self.param_env, ty, span)
    }
}

fn lint_level_for_hir_id(tcx: TyCtxt, mut id: ast::NodeId) -> ast::NodeId {
    // Right now we insert a `with_ignore` node in the dep graph here to
    // ignore the fact that `lint_levels` below depends on the entire crate.
    // For now this'll prevent false positives of recompiling too much when
    // anything changes.
    //
    // Once red/green incremental compilation lands we should be able to
    // remove this because while the crate changes often the lint level map
    // will change rarely.
    tcx.dep_graph.with_ignore(|| {
        let sets = tcx.lint_levels(LOCAL_CRATE);
        loop {
            let hir_id = tcx.hir.definitions().node_to_hir_id(id);
            if sets.lint_level_set(hir_id).is_some() {
                return id
            }
            let next = tcx.hir.get_parent_node(id);
            if next == id {
                bug!("lint traversal reached the root of the crate");
            }
            id = next;
        }
    })
}

mod block;
mod expr;
mod to_ref;
<commit_msg>Prefer `to_value_with_len` over manual expanison of it<commit_after>// Copyright 2015 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

//! This module contains the code to convert from the wacky tcx data
//! structures into the hair. The `builder` is generally ignorant of
//! the tcx etc, and instead goes through the `Cx` for most of its
//! work.
//!

use hair::*;

use rustc_data_structures::indexed_vec::Idx;
use rustc::hir::def_id::{DefId, LOCAL_CRATE};
use rustc::hir::map::blocks::FnLikeNode;
use rustc::middle::region;
use rustc::infer::InferCtxt;
use rustc::ty::subst::Subst;
use rustc::ty::{self, Ty, TyCtxt};
use rustc::ty::subst::{Kind, Substs};
use syntax::ast::{self, LitKind};
use syntax::attr;
use syntax::symbol::Symbol;
use rustc::hir;
use rustc_data_structures::sync::Lrc;
use hair::pattern::parse_float;

#[derive(Clone)]
pub struct Cx<'a, 'gcx: 'a + 'tcx, 'tcx: 'a> {
    tcx: TyCtxt<'a, 'gcx, 'tcx>,
    infcx: &'a InferCtxt<'a, 'gcx, 'tcx>,

    pub root_lint_level: ast::NodeId,
    pub param_env: ty::ParamEnv<'gcx>,

    /// Identity `Substs` for use with const-evaluation.
    pub identity_substs: &'gcx Substs<'gcx>,

    pub region_scope_tree: Lrc<region::ScopeTree>,
    pub tables: &'a ty::TypeckTables<'gcx>,

    /// This is `Constness::Const` if we are compiling a `static`,
    /// `const`, or the body of a `const fn`.
    constness: hir::Constness,

    /// What kind of body is being compiled.
    pub body_owner_kind: hir::BodyOwnerKind,

    /// True if this constant/function needs overflow checks.
    check_overflow: bool,
}

impl<'a, 'gcx, 'tcx> Cx<'a, 'gcx, 'tcx> {
    pub fn new(infcx: &'a InferCtxt<'a, 'gcx, 'tcx>,
               src_id: ast::NodeId) -> Cx<'a, 'gcx, 'tcx> {
        let tcx = infcx.tcx;
        let src_def_id = tcx.hir.local_def_id(src_id);
        let body_owner_kind = tcx.hir.body_owner_kind(src_id);

        let constness = match body_owner_kind {
            hir::BodyOwnerKind::Const |
            hir::BodyOwnerKind::Static(_) => hir::Constness::Const,
            hir::BodyOwnerKind::Fn => {
                let fn_like = FnLikeNode::from_node(infcx.tcx.hir.get(src_id));
                fn_like.map_or(hir::Constness::NotConst, |f| f.constness())
            }
        };

        let attrs = tcx.hir.attrs(src_id);

        // Some functions always have overflow checks enabled,
        // however, they may not get codegen'd, depending on
        // the settings for the crate they are codegened in.
        let mut check_overflow = attr::contains_name(attrs, "rustc_inherit_overflow_checks");

        // Respect -C overflow-checks.
        check_overflow |= tcx.sess.overflow_checks();

        // Constants and const fn's always need overflow checks.
        check_overflow |= constness == hir::Constness::Const;

        let lint_level = lint_level_for_hir_id(tcx, src_id);
        Cx {
            tcx,
            infcx,
            root_lint_level: lint_level,
            param_env: tcx.param_env(src_def_id),
            identity_substs: Substs::identity_for_item(tcx.global_tcx(), src_def_id),
            region_scope_tree: tcx.region_scope_tree(src_def_id),
            tables: tcx.typeck_tables_of(src_def_id),
            constness,
            body_owner_kind,
            check_overflow,
        }
    }

}

impl<'a, 'gcx, 'tcx> Cx<'a, 'gcx, 'tcx> {
    /// Normalizes `ast` into the appropriate `mirror` type.
    pub fn mirror<M: Mirror<'tcx>>(&mut self, ast: M) -> M::Output {
        ast.make_mirror(self)
    }

    pub fn usize_ty(&mut self) -> Ty<'tcx> {
        self.tcx.types.usize
    }

    pub fn usize_literal(&mut self, value: u64) -> Literal<'tcx> {
        Literal::Value {
            value: ty::Const::from_usize(self.tcx, value),
        }
    }

    pub fn bool_ty(&mut self) -> Ty<'tcx> {
        self.tcx.types.bool
    }

    pub fn unit_ty(&mut self) -> Ty<'tcx> {
        self.tcx.mk_nil()
    }

    pub fn true_literal(&mut self) -> Literal<'tcx> {
        Literal::Value {
            value: ty::Const::from_bool(self.tcx, true),
        }
    }

    pub fn false_literal(&mut self) -> Literal<'tcx> {
        Literal::Value {
            value: ty::Const::from_bool(self.tcx, false),
        }
    }

    // FIXME: Combine with rustc_mir::hair::pattern::lit_to_const
    pub fn const_eval_literal(
        &mut self,
        lit: &'tcx ast::LitKind,
        ty: Ty<'tcx>,
        sp: Span,
        neg: bool,
    ) -> Literal<'tcx> {
        trace!("const_eval_literal: {:#?}, {:?}, {:?}, {:?}", lit, ty, sp, neg);

        let parse_float = |num, fty| -> ConstValue<'tcx> {
            parse_float(num, fty, neg).unwrap_or_else(|_| {
                // FIXME(#31407) this is only necessary because float parsing is buggy
                self.tcx.sess.span_fatal(sp, "could not evaluate float literal (see issue #31407)");
            })
        };

        let trunc = |n| {
            let param_ty = self.param_env.and(self.tcx.lift_to_global(&ty).unwrap());
            let bit_width = self.tcx.layout_of(param_ty).unwrap().size.bits();
            trace!("trunc {} with size {} and amt {}", n, bit_width, 128 - bit_width);
            let amt = 128 - bit_width;
            let result = (n << amt) >> amt;
            trace!("trunc result: {}", result);
            ConstValue::Scalar(Scalar::Bits {
                bits: result,
                defined: bit_width as u8,
            })
        };

        use rustc::mir::interpret::*;
        let lit = match *lit {
            LitKind::Str(ref s, _) => {
                let s = s.as_str();
                let id = self.tcx.allocate_bytes(s.as_bytes());
                let value = Scalar::Ptr(id.into()).to_value_with_len(s.len() as u64, self.tcx);
                ConstValue::from_byval_value(value)
            },
            LitKind::ByteStr(ref data) => {
                let id = self.tcx.allocate_bytes(data);
                ConstValue::Scalar(Scalar::Ptr(id.into()))
            },
            LitKind::Byte(n) => ConstValue::Scalar(Scalar::Bits {
                bits: n as u128,
                defined: 8,
            }),
            LitKind::Int(n, _) if neg => {
                let n = n as i128;
                let n = n.overflowing_neg().0;
                trunc(n as u128)
            },
            LitKind::Int(n, _) => trunc(n),
            LitKind::Float(n, fty) => {
                parse_float(n, fty)
            }
            LitKind::FloatUnsuffixed(n) => {
                let fty = match ty.sty {
                    ty::TyFloat(fty) => fty,
                    _ => bug!()
                };
                parse_float(n, fty)
            }
            LitKind::Bool(b) => ConstValue::Scalar(Scalar::Bits {
                bits: b as u128,
                defined: 8,
            }),
            LitKind::Char(c) => ConstValue::Scalar(Scalar::Bits {
                bits: c as u128,
                defined: 32,
            }),
        };
        Literal::Value {
            value: ty::Const::from_const_value(self.tcx, lit, ty)
        }
    }

    pub fn pattern_from_hir(&mut self, p: &hir::Pat) -> Pattern<'tcx> {
        let tcx = self.tcx.global_tcx();
        let p = match tcx.hir.get(p.id) {
            hir::map::NodePat(p) | hir::map::NodeBinding(p) => p,
            node => bug!("pattern became {:?}", node)
        };
        Pattern::from_hir(tcx,
                          self.param_env.and(self.identity_substs),
                          self.tables(),
                          p)
    }

    pub fn trait_method(&mut self,
                        trait_def_id: DefId,
                        method_name: &str,
                        self_ty: Ty<'tcx>,
                        params: &[Kind<'tcx>])
                        -> (Ty<'tcx>, Literal<'tcx>) {
        let method_name = Symbol::intern(method_name);
        let substs = self.tcx.mk_substs_trait(self_ty, params);
        for item in self.tcx.associated_items(trait_def_id) {
            if item.kind == ty::AssociatedKind::Method && item.name == method_name {
                let method_ty = self.tcx.type_of(item.def_id);
                let method_ty = method_ty.subst(self.tcx, substs);
                return (method_ty,
                        Literal::Value {
                            value: ty::Const::zero_sized(self.tcx, method_ty)
                        });
            }
        }

        bug!("found no method `{}` in `{:?}`", method_name, trait_def_id);
    }

    pub fn all_fields(&mut self, adt_def: &ty::AdtDef, variant_index: usize) -> Vec<Field> {
        (0..adt_def.variants[variant_index].fields.len())
            .map(Field::new)
            .collect()
    }

    pub fn needs_drop(&mut self, ty: Ty<'tcx>) -> bool {
        let (ty, param_env) = self.tcx.lift_to_global(&(ty, self.param_env)).unwrap_or_else(|| {
            bug!("MIR: Cx::needs_drop({:?}, {:?}) got \
                  type with inference types/regions",
                 ty, self.param_env);
        });
        ty.needs_drop(self.tcx.global_tcx(), param_env)
    }

    fn lint_level_of(&self, node_id: ast::NodeId) -> LintLevel {
        let hir_id = self.tcx.hir.definitions().node_to_hir_id(node_id);
        let has_lint_level = self.tcx.dep_graph.with_ignore(|| {
            self.tcx.lint_levels(LOCAL_CRATE).lint_level_set(hir_id).is_some()
        });

        if has_lint_level {
            LintLevel::Explicit(node_id)
        } else {
            LintLevel::Inherited
        }
    }

    pub fn tcx(&self) -> TyCtxt<'a, 'gcx, 'tcx> {
        self.tcx
    }

    pub fn tables(&self) -> &'a ty::TypeckTables<'gcx> {
        self.tables
    }

    pub fn check_overflow(&self) -> bool {
        self.check_overflow
    }

    pub fn type_moves_by_default(&self, ty: Ty<'tcx>, span: Span) -> bool {
        self.infcx.type_moves_by_default(self.param_env, ty, span)
    }
}

fn lint_level_for_hir_id(tcx: TyCtxt, mut id: ast::NodeId) -> ast::NodeId {
    // Right now we insert a `with_ignore` node in the dep graph here to
    // ignore the fact that `lint_levels` below depends on the entire crate.
    // For now this'll prevent false positives of recompiling too much when
    // anything changes.
    //
    // Once red/green incremental compilation lands we should be able to
    // remove this because while the crate changes often the lint level map
    // will change rarely.
    tcx.dep_graph.with_ignore(|| {
        let sets = tcx.lint_levels(LOCAL_CRATE);
        loop {
            let hir_id = tcx.hir.definitions().node_to_hir_id(id);
            if sets.lint_level_set(hir_id).is_some() {
                return id
            }
            let next = tcx.hir.get_parent_node(id);
            if next == id {
                bug!("lint traversal reached the root of the crate");
            }
            id = next;
        }
    })
}

mod block;
mod expr;
mod to_ref;
<|endoftext|>