<commit_before>use bridges::{BridgeData, complete_auth};
use config::GoogleConfig;
use crypto;
use email_address::EmailAddress;
use error::BrokerError;
use futures::{Future, future};
use http::{ContextHandle, HandlerResult};
use hyper::{Response, StatusCode};
use hyper::header::{ContentType, Location};
use serde_helpers::UrlDef;
use std::rc::Rc;
use store_cache::{CacheKey, fetch_json_url};
use time::now_utc;
use url::Url;
use url::percent_encoding::{utf8_percent_encode, QUERY_ENCODE_SET};
use validation;
use webfinger::{Link, Relation};

/// The origin of the Google identity provider.
pub const GOOGLE_IDP_ORIGIN: &'static str = "https://accounts.google.com";

/// Normalization to apply to an email address.
#[derive(Serialize,Deserialize)]
pub enum Normalization {
    None,
    Google,
}

impl Normalization {
    /// Apply normalization to an email address.
    pub fn apply(&self, email_addr: EmailAddress) -> EmailAddress {
        match *self {
            Normalization::None => email_addr,
            Normalization::Google => email_addr.normalize_google(),
        }
    }
}

/// Data we store in the session.
#[derive(Serialize,Deserialize)]
pub struct OidcBridgeData {
    pub origin: String,
    pub client_id: String,
    pub nonce: String,
    pub normalization: Normalization,
}

/// OpenID Connect configuration document.
#[derive(Deserialize)]
struct ProviderConfig {
    #[serde(with = "UrlDef")]
    authorization_endpoint: Url,
    #[serde(with = "UrlDef")]
    jwks_uri: Url,
}

/// OpenID Connect key set document.
#[derive(Deserialize)]
struct ProviderKeys {
    #[serde(default)]
    keys: Vec<ProviderKey>,
}

#[derive(Deserialize)]
pub struct ProviderKey {
    #[serde(default)]
    pub kid: String,
    #[serde(rename = "use")]
    #[serde(default)]
    pub use_: String,
    #[serde(default)]
    pub n: String,
    #[serde(default)]
    pub e: String,
}

/// Provide authentication using OpenID Connect.
///
/// Redirect the user agent to the provider authorization endpoint, which we discover by reading
/// the providers configuration document. Included in the request is a nonce which we can later use
/// to definitively match the callback to this request.
///
/// This function handles both Portier providers, which works without registration, as well as
/// the Google provider, for which we have a preregistered `client_id`.
pub fn auth(ctx_handle: &ContextHandle, email_addr: &Rc<EmailAddress>, link: &Link)
    -> HandlerResult {

    let ctx = ctx_handle.borrow();

    // Generate a nonce for the provider.
    let provider_nonce = crypto::nonce();

    // Determine the parameters to use, based on the webfinger link.
    let provider_origin = match validation::parse_oidc_href(&link.href) {
        Some(origin) => origin,
        None => return Box::new(future::err(BrokerError::Provider(
            format!("invalid href: {}", link.href)))),
    };
    let bridge_data = Rc::new(match link.rel {
        Relation::Portier => {
            OidcBridgeData {
                origin: provider_origin,
                client_id: ctx.app.public_url.clone(),
                nonce: provider_nonce,
                normalization: Normalization::None,
            }
        },
        // Delegate to the OpenID Connect bridge for Google, if configured.
        Relation::Google => {
            let client_id = match ctx.app.google {
                Some(GoogleConfig { ref client_id }) => client_id,
                None => return Box::new(future::err(BrokerError::ProviderCancelled)),
            };
            if provider_origin != GOOGLE_IDP_ORIGIN {
                return Box::new(future::err(BrokerError::Provider(
                    format!("invalid href: Google provider only supports {}", GOOGLE_IDP_ORIGIN))));
            }
            OidcBridgeData {
                origin: provider_origin,
                client_id: client_id.clone(),
                nonce: provider_nonce,
                normalization: Normalization::Google,
            }
        },
    });

    // Retrieve the provider's configuration and extract the
    // `authorization_endpoint` from it.
    let bridge_data2 = bridge_data.clone();
    let f = {
        let config_url = build_config_url(&bridge_data2.origin);
        fetch_json_url(&ctx.app, config_url, &CacheKey::OidcConfig {
            origin: bridge_data2.origin.as_str(),
        }).map_err(move |e| BrokerError::Provider(
            format!("could not fetch {}'s configuration: {}", bridge_data2.origin, e)))
    };

    let ctx_handle = ctx_handle.clone();
    let email_addr = email_addr.clone();
    let f = f.and_then(move |provider_config: ProviderConfig| {
        let mut ctx = ctx_handle.borrow_mut();

        // Create the URL to redirect to, properly escaping all parameters.
        let result = Url::parse(&vec![
            provider_config.authorization_endpoint.as_str(),
            "?",
            "client_id=",
            &utf8_percent_encode(&bridge_data.client_id, QUERY_ENCODE_SET).to_string(),
            "&response_type=id_token",
            "&scope=",
            &utf8_percent_encode("openid email", QUERY_ENCODE_SET).to_string(),
            "&redirect_uri=",
            &utf8_percent_encode(&format!("{}/callback", &ctx.app.public_url),
                                 QUERY_ENCODE_SET).to_string(),
            "&state=",
            &utf8_percent_encode(&ctx.session_id, QUERY_ENCODE_SET).to_string(),
            "&nonce=",
            &utf8_percent_encode(&bridge_data.nonce, QUERY_ENCODE_SET).to_string(),
            "&login_hint=",
            &utf8_percent_encode(email_addr.as_str(), QUERY_ENCODE_SET).to_string(),
        ].join(""));
        let auth_url = match result {
            Ok(url) => url,
            Err(_) => {
                let domain = email_addr.domain();
                return future::err(BrokerError::Provider(
                    format!("failed to build valid authorization URL from {}'s 'authorization_endpoint'", domain)));
            },
        };

        // Save session data, committing the session to this provider.
        let bridge_data = Rc::try_unwrap(bridge_data).map_err(|_| ()).expect("lingering oidc bridge data references");
        match ctx.save_session(BridgeData::Oidc(bridge_data)) {
            Ok(true) => {},
            Ok(false) => return future::err(BrokerError::ProviderCancelled),
            Err(e) => return future::err(e),
        }

        let res = Response::new()
            .with_status(StatusCode::SeeOther)
            .with_header(Location::new(auth_url.into_string()));
        future::ok(res)
    });

    Box::new(f)
}

/// Request handler for OpenID Connect callbacks with response mode 'fragment'
///
/// For providers that don't support `response_mode=form_post`, we capture the fragment parameters
/// in javascript and emulate the POST request.
pub fn fragment_callback(ctx_handle: ContextHandle) -> HandlerResult {
    let ctx = ctx_handle.borrow();

    let res = Response::new()
        .with_header(ContentType::html())
        .with_body(ctx.app.templates.fragment_callback.render(&[]));
    Box::new(future::ok(res))
}

/// Request handler for OpenID Connect callbacks.
///
/// Match the returned email address and nonce against our session data, then extract the identity
/// token returned by the provider and verify it. Return an identity token for the relying party if
/// successful, or an error message otherwise.
pub fn callback(ctx_handle: ContextHandle) -> HandlerResult {
    let (bridge_data, id_token) = {
        let mut ctx = ctx_handle.borrow_mut();

        let session_id = try_get_param!(ctx, "state");
        let bridge_data = match ctx.load_session(&session_id) {
            Ok(BridgeData::Oidc(bridge_data)) => Rc::new(bridge_data),
            Ok(_) => return Box::new(future::err(BrokerError::Input("invalid session".to_owned()))),
            Err(e) => return Box::new(future::err(e)),
        };

        let id_token = try_get_param!(ctx, "id_token");
        (bridge_data, id_token)
    };

    // Request the provider's configuration to get the `jwks_uri` values from it.
    let bridge_data2 = bridge_data.clone();
    let f = {
        let ctx = ctx_handle.borrow();
        let config_url = build_config_url(&bridge_data2.origin);
        fetch_json_url(&ctx.app, config_url, &CacheKey::OidcConfig {
            origin: bridge_data2.origin.as_str(),
        }).map_err(move |e| BrokerError::Provider(
            format!("could not fetch {}'s configuration: {}", bridge_data2.origin, e)))
    };

    // Grab the keys from the provider, then verify the signature.
    let ctx_handle2 = ctx_handle.clone();
    let bridge_data2 = bridge_data.clone();
    let f = f.and_then(move |provider_config: ProviderConfig| {
        let ctx = ctx_handle2.borrow();
        fetch_json_url(&ctx.app, provider_config.jwks_uri, &CacheKey::OidcKeySet {
            origin: bridge_data2.origin.as_str(),
        })
            .then(move |result| {
                let key_set: ProviderKeys = result.map_err(|e| BrokerError::Provider(
                    format!("could not fetch {}'s keys: {}", &bridge_data2.origin, e)))?;
                crypto::verify_jws(&id_token, &key_set.keys).map_err(|_| BrokerError::Provider(
                    format!("could not verify the token received from {}", &bridge_data2.origin)))
            })
    });

    let ctx_handle = ctx_handle.clone();
    let f = f.and_then(move |jwt_payload| {
        let ctx = ctx_handle.borrow();
        let data = ctx.session_data.as_ref().expect("session vanished");

        // Normalize the email address according to provider rules.
        let email_addr = bridge_data.normalization.apply(data.email_addr.clone());

        // Extract the token claims.
        let descr = format!("{}'s token payload", email_addr.domain());
        let iss = try_get_json_field!(jwt_payload, "iss", descr);
        let aud = try_get_json_field!(jwt_payload, "aud", descr);
        let token_addr = try_get_json_field!(jwt_payload, "email", descr);
        let exp = try_get_json_field!(jwt_payload, "exp", |v| v.as_i64(), descr);
        let nonce = try_get_json_field!(jwt_payload, "nonce", descr);

        // Normalize the token email address too.
        let token_addr: EmailAddress = match token_addr.parse() {
            Ok(addr) => bridge_data.normalization.apply(addr),
            Err(_) => return future::err(BrokerError::Provider(format!(
                    "failed to parse email from {}", descr))),
        };

        // Verify the token claims.
        check_field!(iss == bridge_data.origin, "iss", descr);
        check_field!(aud == *bridge_data.client_id, "aud", descr);
        check_field!(nonce == *bridge_data.nonce, "nonce", descr);
        check_field!(token_addr == email_addr, "email", descr);

        let now = now_utc().to_timespec().sec;
        check_field!(now < exp, "exp", descr);

        // If everything is okay, build a new identity token and send it
        // to the relying party.
        future::result(complete_auth(&*ctx))
    });

    Box::new(f)
}

