<commit_before>extern crate serde;
#[macro_use] extern crate serde_derive;

use std::collections::HashMap;
use std::fmt;

#[derive(Clone, Debug, PartialEq, Serialize)]
pub enum AMQPType {
    Boolean,
    ShortShortInt,
    ShortShortUInt,
    ShortInt,
    ShortUInt,
    LongInt,
    LongUInt,
    LongLongInt,
    LongLongUInt,
    Float,
    Double,
    DecimalValue,
    ShortString,
    LongString,
    FieldArray,
    TimeStamp,
    FieldTable,
    Void,
}

impl AMQPType {
    pub fn from_id(id: char) -> AMQPType {
        match id {
            't' => AMQPType::Boolean,
            'b' => AMQPType::ShortShortInt,
            'B' => AMQPType::ShortShortUInt,
            'U' => AMQPType::ShortInt,
            'u' => AMQPType::ShortUInt,
            'I' => AMQPType::LongInt,
            'i' => AMQPType::LongUInt,
            'L' => AMQPType::LongLongInt,
            'l' => AMQPType::LongLongUInt,
            'f' => AMQPType::Float,
            'd' => AMQPType::Double,
            'D' => AMQPType::DecimalValue,
            's' => AMQPType::ShortString,
            'S' => AMQPType::LongString,
            'A' => AMQPType::FieldArray,
            'T' => AMQPType::TimeStamp,
            'F' => AMQPType::FieldTable,
            'V' => AMQPType::Void,
            _ => panic!("Unknown type id {}", id), /*FIXME: don't panic*/
        }
    }

    pub fn get_id(&self) -> char {
        match *self {
            AMQPType::Boolean => 't',
            AMQPType::ShortShortInt => 'b',
            AMQPType::ShortShortUInt => 'B',
            AMQPType::ShortInt => 'U',
            AMQPType::ShortUInt => 'u',
            AMQPType::LongInt => 'I',
            AMQPType::LongUInt => 'i',
            AMQPType::LongLongInt => 'L',
            AMQPType::LongLongUInt => 'l',
            AMQPType::Float => 'f',
            AMQPType::Double => 'd',
            AMQPType::DecimalValue => 'D',
            AMQPType::ShortString => 's',
            AMQPType::LongString => 'S',
            AMQPType::FieldArray => 'A',
            AMQPType::TimeStamp => 'T',
            AMQPType::FieldTable => 'F',
            AMQPType::Void => 'V',
        }
    }

    pub fn to_string(&self) -> String {
        format!("{}", self)
    }
}

impl fmt::Display for AMQPType {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "{:?}", self)
    }
}

pub type Boolean        = bool;
pub type ShortShortInt  = i8;
pub type ShortShortUInt = u8;
pub type ShortInt       = i16;
pub type ShortUInt      = u16;
pub type LongInt        = i32;
pub type LongUInt       = u32;
pub type LongLongInt    = i64;
pub type LongLongUInt   = u64;
pub type Float          = f32;
pub type Double         = f64;
pub type DecimalValue   = i64;    /* FIXME: find out what this is exactly */
pub type ShortString    = String; /* TODO: don't allow size >= 255 */
pub type LongString     = String;
pub type FieldArray     = Vec<AMQPValue>;
pub type TimeStamp      = u64;
pub type FieldTable     = HashMap<String, AMQPValue>;
pub type Void           = ();

#[derive(Clone, Debug, PartialEq, Serialize)]
pub enum AMQPValue {
    Boolean(Boolean),
    ShortShortInt(ShortShortInt),
    ShortShortUInt(ShortShortUInt),
    ShortInt(ShortInt),
    ShortUInt(ShortUInt),
    LongInt(LongInt),
    LongUInt(LongUInt),
    LongLongInt(LongLongInt),
    LongLongUInt(LongLongUInt),
    Float(Float),
    Double(Double),
    DecimalValue(DecimalValue),
    ShortString(ShortString),
    LongString(LongString),
    FieldArray(FieldArray),
    TimeStamp(TimeStamp),
    FieldTable(FieldTable),
    Void(Void),
}
<commit_msg>types: correctly define DecimalValue<commit_after>extern crate serde;
#[macro_use] extern crate serde_derive;

use std::collections::HashMap;
use std::fmt;

#[derive(Clone, Debug, PartialEq, Serialize)]
pub enum AMQPType {
    Boolean,
    ShortShortInt,
    ShortShortUInt,
    ShortInt,
    ShortUInt,
    LongInt,
    LongUInt,
    LongLongInt,
    LongLongUInt,
    Float,
    Double,
    DecimalValue,
    ShortString,
    LongString,
    FieldArray,
    TimeStamp,
    FieldTable,
    Void,
}

impl AMQPType {
    pub fn from_id(id: char) -> AMQPType {
        match id {
            't' => AMQPType::Boolean,
            'b' => AMQPType::ShortShortInt,
            'B' => AMQPType::ShortShortUInt,
            'U' => AMQPType::ShortInt,
            'u' => AMQPType::ShortUInt,
            'I' => AMQPType::LongInt,
            'i' => AMQPType::LongUInt,
            'L' => AMQPType::LongLongInt,
            'l' => AMQPType::LongLongUInt,
            'f' => AMQPType::Float,
            'd' => AMQPType::Double,
            'D' => AMQPType::DecimalValue,
            's' => AMQPType::ShortString,
            'S' => AMQPType::LongString,
            'A' => AMQPType::FieldArray,
            'T' => AMQPType::TimeStamp,
            'F' => AMQPType::FieldTable,
            'V' => AMQPType::Void,
            _ => panic!("Unknown type id {}", id), /*FIXME: don't panic*/
        }
    }

    pub fn get_id(&self) -> char {
        match *self {
            AMQPType::Boolean => 't',
            AMQPType::ShortShortInt => 'b',
            AMQPType::ShortShortUInt => 'B',
            AMQPType::ShortInt => 'U',
            AMQPType::ShortUInt => 'u',
            AMQPType::LongInt => 'I',
            AMQPType::LongUInt => 'i',
            AMQPType::LongLongInt => 'L',
            AMQPType::LongLongUInt => 'l',
            AMQPType::Float => 'f',
            AMQPType::Double => 'd',
            AMQPType::DecimalValue => 'D',
            AMQPType::ShortString => 's',
            AMQPType::LongString => 'S',
            AMQPType::FieldArray => 'A',
            AMQPType::TimeStamp => 'T',
            AMQPType::FieldTable => 'F',
            AMQPType::Void => 'V',
        }
    }

    pub fn to_string(&self) -> String {
        format!("{}", self)
    }
}

impl fmt::Display for AMQPType {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "{:?}", self)
    }
}

pub type Boolean        = bool;
pub type ShortShortInt  = i8;
pub type ShortShortUInt = u8;
pub type ShortInt       = i16;
pub type ShortUInt      = u16;
pub type LongInt        = i32;
pub type LongUInt       = u32;
pub type LongLongInt    = i64;
pub type LongLongUInt   = u64;
pub type Float          = f32;
pub type Double         = f64;
pub type ShortString    = String; /* TODO: don't allow size >= 255 */
pub type LongString     = String;
pub type FieldArray     = Vec<AMQPValue>;
pub type TimeStamp      = u64;
pub type FieldTable     = HashMap<String, AMQPValue>;
pub type Void           = ();

#[derive(Clone, Debug, PartialEq, Serialize)]
pub struct DecimalValue {
    pub scale: u8,
    pub value: LongUInt,
}

#[derive(Clone, Debug, PartialEq, Serialize)]
pub enum AMQPValue {
    Boolean(Boolean),
    ShortShortInt(ShortShortInt),
    ShortShortUInt(ShortShortUInt),
    ShortInt(ShortInt),
    ShortUInt(ShortUInt),
    LongInt(LongInt),
    LongUInt(LongUInt),
    LongLongInt(LongLongInt),
    LongLongUInt(LongLongUInt),
    Float(Float),
    Double(Double),
    DecimalValue(DecimalValue),
    ShortString(ShortString),
    LongString(LongString),
    FieldArray(FieldArray),
    TimeStamp(TimeStamp),
    FieldTable(FieldTable),
    Void(Void),
}
<|endoftext|>


<commit_before>//! Handlers for the server.

use std::collections::BTreeMap;

use rustc_serialize::json::{Json, ToJson};
use iron::prelude::*;
use iron::{status, headers, middleware};
use iron::modifiers::Header;
use router::Router;

use urlencoded;

use ::api;
use ::api::optional::Optional;
use ::datastore;
use ::sensors;

#[derive(Debug)]
struct ErrorResponse {
    reason: String,
}

impl ToJson for ErrorResponse {
    /// Serialize an ErrorResponse object into a proper JSON structure.
    fn to_json(&self) -> Json {
        let mut d = BTreeMap::new();
        d.insert("status".to_string(), "error".to_json());
        d.insert("reason".to_string(), self.reason.to_json());
        Json::Object(d)
    }
}

pub struct ReadHandler {
    status: api::Status,
    datastore: datastore::SafeDataStore,
    sensor_specs: sensors::SafeSensorSpecs,
}

impl ReadHandler {
    pub fn new(status: api::Status, datastore: datastore::SafeDataStore, sensor_specs: sensors::SafeSensorSpecs) -> ReadHandler {
        ReadHandler {
            status: status,
            datastore: datastore,
            sensor_specs: sensor_specs,
        }
    }

    fn build_response_json(&self) -> Json {

        // Create a mutable copy of the status struct
        let mut status_copy = self.status.clone();

        // Process registered sensors
        for sensor_spec in self.sensor_specs.lock().unwrap().iter() {
            sensor_spec.get_sensor_value(self.datastore.clone()).map(|value| {
                if status_copy.sensors.is_absent() {
                    status_copy.sensors = Optional::Value(api::Sensors {
                        people_now_present: Optional::Absent,
                        temperature: Optional::Absent,
                    });
                }
                sensor_spec.template.to_sensor(&value, &mut status_copy.sensors.as_mut().unwrap());
            });
        }

        // Update state depending on number of people present
        let people_now_present: Option<i64> = status_copy.sensors.as_ref()
            .and_then(|sensors| sensors.people_now_present.as_ref())
            .map(|people_now_present| people_now_present[0].value)
            .into();
        if let Some(count) = people_now_present {
            status_copy.state.open = Some(count > 0);
            if count == 1 {
                status_copy.state.message = Optional::Value(format!("{} person here right now", count));
            } else if count > 1 {
                status_copy.state.message = Optional::Value(format!("{} people here right now", count));
            }
        }

        // Serialize to JSON
        status_copy.to_json()
    }
}

impl middleware::Handler for ReadHandler {

    /// Return the current status JSON.
    fn handle(&self, req: &mut Request) -> IronResult<Response> {

        println!("{} /{} from {}", req.method, req.url.path[0], req.remote_addr);

        // Get response body
        let body = self.build_response_json().to_string();

        // Create response
        let response = Response::with((status::Ok, body))
            // Set headers
            .set(Header(headers::ContentType("application/json; charset=utf-8".parse().unwrap())))
            .set(Header(headers::CacheControl(vec![headers::CacheDirective::NoCache])))
            .set(Header(headers::AccessControlAllowOrigin::Any));

        Ok(response)
    }

}

pub struct UpdateHandler {
    datastore: datastore::SafeDataStore,
    sensor_specs: sensors::SafeSensorSpecs,
}

error_type! {
    #[derive(Debug)]
    pub enum UpdateHandlerError {
        UnknownSensor(String) {
            desc (sensor) &sensor;
        },
        DataStoreError(datastore::DataStoreError) {
            cause;
        }
    }
}

impl UpdateHandler {
    pub fn new(datastore: datastore::SafeDataStore, sensor_specs: sensors::SafeSensorSpecs) -> UpdateHandler {
        UpdateHandler {
            datastore: datastore,
            sensor_specs: sensor_specs,
        }
    }