/// Build the URL to the OpenID Connect configuration for a provider at the given origin.
fn build_config_url(provider_origin: &str) -> Url {
    format!("{}/.well-known/openid-configuration", provider_origin).parse()
        .expect("could not build the OpenID Connect configuration URL")
}
<commit_msg>Enforce HTTPS in OIDC bridge<commit_after>use bridges::{BridgeData, complete_auth};
use config::GoogleConfig;
use crypto;
use email_address::EmailAddress;
use error::BrokerError;
use futures::{Future, future};
use http::{ContextHandle, HandlerResult};
use hyper::{Response, StatusCode};
use hyper::header::{ContentType, Location};
use serde_helpers::UrlDef;
use std::rc::Rc;
use store_cache::{CacheKey, fetch_json_url};
use time::now_utc;
use url::Url;
use url::percent_encoding::{utf8_percent_encode, QUERY_ENCODE_SET};
use validation;
use webfinger::{Link, Relation};

/// The origin of the Google identity provider.
pub const GOOGLE_IDP_ORIGIN: &'static str = "https://accounts.google.com";

/// Normalization to apply to an email address.
#[derive(Serialize,Deserialize)]
pub enum Normalization {
    None,
    Google,
}

impl Normalization {
    /// Apply normalization to an email address.
    pub fn apply(&self, email_addr: EmailAddress) -> EmailAddress {
        match *self {
            Normalization::None => email_addr,
            Normalization::Google => email_addr.normalize_google(),
        }
    }
}

/// Data we store in the session.
#[derive(Serialize,Deserialize)]
pub struct OidcBridgeData {
    pub origin: String,
    pub client_id: String,
    pub nonce: String,
    pub normalization: Normalization,
}

/// OpenID Connect configuration document.
#[derive(Deserialize)]
struct ProviderConfig {
    #[serde(with = "UrlDef")]
    authorization_endpoint: Url,
    #[serde(with = "UrlDef")]
    jwks_uri: Url,
}

/// OpenID Connect key set document.
#[derive(Deserialize)]
struct ProviderKeys {
    #[serde(default)]
    keys: Vec<ProviderKey>,
}

#[derive(Deserialize)]
pub struct ProviderKey {
    #[serde(default)]
    pub kid: String,
    #[serde(rename = "use")]
    #[serde(default)]
    pub use_: String,
    #[serde(default)]
    pub n: String,
    #[serde(default)]
    pub e: String,
}

/// Provide authentication using OpenID Connect.
///
/// Redirect the user agent to the provider authorization endpoint, which we discover by reading
/// the providers configuration document. Included in the request is a nonce which we can later use
/// to definitively match the callback to this request.
///
/// This function handles both Portier providers, which works without registration, as well as
/// the Google provider, for which we have a preregistered `client_id`.
pub fn auth(ctx_handle: &ContextHandle, email_addr: &Rc<EmailAddress>, link: &Link)
    -> HandlerResult {

    let ctx = ctx_handle.borrow();

    // Generate a nonce for the provider.
    let provider_nonce = crypto::nonce();

    // Determine the parameters to use, based on the webfinger link.
    let provider_origin = match validation::parse_oidc_href(&link.href) {
        Some(origin) => origin,
        None => return Box::new(future::err(BrokerError::Provider(
            format!("invalid href (validation failed): {}", link.href)))),
    };
    let bridge_data = Rc::new(match link.rel {
        Relation::Portier => {
            #[cfg(not(feature = "insecure"))] {
                if link.href.scheme() != "https" {
                    return Box::new(future::err(BrokerError::Provider(
                        format!("invalid href (not HTTPS): {}", link.href))));
                }
            }
            OidcBridgeData {
                origin: provider_origin,
                client_id: ctx.app.public_url.clone(),
                nonce: provider_nonce,
                normalization: Normalization::None,
            }
        },
        // Delegate to the OpenID Connect bridge for Google, if configured.
        Relation::Google => {
            let client_id = match ctx.app.google {
                Some(GoogleConfig { ref client_id }) => client_id,
                None => return Box::new(future::err(BrokerError::ProviderCancelled)),
            };
            if provider_origin != GOOGLE_IDP_ORIGIN {
                return Box::new(future::err(BrokerError::Provider(
                    format!("invalid href: Google provider only supports {}", GOOGLE_IDP_ORIGIN))));
            }
            OidcBridgeData {
                origin: provider_origin,
                client_id: client_id.clone(),
                nonce: provider_nonce,
                normalization: Normalization::Google,
            }
        },
    });

    // Retrieve the provider's configuration.
    let f = fetch_config(ctx_handle, &bridge_data);

    let ctx_handle = ctx_handle.clone();
    let email_addr = email_addr.clone();
    let f = f.and_then(move |provider_config: ProviderConfig| {
        let mut ctx = ctx_handle.borrow_mut();

        // Create the URL to redirect to, properly escaping all parameters.
        let result = Url::parse(&vec![
            provider_config.authorization_endpoint.as_str(),
            "?",
            "client_id=",
            &utf8_percent_encode(&bridge_data.client_id, QUERY_ENCODE_SET).to_string(),
            "&response_type=id_token",
            "&scope=",
            &utf8_percent_encode("openid email", QUERY_ENCODE_SET).to_string(),
            "&redirect_uri=",
            &utf8_percent_encode(&format!("{}/callback", &ctx.app.public_url),
                                 QUERY_ENCODE_SET).to_string(),
            "&state=",
            &utf8_percent_encode(&ctx.session_id, QUERY_ENCODE_SET).to_string(),
            "&nonce=",
            &utf8_percent_encode(&bridge_data.nonce, QUERY_ENCODE_SET).to_string(),
            "&login_hint=",
            &utf8_percent_encode(email_addr.as_str(), QUERY_ENCODE_SET).to_string(),
        ].join(""));
        let auth_url = match result {
            Ok(url) => url,
            Err(_) => {
                let domain = email_addr.domain();
                return future::err(BrokerError::Provider(
                    format!("failed to build valid authorization URL from {}'s 'authorization_endpoint'", domain)));
            },
        };

        // Save session data, committing the session to this provider.
        let bridge_data = Rc::try_unwrap(bridge_data).map_err(|_| ()).expect("lingering oidc bridge data references");
        match ctx.save_session(BridgeData::Oidc(bridge_data)) {
            Ok(true) => {},
            Ok(false) => return future::err(BrokerError::ProviderCancelled),
            Err(e) => return future::err(e),
        }

        let res = Response::new()
            .with_status(StatusCode::SeeOther)
            .with_header(Location::new(auth_url.into_string()));
        future::ok(res)
    });

    Box::new(f)
}

/// Request handler for OpenID Connect callbacks with response mode 'fragment'
///
/// For providers that don't support `response_mode=form_post`, we capture the fragment parameters
/// in javascript and emulate the POST request.
pub fn fragment_callback(ctx_handle: ContextHandle) -> HandlerResult {
    let ctx = ctx_handle.borrow();

    let res = Response::new()
        .with_header(ContentType::html())
        .with_body(ctx.app.templates.fragment_callback.render(&[]));
    Box::new(future::ok(res))
}

/// Request handler for OpenID Connect callbacks.
///
/// Match the returned email address and nonce against our session data, then extract the identity
/// token returned by the provider and verify it. Return an identity token for the relying party if
/// successful, or an error message otherwise.
pub fn callback(ctx_handle: ContextHandle) -> HandlerResult {
    let (bridge_data, id_token) = {
        let mut ctx = ctx_handle.borrow_mut();

        let session_id = try_get_param!(ctx, "state");
        let bridge_data = match ctx.load_session(&session_id) {
            Ok(BridgeData::Oidc(bridge_data)) => Rc::new(bridge_data),
            Ok(_) => return Box::new(future::err(BrokerError::Input("invalid session".to_owned()))),
            Err(e) => return Box::new(future::err(e)),
        };

        let id_token = try_get_param!(ctx, "id_token");
        (bridge_data, id_token)
    };

    // Retrieve the provider's configuration.
    let f = fetch_config(&ctx_handle, &bridge_data);

    // Grab the keys from the provider, then verify the signature.
    let ctx_handle2 = ctx_handle.clone();
    let bridge_data2 = bridge_data.clone();
    let f = f.and_then(move |provider_config: ProviderConfig| {
        let ctx = ctx_handle2.borrow();
        fetch_json_url(&ctx.app, provider_config.jwks_uri, &CacheKey::OidcKeySet {
            origin: bridge_data2.origin.as_str(),
        })
            .then(move |result| {
                let key_set: ProviderKeys = result.map_err(|e| BrokerError::Provider(
                    format!("could not fetch {}'s keys: {}", &bridge_data2.origin, e)))?;
                crypto::verify_jws(&id_token, &key_set.keys).map_err(|_| BrokerError::Provider(
                    format!("could not verify the token received from {}", &bridge_data2.origin)))
            })
    });

    let ctx_handle = ctx_handle.clone();
    let f = f.and_then(move |jwt_payload| {
        let ctx = ctx_handle.borrow();
        let data = ctx.session_data.as_ref().expect("session vanished");

        // Normalize the email address according to provider rules.
        let email_addr = bridge_data.normalization.apply(data.email_addr.clone());

        // Extract the token claims.
        let descr = format!("{}'s token payload", email_addr.domain());
        let iss = try_get_json_field!(jwt_payload, "iss", descr);
        let aud = try_get_json_field!(jwt_payload, "aud", descr);
        let token_addr = try_get_json_field!(jwt_payload, "email", descr);
        let exp = try_get_json_field!(jwt_payload, "exp", |v| v.as_i64(), descr);
        let nonce = try_get_json_field!(jwt_payload, "nonce", descr);

        // Normalize the token email address too.
        let token_addr: EmailAddress = match token_addr.parse() {
            Ok(addr) => bridge_data.normalization.apply(addr),
            Err(_) => return future::err(BrokerError::Provider(format!(
                    "failed to parse email from {}", descr))),
        };

        // Verify the token claims.
        check_field!(iss == bridge_data.origin, "iss", descr);
        check_field!(aud == *bridge_data.client_id, "aud", descr);
        check_field!(nonce == *bridge_data.nonce, "nonce", descr);
        check_field!(token_addr == email_addr, "email", descr);

        let now = now_utc().to_timespec().sec;
        check_field!(now < exp, "exp", descr);

        // If everything is okay, build a new identity token and send it
        // to the relying party.
        future::result(complete_auth(&*ctx))
    });

    Box::new(f)
}

// Retrieve and verify the provider's configuration.
fn fetch_config(ctx_handle: &ContextHandle, bridge_data: &Rc<OidcBridgeData>)
    -> Box<Future<Item=ProviderConfig, Error=BrokerError>> {

    let config_url = format!("{}/.well-known/openid-configuration", bridge_data.origin).parse()
        .expect("could not build the OpenID Connect configuration URL");

    let ctx = ctx_handle.borrow();
    let bridge_data = bridge_data.clone();
    let f = fetch_json_url::<ProviderConfig>(&ctx.app, config_url, &CacheKey::OidcConfig {
        origin: bridge_data.origin.as_str(),
    });

    let f = f.then(move |result| {
        let provider_config = match result {
            Ok(provider_config) => provider_config,
            Err(e) => return future::err(BrokerError::Provider(
                format!("could not fetch {}'s configuration: {}", bridge_data.origin, e))),
        };

        #[cfg(not(feature = "insecure"))] {
            if provider_config.authorization_endpoint.scheme() != "https" {
                return future::err(BrokerError::Provider(
                    format!("{}'s authorization_endpoint is not HTTPS", bridge_data.origin)));
            }
            if provider_config.jwks_uri.scheme() != "https" {
                return future::err(BrokerError::Provider(
                    format!("{}'s jwks_uri is not HTTPS", bridge_data.origin)));
            }
        }

        future::ok(provider_config)
    });

    Box::new(f)
}
<|endoftext|>


<commit_before>use mem;
use parser::ast::ctxt::Context;

use parser::ast::*;
use parser::ast::Stmt::*;
use parser::ast::Expr::*;
use parser::ast::visit::*;

pub fn generate<'a, 'ast>(ctxt: &'a Context<'a, 'ast>, fct: &'ast Function) {
    CodeGenInfo::new(ctxt, fct).generate();
}

struct CodeGenInfo<'a, 'ast: 'a> {
    ctxt: &'a Context<'a, 'ast>,
    fct: &'ast Function,

    stacksize: u32,
    contains_fct_invocation: bool,
}

impl<'a, 'ast> CodeGenInfo<'a, 'ast> {
    fn new(ctxt: &'a Context<'a, 'ast>, fct: &'ast Function) -> CodeGenInfo<'a, 'ast> {
        CodeGenInfo {
            ctxt: ctxt,
            fct: fct,

            stacksize: 0,
            contains_fct_invocation: false,
        }
    }

    fn generate(&mut self) {
        self.visit_fct(self.fct);

        self.ctxt.function(self.fct.id, |fct| {
            fct.stacksize = self.stacksize;
            fct.contains_fct_invocation = self.contains_fct_invocation;
        });
    }

    fn increase_stack(&mut self, id: NodeId) {
        self.ctxt.var(id, |v| {
            let ty_size = v.data_type.size();
            self.stacksize = mem::align(self.stacksize + ty_size, ty_size);
            v.offset = -(self.stacksize as i32);
        });
    }
}

impl<'a, 'ast> Visitor<'ast> for CodeGenInfo<'a, 'ast> {
    fn visit_param(&mut self, p: &'ast Param) {
        self.increase_stack(p.id);
    }

    fn visit_stmt(&mut self, s: &'ast Stmt) {
        if let StmtVar(ref var) = *s {
            self.increase_stack(var.id);
        }

        visit::walk_stmt(self, s);
    }

    fn visit_expr(&mut self, e: &'ast Expr) {
        if let ExprCall(_) = *e {
            self.contains_fct_invocation = true;
        }

        visit::walk_expr(self, e);
    }
}

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

    use parser::ast::ctxt::Context;
    use driver::cmd::Args;
    use error::msg::Msg;
    use parser::ast;
    use parser::Parser;
    use semck;

    pub fn parse<F>(code: &'static str, f: F)
            where F: for<'a, 'ast> FnOnce(&'ast Context<'a, 'ast>) -> () {
        let mut parser = Parser::from_str(code);
        let (ast, interner) = parser.parse().unwrap();
        let map = ast::map::build(&ast, &interner);
        let args : Args = Default::default();

        ast::dump::dump(&ast, &interner);

        let ctxt = Context::new(&args, &interner, &map, &ast);

        semck::check(&ctxt);
        assert!(!ctxt.diag.borrow().has_errors());

        f(&ctxt);
    }

    #[test]
    fn test_invocation_flag() {
        parse("fn f() { g(); } fn g() { }", |ctxt| {
            let fct1 = ctxt.ast.elements[0].to_function().unwrap();
            generate(ctxt, fct1);
            assert_eq!(true, ctxt.function(fct1.id, |fct| fct.contains_fct_invocation));

            let fct2 = ctxt.ast.elements[1].to_function().unwrap();
            generate(ctxt, fct2);
            assert_eq!(false, ctxt.function(fct2.id, |fct| fct.contains_fct_invocation));
        });
    }

    #[test]
    fn test_param_offset() {
        parse("fn f(a: bool, b: int) { var c = 1; }", |ctxt| {
            let fct = ctxt.ast.elements[0].to_function().unwrap();
            generate(ctxt, fct);
            assert_eq!(12, ctxt.function(fct.id, |fct| fct.stacksize));

            ctxt.function(fct.id, |fct| {
                for (varid, offset) in fct.vars.iter().zip(&[-1, -8, -12]) {
                    assert_eq!(*offset, ctxt.var_infos.borrow()[varid.0].offset);
                }
            });
        });
    }

    #[test]
    fn test_var_offset() {
        parse("fn f() { var a = true; var b = false; var c = 2; var d = \"abc\"; }", |ctxt| {
            let fct = ctxt.ast.elements[0].to_function().unwrap();
            generate(ctxt, fct);
            assert_eq!(16, ctxt.function(fct.id, |fct| fct.stacksize));

            ctxt.function(fct.id, |fct| {
                for (varid, offset) in fct.vars.iter().zip(&[-1, -2, -8, -16]) {
                    assert_eq!(*offset, ctxt.var_infos.borrow()[varid.0].offset);
                }
            });
        });
    }
}
<commit_msg>codegen: handle params > 6 correctly<commit_after>use codegen::x64::reg::*;
use mem;
use parser::ast::ctxt::Context;

use parser::ast::*;
use parser::ast::Stmt::*;
use parser::ast::Expr::*;
use parser::ast::visit::*;

pub fn generate<'a, 'ast>(ctxt: &'a Context<'a, 'ast>, fct: &'ast Function) {
    CodeGenInfo::new(ctxt, fct).generate();
}

struct CodeGenInfo<'a, 'ast: 'a> {
    ctxt: &'a Context<'a, 'ast>,
    fct: &'ast Function,

    stacksize: u32,
    param_offset: i32,
    contains_fct_invocation: bool,
}

impl<'a, 'ast> CodeGenInfo<'a, 'ast> {
    fn new(ctxt: &'a Context<'a, 'ast>, fct: &'ast Function) -> CodeGenInfo<'a, 'ast> {
        CodeGenInfo {
            ctxt: ctxt,
            fct: fct,

            stacksize: 0,

            // first param offset to rbp is +16,
            // rbp+0 -> saved rbp
            // rbp+8 -> return address
            param_offset: 16,

            contains_fct_invocation: false,
        }
    }

    fn generate(&mut self) {
        self.visit_fct(self.fct);

        self.ctxt.function(self.fct.id, |fct| {
            fct.stacksize = self.stacksize;
            fct.contains_fct_invocation = self.contains_fct_invocation;
        });
    }

    fn increase_stack(&mut self, id: NodeId) {
        self.ctxt.var(id, |v| {
            let ty_size = v.data_type.size();
            self.stacksize = mem::align(self.stacksize + ty_size, ty_size);
            v.offset = -(self.stacksize as i32);
        });
    }
}

impl<'a, 'ast> Visitor<'ast> for CodeGenInfo<'a, 'ast> {
    fn visit_param(&mut self, p: &'ast Param) {
        // on x64 only the first 6 parameters are stored in registers
        if (p.idx as usize) < PARAM_REGS.len() {
            self.increase_stack(p.id);

        // the rest of the parameters need to be stored in the callers stack
        } else {
            self.ctxt.var(p.id, |v| {
                v.offset = self.param_offset;

                // all params on stack need size of 8
                self.param_offset += 8;
            });
        }
    }

    fn visit_stmt(&mut self, s: &'ast Stmt) {
        if let StmtVar(ref var) = *s {
            self.increase_stack(var.id);
        }

        visit::walk_stmt(self, s);
    }

    fn visit_expr(&mut self, e: &'ast Expr) {
        if let ExprCall(_) = *e {
            self.contains_fct_invocation = true;
        }

        visit::walk_expr(self, e);
    }
}

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

    use parser::ast::ctxt::Context;
    use driver::cmd::Args;
    use error::msg::Msg;
    use parser::ast;
    use parser::Parser;
    use semck;

    pub fn parse<F>(code: &'static str, f: F)
            where F: for<'a, 'ast> FnOnce(&'ast Context<'a, 'ast>) -> () {
        let mut parser = Parser::from_str(code);
        let (ast, interner) = parser.parse().unwrap();
        let map = ast::map::build(&ast, &interner);
        let args : Args = Default::default();

        ast::dump::dump(&ast, &interner);

        let ctxt = Context::new(&args, &interner, &map, &ast);

        semck::check(&ctxt);
        assert!(!ctxt.diag.borrow().has_errors());

        f(&ctxt);
    }

    #[test]
    fn test_invocation_flag() {
        parse("fn f() { g(); } fn g() { }", |ctxt| {
            let fct1 = ctxt.ast.elements[0].to_function().unwrap();
            generate(ctxt, fct1);
            assert_eq!(true, ctxt.function(fct1.id, |fct| fct.contains_fct_invocation));

            let fct2 = ctxt.ast.elements[1].to_function().unwrap();
            generate(ctxt, fct2);
            assert_eq!(false, ctxt.function(fct2.id, |fct| fct.contains_fct_invocation));
        });
    }

    #[test]
    fn test_param_offset() {
        parse("fn f(a: bool, b: int) { var c = 1; }", |ctxt| {
            let fct = ctxt.ast.elements[0].to_function().unwrap();
            generate(ctxt, fct);
            assert_eq!(12, ctxt.function(fct.id, |fct| fct.stacksize));

            ctxt.function(fct.id, |fct| {
                for (varid, offset) in fct.vars.iter().zip(&[-1, -8, -12]) {
                    assert_eq!(*offset, ctxt.var_infos.borrow()[varid.0].offset);
                }
            });
        });
    }

    #[test]
    fn test_var_offset() {
        parse("fn f() { var a = true; var b = false; var c = 2; var d = \"abc\"; }", |ctxt| {
            let fct = ctxt.ast.elements[0].to_function().unwrap();
            generate(ctxt, fct);
            assert_eq!(16, ctxt.function(fct.id, |fct| fct.stacksize));

            ctxt.function(fct.id, |fct| {
                for (varid, offset) in fct.vars.iter().zip(&[-1, -2, -8, -16]) {
                    assert_eq!(*offset, ctxt.var_infos.borrow()[varid.0].offset);
                }
            });
        });
    }
}
<|endoftext|>


<commit_before>pub use crate::wait::NotifyReady;

use std::fmt;

use crate:: {
  error::Error,
  wait::Wait,
};

pub struct Confirmation<T, I=()> {
  kind: ConfirmationKind<T, I>,
}

impl<T, I> Confirmation<T, I> {
  pub(crate) fn new(wait: Wait<T>) -> Self {
    Self { kind: ConfirmationKind::Wait(wait) }
  }

  pub(crate) fn new_error(error: Error) -> Self {
    let (wait, wait_handle) = Wait::new();
    wait_handle.error(error);
    Self::new(wait)
  }

  pub fn as_error(self) -> Result<(), Error> {
    self.try_wait().transpose().map(|_| ())
  }