    /// Update sensor value in the `DataStore`
    fn update_sensor(&self, sensor: &str, value: &str) -> Result<(), UpdateHandlerError> {
        // Validate sensor
        try!(self.sensor_specs.lock().unwrap().iter()
                              .find(|&spec| spec.data_key == sensor)
                              .ok_or(UpdateHandlerError::UnknownSensor(sensor.to_string())));

        // Store data to datastore
        let datastore_ref = self.datastore.clone();
        let mut datastore_lock = datastore_ref.lock().unwrap();
        datastore_lock.store(sensor, value)
                      .map(|_| ()).map_err(|e| {
                          error!("Could not update sensor value in datastore: {:?}", e);
                          UpdateHandlerError::DataStoreError(e)
                      })
    }

    /// Build an OK response with the `HTTP 204 No Content` status code.
    fn ok_response(&self) -> Response {
        Response::with((status::NoContent))
            // Set headers
            .set(Header(headers::ContentType("application/json; charset=utf-8".parse().unwrap())))
            .set(Header(headers::CacheControl(vec![headers::CacheDirective::NoCache])))
            .set(Header(headers::AccessControlAllowOrigin::Any))
    }

    /// Build an error response with the specified `error_code` and the specified `reason` text.
    fn err_response(&self, error_code: status::Status, reason: &str) -> Response {
        let error = ErrorResponse { reason: reason.to_string() };
        Response::with((error_code, error.to_json().to_string()))
            // Set headers
            .set(Header(headers::ContentType("application/json; charset=utf-8".parse().unwrap())))
            .set(Header(headers::CacheControl(vec![headers::CacheDirective::NoCache])))
            .set(Header(headers::AccessControlAllowOrigin::Any))
    }

}

impl middleware::Handler for UpdateHandler {

    /// Update the sensor, return correct status code.
    fn handle(&self, req: &mut Request) -> IronResult<Response> {

        // TODO: create macro for these println! invocations.
        println!("{} /{} from {}", req.method, req.url.path[0], req.remote_addr);

        // Get sensor name
        let sensor_name;
        {
            // TODO: Properly propagate errors
            let params = req.extensions.get::<Router>().unwrap();
            sensor_name = params.find("sensor").unwrap().to_string();
        }

        // Get sensor value
        let sensor_value;
        {
            let params = req.get_ref::<urlencoded::UrlEncodedBody>().unwrap();
            sensor_value = match params.get("value") {
                Some(ref values) =>  match values.len() {
                    1 => values[0].to_string(),
                    _ => return Ok(self.err_response(status::BadRequest, "Too many values specified")),
                },
                None => return Ok(self.err_response(status::BadRequest, "\"value\" parameter not specified")),
            }
        }

        // Update values in datastore
        if let Err(e) = self.update_sensor(&sensor_name, &sensor_value) {
            error!("update_sensor() failed: {:?}", e);
            let response = match e {
                UpdateHandlerError::UnknownSensor(sensor) =>
                    self.err_response(status::BadRequest, &format!("Unknown sensor: {}", sensor)),
                UpdateHandlerError::DataStoreError(_) =>
                    self.err_response(status::InternalServerError, "Updating values in datastore failed"),
            };
            return Ok(response)
        };

        // Create response
        Ok(self.ok_response())
    }

}
<commit_msg>People present sensor should always contain positive values<commit_after>//! Handlers for the server.

use std::collections::BTreeMap;

use rustc_serialize::json::{Json, ToJson};
use iron::prelude::*;
use iron::{status, headers, middleware};
use iron::modifiers::Header;
use router::Router;

use urlencoded;

use ::api;
use ::api::optional::Optional;
use ::datastore;
use ::sensors;

#[derive(Debug)]
struct ErrorResponse {
    reason: String,
}

impl ToJson for ErrorResponse {
    /// Serialize an ErrorResponse object into a proper JSON structure.
    fn to_json(&self) -> Json {
        let mut d = BTreeMap::new();
        d.insert("status".to_string(), "error".to_json());
        d.insert("reason".to_string(), self.reason.to_json());
        Json::Object(d)
    }
}

pub struct ReadHandler {
    status: api::Status,
    datastore: datastore::SafeDataStore,
    sensor_specs: sensors::SafeSensorSpecs,
}

impl ReadHandler {
    pub fn new(status: api::Status, datastore: datastore::SafeDataStore, sensor_specs: sensors::SafeSensorSpecs) -> ReadHandler {
        ReadHandler {
            status: status,
            datastore: datastore,
            sensor_specs: sensor_specs,
        }
    }

    fn build_response_json(&self) -> Json {

        // Create a mutable copy of the status struct
        let mut status_copy = self.status.clone();

        // Process registered sensors
        for sensor_spec in self.sensor_specs.lock().unwrap().iter() {
            sensor_spec.get_sensor_value(self.datastore.clone()).map(|value| {
                if status_copy.sensors.is_absent() {
                    status_copy.sensors = Optional::Value(api::Sensors {
                        people_now_present: Optional::Absent,
                        temperature: Optional::Absent,
                    });
                }
                sensor_spec.template.to_sensor(&value, &mut status_copy.sensors.as_mut().unwrap());
            });
        }

        // Update state depending on number of people present
        let people_now_present: Option<u64> = status_copy.sensors.as_ref()
            .and_then(|sensors| sensors.people_now_present.as_ref())
            .map(|people_now_present| people_now_present[0].value)
            .into();
        if let Some(count) = people_now_present {
            status_copy.state.open = Some(count > 0);
            if count == 1 {
                status_copy.state.message = Optional::Value(format!("{} person here right now", count));
            } else if count > 1 {
                status_copy.state.message = Optional::Value(format!("{} people here right now", count));
            }
        }

        // Serialize to JSON
        status_copy.to_json()
    }
}

impl middleware::Handler for ReadHandler {

    /// Return the current status JSON.
    fn handle(&self, req: &mut Request) -> IronResult<Response> {

        println!("{} /{} from {}", req.method, req.url.path[0], req.remote_addr);

        // Get response body
        let body = self.build_response_json().to_string();

        // Create response
        let response = Response::with((status::Ok, body))
            // Set headers
            .set(Header(headers::ContentType("application/json; charset=utf-8".parse().unwrap())))
            .set(Header(headers::CacheControl(vec![headers::CacheDirective::NoCache])))
            .set(Header(headers::AccessControlAllowOrigin::Any));

        Ok(response)
    }

}

pub struct UpdateHandler {
    datastore: datastore::SafeDataStore,
    sensor_specs: sensors::SafeSensorSpecs,
}

error_type! {
    #[derive(Debug)]
    pub enum UpdateHandlerError {
        UnknownSensor(String) {
            desc (sensor) &sensor;
        },
        DataStoreError(datastore::DataStoreError) {
            cause;
        }
    }
}

impl UpdateHandler {
    pub fn new(datastore: datastore::SafeDataStore, sensor_specs: sensors::SafeSensorSpecs) -> UpdateHandler {
        UpdateHandler {
            datastore: datastore,
            sensor_specs: sensor_specs,
        }
    }

    /// Update sensor value in the `DataStore`
    fn update_sensor(&self, sensor: &str, value: &str) -> Result<(), UpdateHandlerError> {
        // Validate sensor
        try!(self.sensor_specs.lock().unwrap().iter()
                              .find(|&spec| spec.data_key == sensor)
                              .ok_or(UpdateHandlerError::UnknownSensor(sensor.to_string())));

        // Store data to datastore
        let datastore_ref = self.datastore.clone();
        let mut datastore_lock = datastore_ref.lock().unwrap();
        datastore_lock.store(sensor, value)
                      .map(|_| ()).map_err(|e| {
                          error!("Could not update sensor value in datastore: {:?}", e);
                          UpdateHandlerError::DataStoreError(e)
                      })
    }

    /// Build an OK response with the `HTTP 204 No Content` status code.
    fn ok_response(&self) -> Response {
        Response::with((status::NoContent))
            // Set headers
            .set(Header(headers::ContentType("application/json; charset=utf-8".parse().unwrap())))
            .set(Header(headers::CacheControl(vec![headers::CacheDirective::NoCache])))
            .set(Header(headers::AccessControlAllowOrigin::Any))
    }

    /// Build an error response with the specified `error_code` and the specified `reason` text.
    fn err_response(&self, error_code: status::Status, reason: &str) -> Response {
        let error = ErrorResponse { reason: reason.to_string() };
        Response::with((error_code, error.to_json().to_string()))
            // Set headers
            .set(Header(headers::ContentType("application/json; charset=utf-8".parse().unwrap())))
            .set(Header(headers::CacheControl(vec![headers::CacheDirective::NoCache])))
            .set(Header(headers::AccessControlAllowOrigin::Any))
    }

}

impl middleware::Handler for UpdateHandler {

    /// Update the sensor, return correct status code.
    fn handle(&self, req: &mut Request) -> IronResult<Response> {

        // TODO: create macro for these println! invocations.
        println!("{} /{} from {}", req.method, req.url.path[0], req.remote_addr);

        // Get sensor name
        let sensor_name;
        {
            // TODO: Properly propagate errors
            let params = req.extensions.get::<Router>().unwrap();
            sensor_name = params.find("sensor").unwrap().to_string();
        }

        // Get sensor value
        let sensor_value;
        {
            let params = req.get_ref::<urlencoded::UrlEncodedBody>().unwrap();
            sensor_value = match params.get("value") {
                Some(ref values) =>  match values.len() {
                    1 => values[0].to_string(),
                    _ => return Ok(self.err_response(status::BadRequest, "Too many values specified")),
                },
                None => return Ok(self.err_response(status::BadRequest, "\"value\" parameter not specified")),
            }
        }

        // Update values in datastore
        if let Err(e) = self.update_sensor(&sensor_name, &sensor_value) {
            error!("update_sensor() failed: {:?}", e);
            let response = match e {
                UpdateHandlerError::UnknownSensor(sensor) =>
                    self.err_response(status::BadRequest, &format!("Unknown sensor: {}", sensor)),
                UpdateHandlerError::DataStoreError(_) =>
                    self.err_response(status::InternalServerError, "Updating values in datastore failed"),
            };
            return Ok(response)
        };

        // Create response
        Ok(self.ok_response())
    }

}
<|endoftext|>


<commit_before>//! This module provides a token source (`GetToken`) that obtains tokens for service accounts.
//! Service accounts are usually used by software (i.e., non-human actors) to get access to
//! resources. Currently, this module only works with RS256 JWTs, and uses the Google request URL
//! by default (this can be amended by a patch).
//!
//! Resources:
//! - [Using OAuth 2.0 for Server to Server
//! Applications](https://developers.google.com/identity/protocols/OAuth2ServiceAccount)
//! - [JSON Web Tokens](https://jwt.io/)
//!
//! Copyright (c) 2016 Google Inc (lewinb@google.com).
//!

use std::borrow::BorrowMut;
use std::default::Default;
use std::error;
use std::io::{Read, Write};
use std::result;

use authenticator::GetToken;
use types::{StringError, Token};
use storage::{hash_scopes, MemoryStorage, TokenStorage};

use hyper::header;
use url::form_urlencoded;

use base64;
use chrono;
use hyper;
use openssl;
use serde_json;

const GRANT_TYPE: &'static str = "urn:ietf:params:oauth:grant-type:jwt-bearer";
const TOKEN_REQUEST_URL: &'static str = "https://www.googleapis.com/oauth2/v4/token";
const GOOGLE_RS256_HEAD: &'static str = "{\"alg\":\"RS256\",\"typ\":\"JWT\"}";

// Encodes s as Base64
fn encode_base64<T: AsRef<[u8]>>(s: T) -> String {
    base64::encode_mode(s.as_ref(), base64::Base64Mode::UrlSafe)
}

// Calculates the SHA256 hash.
fn hash_sha256(data: &[u8]) -> Vec<u8> {
    let mut hasher = openssl::crypto::hash::Hasher::new(openssl::crypto::hash::Type::SHA256);
    let _ = hasher.write(data);
    hasher.finish()
}

// Signs the hash with key.
fn sign_rsa(key: &openssl::crypto::rsa::RSA, hash: &[u8]) -> String {
    let signature = key.sign(openssl::crypto::hash::Type::SHA256, hash).unwrap();
    let b64_signature = encode_base64(signature);

    b64_signature
}

// Reads an RSA key from pem_pkcs8 (the format of the 'private_key' field in the service account
// key).
fn decode_rsa_key(pem_pkcs8: &str) -> Result<openssl::crypto::rsa::RSA, Box<error::Error>> {
    let private_key = pem_pkcs8.to_string().replace("\\n", "\n");
    let privkey = openssl::crypto::rsa::RSA::private_key_from_pem(&mut private_key.as_bytes());

    match privkey {
        Err(e) => Err(Box::new(e)),
        Ok(key) => Ok(key),
    }
}

/// JSON schema of secret service account key. You can obtain the key from
/// the Cloud Console at https://console.cloud.google.com/.
///
/// You can use `helpers::service_account_key_from_file()` as a quick way to read a JSON client
/// secret into a ServiceAccountKey.
#[derive(Serialize, Deserialize, Debug, Clone)]
pub struct ServiceAccountKey {
    #[serde(rename="type")]
    pub key_type: Option<String>,
    pub project_id: Option<String>,
    pub private_key_id: Option<String>,
    pub private_key: Option<String>,
    pub client_email: Option<String>,
    pub client_id: Option<String>,
    pub auth_uri: Option<String>,
    pub token_uri: Option<String>,
    pub auth_provier_x509_cert_url: Option<String>,
    pub client_x509_cert_url: Option<String>,
}

#[derive(Serialize, Debug)]
struct Claims {
    iss: String,
    aud: String,
    exp: i64,
    iat: i64,
    sub: Option<String>,
    scope: String,
}

struct JWT {
    header: String,
    claims: Claims,
}

impl JWT {
    fn new(claims: Claims) -> JWT {
        JWT {
            header: GOOGLE_RS256_HEAD.to_string(),
            claims: claims,
        }
    }
    // Encodes the first two parts (header and claims) to base64 and assembles them into a form
    // ready to be signed.
    fn encode_claims(&self) -> String {
        let mut head = encode_base64(&self.header);
        let claims = encode_base64(serde_json::to_string(&self.claims).unwrap());

        head.push_str(".");
        head.push_str(&claims);
        head
    }

    fn sign(&self, private_key: &str) -> Result<String, Box<error::Error>> {
        let mut jwt_head = self.encode_claims();

        let key = try!(decode_rsa_key(private_key));
        let hash = hash_sha256(&jwt_head.as_bytes());
        let signature = sign_rsa(&key, &hash);

        jwt_head.push_str(".");
        jwt_head.push_str(&signature);

        Ok(jwt_head)
    }
}

fn init_claims_from_key<'a, I, T>(key: &ServiceAccountKey, scopes: I) -> Claims
    where T: AsRef<str> + 'a,
          I: IntoIterator<Item = &'a T>
{
    let iat = chrono::UTC::now().timestamp();
    let expiry = iat + 3600 - 5; // Max validity is 1h.

    let mut scopes_string = scopes.into_iter().fold(String::new(), |mut acc, sc| {
        acc.push_str(sc.as_ref());
        acc.push_str(" ");
        acc
    });
    scopes_string.pop();

    Claims {
        iss: key.client_email.clone().unwrap(),
        aud: TOKEN_REQUEST_URL.to_string(),
        exp: expiry,
        iat: iat,
        sub: None,
        scope: scopes_string,
    }
}

/// See "Additional claims" at https://developers.google.com/identity/protocols/OAuth2ServiceAccount
#[allow(dead_code)]
fn set_sub_claim(mut claims: Claims, sub: String) -> Claims {
    claims.sub = Some(sub);
    claims
}

/// A token source (`GetToken`) yielding OAuth tokens for services that use ServiceAccount authorization.
/// This token source caches token and automatically renews expired ones.
pub struct ServiceAccountAccess<C> {
    client: C,
    key: ServiceAccountKey,
    cache: MemoryStorage,
}

/// This is the schema of the server's response.
#[derive(Deserialize, Debug)]
struct TokenResponse {
    access_token: Option<String>,
    token_type: Option<String>,
    expires_in: Option<i64>,
}

impl TokenResponse {
    fn to_oauth_token(self) -> Token {
        let expires_ts = chrono::UTC::now().timestamp() + self.expires_in.unwrap_or(0);

        Token {
            access_token: self.access_token.unwrap(),
            token_type: self.token_type.unwrap(),
            refresh_token: String::new(),
            expires_in: self.expires_in,
            expires_in_timestamp: Some(expires_ts),
        }
    }
}

impl<'a, C> ServiceAccountAccess<C>
    where C: BorrowMut<hyper::Client>
{
    /// Returns a new `ServiceAccountAccess` token source.
    #[allow(dead_code)]
    pub fn new(key: ServiceAccountKey, client: C) -> ServiceAccountAccess<C> {
        ServiceAccountAccess {
            client: client,
            key: key,
            cache: MemoryStorage::default(),
        }
    }

    fn request_token(&mut self, scopes: &Vec<&str>) -> result::Result<Token, Box<error::Error>> {
        let signed = try!(JWT::new(init_claims_from_key(&self.key, scopes))
            .sign(self.key.private_key.as_ref().unwrap()));

        let body = form_urlencoded::serialize(vec![("grant_type".to_string(),
                                                    GRANT_TYPE.to_string()),
                                                   ("assertion".to_string(), signed)]);

        let mut response = String::new();
        let mut result = try!(self.client
            .borrow_mut()
            .post(TOKEN_REQUEST_URL)
            .body(&body)
            .header(header::ContentType("application/x-www-form-urlencoded".parse().unwrap()))
            .send());

        try!(result.read_to_string(&mut response));

        let token: Result<TokenResponse, serde_json::error::Error> =
            serde_json::from_str(&response);

        match token {
            Err(e) => return Err(Box::new(e)),
            Ok(token) => {
                if token.access_token.is_none() || token.token_type.is_none() ||
                   token.expires_in.is_none() {
                    Err(Box::new(StringError::new("Token response lacks fields".to_string(),
                                                  Some(&format!("{:?}", token)))))
                } else {
                    Ok(token.to_oauth_token())
                }
            }
        }
    }
}

impl<C: BorrowMut<hyper::Client>> GetToken for ServiceAccountAccess<C> {
    fn token<'b, I, T>(&mut self, scopes: I) -> result::Result<Token, Box<error::Error>>
        where T: AsRef<str> + Ord + 'b,
              I: IntoIterator<Item = &'b T>
    {
        let (hash, scps) = hash_scopes(scopes);

        if let Some(token) = try!(self.cache.get(hash, &scps)) {
            if !token.expired() {
                return Ok(token);
            }
        }

        let token = try!(self.request_token(&scps));
        let _ = self.cache.set(hash, &scps, Some(token.clone()));

        Ok(token)
    }

    fn api_key(&mut self) -> Option<String> {
        None
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use helper::service_account_key_from_file;
    use hyper;
    use authenticator::GetToken;

    // This is a valid but deactivated key.
    const TEST_PRIVATE_KEY_PATH: &'static str = "examples/Sanguine-69411a0c0eea.json";

    // This only works if you have the right key file; if needed, download your own and test it
    // with that.
    // #[test]
    #[allow(dead_code)]
    fn test_service_account_e2e() {
        let key = service_account_key_from_file(&TEST_PRIVATE_KEY_PATH.to_string()).unwrap();
        let mut acc = ServiceAccountAccess::new(key, hyper::client::Client::new());
        println!("{:?}",
                 acc.token(vec![&"https://www.googleapis.com/auth/pubsub"]).unwrap());
    }

    #[test]
    fn test_jwt_initialize_claims() {
        let key = service_account_key_from_file(&TEST_PRIVATE_KEY_PATH.to_string()).unwrap();
        let scopes = vec!["scope1", "scope2", "scope3"];
        let claims = super::init_claims_from_key(&key, &scopes);

        assert_eq!(claims.iss,
                   "oauth2-public-test@sanguine-rhythm-105020.iam.gserviceaccount.com".to_string());
        assert_eq!(claims.scope, "scope1 scope2 scope3".to_string());
        assert_eq!(claims.aud,
                   "https://www.googleapis.com/oauth2/v4/token".to_string());
        assert!(claims.exp > 1000000000);
        assert!(claims.iat < claims.exp);
        assert_eq!(claims.exp - claims.iat, 3595);
    }

    #[test]
    fn test_jwt_sign() {
        let key = service_account_key_from_file(&TEST_PRIVATE_KEY_PATH.to_string()).unwrap();
        let scopes = vec!["scope1", "scope2", "scope3"];
        let claims = super::init_claims_from_key(&key, &scopes);
        let jwt = super::JWT::new(claims);
        let signature = jwt.sign(key.private_key.as_ref().unwrap());

        assert!(signature.is_ok());

        let signature = signature.unwrap();
        assert_eq!(signature.split(".").nth(0).unwrap(),
                   "eyJhbGciOiJSUzI1NiIsInR5cCI6IkpXVCJ9");
    }
}
<commit_msg>feat(service_account): Enable service account oauth with other providers<commit_after>//! This module provides a token source (`GetToken`) that obtains tokens for service accounts.
//! Service accounts are usually used by software (i.e., non-human actors) to get access to
//! resources. Currently, this module only works with RS256 JWTs, and uses the Google request URL
//! by default (this can be amended by a patch).
//!
//! Resources:
//! - [Using OAuth 2.0 for Server to Server
//! Applications](https://developers.google.com/identity/protocols/OAuth2ServiceAccount)
//! - [JSON Web Tokens](https://jwt.io/)
//!
//! Copyright (c) 2016 Google Inc (lewinb@google.com).
//!

use std::borrow::BorrowMut;
use std::default::Default;
use std::error;
use std::io::{Read, Write};
use std::result;

use authenticator::GetToken;
use types::{StringError, Token};
use storage::{hash_scopes, MemoryStorage, TokenStorage};

use hyper::header;
use url::form_urlencoded;

use base64;
use chrono;
use hyper;
use openssl;
use serde_json;

const GRANT_TYPE: &'static str = "urn:ietf:params:oauth:grant-type:jwt-bearer";
const GOOGLE_RS256_HEAD: &'static str = "{\"alg\":\"RS256\",\"typ\":\"JWT\"}";

// Encodes s as Base64
fn encode_base64<T: AsRef<[u8]>>(s: T) -> String {
    base64::encode_mode(s.as_ref(), base64::Base64Mode::UrlSafe)
}

// Calculates the SHA256 hash.
fn hash_sha256(data: &[u8]) -> Vec<u8> {
    let mut hasher = openssl::crypto::hash::Hasher::new(openssl::crypto::hash::Type::SHA256);
    let _ = hasher.write(data);
    hasher.finish()
}

// Signs the hash with key.
fn sign_rsa(key: &openssl::crypto::rsa::RSA, hash: &[u8]) -> String {
    let signature = key.sign(openssl::crypto::hash::Type::SHA256, hash).unwrap();
    let b64_signature = encode_base64(signature);

    b64_signature
}

// Reads an RSA key from pem_pkcs8 (the format of the 'private_key' field in the service account
// key).
fn decode_rsa_key(pem_pkcs8: &str) -> Result<openssl::crypto::rsa::RSA, Box<error::Error>> {
    let private_key = pem_pkcs8.to_string().replace("\\n", "\n");
    let privkey = openssl::crypto::rsa::RSA::private_key_from_pem(&mut private_key.as_bytes());

    match privkey {
        Err(e) => Err(Box::new(e)),
        Ok(key) => Ok(key),
    }
}

/// JSON schema of secret service account key. You can obtain the key from
/// the Cloud Console at https://console.cloud.google.com/.
///
/// You can use `helpers::service_account_key_from_file()` as a quick way to read a JSON client
/// secret into a ServiceAccountKey.
#[derive(Serialize, Deserialize, Debug, Clone)]
pub struct ServiceAccountKey {
    #[serde(rename="type")]
    pub key_type: Option<String>,
    pub project_id: Option<String>,
    pub private_key_id: Option<String>,
    pub private_key: Option<String>,
    pub client_email: Option<String>,
    pub client_id: Option<String>,
    pub auth_uri: Option<String>,
    pub token_uri: Option<String>,
    pub auth_provier_x509_cert_url: Option<String>,
    pub client_x509_cert_url: Option<String>,
}

#[derive(Serialize, Debug)]
struct Claims {
    iss: String,
    aud: String,
    exp: i64,
    iat: i64,
    sub: Option<String>,
    scope: String,
}

struct JWT {
    header: String,
    claims: Claims,
}

impl JWT {
    fn new(claims: Claims) -> JWT {
        JWT {
            header: GOOGLE_RS256_HEAD.to_string(),
            claims: claims,
        }
    }
    // Encodes the first two parts (header and claims) to base64 and assembles them into a form
    // ready to be signed.
    fn encode_claims(&self) -> String {
        let mut head = encode_base64(&self.header);
        let claims = encode_base64(serde_json::to_string(&self.claims).unwrap());

        head.push_str(".");
        head.push_str(&claims);
        head
    }