  pub fn subscribe(&self, task: Box<dyn NotifyReady + Send>) {
    match &self.kind {
      ConfirmationKind::Wait(wait)   => wait.subscribe(task),
      ConfirmationKind::Map(wait, _) => wait.subscribe(task),
    }
  }

  pub fn has_subscriber(&self) -> bool {
    match &self.kind {
      ConfirmationKind::Wait(wait)   => wait.has_subscriber(),
      ConfirmationKind::Map(wait, _) => wait.has_subscriber(),
    }
  }

  pub fn try_wait(&self) -> Option<Result<T, Error>> {
    match &self.kind {
      ConfirmationKind::Wait(wait)   => wait.try_wait(),
      ConfirmationKind::Map(wait, f) => wait.try_wait().map(|res| res.map(f)),
    }
  }

  pub fn wait(self) -> Result<T, Error> {
    match self.kind {
      ConfirmationKind::Wait(wait)   => wait.wait(),
      ConfirmationKind::Map(wait, f) => wait.wait().map(f),
    }
  }
}

impl<T> Confirmation<T> {
  pub(crate) fn map<M>(self, f: Box<dyn Fn(T) -> M + Send + 'static>) -> Confirmation<M, T> {
    Confirmation { kind: ConfirmationKind::Map(Box::new(self), f) }
  }
}

enum ConfirmationKind<T, I> {
  Wait(Wait<T>),
  Map(Box<Confirmation<I>>, Box<dyn Fn(I) -> T + Send + 'static>)
}

impl<T> From<Result<Wait<T>, Error>> for Confirmation<T> {
  fn from(res: Result<Wait<T>, Error>) -> Self {
    match res {
      Ok(wait) => Confirmation::new(wait),
      Err(err) => Confirmation::new_error(err),
    }
  }
}

impl<T> fmt::Debug for Confirmation<T> {
  fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
    write!(f, "Confirmation")
  }
}

#[cfg(feature = "futures")]
pub(crate) mod futures {
  use super::*;

  use ::futures::{
    future::TryFuture,
    task::{Context, Poll, Waker},
  };

  use std::pin::Pin;

  impl<T, I> TryFuture for Confirmation<T, I> {
    type Ok = T;
    type Error = Error;

    fn try_poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<Self::Ok, Self::Error>> {
      if !self.has_subscriber() {
        self.subscribe(Box::new(Watcher(cx.waker().clone())));
      }
      self.try_wait().map(Poll::Ready).unwrap_or(Poll::Pending)
    }
  }

  pub(crate) struct Watcher(pub(crate) Waker);

  impl NotifyReady for Watcher {
    fn notify(&self) {
      self.0.wake_by_ref();
    }
  }
}
<commit_msg>confirmation: implement Future, not TryFuture<commit_after>pub use crate::wait::NotifyReady;

use std::fmt;

use crate:: {
  error::Error,
  wait::Wait,
};

pub struct Confirmation<T, I=()> {
  kind: ConfirmationKind<T, I>,
}

impl<T, I> Confirmation<T, I> {
  pub(crate) fn new(wait: Wait<T>) -> Self {
    Self { kind: ConfirmationKind::Wait(wait) }
  }

  pub(crate) fn new_error(error: Error) -> Self {
    let (wait, wait_handle) = Wait::new();
    wait_handle.error(error);
    Self::new(wait)
  }

  pub fn as_error(self) -> Result<(), Error> {
    self.try_wait().transpose().map(|_| ())
  }

  pub fn subscribe(&self, task: Box<dyn NotifyReady + Send>) {
    match &self.kind {
      ConfirmationKind::Wait(wait)   => wait.subscribe(task),
      ConfirmationKind::Map(wait, _) => wait.subscribe(task),
    }
  }

  pub fn has_subscriber(&self) -> bool {
    match &self.kind {
      ConfirmationKind::Wait(wait)   => wait.has_subscriber(),
      ConfirmationKind::Map(wait, _) => wait.has_subscriber(),
    }
  }

  pub fn try_wait(&self) -> Option<Result<T, Error>> {
    match &self.kind {
      ConfirmationKind::Wait(wait)   => wait.try_wait(),
      ConfirmationKind::Map(wait, f) => wait.try_wait().map(|res| res.map(f)),
    }
  }

  pub fn wait(self) -> Result<T, Error> {
    match self.kind {
      ConfirmationKind::Wait(wait)   => wait.wait(),
      ConfirmationKind::Map(wait, f) => wait.wait().map(f),
    }
  }
}

impl<T> Confirmation<T> {
  pub(crate) fn map<M>(self, f: Box<dyn Fn(T) -> M + Send + 'static>) -> Confirmation<M, T> {
    Confirmation { kind: ConfirmationKind::Map(Box::new(self), f) }
  }
}

enum ConfirmationKind<T, I> {
  Wait(Wait<T>),
  Map(Box<Confirmation<I>>, Box<dyn Fn(I) -> T + Send + 'static>)
}

impl<T> From<Result<Wait<T>, Error>> for Confirmation<T> {
  fn from(res: Result<Wait<T>, Error>) -> Self {
    match res {
      Ok(wait) => Confirmation::new(wait),
      Err(err) => Confirmation::new_error(err),
    }
  }
}

impl<T> fmt::Debug for Confirmation<T> {
  fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
    write!(f, "Confirmation")
  }
}

#[cfg(feature = "futures")]
pub(crate) mod futures {
  use super::*;

  use ::futures::{
    future::Future,
    task::{Context, Poll, Waker},
  };

  use std::pin::Pin;

  impl<T, I> Future for Confirmation<T, I> {
    type Output = Result<T, Error>;

    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
      if !self.has_subscriber() {
        self.subscribe(Box::new(Watcher(cx.waker().clone())));
      }
      self.try_wait().map(Poll::Ready).unwrap_or(Poll::Pending)
    }
  }

  pub(crate) struct Watcher(pub(crate) Waker);

  impl NotifyReady for Watcher {
    fn notify(&self) {
      self.0.wake_by_ref();
    }
  }
}
<|endoftext|>


<commit_before>// iterators2.rs
// In this module, you'll learn some of unique advantages that iterators can offer
// Step 1. Complete the `capitalize_first` function to pass the first two cases
// Step 2. Apply the `capitalize_first` function to a vector of strings, ensuring that it
// Step 3. Apply the `capitalize_first` function again to a list, but try and ensure it returns a single string
// As always, there are hints below!

pub fn capitalize_first(input: &str) -> String {
    let mut c = input.chars();
    match c.next() {
        None => String::new(),
        Some(first) => first.collect()::<String>() + c.as_str(),
    }
}

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

    // Step 1.
    // Tests that verify your `capitalize_first` function implementation
    #[test]
    fn test_success() {
        assert_eq!(capitalize_first("hello"), "Hello");
    }

    #[test]
    fn test_empty() {
        assert_eq!(capitalize_first(""), "");
    }

    // Step 2.
    #[test]
    fn test_iterate_string_vec() {
        let words = vec!["hello", "world"];
        let capitalized_words: Vec<String> = // TODO
        assert_eq!(capitalized_words, ["Hello", "World"]);
    }

    #[test]
    fn test_iterate_into_string() {
        let words = vec!["hello", " ", "world"];
        let capitalized_words = // TODO
        assert_eq!(capitalized_words, "Hello World");
    }
}

// Step 1
// You need to call something on `first` before it can be collected
// Currently it's type is `char`. Have a look at the methods that are available on that type:
// https://doc.rust-lang.org/std/primitive.char.html

// Step 2
// First you'll need to turn the Vec into an iterator
// Then you'll need to apply your function unto each item in the vector
// P.s. Don't forget to collect() at the end!

// Step 3.
// This is very similar to the previous test. The only real change is that you will need to
// alter the type that collect is coerced into. For a bonus you could try doing this with a
// turbofish
<commit_msg>Auto merge of #207 - nkanderson:fix_iterators2, r=komaeda<commit_after>// iterators2.rs
// In this module, you'll learn some of unique advantages that iterators can offer
// Step 1. Complete the `capitalize_first` function to pass the first two cases
// Step 2. Apply the `capitalize_first` function to a vector of strings, ensuring that it returns a vector of strings as well
// Step 3. Apply the `capitalize_first` function again to a list, but try and ensure it returns a single string
// As always, there are hints below!

pub fn capitalize_first(input: &str) -> String {
    let mut c = input.chars();
    match c.next() {
        None => String::new(),
        Some(first) => first.collect::<String>() + c.as_str(),
    }
}

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

    // Step 1.
    // Tests that verify your `capitalize_first` function implementation
    #[test]
    fn test_success() {
        assert_eq!(capitalize_first("hello"), "Hello");
    }

    #[test]
    fn test_empty() {
        assert_eq!(capitalize_first(""), "");
    }

    // Step 2.
    #[test]
    fn test_iterate_string_vec() {
        let words = vec!["hello", "world"];
        let capitalized_words: Vec<String> = // TODO
        assert_eq!(capitalized_words, ["Hello", "World"]);
    }

    #[test]
    fn test_iterate_into_string() {
        let words = vec!["hello", " ", "world"];
        let capitalized_words = // TODO
        assert_eq!(capitalized_words, "Hello World");
    }
}

// Step 1
// You need to call something on `first` before it can be collected
// Currently it's type is `char`. Have a look at the methods that are available on that type:
// https://doc.rust-lang.org/std/primitive.char.html

// Step 2
// First you'll need to turn the Vec into an iterator
// Then you'll need to apply your function unto each item in the vector
// P.s. Don't forget to collect() at the end!

// Step 3.
// This is very similar to the previous test. The only real change is that you will need to
// alter the type that collect is coerced into. For a bonus you could try doing this with a
// turbofish
<|endoftext|>


<commit_before>// Copyright (c) 2014 Robert Clipsham <robert@octarineparrot.com>
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

extern crate libc;

use std::collections::{VecDeque};
use std::cmp;
use std::ffi::CString;
use std::old_io::{IoResult, IoError};
use std::iter::repeat;
use std::mem;
use std::sync::Arc;

use bindings::bpf;
use old_packet::Packet;
use old_packet::ethernet::{EthernetHeader, MutableEthernetHeader};
use datalink::DataLinkChannelType;
use datalink::DataLinkChannelType::{Layer2, Layer3};
use internal;
use util::NetworkInterface;