    fn sign(&self, private_key: &str) -> Result<String, Box<error::Error>> {
        let mut jwt_head = self.encode_claims();

        let key = try!(decode_rsa_key(private_key));
        let hash = hash_sha256(&jwt_head.as_bytes());
        let signature = sign_rsa(&key, &hash);

        jwt_head.push_str(".");
        jwt_head.push_str(&signature);

        Ok(jwt_head)
    }
}

fn init_claims_from_key<'a, I, T>(key: &ServiceAccountKey, scopes: I) -> Claims
    where T: AsRef<str> + 'a,
          I: IntoIterator<Item = &'a T>
{
    let iat = chrono::UTC::now().timestamp();
    let expiry = iat + 3600 - 5; // Max validity is 1h.

    let mut scopes_string = scopes.into_iter().fold(String::new(), |mut acc, sc| {
        acc.push_str(sc.as_ref());
        acc.push_str(" ");
        acc
    });
    scopes_string.pop();

    Claims {
        iss: key.client_email.clone().unwrap(),
        aud: key.token_uri.clone().unwrap(),
        exp: expiry,
        iat: iat,
        sub: None,
        scope: scopes_string,
    }
}

/// See "Additional claims" at https://developers.google.com/identity/protocols/OAuth2ServiceAccount
#[allow(dead_code)]
fn set_sub_claim(mut claims: Claims, sub: String) -> Claims {
    claims.sub = Some(sub);
    claims
}

/// A token source (`GetToken`) yielding OAuth tokens for services that use ServiceAccount authorization.
/// This token source caches token and automatically renews expired ones.
pub struct ServiceAccountAccess<C> {
    client: C,
    key: ServiceAccountKey,
    cache: MemoryStorage,
}

/// This is the schema of the server's response.
#[derive(Deserialize, Debug)]
struct TokenResponse {
    access_token: Option<String>,
    token_type: Option<String>,
    expires_in: Option<i64>,
}

impl TokenResponse {
    fn to_oauth_token(self) -> Token {
        let expires_ts = chrono::UTC::now().timestamp() + self.expires_in.unwrap_or(0);

        Token {
            access_token: self.access_token.unwrap(),
            token_type: self.token_type.unwrap(),
            refresh_token: String::new(),
            expires_in: self.expires_in,
            expires_in_timestamp: Some(expires_ts),
        }
    }
}

impl<'a, C> ServiceAccountAccess<C>
    where C: BorrowMut<hyper::Client>
{
    /// Returns a new `ServiceAccountAccess` token source.
    #[allow(dead_code)]
    pub fn new(key: ServiceAccountKey, client: C) -> ServiceAccountAccess<C> {
        ServiceAccountAccess {
            client: client,
            key: key,
            cache: MemoryStorage::default(),
        }
    }

    fn request_token(&mut self, scopes: &Vec<&str>) -> result::Result<Token, Box<error::Error>> {
        let signed = try!(JWT::new(init_claims_from_key(&self.key, scopes))
            .sign(self.key.private_key.as_ref().unwrap()));

        let body = form_urlencoded::serialize(vec![("grant_type".to_string(),
                                                    GRANT_TYPE.to_string()),
                                                   ("assertion".to_string(), signed)]);

        let mut response = String::new();
        let mut result = try!(self.client
            .borrow_mut()
            .post(self.key.token_uri.as_ref().unwrap())
            .body(&body)
            .header(header::ContentType("application/x-www-form-urlencoded".parse().unwrap()))
            .send());

        try!(result.read_to_string(&mut response));

        let token: Result<TokenResponse, serde_json::error::Error> =
            serde_json::from_str(&response);

        match token {
            Err(e) => return Err(Box::new(e)),
            Ok(token) => {
                if token.access_token.is_none() || token.token_type.is_none() ||
                   token.expires_in.is_none() {
                    Err(Box::new(StringError::new("Token response lacks fields".to_string(),
                                                  Some(&format!("{:?}", token)))))
                } else {
                    Ok(token.to_oauth_token())
                }
            }
        }
    }
}

impl<C: BorrowMut<hyper::Client>> GetToken for ServiceAccountAccess<C> {
    fn token<'b, I, T>(&mut self, scopes: I) -> result::Result<Token, Box<error::Error>>
        where T: AsRef<str> + Ord + 'b,
              I: IntoIterator<Item = &'b T>
    {
        let (hash, scps) = hash_scopes(scopes);

        if let Some(token) = try!(self.cache.get(hash, &scps)) {
            if !token.expired() {
                return Ok(token);
            }
        }

        let token = try!(self.request_token(&scps));
        let _ = self.cache.set(hash, &scps, Some(token.clone()));

        Ok(token)
    }

    fn api_key(&mut self) -> Option<String> {
        None
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use helper::service_account_key_from_file;
    use hyper;
    use authenticator::GetToken;

    // This is a valid but deactivated key.
    const TEST_PRIVATE_KEY_PATH: &'static str = "examples/Sanguine-69411a0c0eea.json";

    // This only works if you have the right key file; if needed, download your own and test it
    // with that.
    // #[test]
    #[allow(dead_code)]
    fn test_service_account_e2e() {
        let key = service_account_key_from_file(&TEST_PRIVATE_KEY_PATH.to_string()).unwrap();
        let mut acc = ServiceAccountAccess::new(key, hyper::client::Client::new());
        println!("{:?}",
                 acc.token(vec![&"https://www.googleapis.com/auth/pubsub"]).unwrap());
    }

    #[test]
    fn test_jwt_initialize_claims() {
        let key = service_account_key_from_file(&TEST_PRIVATE_KEY_PATH.to_string()).unwrap();
        let scopes = vec!["scope1", "scope2", "scope3"];
        let claims = super::init_claims_from_key(&key, &scopes);

        assert_eq!(claims.iss,
                   "oauth2-public-test@sanguine-rhythm-105020.iam.gserviceaccount.com".to_string());
        assert_eq!(claims.scope, "scope1 scope2 scope3".to_string());
        assert_eq!(claims.aud,
                   "https://accounts.google.com/o/oauth2/token".to_string());
        assert!(claims.exp > 1000000000);
        assert!(claims.iat < claims.exp);
        assert_eq!(claims.exp - claims.iat, 3595);
    }

    #[test]
    fn test_jwt_sign() {
        let key = service_account_key_from_file(&TEST_PRIVATE_KEY_PATH.to_string()).unwrap();
        let scopes = vec!["scope1", "scope2", "scope3"];
        let claims = super::init_claims_from_key(&key, &scopes);
        let jwt = super::JWT::new(claims);
        let signature = jwt.sign(key.private_key.as_ref().unwrap());

        assert!(signature.is_ok());

        let signature = signature.unwrap();
        assert_eq!(signature.split(".").nth(0).unwrap(),
                   "eyJhbGciOiJSUzI1NiIsInR5cCI6IkpXVCJ9");
    }
}
<|endoftext|>


<commit_before>use std::rc::Rc;
use std::ops::Deref;
use std::path::PathBuf;
use std::collections::HashSet;
use cargo::core::Workspace;
use cargo::sources::PathSource;
use syntex_syntax::attr;
use syntex_syntax::visit::{self, Visitor};
use syntex_syntax::codemap::{CodeMap, Span, FilePathMapping};
use syntex_syntax::ast::{NodeId, Mac, Attribute, Mod, StructField, Block, Item, ItemKind};
use syntex_syntax::parse::{self, ParseSess};
use syntex_syntax::errors::Handler;
use syntex_syntax::errors::emitter::ColorConfig;
use config::Config;

struct IgnoredLines<'a> {
    lines: Vec<(PathBuf, usize)>,
    covered: &'a HashSet<PathBuf>,
    codemap: &'a CodeMap,
    config: &'a Config,
}

/// Returns a list of files and line numbers to ignore (not indexes!)
pub fn get_lines_to_ignore(project: &Workspace, config: &Config) -> Vec<(PathBuf, usize)> {
    let mut result: Vec<(PathBuf, usize)> = Vec::new();

    let pkg = project.current().unwrap();
    let mut src = PathSource::new(pkg.root(), pkg.package_id().source_id(), project.config());
    if let Ok(package) = src.root_package() {

        let codemap = Rc::new(CodeMap::new(FilePathMapping::empty()));
        let handler = Handler::with_tty_emitter(ColorConfig::Auto, false, false, Some(codemap.clone()));
        let parse_session = ParseSess::with_span_handler(handler, codemap.clone());

        let mut done_files: HashSet<PathBuf> = HashSet::new();
        for target in package.targets() {
            let file = target.src_path();
            if !(config.ignore_tests && file.starts_with(project.root().join("tests"))) {
                let mut parser = parse::new_parser_from_file(&parse_session, file);
                parser.cfg_mods = false;
                if let Ok(krate) = parser.parse_crate_mod() {

                    let mut lines = {
                        let mut visitor = IgnoredLines::from_session(&parse_session, &done_files, config);
                        visitor.visit_mod(&krate.module, krate.span, &krate.attrs, NodeId::new(0));
                        visitor.lines
                    };

                    for l in &lines {
                        done_files.insert(l.0.clone());
                    }
                    result.append(&mut lines);
                }
            }
        }
    }
    result
}

impl<'a> IgnoredLines<'a> {
    /// Construct a new ignored lines object for the given project
    fn from_session(session: &'a ParseSess,
                    covered: &'a HashSet<PathBuf>,
                    config: &'a Config) -> IgnoredLines<'a> {
        IgnoredLines {
            lines: vec![],
            covered: covered,
            codemap: session.codemap(),
            config: config,
        }
    }

    /// Add lines to the line ignore list
    fn ignore_lines(&mut self, span: Span) {
        if let Ok(s) = self.codemap.span_to_lines(span) {
            for line in &s.lines {
                let pb = PathBuf::from(self.codemap.span_to_filename(span) as String);
                // Line number is index+1
                self.lines.push((pb, line.line_index + 1));
            }
        }
    }

    /// Looks for #[cfg(test)] attribute.
    fn contains_cfg_test(&mut self, attrs: &[Attribute]) -> bool {
        attrs.iter()
             .filter(|x| x.path == "cfg")
             .filter_map(|x| x.meta_item_list())
             .flat_map(|x| x)
             .any(|x| {
                 if let Some(w) = x.word() {
                    w.name().as_str() == "test"
                 } else {
                     false
                 }
             })
    }

    /// This function finds ignorable lines within actual coverable code.
    /// As opposed to other functions which find isolated lines that aren't
    /// executed or lines filtered by the user. These lines are things like
    /// close braces that are within coverable code but not coverable.
    fn find_ignorable_lines(&mut self, span: Span) {
        if let Ok(l) = self.codemap.span_to_lines(span) {
            for line in &l.lines {
                let pb = PathBuf::from(self.codemap.span_to_filename(span) as String);
                if let Some(s) = l.file.get_line(line.line_index) {
                    // Is this one of those pointless {, } or }; only lines?
                    if !s.chars().any(|x| !"{}[];\t ,".contains(x)) {
                        self.lines.push((pb, line.line_index + 1));
                    }
                }
            }
        }
    }
}

impl<'v, 'a> Visitor<'v> for IgnoredLines<'a> {

    fn visit_item(&mut self, i: &'v Item) {
        match i.node {
            ItemKind::ExternCrate(..) => self.ignore_lines(i.span),
            ItemKind::Fn(_, _, _, _, _, ref block) => {
                if attr::contains_name(&i.attrs, "test") && self.config.ignore_tests {
                    self.ignore_lines(i.span);
                    self.ignore_lines(block.deref().span);
                }
            },
            _ => {},
        }
        visit::walk_item(self, i);
    }

    fn visit_mod(&mut self, m: &'v Mod, s: Span, _attrs: &[Attribute], _n: NodeId) {
        // If mod is cfg(test) and --ignore-tests ignore contents!
        if let Ok(fl) = self.codemap.span_to_lines(s) {
            if self.config.ignore_tests && self.contains_cfg_test(_attrs) {
                self.ignore_lines(s);
                if fl.lines.len() == 1 {
                    // Ignore the file
                    self.ignore_lines(m.inner);
                }
            }
            else {
                if fl.lines.len() == 1 {
                    // mod imports show up as coverable. Ignore
                    self.ignore_lines(s);
                }
                let mod_path = PathBuf::from(self.codemap.span_to_filename(m.inner));
                if !self.covered.contains(&mod_path) {
                    visit::walk_mod(self, m);
                }
            }
        }
    }

    fn visit_mac(&mut self, mac: &Mac) {
        // Use this to ignore unreachable lines
        let mac_text = &format!("{}", mac.node.path)[..];
        // TODO unimplemented should have extra logic to exclude the
        // function from coverage
        match mac_text {
            "unimplemented" => self.ignore_lines(mac.span),
            "unreachable" => self.ignore_lines(mac.span),
            _ => {},
        }
        visit::walk_mac(self, mac);
    }

    /// Ignores attributes which may get identified as coverable lines.
    fn visit_attribute(&mut self, attr: &Attribute) {
        if attr.check_name("derive") {
            self.ignore_lines(attr.span);
        }
    }

    /// Struct fields are mistakenly identified as instructions and uncoverable.
    fn visit_struct_field(&mut self, s: &'v StructField) {
        self.ignore_lines(s.span);
        visit::walk_struct_field(self, s);
    }

    fn visit_block(&mut self, b: &'v Block) {
        self.find_ignorable_lines(b.span);
        visit::walk_block(self, b);
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use syntex_syntax::parse::filemap_to_parser;
    use syntex_syntax::parse::parser::Parser;

    struct TestContext {
        conf: Config,
        codemap: Rc<CodeMap>,
        parse_session: ParseSess
    }

    impl TestContext {
        fn generate_parser(&self, filename: &str, src_string: &str) -> Parser {
            let filemap = self.codemap.new_filemap(filename.to_string(),
                                                   src_string.to_string());
            filemap_to_parser(&self.parse_session, filemap)
        }
    }

    impl Default for TestContext {
        fn default() -> TestContext {
            let codemap = Rc::new(CodeMap::new(FilePathMapping::empty()));
            let handler = Handler::with_tty_emitter(ColorConfig::Auto, false, false, Some(codemap.clone()));
            let parse_session = ParseSess::with_span_handler(handler, codemap.clone());
            TestContext {
                conf: Config::default(),
                codemap: codemap,
                parse_session: parse_session
            }
        }
    }

    fn parse_crate(ctx: &TestContext, parser: &mut Parser) -> Vec<usize>  {
        let krate = parser.parse_crate_mod();
        assert!(krate.is_ok());
        let krate = krate.unwrap();
        let unused: HashSet<PathBuf> = HashSet::new();
        let mut visitor = IgnoredLines::from_session(&ctx.parse_session, &unused, &ctx.conf);
        visitor.visit_mod(&krate.module, krate.span, &krate.attrs, NodeId::new(0));
        visitor.lines.iter().map(|x| x.1).collect::<Vec<_>>()
    }

    #[test]
    fn filter_struct_members() {
        let ctx = TestContext::default();
        let mut parser = ctx.generate_parser("struct_test.rs", "#[derive(Debug)]\npub struct Struct {\npub i: i32,\nj:String,\n}");
        let lines = parse_crate(&ctx, &mut parser);

        assert_eq!(lines.len(), 3);
        assert!(lines.contains(&1));
        assert!(lines.contains(&3));
        assert!(lines.contains(&4));
    }

    #[test]
    fn filter_mods() {
        let ctx = TestContext::default();
        let mut parser = ctx.generate_parser("test.rs", "mod foo {\nfn double(x:i32)->i32 {\n x*2\n}\n}");
        let lines = parse_crate(&ctx, &mut parser);
        assert!(!lines.contains(&3));

        let mut parser = ctx.generate_parser("test.rs", "mod foo{}");
        let lines = parse_crate(&ctx, &mut parser);
        assert!(lines.contains(&1));
    }

    #[test]
    fn filter_macros() {
        let ctx = TestContext::default();
        let mut parser = ctx.generate_parser("test.rs", "\n\nfn unused() {\nunimplemented!();\n}");

        let lines = parse_crate(&ctx, &mut parser);
        // Braces should be ignored so number could be higher
        assert!(lines.len() >= 1);
        assert!(lines.contains(&4));

        let mut parser = ctx.generate_parser("test.rs", "fn unused() {\nunreachable!();\n}");
        let lines = parse_crate(&ctx, &mut parser);
        assert!(lines.len() >= 1);
        assert!(lines.contains(&2));

        let mut parser = ctx.generate_parser("test.rs", "fn unused() {\nprintln!(\"text\");\n}");
        let lines = parse_crate(&ctx, &mut parser);
        assert!(!lines.contains(&2));
    }

    #[test]
    fn filter_tests() {
        let ctx = TestContext::default();
        let src_string = "#[cfg(test)]\nmod tests {\n fn boo(){\nassert!(true);\n}\n}";
        let mut parser = ctx.generate_parser("test.rs", src_string);
        let lines = parse_crate(&ctx, &mut parser);
        assert!(!lines.contains(&4));

        let mut ctx = TestContext::default();
        ctx.conf.ignore_tests = true;
        let mut parser = ctx.generate_parser("test.rs", src_string);
        let lines = parse_crate(&ctx, &mut parser);
        assert!(lines.contains(&4));

        let ctx = TestContext::default();
        let src_string = "#[test]\nfn mytest() { \n assert!(true);\n}";
        let mut parser = ctx.generate_parser("test.rs", src_string);
        let lines = parse_crate(&ctx, &mut parser);
        assert!(!lines.contains(&2));
        assert!(!lines.contains(&3));

        let mut ctx = TestContext::default();
        ctx.conf.ignore_tests = true;
        let mut parser = ctx.generate_parser("test.rs", src_string);
        let lines = parse_crate(&ctx, &mut parser);
        assert!(lines.contains(&2));
        assert!(lines.contains(&3));

    }
}
<commit_msg>Added coverable ability to IgnoredLines<commit_after>use std::rc::Rc;
use std::ops::Deref;
use std::path::PathBuf;
use std::collections::{HashSet, HashMap};
use cargo::core::Workspace;
use cargo::sources::PathSource;
use syntex_syntax::attr;
use syntex_syntax::visit::{self, Visitor, FnKind};
use syntex_syntax::codemap::{CodeMap, Span, FilePathMapping};
use syntex_syntax::ast::{NodeId, Mac, Attribute, Mod, StructField, Block, Item, ItemKind, FnDecl};
use syntex_syntax::parse::{self, ParseSess};
use syntex_syntax::errors::Handler;
use syntex_syntax::errors::emitter::ColorConfig;
use config::Config;

/// Represents the results of analysis of a single file. Does not store the file
/// in question as this is expected to be maintained by the user.
struct LineAnalysis {
    /// This represents lines that should be ignored in coverage
    /// but may be identifed as coverable in the DWARF tables
    pub ignore: HashSet<usize>,
    /// This represents lines that should be included in coverage
    /// But may be ignored.
    pub cover: HashSet<usize>,
}

impl LineAnalysis {
    fn new() -> LineAnalysis {
        LineAnalysis {
            ignore: HashSet::new(),
            cover: HashSet::new()
        }
    }
}

struct IgnoredLines<'a> {
    lines: Vec<(PathBuf, usize)>,
    coverable: Vec<(PathBuf, usize)>,
    covered: &'a HashSet<PathBuf>,
    codemap: &'a CodeMap,
    config: &'a Config,
}

/// Returns a list of files and line numbers to ignore (not indexes!)
pub fn get_lines_to_ignore(project: &Workspace, config: &Config) -> Vec<(PathBuf, usize)> {
    let mut result: Vec<(PathBuf, usize)> = Vec::new();
    let pkg = project.current().unwrap();
    let mut src = PathSource::new(pkg.root(), pkg.package_id().source_id(), project.config());
    if let Ok(package) = src.root_package() {

        let codemap = Rc::new(CodeMap::new(FilePathMapping::empty()));
        let handler = Handler::with_tty_emitter(ColorConfig::Auto, false, false, Some(codemap.clone()));
        let parse_session = ParseSess::with_span_handler(handler, codemap.clone());

        let mut done_files: HashSet<PathBuf> = HashSet::new();
        for target in package.targets() {
            let file = target.src_path();
            if !(config.ignore_tests && file.starts_with(project.root().join("tests"))) {
                let mut parser = parse::new_parser_from_file(&parse_session, file);
                parser.cfg_mods = false;
                if let Ok(krate) = parser.parse_crate_mod() {

                    let mut lines = {
                        let mut visitor = IgnoredLines::from_session(&parse_session, &done_files, config);
                        visitor.visit_mod(&krate.module, krate.span, &krate.attrs, NodeId::new(0));
                        visitor.lines
                    };

                    for l in &lines {
                        done_files.insert(l.0.clone());
                    }
                    result.append(&mut lines);
                }
            }
        }
    }
    result
}

fn add_lines(codemap: &CodeMap, lines: &mut Vec<(PathBuf, usize)>, s: Span) {
    if let Ok(ls) = codemap.span_to_lines(s) {
        for line in &ls.lines {
            let pb = PathBuf::from(codemap.span_to_filename(s) as String);
            // Line number is index+1
            lines.push((pb, line.line_index + 1));
        }
    }
}

impl<'a> IgnoredLines<'a> {
    /// Construct a new ignored lines object for the given project
    fn from_session(session: &'a ParseSess,
                    covered: &'a HashSet<PathBuf>,
                    config: &'a Config) -> IgnoredLines<'a> {
        IgnoredLines {
            lines: vec![],
            coverable: vec![],
            covered: covered,
            codemap: session.codemap(),
            config: config,
        }
    }

    /// Add lines to the line ignore list
    fn ignore_lines(&mut self, span: Span) {
        add_lines(self.codemap, &mut self.lines, span);
    }

    fn cover_lines(&mut self, span: Span) {
        add_lines(self.codemap, &mut self.coverable, span);
    }

    /// Looks for #[cfg(test)] attribute.
    fn contains_cfg_test(&mut self, attrs: &[Attribute]) -> bool {
        attrs.iter()
             .filter(|x| x.path == "cfg")
             .filter_map(|x| x.meta_item_list())
             .flat_map(|x| x)
             .any(|x| {
                 if let Some(w) = x.word() {
                    w.name().as_str() == "test"
                 } else {
                     false
                 }
             })
    }

    /// This function finds ignorable lines within actual coverable code.
    /// As opposed to other functions which find isolated lines that aren't
    /// executed or lines filtered by the user. These lines are things like
    /// close braces that are within coverable code but not coverable.
    fn find_ignorable_lines(&mut self, span: Span) {
        if let Ok(l) = self.codemap.span_to_lines(span) {
            for line in &l.lines {
                let pb = PathBuf::from(self.codemap.span_to_filename(span) as String);
                if let Some(s) = l.file.get_line(line.line_index) {
                    // Is this one of those pointless {, } or }; only lines?
                    if !s.chars().any(|x| !"{}[];\t ,".contains(x)) {
                        self.lines.push((pb, line.line_index + 1));
                    }
                }
            }
        }
    }
}

impl<'v, 'a> Visitor<'v> for IgnoredLines<'a> {