// NOTE buffer must be word aligned.
pub fn datalink_channel(network_interface: &NetworkInterface,
                        write_buffer_size: usize,
                        read_buffer_size: usize,
                        channel_type: DataLinkChannelType)
    -> IoResult<(DataLinkSenderImpl, DataLinkReceiverImpl)> {
    #[cfg(target_os = "freebsd")]
    fn get_fd() -> libc::c_int {
        unsafe {
            libc::open(CString::from_slice("/dev/bpf").as_ptr(), libc::O_RDWR, 0)
        }
    }

    #[cfg(target_os = "macos")]
    fn get_fd() -> libc::c_int {
        // FIXME This is an arbitrary number of attempts
        for i in range(0, 1_000is) {
            let fd = unsafe {
                let file_name = format!("/dev/bpf{}", i);
                libc::open(CString::from_slice(file_name.as_bytes()).as_ptr(), libc::O_RDWR, 0)
            };
            if fd != -1 {
                return fd;
            }
        }
        return -1;
    }

    #[cfg(target_os = "freebsd")]
    fn set_feedback(fd: libc::c_int) -> Result<(), IoError> {
        let one: libc::c_uint = 1;
        if unsafe { bpf::ioctl(fd, bpf::BIOCFEEDBACK, &one) } == -1 {
            let err = IoError::last_error();
            unsafe { libc::close(fd); }
            return Err(err);
        }
        Ok(())
    }

    #[cfg(target_os = "macos")]
    fn set_feedback(_fd: libc::c_int) -> Result<(), IoError> {
        Ok(())
    }

    match channel_type {
        Layer2 => (),
        Layer3(_) => unimplemented!(),
    }

    let fd = get_fd();
    if fd == -1 {
        return Err(IoError::last_error());
    }
    let mut iface: bpf::ifreq = unsafe { mem::zeroed() };
    let mut i = 0;
    for c in network_interface.name.as_slice().bytes() {
        iface.ifr_name[i] = c as i8;
        i += 1;
    }

    let buflen = read_buffer_size as libc::c_uint;
    // NOTE Buffer length must be set before binding to an interface
    //      otherwise this will return Invalid Argument
    if unsafe { bpf::ioctl(fd, bpf::BIOCSBLEN, &buflen) } == -1 {
        let err = IoError::last_error();
        unsafe { libc::close(fd); }
        return Err(err);
    }

    // Set the interface to use
    if unsafe { bpf::ioctl(fd, bpf::BIOCSETIF, &iface) } == -1 {
        let err = IoError::last_error();
        unsafe { libc::close(fd); }
        return Err(err);
    }

    // Return from read as soon as packets are available - don't wait to fill the buffer
    let one: libc::c_uint = 1;
    if unsafe { bpf::ioctl(fd, bpf::BIOCIMMEDIATE, &one) } == -1 {
        let err = IoError::last_error();
        unsafe { libc::close(fd); }
        return Err(err);
    }

    let mut header_size = 0;

    // Get the device type
    let mut dlt: libc::c_uint = 0;
    if unsafe { bpf::ioctl(fd, bpf::BIOCGDLT, &mut dlt) } == -1 {
        let err = IoError::last_error();
        unsafe { libc::close(fd); }
        return Err(err);
    }

    // The loopback IoError::last_error()device does weird things
    // FIXME This should really just be another L2 packet header type
    if dlt == bpf::DLT_NULL {
        header_size = 4;

        // Allow packets to be read back after they are written
        match set_feedback(fd) {
            Err(e) => return Err(e),
            _ => ()
        }

    } else {
        // Don't fill in source MAC
        if unsafe { bpf::ioctl(fd, bpf::BIOCSHDRCMPLT, &one) } == -1 {
            let err = IoError::last_error();
            unsafe { libc::close(fd); }
            return Err(err);
        }
    }

    let fd = Arc::new(internal::FileDesc { fd: fd });
    let sender = DataLinkSenderImpl {
        fd: fd.clone(),
        write_buffer: repeat(0u8).take(write_buffer_size).collect(),
        header_size: header_size,
    };
    let receiver = DataLinkReceiverImpl {
        fd: fd,
        read_buffer: repeat(0u8).take(read_buffer_size).collect(),
        header_size: header_size,
    };

    Ok((sender, receiver))
}

pub struct DataLinkSenderImpl {
    fd: Arc<internal::FileDesc>,
    write_buffer: Vec<u8>,
    header_size: usize,
}

impl DataLinkSenderImpl {
    pub fn build_and_send<F>(&mut self, num_packets: usize, packet_size: usize,
                          func: &mut F) -> Option<IoResult<()>>
        where F : FnMut(MutableEthernetHeader)
    {
        let len = num_packets * (packet_size + self.header_size);
        if len >= self.write_buffer.len() {
            None
        } else {
            let min = cmp::min(self.write_buffer.len(), len);
            for chunk in self.write_buffer[..min]
                                          .chunks_mut(packet_size + self.header_size) {
                // If we're sending on the loopback device, the first 4 bytes must be set to
                // AF_INET
                if self.header_size == 4 {
                    unsafe {
                        *(chunk.as_mut_ptr() as *mut u32) = libc::AF_INET as u32;
                    }
                }
                {
                    let eh = MutableEthernetHeader::new(&mut chunk[self.header_size..]);
                    func(eh);
                }
                match unsafe { libc::write(self.fd.fd,
                                           chunk.as_ptr() as *const libc::c_void,
                                           chunk.len() as libc::size_t) } {
                    len if len == -1 => return Some(Err(IoError::last_error())),
                    _ => ()
                }
            }
            Some(Ok(()))
        }
    }

    pub fn send_to(&mut self, packet: EthernetHeader, _dst: Option<NetworkInterface>)
        -> Option<IoResult<()>> {
        match unsafe { libc::write(self.fd.fd,
                                   packet.packet().as_ptr() as *const libc::c_void,
                                   packet.packet().len() as libc::size_t) } {
            len if len == -1 => Some(Err(IoError::last_error())),
            _ => Some(Ok(()))
        }
    }
}

pub struct DataLinkReceiverImpl {
    fd: Arc<internal::FileDesc>,
    read_buffer: Vec<u8>,
    header_size: usize,
}

impl DataLinkReceiverImpl {
    pub fn iter<'a>(&'a mut self) -> DataLinkChannelIteratorImpl<'a> {
        let buflen = self.read_buffer.len();
        DataLinkChannelIteratorImpl {
            pc: self,
            // Enough room for minimally sized packets without reallocating
            packets: VecDeque::with_capacity(buflen / 64)
        }
    }
}

pub struct DataLinkChannelIteratorImpl<'a> {
    pc: &'a mut DataLinkReceiverImpl,
    packets: VecDeque<(usize, usize)>,
}

impl<'a> DataLinkChannelIteratorImpl<'a> {
    pub fn next<'c>(&'c mut self) -> IoResult<EthernetHeader<'c>> {
        if self.packets.is_empty() {
            let buflen = match unsafe {
                libc::read(self.pc.fd.fd,
                           self.pc.read_buffer.as_ptr() as *mut libc::c_void,
                           self.pc.read_buffer.len() as libc::size_t)
            } {
                len if len > 0 => len,
                _ => return Err(IoError::last_error())
            };
            let mut ptr = self.pc.read_buffer.as_mut_ptr();
            let end = unsafe { self.pc.read_buffer.as_ptr().offset(buflen as isize) };
            while (ptr as *const u8) < end {
                unsafe {
                    let packet: *const bpf::bpf_hdr = mem::transmute(ptr);
                    let start = ptr as isize +
                                (*packet).bh_hdrlen as isize -
                                self.pc.read_buffer.as_ptr() as isize;
                    self.packets.push_back((start as usize + self.pc.header_size,
                                      (*packet).bh_caplen as usize - self.pc.header_size));
                    let offset = (*packet).bh_hdrlen as isize + (*packet).bh_caplen as isize;
                    ptr = ptr.offset(bpf::BPF_WORDALIGN(offset));
                }
            }
        }
        let (start, len) = self.packets.pop_front().unwrap();
        Ok(EthernetHeader::new(&self.pc.read_buffer[start .. start + len]))
    }
}

<commit_msg>Use of `range` is no longer permitted, use `x..y` instead.<commit_after>// Copyright (c) 2014 Robert Clipsham <robert@octarineparrot.com>
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

extern crate libc;

use std::collections::{VecDeque};
use std::cmp;
use std::ffi::CString;
use std::old_io::{IoResult, IoError};
use std::iter::repeat;
use std::mem;
use std::sync::Arc;

use bindings::bpf;
use old_packet::Packet;
use old_packet::ethernet::{EthernetHeader, MutableEthernetHeader};
use datalink::DataLinkChannelType;
use datalink::DataLinkChannelType::{Layer2, Layer3};
use internal;
use util::NetworkInterface;

// NOTE buffer must be word aligned.
pub fn datalink_channel(network_interface: &NetworkInterface,
                        write_buffer_size: usize,
                        read_buffer_size: usize,
                        channel_type: DataLinkChannelType)
    -> IoResult<(DataLinkSenderImpl, DataLinkReceiverImpl)> {
    #[cfg(target_os = "freebsd")]
    fn get_fd() -> libc::c_int {
        unsafe {
            libc::open(CString::from_slice("/dev/bpf").as_ptr(), libc::O_RDWR, 0)
        }
    }

    #[cfg(target_os = "macos")]
    fn get_fd() -> libc::c_int {
        // FIXME This is an arbitrary number of attempts
        for i in (0..1_000is) {
            let fd = unsafe {
                let file_name = format!("/dev/bpf{}", i);
                libc::open(CString::from_slice(file_name.as_bytes()).as_ptr(), libc::O_RDWR, 0)
            };
            if fd != -1 {
                return fd;
            }
        }
        return -1;
    }

    #[cfg(target_os = "freebsd")]
    fn set_feedback(fd: libc::c_int) -> Result<(), IoError> {
        let one: libc::c_uint = 1;
        if unsafe { bpf::ioctl(fd, bpf::BIOCFEEDBACK, &one) } == -1 {
            let err = IoError::last_error();
            unsafe { libc::close(fd); }
            return Err(err);
        }
        Ok(())
    }

    #[cfg(target_os = "macos")]
    fn set_feedback(_fd: libc::c_int) -> Result<(), IoError> {
        Ok(())
    }

    match channel_type {
        Layer2 => (),
        Layer3(_) => unimplemented!(),
    }

    let fd = get_fd();
    if fd == -1 {
        return Err(IoError::last_error());
    }
    let mut iface: bpf::ifreq = unsafe { mem::zeroed() };
    let mut i = 0;
    for c in network_interface.name.as_slice().bytes() {
        iface.ifr_name[i] = c as i8;
        i += 1;
    }

    let buflen = read_buffer_size as libc::c_uint;
    // NOTE Buffer length must be set before binding to an interface
    //      otherwise this will return Invalid Argument
    if unsafe { bpf::ioctl(fd, bpf::BIOCSBLEN, &buflen) } == -1 {
        let err = IoError::last_error();
        unsafe { libc::close(fd); }
        return Err(err);
    }

    // Set the interface to use
    if unsafe { bpf::ioctl(fd, bpf::BIOCSETIF, &iface) } == -1 {
        let err = IoError::last_error();
        unsafe { libc::close(fd); }
        return Err(err);
    }

    // Return from read as soon as packets are available - don't wait to fill the buffer
    let one: libc::c_uint = 1;
    if unsafe { bpf::ioctl(fd, bpf::BIOCIMMEDIATE, &one) } == -1 {
        let err = IoError::last_error();
        unsafe { libc::close(fd); }
        return Err(err);
    }

    let mut header_size = 0;

    // Get the device type
    let mut dlt: libc::c_uint = 0;
    if unsafe { bpf::ioctl(fd, bpf::BIOCGDLT, &mut dlt) } == -1 {
        let err = IoError::last_error();
        unsafe { libc::close(fd); }
        return Err(err);
    }

    // The loopback IoError::last_error()device does weird things
    // FIXME This should really just be another L2 packet header type
    if dlt == bpf::DLT_NULL {
        header_size = 4;

        // Allow packets to be read back after they are written
        match set_feedback(fd) {
            Err(e) => return Err(e),
            _ => ()
        }

    } else {
        // Don't fill in source MAC
        if unsafe { bpf::ioctl(fd, bpf::BIOCSHDRCMPLT, &one) } == -1 {
            let err = IoError::last_error();
            unsafe { libc::close(fd); }
            return Err(err);
        }
    }

    let fd = Arc::new(internal::FileDesc { fd: fd });
    let sender = DataLinkSenderImpl {
        fd: fd.clone(),
        write_buffer: repeat(0u8).take(write_buffer_size).collect(),
        header_size: header_size,
    };
    let receiver = DataLinkReceiverImpl {
        fd: fd,
        read_buffer: repeat(0u8).take(read_buffer_size).collect(),
        header_size: header_size,
    };

    Ok((sender, receiver))
}

pub struct DataLinkSenderImpl {
    fd: Arc<internal::FileDesc>,
    write_buffer: Vec<u8>,
    header_size: usize,
}

impl DataLinkSenderImpl {
    pub fn build_and_send<F>(&mut self, num_packets: usize, packet_size: usize,
                          func: &mut F) -> Option<IoResult<()>>
        where F : FnMut(MutableEthernetHeader)
    {
        let len = num_packets * (packet_size + self.header_size);
        if len >= self.write_buffer.len() {
            None
        } else {
            let min = cmp::min(self.write_buffer.len(), len);
            for chunk in self.write_buffer[..min]
                                          .chunks_mut(packet_size + self.header_size) {
                // If we're sending on the loopback device, the first 4 bytes must be set to
                // AF_INET
                if self.header_size == 4 {
                    unsafe {
                        *(chunk.as_mut_ptr() as *mut u32) = libc::AF_INET as u32;
                    }
                }
                {
                    let eh = MutableEthernetHeader::new(&mut chunk[self.header_size..]);
                    func(eh);
                }
                match unsafe { libc::write(self.fd.fd,
                                           chunk.as_ptr() as *const libc::c_void,
                                           chunk.len() as libc::size_t) } {
                    len if len == -1 => return Some(Err(IoError::last_error())),
                    _ => ()
                }
            }
            Some(Ok(()))
        }
    }

    pub fn send_to(&mut self, packet: EthernetHeader, _dst: Option<NetworkInterface>)
        -> Option<IoResult<()>> {
        match unsafe { libc::write(self.fd.fd,
                                   packet.packet().as_ptr() as *const libc::c_void,
                                   packet.packet().len() as libc::size_t) } {
            len if len == -1 => Some(Err(IoError::last_error())),
            _ => Some(Ok(()))
        }
    }
}

pub struct DataLinkReceiverImpl {
    fd: Arc<internal::FileDesc>,
    read_buffer: Vec<u8>,
    header_size: usize,
}

impl DataLinkReceiverImpl {
    pub fn iter<'a>(&'a mut self) -> DataLinkChannelIteratorImpl<'a> {
        let buflen = self.read_buffer.len();
        DataLinkChannelIteratorImpl {
            pc: self,
            // Enough room for minimally sized packets without reallocating
            packets: VecDeque::with_capacity(buflen / 64)
        }
    }
}

pub struct DataLinkChannelIteratorImpl<'a> {
    pc: &'a mut DataLinkReceiverImpl,
    packets: VecDeque<(usize, usize)>,
}

impl<'a> DataLinkChannelIteratorImpl<'a> {
    pub fn next<'c>(&'c mut self) -> IoResult<EthernetHeader<'c>> {
        if self.packets.is_empty() {
            let buflen = match unsafe {
                libc::read(self.pc.fd.fd,
                           self.pc.read_buffer.as_ptr() as *mut libc::c_void,
                           self.pc.read_buffer.len() as libc::size_t)
            } {
                len if len > 0 => len,
                _ => return Err(IoError::last_error())
            };
            let mut ptr = self.pc.read_buffer.as_mut_ptr();
            let end = unsafe { self.pc.read_buffer.as_ptr().offset(buflen as isize) };
            while (ptr as *const u8) < end {
                unsafe {
                    let packet: *const bpf::bpf_hdr = mem::transmute(ptr);
                    let start = ptr as isize +
                                (*packet).bh_hdrlen as isize -
                                self.pc.read_buffer.as_ptr() as isize;
                    self.packets.push_back((start as usize + self.pc.header_size,
                                      (*packet).bh_caplen as usize - self.pc.header_size));
                    let offset = (*packet).bh_hdrlen as isize + (*packet).bh_caplen as isize;
                    ptr = ptr.offset(bpf::BPF_WORDALIGN(offset));
                }
            }
        }
        let (start, len) = self.packets.pop_front().unwrap();
        Ok(EthernetHeader::new(&self.pc.read_buffer[start .. start + len]))
    }
}

<|endoftext|>


<commit_before>//
//  SOS: the Stupid Operating System
//  by Hawk Weisman (hi@hawkweisman.me)
//
//  Copyright (c) 2015 Hawk Weisman
//  Released under the terms of the MIT license. See `LICENSE` in the root
//  directory of this repository for more information.
//
//! Architecture-specific memory management.
use core::ptr::Unique;
use core::convert;

use ::memory::{VAddr, Addr};
use ::memory::paging::{Page, Mapper};

use alloc::{Allocator};

pub mod table;
pub mod entry;

use self::table::*;
use self::entry::Flags;

use core::ops;

pub const PAGE_SHIFT: u8 = 12;
/// The size of a page (4mb)
//  TODO: can we possibly rewrite this so that we can handle pages
//        in excess of 4 megs?
pub const PAGE_SIZE: u64 = 1 << PAGE_SHIFT; // 4096

extern {
    // TODO: It would be really nice if there was a less ugly way of doing
    // this... (read: after the Revolution when we add memory regions to the
    // heap programmatically.)
    pub static mut HEAP_BASE: u8;
    pub static mut HEAP_TOP: u8;
    // Of course, we will still need to export the kernel stack addresses like
    // this, but it would be nice if they could be, i dont know, not mut u8s
    // pointers, like God intended.
    pub static mut STACK_BASE: u8;
    pub static mut STACK_TOP: u8;
}

/// A frame (physical page)
#[derive(Copy, Clone, Debug, PartialEq, Eq, PartialOrd, Ord)]
pub struct Frame { pub number: u64 }

impl ops::Add<u64> for Frame {
    type Output = Frame;

    #[inline]
    fn add(self, amount: u64) -> Frame {
        Frame { number: self.number + amount }
    }
}

impl ops::Add<usize> for Frame {
    type Output = Frame;

    #[inline]
    fn add(self, amount: usize) -> Frame {
        Frame { number: self.number + amount as u64 }
    }
}

impl Frame {

    /// Returns the physical address where this frame starts.
    #[inline]
    pub const fn base_addr(&self) -> PAddr {
        PAddr(self.number << PAGE_SHIFT)
    }

    /// Returns a new frame containing `addr`
    #[inline]
    pub const fn containing(addr: PAddr) -> Frame {
        Frame { number: addr.0 / PAGE_SIZE }
    }
}

/// A physical (linear) memory address is a 64-bit unsigned integer
#[derive(Copy, Clone, Debug, Eq, Ord, PartialEq, PartialOrd)]
pub struct PAddr(u64);

impl Addr<u64> for PAddr { }

impl_addr! { PAddr, u64 }

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

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

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

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

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

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

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

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

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

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

    fn map_to_any<A>(&mut self, page: Page, flags: Flags, alloc: &mut A)
    where A: Allocator {
        // TODO: this is Definitely Wrong; our malloc just gives us
        //       pointers instead of allocating as frames that we coerce to
        //       pointers. might want to rewrite that.
        let frame = unsafe {
            alloc.allocate(PAGE_SIZE as usize, PAGE_SIZE as usize)
            // also, "PAGE_SIZE, PAGE_SIZE" is Almost Certainly the wrong size
            // and alignment for the allocation request - I think i left it that
            // way because i couldn't figure it out at the time and am an idiot.
            //      -- eliza
                    .expect("Couldn't map page, out of frames!")
        };
        unimplemented!()
        //self.map_to(page, frame, flags, alloc);
    }

}

impl ActivePML4 {

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

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

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

}
<commit_msg>[x86_64] frames can be turned into ptrs this is so we can reuse our heap allocator as the rust allocator (I hope)<commit_after>//
//  SOS: the Stupid Operating System
//  by Hawk Weisman (hi@hawkweisman.me)
//
//  Copyright (c) 2015 Hawk Weisman
//  Released under the terms of the MIT license. See `LICENSE` in the root
//  directory of this repository for more information.
//
//! Architecture-specific memory management.
use core::ptr::Unique;
use core::mem;

use ::memory::{VAddr, Addr};
use ::memory::paging::{Page, Mapper};

use alloc::{Allocator};

pub mod table;
pub mod entry;

use self::table::*;
use self::entry::Flags;

use core::ops;

pub const PAGE_SHIFT: u8 = 12;
/// The size of a page (4mb)
//  TODO: can we possibly rewrite this so that we can handle pages
//        in excess of 4 megs?
pub const PAGE_SIZE: u64 = 1 << PAGE_SHIFT; // 4096

extern {
    // TODO: It would be really nice if there was a less ugly way of doing
    // this... (read: after the Revolution when we add memory regions to the
    // heap programmatically.)
    pub static mut HEAP_BASE: u8;
    pub static mut HEAP_TOP: u8;
    // Of course, we will still need to export the kernel stack addresses like
    // this, but it would be nice if they could be, i dont know, not mut u8s
    // pointers, like God intended.
    pub static mut STACK_BASE: u8;
    pub static mut STACK_TOP: u8;
}

/// A frame (physical page)
#[derive(Copy, Clone, Debug, PartialEq, Eq, PartialOrd, Ord)]
pub struct Frame { pub number: u64 }

impl ops::Add<u64> for Frame {
    type Output = Frame;

    #[inline]
    fn add(self, amount: u64) -> Frame {
        Frame { number: self.number + amount }
    }
}

impl ops::Add<usize> for Frame {
    type Output = Frame;

    #[inline]
    fn add(self, amount: usize) -> Frame {
        Frame { number: self.number + amount as u64 }
    }
}

impl Frame {

    /// Returns the physical address where this frame starts.
    #[inline]
    pub const fn base_addr(&self) -> PAddr {
        PAddr(self.number << PAGE_SHIFT)
    }

    /// Returns a new frame containing `addr`
    #[inline]
    pub const fn containing(addr: PAddr) -> Frame {
        Frame { number: addr.0 / PAGE_SIZE }
    }

    /// Convert the frame into a raw pointer to the frame's base address
    #[inline]
    pub unsafe fn as_ptr<T>(&self) -> *const T {
        mem::transmute(self.base_addr())
    }