    fn visit_item(&mut self, i: &'v Item) {
        match i.node {
            ItemKind::ExternCrate(..) => self.ignore_lines(i.span),
            ItemKind::Fn(_, _, _, _, _, ref block) => {
                if attr::contains_name(&i.attrs, "test") && self.config.ignore_tests {
                    self.ignore_lines(i.span);
                    self.ignore_lines(block.deref().span);
                }
            },
            _ => {},
        }
        visit::walk_item(self, i);
    }

    fn visit_mod(&mut self, m: &'v Mod, s: Span, _attrs: &[Attribute], _n: NodeId) {
        // If mod is cfg(test) and --ignore-tests ignore contents!
        if let Ok(fl) = self.codemap.span_to_lines(s) {
            if self.config.ignore_tests && self.contains_cfg_test(_attrs) {
                self.ignore_lines(s);
                if fl.lines.len() == 1 {
                    // Ignore the file
                    self.ignore_lines(m.inner);
                }
            }
            else {
                if fl.lines.len() == 1 {
                    // mod imports show up as coverable. Ignore
                    self.ignore_lines(s);
                }
                let mod_path = PathBuf::from(self.codemap.span_to_filename(m.inner));
                if !self.covered.contains(&mod_path) {
                    visit::walk_mod(self, m);
                }
            }
        }
    }

    fn visit_fn(&mut self, fk: FnKind, fd: &FnDecl, s: Span, _: NodeId) {
        match fk {
            FnKind::ItemFn(_, g, _,_,_,_,_) => {
                if !g.ty_params.is_empty() {
                    self.cover_lines(s);
                }
            },
            FnKind::Method(_, sig, _, _) => {
                if !sig.generics.ty_params.is_empty() {
                    self.cover_lines(s);
                }
            },
            _ => {},
        }
        visit::walk_fn(self, fk, fd, s);
    }

    fn visit_mac(&mut self, mac: &Mac) {
        // Use this to ignore unreachable lines
        let mac_text = &format!("{}", mac.node.path)[..];
        // TODO unimplemented should have extra logic to exclude the
        // function from coverage
        match mac_text {
            "unimplemented" => self.ignore_lines(mac.span),
            "unreachable" => self.ignore_lines(mac.span),
            _ => {},
        }
        visit::walk_mac(self, mac);
    }

    /// Ignores attributes which may get identified as coverable lines.
    fn visit_attribute(&mut self, attr: &Attribute) {
        if attr.check_name("derive") {
            self.ignore_lines(attr.span);
        }
    }

    /// Struct fields are mistakenly identified as instructions and uncoverable.
    fn visit_struct_field(&mut self, s: &'v StructField) {
        self.ignore_lines(s.span);
        visit::walk_struct_field(self, s);
    }

    fn visit_block(&mut self, b: &'v Block) {
        self.find_ignorable_lines(b.span);
        visit::walk_block(self, b);
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use syntex_syntax::parse::filemap_to_parser;
    use syntex_syntax::parse::parser::Parser;

    struct TestContext {
        conf: Config,
        codemap: Rc<CodeMap>,
        parse_session: ParseSess
    }

    impl TestContext {
        fn generate_parser(&self, filename: &str, src_string: &str) -> Parser {
            let filemap = self.codemap.new_filemap(filename.to_string(),
                                                   src_string.to_string());
            filemap_to_parser(&self.parse_session, filemap)
        }
    }

    impl Default for TestContext {
        fn default() -> TestContext {
            let codemap = Rc::new(CodeMap::new(FilePathMapping::empty()));
            let handler = Handler::with_tty_emitter(ColorConfig::Auto, false, false, Some(codemap.clone()));
            let parse_session = ParseSess::with_span_handler(handler, codemap.clone());
            TestContext {
                conf: Config::default(),
                codemap: codemap,
                parse_session: parse_session
            }
        }
    }

    fn parse_crate(ctx: &TestContext, parser: &mut Parser) -> Vec<usize>  {
        let krate = parser.parse_crate_mod();
        assert!(krate.is_ok());
        let krate = krate.unwrap();
        let unused: HashSet<PathBuf> = HashSet::new();
        let mut visitor = IgnoredLines::from_session(&ctx.parse_session, &unused, &ctx.conf);
        visitor.visit_mod(&krate.module, krate.span, &krate.attrs, NodeId::new(0));
        visitor.lines.iter().map(|x| x.1).collect::<Vec<_>>()
    }

    #[test]
    fn filter_struct_members() {
        let ctx = TestContext::default();
        let mut parser = ctx.generate_parser("struct_test.rs", "#[derive(Debug)]\npub struct Struct {\npub i: i32,\nj:String,\n}");
        let lines = parse_crate(&ctx, &mut parser);

        assert_eq!(lines.len(), 3);
        assert!(lines.contains(&1));
        assert!(lines.contains(&3));
        assert!(lines.contains(&4));
    }

    #[test]
    fn filter_mods() {
        let ctx = TestContext::default();
        let mut parser = ctx.generate_parser("test.rs", "mod foo {\nfn double(x:i32)->i32 {\n x*2\n}\n}");
        let lines = parse_crate(&ctx, &mut parser);
        assert!(!lines.contains(&3));

        let mut parser = ctx.generate_parser("test.rs", "mod foo{}");
        let lines = parse_crate(&ctx, &mut parser);
        assert!(lines.contains(&1));
    }

    #[test]
    fn filter_macros() {
        let ctx = TestContext::default();
        let mut parser = ctx.generate_parser("test.rs", "\n\nfn unused() {\nunimplemented!();\n}");

        let lines = parse_crate(&ctx, &mut parser);
        // Braces should be ignored so number could be higher
        assert!(lines.len() >= 1);
        assert!(lines.contains(&4));

        let mut parser = ctx.generate_parser("test.rs", "fn unused() {\nunreachable!();\n}");
        let lines = parse_crate(&ctx, &mut parser);
        assert!(lines.len() >= 1);
        assert!(lines.contains(&2));

        let mut parser = ctx.generate_parser("test.rs", "fn unused() {\nprintln!(\"text\");\n}");
        let lines = parse_crate(&ctx, &mut parser);
        assert!(!lines.contains(&2));
    }

    #[test]
    fn filter_tests() {
        let ctx = TestContext::default();
        let src_string = "#[cfg(test)]\nmod tests {\n fn boo(){\nassert!(true);\n}\n}";
        let mut parser = ctx.generate_parser("test.rs", src_string);
        let lines = parse_crate(&ctx, &mut parser);
        assert!(!lines.contains(&4));

        let mut ctx = TestContext::default();
        ctx.conf.ignore_tests = true;
        let mut parser = ctx.generate_parser("test.rs", src_string);
        let lines = parse_crate(&ctx, &mut parser);
        assert!(lines.contains(&4));

        let ctx = TestContext::default();
        let src_string = "#[test]\nfn mytest() { \n assert!(true);\n}";
        let mut parser = ctx.generate_parser("test.rs", src_string);
        let lines = parse_crate(&ctx, &mut parser);
        assert!(!lines.contains(&2));
        assert!(!lines.contains(&3));

        let mut ctx = TestContext::default();
        ctx.conf.ignore_tests = true;
        let mut parser = ctx.generate_parser("test.rs", src_string);
        let lines = parse_crate(&ctx, &mut parser);
        assert!(lines.contains(&2));
        assert!(lines.contains(&3));

    }
}
<|endoftext|>


<commit_before>#![feature(plugin)]
#![plugin(clippy)]

#![warn(while_let_loop, empty_loop, while_let_on_iterator)]
#![allow(dead_code, never_loop, unused, cyclomatic_complexity)]

fn main() {
    let y = Some(true);
    loop {
        if let Some(_x) = y {
            let _v = 1;
        } else {
            break
        }
    }
    loop { // no error, break is not in else clause
        if let Some(_x) = y {
            let _v = 1;
        }
        break;
    }
    loop {
        match y {
            Some(_x) => true,
            None => break
        };
    }
    loop {
        let x = match y {
            Some(x) => x,
            None => break
        };
        let _x = x;
        let _str = "foo";
    }
    loop {
        let x = match y {
            Some(x) => x,
            None => break,
        };
        { let _a = "bar"; };
        { let _b = "foobar"; }
    }
    loop { // no error, else branch does something other than break
        match y {
            Some(_x) => true,
            _ => {
                let _z = 1;
                break;
            }
        };
    }
    while let Some(x) = y { // no error, obviously
        println!("{}", x);
    }

    // #675, this used to have a wrong suggestion
    loop {
        let (e, l) = match "".split_whitespace().next() {
            Some(word) => (word.is_empty(), word.len()),
            None => break
        };

        let _ = (e, l);
    }

    let mut iter = 1..20;
    while let Option::Some(x) = iter.next() {
        println!("{}", x);
    }

    let mut iter = 1..20;
    while let Some(x) = iter.next() {
        println!("{}", x);
    }

    let mut iter = 1..20;
    while let Some(_) = iter.next() {}

    let mut iter = 1..20;
    while let None = iter.next() {} // this is fine (if nonsensical)

    let mut iter = 1..20;
    if let Some(x) = iter.next() { // also fine
        println!("{}", x)
    }

    // the following shouldn't warn because it can't be written with a for loop
    let mut iter = 1u32..20;
    while let Some(x) = iter.next() {
        println!("next: {:?}", iter.next())
    }

    // neither can this
    let mut iter = 1u32..20;
    while let Some(x) = iter.next() {
        println!("next: {:?}", iter.next());
    }

    // or this
    let mut iter = 1u32..20;
    while let Some(x) = iter.next() {break;}
    println!("Remaining iter {:?}", iter);

    // or this
    let mut iter = 1u32..20;
    while let Some(x) = iter.next() {
        iter = 1..20;
    }
}

// regression test (#360)
// this should not panic
// it's okay if further iterations of the lint
// cause this function to trigger it
fn no_panic<T>(slice: &[T]) {
    let mut iter = slice.iter();
    loop {
        let _ = match iter.next() {
            Some(ele) => ele,
            None => break
        };
        loop {}
    }
}

fn issue1017() {
    let r: Result<u32, u32> = Ok(42);
    let mut len = 1337;

    loop {
        match r {
            Err(_) => len = 0,
            Ok(length) => {
                len = length;
                break
            }
        }
    }
}

// Issue #1188
fn refutable() {
    let a = [42, 1337];
    let mut b = a.iter();

    // consume all the 42s
    while let Some(&42) = b.next() {
    }

    let a = [(1, 2, 3)];
    let mut b = a.iter();

    while let Some(&(1, 2, 3)) = b.next() {
    }

    let a = [Some(42)];
    let mut b = a.iter();

    while let Some(&None) = b.next() {
    }

    /* This gives “refutable pattern in `for` loop binding: `&_` not covered”
    for &42 in b {}
    for &(1, 2, 3) in b {}
    for &Option::None in b.next() {}
    // */

    let mut y = a.iter();
    loop { // x is reused, so don't lint here
        while let Some(v) = y.next() {
        }
    }

    let mut y = a.iter();
    for _ in 0..2 {
        while let Some(v) = y.next() { // y is reused, don't lint
        }
    }

    loop {
        let mut y = a.iter();
        while let Some(v) = y.next() { // use a for loop here
        }
    }

    //should not trigger while_let_loop lint because break passes an expression
    let a = Some(10);
    let b = loop {
        if let Some(c) = a {
            break Some(c);
        } else {
            break None;
        }
    };
}
<commit_msg>format comment<commit_after>#![feature(plugin)]
#![plugin(clippy)]

#![warn(while_let_loop, empty_loop, while_let_on_iterator)]
#![allow(dead_code, never_loop, unused, cyclomatic_complexity)]

fn main() {
    let y = Some(true);
    loop {
        if let Some(_x) = y {
            let _v = 1;
        } else {
            break
        }
    }
    loop { // no error, break is not in else clause
        if let Some(_x) = y {
            let _v = 1;
        }
        break;
    }
    loop {
        match y {
            Some(_x) => true,
            None => break
        };
    }
    loop {
        let x = match y {
            Some(x) => x,
            None => break
        };
        let _x = x;
        let _str = "foo";
    }
    loop {
        let x = match y {
            Some(x) => x,
            None => break,
        };
        { let _a = "bar"; };
        { let _b = "foobar"; }
    }
    loop { // no error, else branch does something other than break
        match y {
            Some(_x) => true,
            _ => {
                let _z = 1;
                break;
            }
        };
    }
    while let Some(x) = y { // no error, obviously
        println!("{}", x);
    }

    // #675, this used to have a wrong suggestion
    loop {
        let (e, l) = match "".split_whitespace().next() {
            Some(word) => (word.is_empty(), word.len()),
            None => break
        };

        let _ = (e, l);
    }

    let mut iter = 1..20;
    while let Option::Some(x) = iter.next() {
        println!("{}", x);
    }

    let mut iter = 1..20;
    while let Some(x) = iter.next() {
        println!("{}", x);
    }

    let mut iter = 1..20;
    while let Some(_) = iter.next() {}

    let mut iter = 1..20;
    while let None = iter.next() {} // this is fine (if nonsensical)

    let mut iter = 1..20;
    if let Some(x) = iter.next() { // also fine
        println!("{}", x)
    }

    // the following shouldn't warn because it can't be written with a for loop
    let mut iter = 1u32..20;
    while let Some(x) = iter.next() {
        println!("next: {:?}", iter.next())
    }

    // neither can this
    let mut iter = 1u32..20;
    while let Some(x) = iter.next() {
        println!("next: {:?}", iter.next());
    }

    // or this
    let mut iter = 1u32..20;
    while let Some(x) = iter.next() {break;}
    println!("Remaining iter {:?}", iter);

    // or this
    let mut iter = 1u32..20;
    while let Some(x) = iter.next() {
        iter = 1..20;
    }
}

// regression test (#360)
// this should not panic
// it's okay if further iterations of the lint
// cause this function to trigger it
fn no_panic<T>(slice: &[T]) {
    let mut iter = slice.iter();
    loop {
        let _ = match iter.next() {
            Some(ele) => ele,
            None => break
        };
        loop {}
    }
}

fn issue1017() {
    let r: Result<u32, u32> = Ok(42);
    let mut len = 1337;

    loop {
        match r {
            Err(_) => len = 0,
            Ok(length) => {
                len = length;
                break
            }
        }
    }
}

// Issue #1188
fn refutable() {
    let a = [42, 1337];
    let mut b = a.iter();

    // consume all the 42s
    while let Some(&42) = b.next() {
    }

    let a = [(1, 2, 3)];
    let mut b = a.iter();

    while let Some(&(1, 2, 3)) = b.next() {
    }

    let a = [Some(42)];
    let mut b = a.iter();

    while let Some(&None) = b.next() {
    }

    /* This gives “refutable pattern in `for` loop binding: `&_` not covered”
    for &42 in b {}
    for &(1, 2, 3) in b {}
    for &Option::None in b.next() {}
    // */

    let mut y = a.iter();
    loop { // x is reused, so don't lint here
        while let Some(v) = y.next() {
        }
    }

    let mut y = a.iter();
    for _ in 0..2 {
        while let Some(v) = y.next() { // y is reused, don't lint
        }
    }

    loop {
        let mut y = a.iter();
        while let Some(v) = y.next() { // use a for loop here
        }
    }

    // should not trigger while_let_loop lint because break passes an expression
    let a = Some(10);
    let b = loop {
        if let Some(c) = a {
            break Some(c);
        } else {
            break None;
        }
    };
}
<|endoftext|>


<commit_before>use ntapi::ntioapi::{IO_STATUS_BLOCK_u, IO_STATUS_BLOCK};
use ntapi::ntioapi::{NtCancelIoFileEx, NtDeviceIoControlFile};
use ntapi::ntrtl::RtlNtStatusToDosError;
use std::fmt;
use std::fs::File;
use std::io;
use std::mem::size_of;
use std::os::windows::io::AsRawHandle;
use std::ptr::null_mut;
use winapi::shared::ntdef::{HANDLE, LARGE_INTEGER, NTSTATUS, PVOID, ULONG};
use winapi::shared::ntstatus::{STATUS_NOT_FOUND, STATUS_PENDING, STATUS_SUCCESS};

const IOCTL_AFD_POLL: ULONG = 0x00012024;

/// Winsock2 AFD driver instance.
///
/// All operations are unsafe due to IO_STATUS_BLOCK parameter are being used by Afd driver during STATUS_PENDING before I/O Completion Port returns its result.
#[derive(Debug)]
pub struct Afd {
    fd: File,
}

#[repr(C)]
#[derive(Debug)]
pub struct AfdPollHandleInfo {
    pub handle: HANDLE,
    pub events: ULONG,
    pub status: NTSTATUS,
}

unsafe impl Send for AfdPollHandleInfo {}

#[repr(C)]
pub struct AfdPollInfo {
    pub timeout: LARGE_INTEGER,
    // Can have only value 1.
    pub number_of_handles: ULONG,
    pub exclusive: ULONG,
    pub handles: [AfdPollHandleInfo; 1],
}

impl fmt::Debug for AfdPollInfo {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("AfdPollInfo").finish()
    }
}

impl Afd {
    /// Poll `Afd` instance with `AfdPollInfo`.
    ///
    /// # Unsafety
    ///
    /// This function is unsafe due to memory of `IO_STATUS_BLOCK` still being used by `Afd` instance while `Ok(false)` (`STATUS_PENDING`).
    /// `iosb` needs to be untouched after the call while operation is in effective at ALL TIME except for `cancel` method.
    /// So be careful not to `poll` twice while polling.
    /// User should deallocate there overlapped value when error to prevent memory leak.
    pub unsafe fn poll(
        &self,
        info: &mut AfdPollInfo,
        iosb: *mut IO_STATUS_BLOCK,
        overlapped: PVOID,
    ) -> io::Result<bool> {
        let info_ptr: PVOID = info as *mut _ as PVOID;
        (*iosb).u.Status = STATUS_PENDING;
        let status = NtDeviceIoControlFile(
            self.fd.as_raw_handle(),
            null_mut(),
            None,
            overlapped,
            iosb,
            IOCTL_AFD_POLL,
            info_ptr,
            size_of::<AfdPollInfo>() as u32,
            info_ptr,
            size_of::<AfdPollInfo>() as u32,
        );
        match status {
            STATUS_SUCCESS => Ok(true),
            STATUS_PENDING => Ok(false),
            _ => Err(io::Error::from_raw_os_error(
                RtlNtStatusToDosError(status) as i32
            )),
        }
    }

    /// Cancel previous polled request of `Afd`.
    ///
    /// iosb needs to be used by `poll` first for valid `cancel`.
    ///
    /// # Unsafety
    ///
    /// This function is unsafe due to memory of `IO_STATUS_BLOCK` still being used by `Afd` instance while `Ok(false)` (`STATUS_PENDING`).
    /// Use it only with request is still being polled so that you have valid `IO_STATUS_BLOCK` to use.
    /// User should NOT deallocate there overlapped value after the `cancel` to prevent double free.
    pub unsafe fn cancel(&self, iosb: *mut IO_STATUS_BLOCK) -> io::Result<()> {
        if (*iosb).u.Status != STATUS_PENDING {
            return Ok(());
        }

        let mut cancel_iosb = IO_STATUS_BLOCK {
            u: IO_STATUS_BLOCK_u { Status: 0 },
            Information: 0,
        };
        let status = NtCancelIoFileEx(self.fd.as_raw_handle(), iosb, &mut cancel_iosb);
        if status == STATUS_SUCCESS || status == STATUS_NOT_FOUND {
            return Ok(());
        }
        Err(io::Error::from_raw_os_error(
            RtlNtStatusToDosError(status) as i32
        ))
    }
}

cfg_io_source! {
    use std::mem::zeroed;
    use std::os::windows::io::{FromRawHandle, RawHandle};
    use std::sync::atomic::{AtomicUsize, Ordering};

    use miow::iocp::CompletionPort;
    use ntapi::ntioapi::{NtCreateFile, FILE_OPEN};
    use winapi::shared::ntdef::{OBJECT_ATTRIBUTES, UNICODE_STRING, USHORT, WCHAR};
    use winapi::um::handleapi::INVALID_HANDLE_VALUE;
    use winapi::um::winbase::{SetFileCompletionNotificationModes, FILE_SKIP_SET_EVENT_ON_HANDLE};
    use winapi::um::winnt::{SYNCHRONIZE, FILE_SHARE_READ, FILE_SHARE_WRITE};

    const AFD_HELPER_ATTRIBUTES: OBJECT_ATTRIBUTES = OBJECT_ATTRIBUTES {
        Length: size_of::<OBJECT_ATTRIBUTES>() as ULONG,
        RootDirectory: null_mut(),
        ObjectName: &AFD_OBJ_NAME as *const _ as *mut _,
        Attributes: 0,
        SecurityDescriptor: null_mut(),
        SecurityQualityOfService: null_mut(),
    };

    const AFD_OBJ_NAME: UNICODE_STRING = UNICODE_STRING {
        Length: (AFD_HELPER_NAME.len() * size_of::<WCHAR>()) as USHORT,
        MaximumLength: (AFD_HELPER_NAME.len() * size_of::<WCHAR>()) as USHORT,
        Buffer: AFD_HELPER_NAME.as_ptr() as *mut _,
    };

    const AFD_HELPER_NAME: &[WCHAR] = &[
        '\\' as _,
        'D' as _,
        'e' as _,
        'v' as _,
        'i' as _,
        'c' as _,
        'e' as _,
        '\\' as _,
        'A' as _,
        'f' as _,
        'd' as _,
        '\\' as _,
        'M' as _,
        'i' as _,
        'o' as _
    ];

    static NEXT_TOKEN: AtomicUsize = AtomicUsize::new(0);

    impl AfdPollInfo {
        pub fn zeroed() -> AfdPollInfo {
            unsafe { zeroed() }
        }
    }

    impl Afd {
        /// Create new Afd instance.
        pub fn new(cp: &CompletionPort) -> io::Result<Afd> {
            let mut afd_helper_handle: HANDLE = INVALID_HANDLE_VALUE;
            let mut iosb = IO_STATUS_BLOCK {
                u: IO_STATUS_BLOCK_u { Status: 0 },
                Information: 0,
            };

            unsafe {
                let status = NtCreateFile(
                    &mut afd_helper_handle as *mut _,
                    SYNCHRONIZE,
                    &AFD_HELPER_ATTRIBUTES as *const _ as *mut _,
                    &mut iosb,
                    null_mut(),
                    0 as ULONG,
                    FILE_SHARE_READ | FILE_SHARE_WRITE,
                    FILE_OPEN,
                    0 as ULONG,
                    null_mut(),
                    0 as ULONG,
                );
                if status != STATUS_SUCCESS {
                    return Err(io::Error::from_raw_os_error(
                        RtlNtStatusToDosError(status) as i32
                    ));
                }
                let fd = File::from_raw_handle(afd_helper_handle as RawHandle);
                // Increment by 2 to reserve space for other types of handles.
                // Non-AFD types (currently only NamedPipe), use odd numbered
                // tokens. This allows the selector to differentate between them
                // and dispatch events accordingly.
                let token = NEXT_TOKEN.fetch_add(2, Ordering::Relaxed) + 2;
                let afd = Afd { fd };
                cp.add_handle(token, &afd.fd)?;
                match SetFileCompletionNotificationModes(
                    afd_helper_handle,
                    FILE_SKIP_SET_EVENT_ON_HANDLE,
                ) {
                    0 => Err(io::Error::last_os_error()),
                    _ => Ok(afd),
                }
            }
        }
    }
}

pub const POLL_RECEIVE: u32 = 0b000_000_001;
pub const POLL_RECEIVE_EXPEDITED: u32 = 0b000_000_010;
pub const POLL_SEND: u32 = 0b000_000_100;
pub const POLL_DISCONNECT: u32 = 0b000_001_000;
pub const POLL_ABORT: u32 = 0b000_010_000;
pub const POLL_LOCAL_CLOSE: u32 = 0b000_100_000;
// Not used as it indicated in each event where a connection is connected, not
// just the first time a connection is established.
// Also see https://github.com/piscisaureus/wepoll/commit/8b7b340610f88af3d83f40fb728e7b850b090ece.
pub const POLL_CONNECT: u32 = 0b001_000_000;
pub const POLL_ACCEPT: u32 = 0b010_000_000;
pub const POLL_CONNECT_FAIL: u32 = 0b100_000_000;

pub const KNOWN_EVENTS: u32 = POLL_RECEIVE
    | POLL_RECEIVE_EXPEDITED
    | POLL_SEND
    | POLL_DISCONNECT
    | POLL_ABORT
    | POLL_LOCAL_CLOSE
    | POLL_ACCEPT
    | POLL_CONNECT_FAIL;
<commit_msg>Improve error message when opening AFD device<commit_after>use ntapi::ntioapi::{IO_STATUS_BLOCK_u, IO_STATUS_BLOCK};
use ntapi::ntioapi::{NtCancelIoFileEx, NtDeviceIoControlFile};
use ntapi::ntrtl::RtlNtStatusToDosError;
use std::fmt;
use std::fs::File;
use std::io;
use std::mem::size_of;
use std::os::windows::io::AsRawHandle;
use std::ptr::null_mut;
use winapi::shared::ntdef::{HANDLE, LARGE_INTEGER, NTSTATUS, PVOID, ULONG};
use winapi::shared::ntstatus::{STATUS_NOT_FOUND, STATUS_PENDING, STATUS_SUCCESS};

const IOCTL_AFD_POLL: ULONG = 0x00012024;

/// Winsock2 AFD driver instance.
///
/// All operations are unsafe due to IO_STATUS_BLOCK parameter are being used by Afd driver during STATUS_PENDING before I/O Completion Port returns its result.
#[derive(Debug)]
pub struct Afd {
    fd: File,
}

#[repr(C)]
#[derive(Debug)]
pub struct AfdPollHandleInfo {
    pub handle: HANDLE,
    pub events: ULONG,
    pub status: NTSTATUS,
}

unsafe impl Send for AfdPollHandleInfo {}

#[repr(C)]
pub struct AfdPollInfo {
    pub timeout: LARGE_INTEGER,
    // Can have only value 1.
    pub number_of_handles: ULONG,
    pub exclusive: ULONG,
    pub handles: [AfdPollHandleInfo; 1],
}

impl fmt::Debug for AfdPollInfo {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("AfdPollInfo").finish()
    }
}

impl Afd {
    /// Poll `Afd` instance with `AfdPollInfo`.
    ///
    /// # Unsafety
    ///
    /// This function is unsafe due to memory of `IO_STATUS_BLOCK` still being used by `Afd` instance while `Ok(false)` (`STATUS_PENDING`).
    /// `iosb` needs to be untouched after the call while operation is in effective at ALL TIME except for `cancel` method.
    /// So be careful not to `poll` twice while polling.
    /// User should deallocate there overlapped value when error to prevent memory leak.
    pub unsafe fn poll(
        &self,
        info: &mut AfdPollInfo,
        iosb: *mut IO_STATUS_BLOCK,
        overlapped: PVOID,
    ) -> io::Result<bool> {
        let info_ptr: PVOID = info as *mut _ as PVOID;
        (*iosb).u.Status = STATUS_PENDING;
        let status = NtDeviceIoControlFile(
            self.fd.as_raw_handle(),
            null_mut(),
            None,
            overlapped,
            iosb,
            IOCTL_AFD_POLL,
            info_ptr,
            size_of::<AfdPollInfo>() as u32,
            info_ptr,
            size_of::<AfdPollInfo>() as u32,
        );
        match status {
            STATUS_SUCCESS => Ok(true),
            STATUS_PENDING => Ok(false),
            _ => Err(io::Error::from_raw_os_error(
                RtlNtStatusToDosError(status) as i32
            )),
        }
    }