    /// Convert the frame into a raw mutable pointer to the frame's base address
    #[inline]
    pub unsafe fn as_mut_ptr<T>(&self) -> *mut T {
        *self.base_addr() as *mut u8 as *mut T
    }
}

/// A physical (linear) memory address is a 64-bit unsigned integer
#[derive(Copy, Clone, Debug, Eq, Ord, PartialEq, PartialOrd)]
pub struct PAddr(u64);

impl Addr<u64> for PAddr { }

impl_addr! { PAddr, u64 }

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

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

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

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

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

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

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

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

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

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

    fn map_to_any<A>(&mut self, page: Page, flags: Flags, alloc: &mut A)
    where A: Allocator {
        // TODO: this is Definitely Wrong; our malloc just gives us
        //       pointers instead of allocating as frames that we coerce to
        //       pointers. might want to rewrite that.
        let frame = unsafe {
            alloc.allocate(PAGE_SIZE as usize, PAGE_SIZE as usize)
            // also, "PAGE_SIZE, PAGE_SIZE" is Almost Certainly the wrong size
            // and alignment for the allocation request - I think i left it that
            // way because i couldn't figure it out at the time and am an idiot.
            //      -- eliza
                    .expect("Couldn't map page, out of frames!")
        };
        unimplemented!()
        //self.map_to(page, frame, flags, alloc);
    }

}

impl ActivePML4 {

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

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

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

}
<|endoftext|>


<commit_before>use super::Operation;
use buffer::Buffer;

/// A collection of operations run as a single/atomic operation.
///
/// Useful for composing smaller, related actions into a larger action, from a history/undo
/// standpoint. A common example of this is character-by-character insertions, which can
/// be undone as a whole, or word by word, instead of one character at a time.
///
/// Because this type implements the Operation trait, it can be placed into the history like
/// any other operation. It's a simple grouping type; it relies on its constituent operations
/// to handle all of their undo/redo implementation details. It exposes two methods on the
/// buffer type to signal the start and end of a group.
pub struct OperationGroup {
    operations: Vec<Box<Operation>>,
}

impl Operation for OperationGroup {
    /// Runs all of the group's individual operations, in order.
    fn run(&mut self, buffer: &mut Buffer) {
        for operation in &mut self.operations {
            operation.run(buffer);
        }
    }

    /// Reverses all of the group's individual operations, in reverse order.
    fn reverse(&mut self, buffer: &mut Buffer) {
        for operation in &mut self.operations.iter_mut().rev() {
            operation.reverse(buffer);
        }
    }

    /// Build a new operation group by manually cloning all of the groups individual operations.
    /// We can't derive this because operations are unsized and need some hand holding.
    fn clone_operation(&self) -> Box<Operation> {
        Box::new(OperationGroup{
            operations: self.operations.iter().map(|o| (*o).clone_operation()).collect()
        })
    }
}

impl OperationGroup {
    /// Adds an operation to the group.
    pub fn add(&mut self, operation: Box<Operation>) {
        self.operations.push(operation);
    }
}

pub fn new() -> OperationGroup {
    OperationGroup{ operations: Vec::new() }
}

impl Buffer {
    /// Tells the buffer to start tracking operations as a single unit, until
    /// end_operation_group is called. Any calls to insert or delete occurring within
    /// these will be undone/applied together when calling undo/redo, respectively.
    pub fn start_operation_group(&mut self) {
        // Create an operation group, if one doesn't already exist.
        match self.operation_group {
            Some(_) => (),
            None => {
                self.operation_group = Some(new());
            }
        }
    }

    /// Tells the buffer to stop tracking operations as a single unit, since
    /// start_operation_group was called. Any calls to insert or delete occurring within
    /// these will be undone/applied together when calling undo/redo, respectively.
    pub fn end_operation_group(&mut self) {
        // Push an open operation group on to the history stack, if one exists.
        match self.operation_group.take() {
            Some(group) => self.history.add(Box::new(group)),
            None => (),
        }
    }
}

#[cfg(test)]
mod tests {
    use super::new;
    use buffer::operations::insert;
    use buffer::Position;
    use buffer::operation::Operation;

    #[test]
    fn run_and_reverse_call_themselves_on_all_operations() {
        let mut group = new();
        let mut buffer = ::buffer::new();

        // Push two insert operations into the group.
        let first = Box::new(insert::new("something".to_string(), Position{ line: 0, offset: 0 }));
        let second = Box::new(insert::new(" else".to_string(), Position{ line: 0, offset: 9 }));
        group.add(first);
        group.add(second);

        // Run the operation group.
        group.run(&mut buffer);

        assert_eq!(buffer.data(), "something else");

        // Reverse the operation group.
        group.reverse(&mut buffer);

        assert_eq!(buffer.data(), "");
    }
}
<commit_msg>Update end_operation_group to drop empty operation groups.<commit_after>use super::Operation;
use buffer::Buffer;

/// A collection of operations run as a single/atomic operation.
///
/// Useful for composing smaller, related actions into a larger action, from a history/undo
/// standpoint. A common example of this is character-by-character insertions, which can
/// be undone as a whole, or word by word, instead of one character at a time.
///
/// Because this type implements the Operation trait, it can be placed into the history like
/// any other operation. It's a simple grouping type; it relies on its constituent operations
/// to handle all of their undo/redo implementation details. It exposes two methods on the
/// buffer type to signal the start and end of a group.
pub struct OperationGroup {
    operations: Vec<Box<Operation>>,
}

impl Operation for OperationGroup {
    /// Runs all of the group's individual operations, in order.
    fn run(&mut self, buffer: &mut Buffer) {
        for operation in &mut self.operations {
            operation.run(buffer);
        }
    }

    /// Reverses all of the group's individual operations, in reverse order.
    fn reverse(&mut self, buffer: &mut Buffer) {
        for operation in &mut self.operations.iter_mut().rev() {
            operation.reverse(buffer);
        }
    }

    /// Build a new operation group by manually cloning all of the groups individual operations.
    /// We can't derive this because operations are unsized and need some hand holding.
    fn clone_operation(&self) -> Box<Operation> {
        Box::new(OperationGroup{
            operations: self.operations.iter().map(|o| (*o).clone_operation()).collect()
        })
    }
}

impl OperationGroup {
    /// Adds an operation to the group.
    pub fn add(&mut self, operation: Box<Operation>) {
        self.operations.push(operation);
    }
}

pub fn new() -> OperationGroup {
    OperationGroup{ operations: Vec::new() }
}

impl Buffer {
    /// Tells the buffer to start tracking operations as a single unit, until
    /// end_operation_group is called. Any calls to insert or delete occurring within
    /// these will be undone/applied together when calling undo/redo, respectively.
    pub fn start_operation_group(&mut self) {
        // Create an operation group, if one doesn't already exist.
        match self.operation_group {
            Some(_) => (),
            None => {
                self.operation_group = Some(new());
            }
        }
    }

    /// Tells the buffer to stop tracking operations as a single unit, since
    /// start_operation_group was called. Any calls to insert or delete occurring within
    /// these will be undone/applied together when calling undo/redo, respectively.
    pub fn end_operation_group(&mut self) {
        // Push an open operation group on to the history stack, if one exists.
        match self.operation_group.take() {
            Some(group) => {
                if !group.operations.is_empty() {
                    self.history.add(Box::new(group))
                }
            },
            None => (),
        }
    }
}

#[cfg(test)]
mod tests {
    use super::new;
    use buffer::operations::insert;
    use buffer::Position;
    use buffer::operation::Operation;

    #[test]
    fn run_and_reverse_call_themselves_on_all_operations() {
        let mut group = new();
        let mut buffer = ::buffer::new();

        // Push two insert operations into the group.
        let first = Box::new(insert::new("something".to_string(), Position{ line: 0, offset: 0 }));
        let second = Box::new(insert::new(" else".to_string(), Position{ line: 0, offset: 9 }));
        group.add(first);
        group.add(second);

        // Run the operation group.
        group.run(&mut buffer);

        assert_eq!(buffer.data(), "something else");

        // Reverse the operation group.
        group.reverse(&mut buffer);

        assert_eq!(buffer.data(), "");
    }

    #[test]
    fn end_operation_group_drops_group_if_empty() {
        let mut buffer = ::buffer::new();
        buffer.insert("amp");

        // Create an empty operation group that
        // shouldn't be added to the buffer history.
        buffer.start_operation_group();
        buffer.end_operation_group();

        // Undo the last change, which should be the initial
        // insert, if the empty operation group was ignored.
        buffer.undo();
        assert_eq!(buffer.data(), "");
    }
}
<|endoftext|>


<commit_before>use crate::controllers::frontend_prelude::*;

use crate::models::{CrateOwner, OwnerKind, User};
use crate::schema::{crate_owners, crates, users};
use crate::sql::lower;
use crate::views::EncodablePublicUser;

/// Handles the `GET /users/:user_id` route.
pub fn show(req: &mut dyn RequestExt) -> EndpointResult {
    use self::users::dsl::{gh_login, id, users};

    let name = lower(&req.params()["user_id"]);
    let conn = req.db_read_prefer_primary()?;
    let user: User = users
        .filter(lower(gh_login).eq(name))
        .order(id.desc())
        .first(&*conn)?;

    Ok(req.json(&json!({ "user": EncodablePublicUser::from(user) })))
}

/// Handles the `GET /users/:user_id/stats` route.
pub fn stats(req: &mut dyn RequestExt) -> EndpointResult {
    use diesel::dsl::sum;

    let user_id = &req.params()["user_id"]
        .parse::<i32>()
        .map_err(|err| err.chain(bad_request("invalid user_id")))?;
    let conn = req.db_write()?;

    let data: i64 = CrateOwner::by_owner_kind(OwnerKind::User)
        .inner_join(crates::table)
        .filter(crate_owners::owner_id.eq(user_id))
        .select(sum(crates::downloads))
        .first::<Option<i64>>(&*conn)?
        .unwrap_or(0);

    Ok(req.json(&json!({ "total_downloads": data })))
}
<commit_msg>Use `db_read_prefer_primary()` for `GET /users/:user_id/stats` endpoint<commit_after>use crate::controllers::frontend_prelude::*;

use crate::models::{CrateOwner, OwnerKind, User};
use crate::schema::{crate_owners, crates, users};
use crate::sql::lower;
use crate::views::EncodablePublicUser;

/// Handles the `GET /users/:user_id` route.
pub fn show(req: &mut dyn RequestExt) -> EndpointResult {
    use self::users::dsl::{gh_login, id, users};

    let name = lower(&req.params()["user_id"]);
    let conn = req.db_read_prefer_primary()?;
    let user: User = users
        .filter(lower(gh_login).eq(name))
        .order(id.desc())
        .first(&*conn)?;

    Ok(req.json(&json!({ "user": EncodablePublicUser::from(user) })))
}