    /// Cancel previous polled request of `Afd`.
    ///
    /// iosb needs to be used by `poll` first for valid `cancel`.
    ///
    /// # Unsafety
    ///
    /// This function is unsafe due to memory of `IO_STATUS_BLOCK` still being used by `Afd` instance while `Ok(false)` (`STATUS_PENDING`).
    /// Use it only with request is still being polled so that you have valid `IO_STATUS_BLOCK` to use.
    /// User should NOT deallocate there overlapped value after the `cancel` to prevent double free.
    pub unsafe fn cancel(&self, iosb: *mut IO_STATUS_BLOCK) -> io::Result<()> {
        if (*iosb).u.Status != STATUS_PENDING {
            return Ok(());
        }

        let mut cancel_iosb = IO_STATUS_BLOCK {
            u: IO_STATUS_BLOCK_u { Status: 0 },
            Information: 0,
        };
        let status = NtCancelIoFileEx(self.fd.as_raw_handle(), iosb, &mut cancel_iosb);
        if status == STATUS_SUCCESS || status == STATUS_NOT_FOUND {
            return Ok(());
        }
        Err(io::Error::from_raw_os_error(
            RtlNtStatusToDosError(status) as i32
        ))
    }
}

cfg_io_source! {
    use std::mem::zeroed;
    use std::os::windows::io::{FromRawHandle, RawHandle};
    use std::sync::atomic::{AtomicUsize, Ordering};

    use miow::iocp::CompletionPort;
    use ntapi::ntioapi::{NtCreateFile, FILE_OPEN};
    use winapi::shared::ntdef::{OBJECT_ATTRIBUTES, UNICODE_STRING, USHORT, WCHAR};
    use winapi::um::handleapi::INVALID_HANDLE_VALUE;
    use winapi::um::winbase::{SetFileCompletionNotificationModes, FILE_SKIP_SET_EVENT_ON_HANDLE};
    use winapi::um::winnt::{SYNCHRONIZE, FILE_SHARE_READ, FILE_SHARE_WRITE};

    const AFD_HELPER_ATTRIBUTES: OBJECT_ATTRIBUTES = OBJECT_ATTRIBUTES {
        Length: size_of::<OBJECT_ATTRIBUTES>() as ULONG,
        RootDirectory: null_mut(),
        ObjectName: &AFD_OBJ_NAME as *const _ as *mut _,
        Attributes: 0,
        SecurityDescriptor: null_mut(),
        SecurityQualityOfService: null_mut(),
    };

    const AFD_OBJ_NAME: UNICODE_STRING = UNICODE_STRING {
        Length: (AFD_HELPER_NAME.len() * size_of::<WCHAR>()) as USHORT,
        MaximumLength: (AFD_HELPER_NAME.len() * size_of::<WCHAR>()) as USHORT,
        Buffer: AFD_HELPER_NAME.as_ptr() as *mut _,
    };

    const AFD_HELPER_NAME: &[WCHAR] = &[
        '\\' as _,
        'D' as _,
        'e' as _,
        'v' as _,
        'i' as _,
        'c' as _,
        'e' as _,
        '\\' as _,
        'A' as _,
        'f' as _,
        'd' as _,
        '\\' as _,
        'M' as _,
        'i' as _,
        'o' as _
    ];

    static NEXT_TOKEN: AtomicUsize = AtomicUsize::new(0);

    impl AfdPollInfo {
        pub fn zeroed() -> AfdPollInfo {
            unsafe { zeroed() }
        }
    }

    impl Afd {
        /// Create new Afd instance.
        pub fn new(cp: &CompletionPort) -> io::Result<Afd> {
            let mut afd_helper_handle: HANDLE = INVALID_HANDLE_VALUE;
            let mut iosb = IO_STATUS_BLOCK {
                u: IO_STATUS_BLOCK_u { Status: 0 },
                Information: 0,
            };

            unsafe {
                let status = NtCreateFile(
                    &mut afd_helper_handle as *mut _,
                    SYNCHRONIZE,
                    &AFD_HELPER_ATTRIBUTES as *const _ as *mut _,
                    &mut iosb,
                    null_mut(),
                    0 as ULONG,
                    FILE_SHARE_READ | FILE_SHARE_WRITE,
                    FILE_OPEN,
                    0 as ULONG,
                    null_mut(),
                    0 as ULONG,
                );
                if status != STATUS_SUCCESS {
                    let raw_err = io::Error::from_raw_os_error(
                        RtlNtStatusToDosError(status) as i32
                    );
                    let msg = format!("Failed to open \\Device\\Afd\\Mio: {}", raw_err);
                    return Err(io::Error::new(raw_err.kind(), msg));
                }
                let fd = File::from_raw_handle(afd_helper_handle as RawHandle);
                // Increment by 2 to reserve space for other types of handles.
                // Non-AFD types (currently only NamedPipe), use odd numbered
                // tokens. This allows the selector to differentate between them
                // and dispatch events accordingly.
                let token = NEXT_TOKEN.fetch_add(2, Ordering::Relaxed) + 2;
                let afd = Afd { fd };
                cp.add_handle(token, &afd.fd)?;
                match SetFileCompletionNotificationModes(
                    afd_helper_handle,
                    FILE_SKIP_SET_EVENT_ON_HANDLE,
                ) {
                    0 => Err(io::Error::last_os_error()),
                    _ => Ok(afd),
                }
            }
        }
    }
}

pub const POLL_RECEIVE: u32 = 0b000_000_001;
pub const POLL_RECEIVE_EXPEDITED: u32 = 0b000_000_010;
pub const POLL_SEND: u32 = 0b000_000_100;
pub const POLL_DISCONNECT: u32 = 0b000_001_000;
pub const POLL_ABORT: u32 = 0b000_010_000;
pub const POLL_LOCAL_CLOSE: u32 = 0b000_100_000;
// Not used as it indicated in each event where a connection is connected, not
// just the first time a connection is established.
// Also see https://github.com/piscisaureus/wepoll/commit/8b7b340610f88af3d83f40fb728e7b850b090ece.
pub const POLL_CONNECT: u32 = 0b001_000_000;
pub const POLL_ACCEPT: u32 = 0b010_000_000;
pub const POLL_CONNECT_FAIL: u32 = 0b100_000_000;

pub const KNOWN_EVENTS: u32 = POLL_RECEIVE
    | POLL_RECEIVE_EXPEDITED
    | POLL_SEND
    | POLL_DISCONNECT
    | POLL_ABORT
    | POLL_LOCAL_CLOSE
    | POLL_ACCEPT
    | POLL_CONNECT_FAIL;
<|endoftext|>