/// Handles the `GET /users/:user_id/stats` route.
pub fn stats(req: &mut dyn RequestExt) -> EndpointResult {
    use diesel::dsl::sum;

    let user_id = &req.params()["user_id"]
        .parse::<i32>()
        .map_err(|err| err.chain(bad_request("invalid user_id")))?;
    let conn = req.db_read_prefer_primary()?;

    let data: i64 = CrateOwner::by_owner_kind(OwnerKind::User)
        .inner_join(crates::table)
        .filter(crate_owners::owner_id.eq(user_id))
        .select(sum(crates::downloads))
        .first::<Option<i64>>(&*conn)?
        .unwrap_or(0);

    Ok(req.json(&json!({ "total_downloads": data })))
}
<|endoftext|>


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

//! The binomial distribution.

use Rng;
use distributions::Distribution;
use distributions::log_gamma::log_gamma;
use std::f64::consts::PI;

/// The binomial distribution `Binomial(n, p)`.
///
/// This distribution has density function:
/// `f(k) = n!/(k! (n-k)!) p^k (1-p)^(n-k)` for `k >= 0`.
///
/// # Example
///
/// ```
/// use rand::distributions::{Binomial, Distribution};
///
/// let bin = Binomial::new(20, 0.3);
/// let v = bin.sample(&mut rand::thread_rng());
/// println!("{} is from a binomial distribution", v);
/// ```
#[derive(Clone, Copy, Debug)]
pub struct Binomial {
    /// Number of trials.
    n: u64,
    /// Probability of success.
    p: f64,
}

impl Binomial {
    /// Construct a new `Binomial` with the given shape parameters `n` (number
    /// of trials) and `p` (probability of success).
    ///
    /// Panics if `p <= 0` or `p >= 1`.
    pub fn new(n: u64, p: f64) -> Binomial {
        assert!(p > 0.0, "Binomial::new called with p <= 0");
        assert!(p < 1.0, "Binomial::new called with p >= 1");
        Binomial { n, p }
    }
}

impl Distribution<u64> for Binomial {
    fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> u64 {
        // binomial distribution is symmetrical with respect to p -> 1-p, k -> n-k
        // switch p so that it is less than 0.5 - this allows for lower expected values
        // we will just invert the result at the end
        let p = if self.p <= 0.5 {
            self.p
        } else {
            1.0 - self.p
        };

        // expected value of the sample
        let expected = self.n as f64 * p;

        let result =
            // for low expected values we just simulate n drawings
            if expected < 25.0 {
                let mut lresult = 0.0;
                for _ in 0 .. self.n {
                    if rng.gen_bool(p) {
                        lresult += 1.0;
                    }
                }
                lresult
            }
            // high expected value - do the rejection method
            else {
                // prepare some cached values
                let float_n = self.n as f64;
                let ln_fact_n = log_gamma(float_n + 1.0);
                let pc = 1.0 - p;
                let log_p = p.ln();
                let log_pc = pc.ln();
                let sq = (expected * (2.0 * pc)).sqrt();

                let mut lresult;

                loop {
                    let mut comp_dev: f64;
                    // we use the lorentzian distribution as the comparison distribution
                    // f(x) ~ 1/(1+x/^2)
                    loop {
                        // draw from the lorentzian distribution
                        comp_dev = (PI*rng.gen::<f64>()).tan();
                        // shift the peak of the comparison ditribution
                        lresult = expected + sq * comp_dev;
                        // repeat the drawing until we are in the range of possible values
                        if lresult >= 0.0 && lresult < float_n + 1.0 {
                            break;
                        }
                    }

                    // the result should be discrete
                    lresult = lresult.floor();

                    let log_binomial_dist = ln_fact_n - log_gamma(lresult+1.0) -
                        log_gamma(float_n - lresult + 1.0) + lresult*log_p + (float_n - lresult)*log_pc;
                    // this is the binomial probability divided by the comparison probability
                    // we will generate a uniform random value and if it is larger than this,
                    // we interpret it as a value falling out of the distribution and repeat
                    let comparison_coeff = (log_binomial_dist.exp() * sq) * (1.2 * (1.0 + comp_dev*comp_dev));

                    if comparison_coeff >= rng.gen() {
                        break;
                    }
                }

                lresult
            };

        // invert the result for p < 0.5
        if p != self.p {
            self.n - result as u64
        } else {
            result as u64
        }
    }
}

#[cfg(test)]
mod test {
    use Rng;
    use distributions::Distribution;
    use super::Binomial;

    fn test_binomial_mean_and_variance<R: Rng>(n: u64, p: f64, rng: &mut R) {
        let binomial = Binomial::new(n, p);

        let expected_mean = n as f64 * p;
        let expected_variance = n as f64 * p * (1.0 - p);

        let mut results = [0.0; 1000];
        for i in results.iter_mut() { *i = binomial.sample(rng) as f64; }

        let mean = results.iter().sum::<f64>() / results.len() as f64;
        assert!((mean as f64 - expected_mean).abs() < expected_mean / 50.0);

        let variance =
            results.iter().map(|x| (x - mean) * (x - mean)).sum::<f64>()
            / results.len() as f64;
        assert!((variance - expected_variance).abs() < expected_variance / 10.0);
    }

    #[test]
    fn test_binomial() {
        let mut rng = ::test::rng(123);
        test_binomial_mean_and_variance(150, 0.1, &mut rng);
        test_binomial_mean_and_variance(70, 0.6, &mut rng);
        test_binomial_mean_and_variance(40, 0.5, &mut rng);
        test_binomial_mean_and_variance(20, 0.7, &mut rng);
        test_binomial_mean_and_variance(20, 0.5, &mut rng);
    }

    #[test]
    #[should_panic]
    fn test_binomial_invalid_lambda_zero() {
        Binomial::new(20, 0.0);
    }

    #[test]
    #[should_panic]
    fn test_binomial_invalid_lambda_neg() {
        Binomial::new(20, -10.0);
    }
}
<commit_msg>Binomial distribution: optimisations<commit_after>// Copyright 2016-2017 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// https://rust-lang.org/COPYRIGHT.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// https://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or https://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

//! The binomial distribution.

use Rng;
use distributions::{Distribution, Bernoulli};
use distributions::log_gamma::log_gamma;
use std::f64::consts::PI;

/// The binomial distribution `Binomial(n, p)`.
///
/// This distribution has density function:
/// `f(k) = n!/(k! (n-k)!) p^k (1-p)^(n-k)` for `k >= 0`.
///
/// # Example
///
/// ```
/// use rand::distributions::{Binomial, Distribution};
///
/// let bin = Binomial::new(20, 0.3);
/// let v = bin.sample(&mut rand::thread_rng());
/// println!("{} is from a binomial distribution", v);
/// ```
#[derive(Clone, Copy, Debug)]
pub struct Binomial {
    /// Number of trials.
    n: u64,
    /// Probability of success.
    p: f64,
}

impl Binomial {
    /// Construct a new `Binomial` with the given shape parameters `n` (number
    /// of trials) and `p` (probability of success).
    ///
    /// Panics if `p <= 0` or `p >= 1`.
    pub fn new(n: u64, p: f64) -> Binomial {
        assert!(p > 0.0, "Binomial::new called with p <= 0");
        assert!(p < 1.0, "Binomial::new called with p >= 1");
        Binomial { n, p }
    }
}

impl Distribution<u64> for Binomial {
    fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> u64 {
        // For low n, it is faster to sample directly. For both methods,
        // performance is independent of p. On Intel Haswell CPU this method
        // appears to be faster for approx n < 300.
        if self.n < 300 {
            let mut result = 0;
            let d = Bernoulli::new(self.p);
            for _ in 0 .. self.n {
                result += rng.sample(d) as u32;
            }
            return result as u64;
        }

        // binomial distribution is symmetrical with respect to p -> 1-p, k -> n-k
        // switch p so that it is less than 0.5 - this allows for lower expected values
        // we will just invert the result at the end
        let p = if self.p <= 0.5 {
            self.p
        } else {
            1.0 - self.p
        };

        // prepare some cached values
        let float_n = self.n as f64;
        let ln_fact_n = log_gamma(float_n + 1.0);
        let pc = 1.0 - p;
        let log_p = p.ln();
        let log_pc = pc.ln();
        let expected = self.n as f64 * p;
        let sq = (expected * (2.0 * pc)).sqrt();

        let mut lresult;

        loop {
            let mut comp_dev: f64;
            // we use the lorentzian distribution as the comparison distribution
            // f(x) ~ 1/(1+x/^2)
            loop {
                // draw from the lorentzian distribution
                comp_dev = (PI*rng.gen::<f64>()).tan();
                // shift the peak of the comparison ditribution
                lresult = expected + sq * comp_dev;
                // repeat the drawing until we are in the range of possible values
                if lresult >= 0.0 && lresult < float_n + 1.0 {
                    break;
                }
            }

            // the result should be discrete
            lresult = lresult.floor();

            let log_binomial_dist = ln_fact_n - log_gamma(lresult+1.0) -
                log_gamma(float_n - lresult + 1.0) + lresult*log_p + (float_n - lresult)*log_pc;
            // this is the binomial probability divided by the comparison probability
            // we will generate a uniform random value and if it is larger than this,
            // we interpret it as a value falling out of the distribution and repeat
            let comparison_coeff = (log_binomial_dist.exp() * sq) * (1.2 * (1.0 + comp_dev*comp_dev));

            if comparison_coeff >= rng.gen() {
                break;
            }
        }

        // invert the result for p < 0.5
        if p != self.p {
            self.n - lresult as u64
        } else {
            lresult as u64
        }
    }
}

#[cfg(test)]
mod test {
    use Rng;
    use distributions::Distribution;
    use super::Binomial;

    fn test_binomial_mean_and_variance<R: Rng>(n: u64, p: f64, rng: &mut R) {
        let binomial = Binomial::new(n, p);

        let expected_mean = n as f64 * p;
        let expected_variance = n as f64 * p * (1.0 - p);

        let mut results = [0.0; 1000];
        for i in results.iter_mut() { *i = binomial.sample(rng) as f64; }

        let mean = results.iter().sum::<f64>() / results.len() as f64;
        assert!((mean as f64 - expected_mean).abs() < expected_mean / 50.0);

        let variance =
            results.iter().map(|x| (x - mean) * (x - mean)).sum::<f64>()
            / results.len() as f64;
        assert!((variance - expected_variance).abs() < expected_variance / 10.0);
    }

    #[test]
    fn test_binomial() {
        let mut rng = ::test::rng(123);
        test_binomial_mean_and_variance(150, 0.1, &mut rng);
        test_binomial_mean_and_variance(70, 0.6, &mut rng);
        test_binomial_mean_and_variance(40, 0.5, &mut rng);
        test_binomial_mean_and_variance(20, 0.7, &mut rng);
        test_binomial_mean_and_variance(20, 0.5, &mut rng);
    }

    #[test]
    #[should_panic]
    fn test_binomial_invalid_lambda_zero() {
        Binomial::new(20, 0.0);
    }

    #[test]
    #[should_panic]
    fn test_binomial_invalid_lambda_neg() {
        Binomial::new(20, -10.0);
    }
}
<|endoftext|>