<commit_before>use chrono::prelude::*;
use chrono::Duration;

use rusoto_core;

use rusoto_core::{AwsCredentials, CredentialsError,
    ProvideAwsCredentials, DispatchSignedRequest};
use ::{AssumeRoleRequest, AssumeRoleResponse, AssumeRoleError,
    AssumeRoleWithSAMLRequest, AssumeRoleWithSAMLResponse, AssumeRoleWithSAMLError,
    AssumeRoleWithWebIdentityRequest, AssumeRoleWithWebIdentityResponse, AssumeRoleWithWebIdentityError,
    DecodeAuthorizationMessageRequest, DecodeAuthorizationMessageResponse, DecodeAuthorizationMessageError,
    GetCallerIdentityRequest, GetCallerIdentityResponse, GetCallerIdentityError,
    GetFederationTokenRequest, GetFederationTokenResponse, GetFederationTokenError,
    GetSessionTokenRequest, GetSessionTokenResponse, GetSessionTokenError, Sts,
    StsClient};

pub const DEFAULT_DURATION_SECONDS: i32 = 3600;
pub const DEFAULT_ROLE_DURATION_SECONDS: i32 = 900;

/// Trait for conversions from STS Credentials to AWS Credentials.
pub trait NewAwsCredsForStsCreds {
    /// Creates an [AwsCredentials](../struct.AwsCredentials.html) from a [Credentials](struct.Credentials.html)
    /// Returns a [CredentialsError](../struct.CredentialsError.html) in case of an error.
    fn new_for_credentials(sts_creds: ::Credentials) -> Result<AwsCredentials, CredentialsError>;
}

impl NewAwsCredsForStsCreds for AwsCredentials {
    fn new_for_credentials(sts_creds: ::Credentials) -> Result<AwsCredentials, CredentialsError> {
        let expires_at = try!(sts_creds.expiration.parse::<DateTime<Utc>>().map_err(CredentialsError::from));

        Ok(AwsCredentials::new(
            sts_creds.access_key_id,
            sts_creds.secret_access_key,
            Some(sts_creds.session_token),
            expires_at))
    }
}

// Trait that defines the STS Client API without any type parameters or assumptions about implementation.
// This is an internal type used to box the [StsClient](struct.StsClient.html) provided in the session token providers' constructors.
pub trait StsSessionCredentialsClient {
    fn assume_role(&self, input: &AssumeRoleRequest) -> Result<AssumeRoleResponse, AssumeRoleError>;

    fn assume_role_with_saml(&self, input: &AssumeRoleWithSAMLRequest) -> Result<AssumeRoleWithSAMLResponse, AssumeRoleWithSAMLError>;

    fn assume_role_with_web_identity(&self, input: &AssumeRoleWithWebIdentityRequest) -> Result<AssumeRoleWithWebIdentityResponse, AssumeRoleWithWebIdentityError>;

    fn decode_authorization_message(&self, input: &DecodeAuthorizationMessageRequest) -> Result<DecodeAuthorizationMessageResponse, DecodeAuthorizationMessageError>;

    fn get_caller_identity(&self, input: &GetCallerIdentityRequest) -> Result<GetCallerIdentityResponse, GetCallerIdentityError> ;

    fn get_federation_token(&self, input: &GetFederationTokenRequest) -> Result<GetFederationTokenResponse, GetFederationTokenError>;

    fn get_session_token(&self, input: &GetSessionTokenRequest) -> Result<GetSessionTokenResponse, GetSessionTokenError>;
}

impl<T> StsSessionCredentialsClient for T where T: Sts {
    fn assume_role(&self, input: &AssumeRoleRequest) -> Result<AssumeRoleResponse, AssumeRoleError> {
        T::assume_role(self, input)
    }

    fn assume_role_with_saml(&self, input: &AssumeRoleWithSAMLRequest) -> Result<AssumeRoleWithSAMLResponse, AssumeRoleWithSAMLError> {
        T::assume_role_with_saml(self, input)
    }

    fn assume_role_with_web_identity(&self, input: &AssumeRoleWithWebIdentityRequest) -> Result<AssumeRoleWithWebIdentityResponse, AssumeRoleWithWebIdentityError> {
        T::assume_role_with_web_identity(self, input)
    }

    fn decode_authorization_message(&self, input: &DecodeAuthorizationMessageRequest) -> Result<DecodeAuthorizationMessageResponse, DecodeAuthorizationMessageError> {
        T::decode_authorization_message(self, input)
    }

    fn get_caller_identity(&self, input: &GetCallerIdentityRequest) -> Result<GetCallerIdentityResponse, GetCallerIdentityError> {
        T::get_caller_identity(self, input)
    }

    fn get_federation_token(&self, input: &GetFederationTokenRequest) -> Result<GetFederationTokenResponse, GetFederationTokenError> {
        T::get_federation_token(self, input)
    }

    fn get_session_token(&self, input: &GetSessionTokenRequest) -> Result<GetSessionTokenResponse, GetSessionTokenError> {
        T::get_session_token(self, input)
    }
}

/// [AwsCredentials](../struct.AwsCredentials.html) provider that calls
/// `GetSessionToken` using the provided [StsClient](struct.StsClient.html).
/// To use with MFA, pass in the MFA serial number then set the MFA code.
/// You will need to ensure the provider has a valid code each time you
/// acquire a new STS token.
pub struct StsSessionCredentialsProvider {
    sts_client: Box<StsSessionCredentialsClient>,
    session_duration: Duration,
    mfa_serial: Option<String>,
    mfa_code: Option<String>,
}

impl StsSessionCredentialsProvider {
    /// Creates a new `StsSessionCredentialsProvider` with the given
    /// [StsClient](struct.StsClient.html) and session parameters.
    ///
    /// * `sts_client` - The [StsClient](struct.StsClient.html) to use to acquire session tokens.
    /// * `duration` - The duration of the session tokens. Default 1 hour.
    /// * `mfa_serial` - Optional MFA hardware device serial number or virtual device ARN. Set the MFA code with `set_mfa_code`.
    pub fn new<P,D>(sts_client: StsClient<P,D>,
            duration: Option<Duration>,
            mfa_serial: Option<String>,
            ) -> StsSessionCredentialsProvider
            where P: ProvideAwsCredentials + 'static, D: DispatchSignedRequest + 'static {
        StsSessionCredentialsProvider {
            sts_client: Box::new(sts_client),
            session_duration: duration.unwrap_or(Duration::seconds(DEFAULT_DURATION_SECONDS as i64)),
            mfa_serial: mfa_serial,
            mfa_code: None,
        }
    }

    /// Set the MFA code for use when acquiring session tokens.
    pub fn set_mfa_code<S>(&mut self, code: S) where S: Into<String> {
        self.mfa_code = Some(code.into());
    }

    /// Clear the MFA code.
    pub fn clear_mfa_code(&mut self) {
        self.mfa_code = None;
    }

    /// Calls `GetSessionToken` to get a session token from the STS Api.
    /// Optionally uses MFA if the MFA serial number and code are set.
    pub fn get_session_token(&self) -> Result<AwsCredentials, CredentialsError> {
        match self.sts_client.get_session_token(
            &GetSessionTokenRequest {
                serial_number: self.mfa_serial.clone(),
                token_code: self.mfa_code.clone(),
                duration_seconds: Some(self.session_duration.num_seconds() as i64),
                ..Default::default()
            }) {
            Ok(resp) => {
                let creds = try!(resp.credentials.ok_or(CredentialsError::new("no credentials in response")));

                AwsCredentials::new_for_credentials(creds)
            },
            err =>
                Err(CredentialsError::new(format!("StsProvider get_session_token error: {:?}", err)))
        }
    }
}

impl ProvideAwsCredentials for StsSessionCredentialsProvider {
    fn credentials(&self) -> Result<AwsCredentials, CredentialsError> {
        self.get_session_token()
    }
}

/// [AwsCredentials](../struct.AwsCredentials.html) provider that calls
/// `AssumeRole` using the provided [StsClient](struct.StsClient.html).
/// To use with MFA, pass in the MFA serial number then set the MFA code.
/// You will need to ensure the provider has a valid code each time you
/// acquire a new STS token.
pub struct StsAssumeRoleSessionCredentialsProvider {
    sts_client: Box<StsSessionCredentialsClient>,
    role_arn: String,
    session_name: String,
    external_id: Option<String>,
    session_duration: Duration,
    scope_down_policy: Option<String>,
    mfa_serial: Option<String>,
    mfa_code: Option<String>,
}

impl StsAssumeRoleSessionCredentialsProvider {
    /// Creates a new `StsAssumeRoleSessionCredentialsProvider` with the given
    /// [StsClient](struct.StsClient.html) and session parameters.
    ///
    /// * `sts_client` - [StsClient](struct.StsClient.html) to use to acquire session tokens.
    /// * `role_arn` - The ARN of the role to assume.
    /// * `session_name` - An identifier for the assumed role session. Minimum length of 2. Maximum length of 64. Pattern: `[\w+=,.@-]*`
    /// * `external_id` -
    /// * `session_duration` - Duration of session tokens. Default 1 hour.
    /// * `scope_down_policy` - Optional inline IAM policy in JSON format to further restrict the access granted to the negotiated session.
    /// * `mfa_serial` - Optional MFA hardware device serial number or virtual device ARN. Use `set_mfa_code` to set the MFA code.
    pub fn new<P,D>(sts_client: StsClient<P,D>,
            role_arn: String,
            session_name: String,
            external_id: Option<String>,
            session_duration: Option<Duration>,
            scope_down_policy: Option<String>,
            mfa_serial: Option<String>
            )
            -> StsAssumeRoleSessionCredentialsProvider
            where P: ProvideAwsCredentials + 'static, D: DispatchSignedRequest + 'static {
        StsAssumeRoleSessionCredentialsProvider {
            sts_client: Box::new(sts_client),
            role_arn: role_arn,
            session_name: session_name,
            external_id: external_id,
            session_duration: session_duration.unwrap_or(Duration::seconds(DEFAULT_ROLE_DURATION_SECONDS as i64)),
            scope_down_policy: scope_down_policy,
            mfa_serial: mfa_serial,
            mfa_code: None,
        }
    }

    /// Set the MFA code for use when acquiring session tokens.
    pub fn set_mfa_code<S>(&mut self, code: S) where S: Into<String> {
        self.mfa_code = Some(code.into());
    }

    /// Clear the MFA code.
    pub fn clear_mfa_code(&mut self) {
        self.mfa_code = None;
    }

    /// Calls `AssumeRole` to get a session token from the STS Api.
    /// Optionally uses MFA if the MFA serial number and code are set.
    pub fn assume_role(&self) -> Result<AwsCredentials, CredentialsError> {
        match self.sts_client.assume_role(&AssumeRoleRequest {
            role_arn: self.role_arn.clone(),
            role_session_name: self.session_name.clone(),
            duration_seconds: Some(self.session_duration.num_seconds() as i64),
            external_id: self.external_id.clone(),
            policy: self.scope_down_policy.clone(),
            serial_number: self.mfa_serial.clone(),
            token_code: self.mfa_code.clone(),
            ..Default::default()
        }) {
            Err(err) =>
                Err(CredentialsError::new(format!("Sts AssumeRoleError: {:?}", err))),
            Ok(resp) => {
                let creds = try!(resp.credentials.ok_or(CredentialsError::new("no credentials in response")));

                AwsCredentials::new_for_credentials(creds)
            }
        }
    }
}

impl ProvideAwsCredentials for StsAssumeRoleSessionCredentialsProvider {
    fn credentials(&self) -> Result<AwsCredentials, CredentialsError> {
        self.assume_role()
    }
}

/// [AwsCredentials](../struct.AwsCredentials.html) provider that calls
/// `AssumeRoleWithWebIdentity` using the provided [StsClient](struct.StsClient.html).
pub struct StsWebIdentityFederationSessionCredentialsProvider {
    sts_client: Box<StsSessionCredentialsClient>,
    wif_token: String,
    wif_provider: Option<String>,
    role_arn: String,
    session_name: String,
    session_duration: Duration,
    scope_down_policy: Option<String>,
}

impl StsWebIdentityFederationSessionCredentialsProvider {
    /// Creates a new `StsWebIdentityFederationSessionCredentialsProvider` with the given
    /// [StsClient](struct.StsClient.html) and session parameters.
    ///
    /// * `sts_client` - The [StsClient](struct.StsClient.html) to use to acquire session tokens.
    /// * `wif_token` - The OAuth 2.0 access token or OpenID Connect ID token that is provided by the identity provider.
    /// * `wif_provider` - The fully qualified host component of the domain name of the identity provider. Only for OAuth 2.0 access tokens. Do not include URL schemes and port numbers.
    /// * `role_arn` - The ARN of the role to assume.
    /// * `session_name` - An identifier for the assumed role session. Minimum length of 2. Maximum length of 64. Pattern: `[\w+=,.@-]*`
    /// * `session_duration` - Duration of session tokens. Default 1 hour.
    /// * `scope_down_policy` - Optional inline IAM policy in JSON format to further restrict the access granted to the negotiated session.
    pub fn new<P,D>(sts_client: StsClient<P,D>,
            wif_token: String,
            wif_provider: Option<String>,
            role_arn: String,
            session_name: String,
            session_duration: Option<Duration>,
            scope_down_policy: Option<String>
            ) -> StsWebIdentityFederationSessionCredentialsProvider
            where P: ProvideAwsCredentials + 'static, D: DispatchSignedRequest + 'static {
        StsWebIdentityFederationSessionCredentialsProvider {
            sts_client: Box::new(sts_client),
            wif_token: wif_token,
            wif_provider: wif_provider,
            role_arn: role_arn,
            session_name: session_name,
            session_duration: session_duration.unwrap_or(Duration::seconds(DEFAULT_DURATION_SECONDS as i64)),
            scope_down_policy: scope_down_policy
        }
    }

    /// Calls `AssumeRoleWithWebIdentity` to get a session token from the STS Api.
    pub fn assume_role_with_web_identity(&self) -> Result<AwsCredentials, CredentialsError> {
        match self.sts_client.assume_role_with_web_identity(&AssumeRoleWithWebIdentityRequest {
            web_identity_token: self.wif_token.clone(),
            provider_id: self.wif_provider.clone(),
            role_arn: self.role_arn.clone(),
            role_session_name: self.session_name.clone(),
            duration_seconds: Some(self.session_duration.num_seconds() as i64),
            policy: self.scope_down_policy.clone(),
            ..Default::default()
        }) {
            Err(err) =>
                Err(CredentialsError::new(format!("Sts AssumeRoleWithWebIdentityError: {:?}", err))),
            Ok(resp) => {
                let creds = try!(resp.credentials.ok_or(CredentialsError::new("no credentials in response")));

                let mut aws_creds = try!(AwsCredentials::new_for_credentials(creds));

                if let Some(subject_from_wif) = resp.subject_from_web_identity_token {
                    aws_creds.claims_mut().insert(rusoto_core::credential::claims::SUBJECT.to_owned(), subject_from_wif);
                }

                if let Some(audience) = resp.audience {
                    aws_creds.claims_mut().insert(rusoto_core::credential::claims::AUDIENCE.to_owned(), audience);
                }

                if let Some(issuer) = resp.provider {
                    aws_creds.claims_mut().insert(rusoto_core::credential::claims::ISSUER.to_owned(), issuer);
                }

                Ok(aws_creds)
            }
        }
    }
}

impl ProvideAwsCredentials for StsWebIdentityFederationSessionCredentialsProvider {
    fn credentials(&self) -> Result<AwsCredentials, CredentialsError> {
        self.assume_role_with_web_identity()
    }
}<commit_msg>Fix broken links sts docs<commit_after>use chrono::prelude::*;
use chrono::Duration;

use rusoto_core;

use rusoto_core::{AwsCredentials, CredentialsError,
    ProvideAwsCredentials, DispatchSignedRequest};
use ::{AssumeRoleRequest, AssumeRoleResponse, AssumeRoleError,
    AssumeRoleWithSAMLRequest, AssumeRoleWithSAMLResponse, AssumeRoleWithSAMLError,
    AssumeRoleWithWebIdentityRequest, AssumeRoleWithWebIdentityResponse, AssumeRoleWithWebIdentityError,
    DecodeAuthorizationMessageRequest, DecodeAuthorizationMessageResponse, DecodeAuthorizationMessageError,
    GetCallerIdentityRequest, GetCallerIdentityResponse, GetCallerIdentityError,
    GetFederationTokenRequest, GetFederationTokenResponse, GetFederationTokenError,
    GetSessionTokenRequest, GetSessionTokenResponse, GetSessionTokenError, Sts,
    StsClient};

pub const DEFAULT_DURATION_SECONDS: i32 = 3600;
pub const DEFAULT_ROLE_DURATION_SECONDS: i32 = 900;

/// Trait for conversions from STS Credentials to AWS Credentials.
pub trait NewAwsCredsForStsCreds {
    /// Creates an [AwsCredentials](../rusoto_credential/struct.AwsCredentials.html) from a [Credentials](struct.Credentials.html)
    /// Returns a [CredentialsError](../rusoto_credential/struct.CredentialsError.html) in case of an error.
    fn new_for_credentials(sts_creds: ::Credentials) -> Result<AwsCredentials, CredentialsError>;
}

impl NewAwsCredsForStsCreds for AwsCredentials {
    fn new_for_credentials(sts_creds: ::Credentials) -> Result<AwsCredentials, CredentialsError> {
        let expires_at = try!(sts_creds.expiration.parse::<DateTime<Utc>>().map_err(CredentialsError::from));

        Ok(AwsCredentials::new(
            sts_creds.access_key_id,
            sts_creds.secret_access_key,
            Some(sts_creds.session_token),
            expires_at))
    }
}

// Trait that defines the STS Client API without any type parameters or assumptions about implementation.
// This is an internal type used to box the [StsClient](struct.StsClient.html) provided in the session token providers' constructors.
pub trait StsSessionCredentialsClient {
    fn assume_role(&self, input: &AssumeRoleRequest) -> Result<AssumeRoleResponse, AssumeRoleError>;

    fn assume_role_with_saml(&self, input: &AssumeRoleWithSAMLRequest) -> Result<AssumeRoleWithSAMLResponse, AssumeRoleWithSAMLError>;

    fn assume_role_with_web_identity(&self, input: &AssumeRoleWithWebIdentityRequest) -> Result<AssumeRoleWithWebIdentityResponse, AssumeRoleWithWebIdentityError>;

    fn decode_authorization_message(&self, input: &DecodeAuthorizationMessageRequest) -> Result<DecodeAuthorizationMessageResponse, DecodeAuthorizationMessageError>;

    fn get_caller_identity(&self, input: &GetCallerIdentityRequest) -> Result<GetCallerIdentityResponse, GetCallerIdentityError> ;

    fn get_federation_token(&self, input: &GetFederationTokenRequest) -> Result<GetFederationTokenResponse, GetFederationTokenError>;

    fn get_session_token(&self, input: &GetSessionTokenRequest) -> Result<GetSessionTokenResponse, GetSessionTokenError>;
}

impl<T> StsSessionCredentialsClient for T where T: Sts {
    fn assume_role(&self, input: &AssumeRoleRequest) -> Result<AssumeRoleResponse, AssumeRoleError> {
        T::assume_role(self, input)
    }

    fn assume_role_with_saml(&self, input: &AssumeRoleWithSAMLRequest) -> Result<AssumeRoleWithSAMLResponse, AssumeRoleWithSAMLError> {
        T::assume_role_with_saml(self, input)
    }

    fn assume_role_with_web_identity(&self, input: &AssumeRoleWithWebIdentityRequest) -> Result<AssumeRoleWithWebIdentityResponse, AssumeRoleWithWebIdentityError> {
        T::assume_role_with_web_identity(self, input)
    }

    fn decode_authorization_message(&self, input: &DecodeAuthorizationMessageRequest) -> Result<DecodeAuthorizationMessageResponse, DecodeAuthorizationMessageError> {
        T::decode_authorization_message(self, input)
    }

    fn get_caller_identity(&self, input: &GetCallerIdentityRequest) -> Result<GetCallerIdentityResponse, GetCallerIdentityError> {
        T::get_caller_identity(self, input)
    }

    fn get_federation_token(&self, input: &GetFederationTokenRequest) -> Result<GetFederationTokenResponse, GetFederationTokenError> {
        T::get_federation_token(self, input)
    }

    fn get_session_token(&self, input: &GetSessionTokenRequest) -> Result<GetSessionTokenResponse, GetSessionTokenError> {
        T::get_session_token(self, input)
    }
}

/// [AwsCredentials](../rusoto_credential/struct.AwsCredentials.html) provider that calls
/// `GetSessionToken` using the provided [StsClient](struct.StsClient.html).
/// To use with MFA, pass in the MFA serial number then set the MFA code.
/// You will need to ensure the provider has a valid code each time you
/// acquire a new STS token.
pub struct StsSessionCredentialsProvider {
    sts_client: Box<StsSessionCredentialsClient>,
    session_duration: Duration,
    mfa_serial: Option<String>,
    mfa_code: Option<String>,
}

impl StsSessionCredentialsProvider {
    /// Creates a new `StsSessionCredentialsProvider` with the given
    /// [StsClient](struct.StsClient.html) and session parameters.
    ///
    /// * `sts_client` - The [StsClient](struct.StsClient.html) to use to acquire session tokens.
    /// * `duration` - The duration of the session tokens. Default 1 hour.
    /// * `mfa_serial` - Optional MFA hardware device serial number or virtual device ARN. Set the MFA code with `set_mfa_code`.
    pub fn new<P,D>(sts_client: StsClient<P,D>,
            duration: Option<Duration>,
            mfa_serial: Option<String>,
            ) -> StsSessionCredentialsProvider
            where P: ProvideAwsCredentials + 'static, D: DispatchSignedRequest + 'static {
        StsSessionCredentialsProvider {
            sts_client: Box::new(sts_client),
            session_duration: duration.unwrap_or(Duration::seconds(DEFAULT_DURATION_SECONDS as i64)),
            mfa_serial: mfa_serial,
            mfa_code: None,
        }
    }

    /// Set the MFA code for use when acquiring session tokens.
    pub fn set_mfa_code<S>(&mut self, code: S) where S: Into<String> {
        self.mfa_code = Some(code.into());
    }

    /// Clear the MFA code.
    pub fn clear_mfa_code(&mut self) {
        self.mfa_code = None;
    }

    /// Calls `GetSessionToken` to get a session token from the STS Api.
    /// Optionally uses MFA if the MFA serial number and code are set.
    pub fn get_session_token(&self) -> Result<AwsCredentials, CredentialsError> {
        match self.sts_client.get_session_token(
            &GetSessionTokenRequest {
                serial_number: self.mfa_serial.clone(),
                token_code: self.mfa_code.clone(),
                duration_seconds: Some(self.session_duration.num_seconds() as i64),
                ..Default::default()
            }) {
            Ok(resp) => {
                let creds = try!(resp.credentials.ok_or(CredentialsError::new("no credentials in response")));

                AwsCredentials::new_for_credentials(creds)
            },
            err =>
                Err(CredentialsError::new(format!("StsProvider get_session_token error: {:?}", err)))
        }
    }
}

impl ProvideAwsCredentials for StsSessionCredentialsProvider {
    fn credentials(&self) -> Result<AwsCredentials, CredentialsError> {
        self.get_session_token()
    }
}

/// [AwsCredentials](../rusoto_credential/struct.AwsCredentials.html) provider that calls
/// `AssumeRole` using the provided [StsClient](struct.StsClient.html).
/// To use with MFA, pass in the MFA serial number then set the MFA code.
/// You will need to ensure the provider has a valid code each time you
/// acquire a new STS token.
pub struct StsAssumeRoleSessionCredentialsProvider {
    sts_client: Box<StsSessionCredentialsClient>,
    role_arn: String,
    session_name: String,
    external_id: Option<String>,
    session_duration: Duration,
    scope_down_policy: Option<String>,
    mfa_serial: Option<String>,
    mfa_code: Option<String>,
}

impl StsAssumeRoleSessionCredentialsProvider {
    /// Creates a new `StsAssumeRoleSessionCredentialsProvider` with the given
    /// [StsClient](struct.StsClient.html) and session parameters.
    ///
    /// * `sts_client` - [StsClient](struct.StsClient.html) to use to acquire session tokens.
    /// * `role_arn` - The ARN of the role to assume.
    /// * `session_name` - An identifier for the assumed role session. Minimum length of 2. Maximum length of 64. Pattern: `[\w+=,.@-]*`
    /// * `external_id` -
    /// * `session_duration` - Duration of session tokens. Default 1 hour.
    /// * `scope_down_policy` - Optional inline IAM policy in JSON format to further restrict the access granted to the negotiated session.
    /// * `mfa_serial` - Optional MFA hardware device serial number or virtual device ARN. Use `set_mfa_code` to set the MFA code.
    pub fn new<P,D>(sts_client: StsClient<P,D>,
            role_arn: String,
            session_name: String,
            external_id: Option<String>,
            session_duration: Option<Duration>,
            scope_down_policy: Option<String>,
            mfa_serial: Option<String>
            )
            -> StsAssumeRoleSessionCredentialsProvider
            where P: ProvideAwsCredentials + 'static, D: DispatchSignedRequest + 'static {
        StsAssumeRoleSessionCredentialsProvider {
            sts_client: Box::new(sts_client),
            role_arn: role_arn,
            session_name: session_name,
            external_id: external_id,
            session_duration: session_duration.unwrap_or(Duration::seconds(DEFAULT_ROLE_DURATION_SECONDS as i64)),
            scope_down_policy: scope_down_policy,
            mfa_serial: mfa_serial,
            mfa_code: None,
        }
    }

    /// Set the MFA code for use when acquiring session tokens.
    pub fn set_mfa_code<S>(&mut self, code: S) where S: Into<String> {
        self.mfa_code = Some(code.into());
    }

    /// Clear the MFA code.
    pub fn clear_mfa_code(&mut self) {
        self.mfa_code = None;
    }

    /// Calls `AssumeRole` to get a session token from the STS Api.
    /// Optionally uses MFA if the MFA serial number and code are set.
    pub fn assume_role(&self) -> Result<AwsCredentials, CredentialsError> {
        match self.sts_client.assume_role(&AssumeRoleRequest {
            role_arn: self.role_arn.clone(),
            role_session_name: self.session_name.clone(),
            duration_seconds: Some(self.session_duration.num_seconds() as i64),
            external_id: self.external_id.clone(),
            policy: self.scope_down_policy.clone(),
            serial_number: self.mfa_serial.clone(),
            token_code: self.mfa_code.clone(),
            ..Default::default()
        }) {
            Err(err) =>
                Err(CredentialsError::new(format!("Sts AssumeRoleError: {:?}", err))),
            Ok(resp) => {
                let creds = try!(resp.credentials.ok_or(CredentialsError::new("no credentials in response")));

                AwsCredentials::new_for_credentials(creds)
            }
        }
    }
}

impl ProvideAwsCredentials for StsAssumeRoleSessionCredentialsProvider {
    fn credentials(&self) -> Result<AwsCredentials, CredentialsError> {
        self.assume_role()
    }
}

/// [AwsCredentials](../rusoto_credential/struct.AwsCredentials.html) provider that calls
/// `AssumeRoleWithWebIdentity` using the provided [StsClient](struct.StsClient.html).
pub struct StsWebIdentityFederationSessionCredentialsProvider {
    sts_client: Box<StsSessionCredentialsClient>,
    wif_token: String,
    wif_provider: Option<String>,
    role_arn: String,
    session_name: String,
    session_duration: Duration,
    scope_down_policy: Option<String>,
}

impl StsWebIdentityFederationSessionCredentialsProvider {
    /// Creates a new `StsWebIdentityFederationSessionCredentialsProvider` with the given
    /// [StsClient](struct.StsClient.html) and session parameters.
    ///
    /// * `sts_client` - The [StsClient](struct.StsClient.html) to use to acquire session tokens.
    /// * `wif_token` - The OAuth 2.0 access token or OpenID Connect ID token that is provided by the identity provider.
    /// * `wif_provider` - The fully qualified host component of the domain name of the identity provider. Only for OAuth 2.0 access tokens. Do not include URL schemes and port numbers.
    /// * `role_arn` - The ARN of the role to assume.
    /// * `session_name` - An identifier for the assumed role session. Minimum length of 2. Maximum length of 64. Pattern: `[\w+=,.@-]*`
    /// * `session_duration` - Duration of session tokens. Default 1 hour.
    /// * `scope_down_policy` - Optional inline IAM policy in JSON format to further restrict the access granted to the negotiated session.
    pub fn new<P,D>(sts_client: StsClient<P,D>,
            wif_token: String,
            wif_provider: Option<String>,
            role_arn: String,
            session_name: String,
            session_duration: Option<Duration>,
            scope_down_policy: Option<String>
            ) -> StsWebIdentityFederationSessionCredentialsProvider
            where P: ProvideAwsCredentials + 'static, D: DispatchSignedRequest + 'static {
        StsWebIdentityFederationSessionCredentialsProvider {
            sts_client: Box::new(sts_client),
            wif_token: wif_token,
            wif_provider: wif_provider,
            role_arn: role_arn,
            session_name: session_name,
            session_duration: session_duration.unwrap_or(Duration::seconds(DEFAULT_DURATION_SECONDS as i64)),
            scope_down_policy: scope_down_policy
        }
    }

    /// Calls `AssumeRoleWithWebIdentity` to get a session token from the STS Api.
    pub fn assume_role_with_web_identity(&self) -> Result<AwsCredentials, CredentialsError> {
        match self.sts_client.assume_role_with_web_identity(&AssumeRoleWithWebIdentityRequest {
            web_identity_token: self.wif_token.clone(),
            provider_id: self.wif_provider.clone(),
            role_arn: self.role_arn.clone(),
            role_session_name: self.session_name.clone(),
            duration_seconds: Some(self.session_duration.num_seconds() as i64),
            policy: self.scope_down_policy.clone(),
            ..Default::default()
        }) {
            Err(err) =>
                Err(CredentialsError::new(format!("Sts AssumeRoleWithWebIdentityError: {:?}", err))),
            Ok(resp) => {
                let creds = try!(resp.credentials.ok_or(CredentialsError::new("no credentials in response")));

                let mut aws_creds = try!(AwsCredentials::new_for_credentials(creds));

                if let Some(subject_from_wif) = resp.subject_from_web_identity_token {
                    aws_creds.claims_mut().insert(rusoto_core::credential::claims::SUBJECT.to_owned(), subject_from_wif);
                }

                if let Some(audience) = resp.audience {
                    aws_creds.claims_mut().insert(rusoto_core::credential::claims::AUDIENCE.to_owned(), audience);
                }

                if let Some(issuer) = resp.provider {
                    aws_creds.claims_mut().insert(rusoto_core::credential::claims::ISSUER.to_owned(), issuer);
                }

                Ok(aws_creds)
            }
        }
    }
}

impl ProvideAwsCredentials for StsWebIdentityFederationSessionCredentialsProvider {
    fn credentials(&self) -> Result<AwsCredentials, CredentialsError> {
        self.assume_role_with_web_identity()
    }
}<|endoftext|>


<commit_before>#![crate_name = "cksum"]
#![feature(macro_rules)]

/*
 * This file is part of the uutils coreutils package.
 *
 * (c) Michael Gehring <mg@ebfe.org>
 *
 * For the full copyright and license information, please view the LICENSE
 * file that was distributed with this source code.
 */

extern crate getopts;

use std::io::{EndOfFile, File, IoError, IoResult, print};
use std::io::stdio::stdin_raw;
use std::mem;

use crc_table::CRC_TABLE;

#[path="../common/util.rs"]
mod util;

mod crc_table;

static NAME: &'static str = "cksum";
static VERSION: &'static str = "1.0.0";

#[inline]
fn crc_update(crc: u32, input: u8) -> u32 {
    (crc << 8) ^ CRC_TABLE[((crc >> 24) as uint ^ input as uint) & 0xFF]
}

#[inline]
fn crc_final(mut crc: u32, mut length: uint) -> u32 {
    while length != 0 {
        crc = crc_update(crc, length as u8);
        length >>= 8;
    }

    !crc
}

fn cksum(fname: &str) -> IoResult<(u32, uint)> {
    let mut crc = 0u32;
    let mut size = 0u;

    let mut rd = try!(open_file(fname));
    let mut bytes: [u8, ..1024 * 1024] = unsafe { mem::uninitialized() };
    loop {
        match rd.read(bytes) {
            Ok(num_bytes) => {
                for &b in bytes.iter().take(num_bytes) {
                    crc = crc_update(crc, b);
                }
                size += num_bytes;
            }
            Err(err) =>  {
                return match err {
                    IoError{kind: k, ..} if k == EndOfFile => Ok((crc_final(crc, size), size)),
                    _ => Err(err),
                }
            }
        }
    }
}

fn open_file(name: &str) -> IoResult<Box<Reader>> {
    match name {
        "-" => Ok(box stdin_raw() as Box<Reader>),
        _   => {
            let f = try!(File::open(&Path::new(name)));
            Ok(box f as Box<Reader>)
        }
    }
}

pub fn uumain(args: Vec<String>) -> int {
    let opts = [
        getopts::optflag("h", "help", "display this help and exit"),
        getopts::optflag("V", "version", "output version information and exit"),
    ];

    let matches = match getopts::getopts(args.tail(), opts) {
        Ok(m) => m,
        Err(err) => fail!("{}", err),
    };

    if matches.opt_present("help") {
        println!("{} {}", NAME, VERSION);
        println!("");
        println!("Usage:");
        println!("  {} [OPTIONS] [FILE]...", NAME);
        println!("");
        print(getopts::usage("Print CRC and size for each file.", opts.as_slice()).as_slice());
        return 0;
    }

    if matches.opt_present("version") {
        println!("{} {}", NAME, VERSION);
        return 0;
    }

    let files = matches.free;

    if files.is_empty() {
        match cksum("-") {
            Ok((crc, size)) => println!("{} {}", crc, size),
            Err(err) => {
                show_error!("{}", err);
                return 2;
            }
        }
        return 0;
    }

    let mut exit_code = 0;
    for fname in files.iter() {
        match cksum(fname.as_slice()) {
            Ok((crc, size)) => println!("{} {} {}", crc, size, fname),
            Err(err) => {
                show_error!("'{}' {}", fname, err);
                exit_code = 2;
            }
        }
    }

    exit_code
}
<commit_msg>cksum: improve performance slightly<commit_after>#![crate_name = "cksum"]
#![feature(macro_rules)]

/*
 * This file is part of the uutils coreutils package.
 *
 * (c) Michael Gehring <mg@ebfe.org>
 *
 * For the full copyright and license information, please view the LICENSE
 * file that was distributed with this source code.
 */

extern crate getopts;

use std::io::{EndOfFile, File, IoError, IoResult, print};
use std::io::stdio::stdin_raw;
use std::mem;

use crc_table::CRC_TABLE;

#[path="../common/util.rs"]
mod util;

mod crc_table;

static NAME: &'static str = "cksum";
static VERSION: &'static str = "1.0.0";

#[inline]
fn crc_update(crc: u32, input: u8) -> u32 {
    (crc << 8) ^ CRC_TABLE[((crc >> 24) as uint ^ input as uint) & 0xFF]
}

#[inline]
fn crc_final(mut crc: u32, mut length: uint) -> u32 {
    while length != 0 {
        crc = crc_update(crc, length as u8);
        length >>= 8;
    }

    !crc
}

#[inline]
fn cksum(fname: &str) -> IoResult<(u32, uint)> {
    let mut crc = 0u32;
    let mut size = 0u;

    let mut stdin_buf;
    let mut file_buf;
    let rd = match fname {
        "-" => {
            stdin_buf = stdin_raw();
            &mut stdin_buf as &mut Reader
        }
        _ => {
            file_buf = try!(File::open(&Path::new(fname)));
            &mut file_buf as &mut Reader
        }
    };

    let mut bytes: [u8, ..1024 * 1024] = unsafe { mem::uninitialized() };
    loop {
        match rd.read(bytes) {
            Ok(num_bytes) => {
                for &b in bytes.slice_to(num_bytes).iter() {
                    crc = crc_update(crc, b);
                }
                size += num_bytes;
            }
            Err(IoError { kind: EndOfFile, .. }) => return Ok((crc_final(crc, size), size)),
            Err(err) => return Err(err)
        }
    }
}

pub fn uumain(args: Vec<String>) -> int {
    let opts = [
        getopts::optflag("h", "help", "display this help and exit"),
        getopts::optflag("V", "version", "output version information and exit"),
    ];

    let matches = match getopts::getopts(args.tail(), opts) {
        Ok(m) => m,
        Err(err) => fail!("{}", err),
    };

    if matches.opt_present("help") {
        println!("{} {}", NAME, VERSION);
        println!("");
        println!("Usage:");
        println!("  {} [OPTIONS] [FILE]...", NAME);
        println!("");
        print(getopts::usage("Print CRC and size for each file.", opts.as_slice()).as_slice());
        return 0;
    }

    if matches.opt_present("version") {
        println!("{} {}", NAME, VERSION);
        return 0;
    }

    let files = matches.free;

    if files.is_empty() {
        match cksum("-") {
            Ok((crc, size)) => println!("{} {}", crc, size),
            Err(err) => {
                show_error!("{}", err);
                return 2;
            }
        }
        return 0;
    }

    let mut exit_code = 0;
    for fname in files.iter() {
        match cksum(fname.as_slice()) {
            Ok((crc, size)) => println!("{} {} {}", crc, size, fname),
            Err(err) => {
                show_error!("'{}' {}", fname, err);
                exit_code = 2;
            }
        }
    }

    exit_code
}
<|endoftext|>


<commit_before>use utils::ring_buffer::BiasedRingBuffer;
use utils::Baseable;

use compressor::{Compressor, NaiveCompressor};

use page_manager::{PageId, Block, BlockIter, BlockId, RamPageCache, BlockManager};

use index::posting::{Posting, DocId, PostingIterator};

pub type UsedCompressor = NaiveCompressor;

pub struct Listing {
    page_list: Vec<PageId>,
    block_biases: Vec<Posting>,
    block_counter: BlockId,
    block_start: Posting,
    block_end: Posting,
    posting_buffer: BiasedRingBuffer<Posting>,
}

impl Listing {
    pub fn new() -> Self {
        Listing {
            page_list: Vec::new(),
            block_biases: Vec::new(),
            block_counter: BlockId::first(),
            posting_buffer: BiasedRingBuffer::new(),
            block_start: Posting(DocId(0)),
            block_end: Posting(DocId(0)),
        }
    }

    pub fn add(&mut self, postings: &[Posting], page_cache: &mut RamPageCache) {
        for (i, posting) in postings.iter().enumerate() {
            self.block_end = *posting;
            self.posting_buffer.push_back(*posting);
            if i % 8 == 0 {
                self.compress_and_ship(page_cache, false);
            }
        }
        self.compress_and_ship(page_cache, false);
    }

    pub fn commit(&mut self, page_cache: &mut RamPageCache) {
        self.compress_and_ship(page_cache, true);
        page_cache.flush_page(*self.page_list.last().unwrap());
    }

    pub fn last_block_id(&self) -> BlockId {
        let mut r = self.block_counter;
        r.dec();
        r
    }

    pub fn posting_iter<'a>(&self, cache: &'a RamPageCache) -> PostingIterator<'a> {
        let block_iter = BlockIter::new(cache, self.page_list.clone(), self.last_block_id());
        PostingIterator::new(block_iter, self.block_biases.clone())
    }

    fn compress_and_ship(&mut self, page_cache: &mut RamPageCache, force: bool) {
        while let Some(block) = UsedCompressor::compress(&mut self.posting_buffer) {
            self.ship(page_cache, block);
        }
        if force && self.posting_buffer.count() > 0 {
            let block = UsedCompressor::force_compress(&mut self.posting_buffer);
            self.ship(page_cache, block);
        }
    }

    fn ship(&mut self, page_cache: &mut RamPageCache, block: Block) {
        // If the block is on a new page
        if self.block_counter == BlockId::first() {
            // Push it on a new page and store the page ide
            self.page_list.push(page_cache.store_block(block));
        } else {
            // Otherwise store it on an existing page
            page_cache.store_in_place(*self.page_list.last().unwrap(), self.block_counter, block)
        }
        // Save with what doc_id the block just stored block starts
        self.block_biases.push(self.block_start);
        // Count up the block
        self.block_counter.inc();
        if let Some(posting) = self.posting_buffer.peek_front() {
            self.block_start = *posting;
        } else {
            self.block_start = self.block_end;
        }
        self.posting_buffer.base_on(self.block_start);
    }
}

#[cfg(test)]
mod tests {

    use super::Listing;

    use test_utils::create_test_dir;

    use index::posting::{Posting, DocId};
    use page_manager::{FsPageManager, RamPageCache, BlockId};

    fn new_cache(name: &str) -> RamPageCache {
        let path = &create_test_dir(format!("listing/{}", name).as_str());
        let pmgr = FsPageManager::new(&path.join("pages.bin"));
        RamPageCache::new(pmgr)
    }

    #[test]
    fn basic_add() {
        let mut cache = new_cache("basic_add");
        let mut listing = Listing::new();
        listing.add(&[Posting(DocId(0))], &mut cache);
        assert_eq!(listing.page_list.len(), 0);
        assert_eq!(listing.posting_buffer.count(), 1);
    }

    #[test]
    fn commit() {
        let mut cache = new_cache("commit");
        let mut listing = Listing::new();
        listing.add(&[Posting(DocId(0))], &mut cache);
        assert_eq!(listing.page_list.len(), 0);
        assert_eq!(listing.posting_buffer.count(), 1);
        listing.commit(&mut cache);
        assert_eq!(listing.page_list.len(), 1);
        assert_eq!(listing.posting_buffer.count(), 0);
        assert_eq!(listing.last_block_id(), BlockId::first());
    }

    #[test]
    fn add() {
        let mut cache = new_cache("add");
        let mut listing = Listing::new();
        listing.add(&[Posting(DocId(0))], &mut cache);
        assert_eq!(listing.page_list.len(), 0);
        assert_eq!(listing.posting_buffer.count(), 1);
        for i in 0..100 {
            listing.add(&[Posting(DocId(i))], &mut cache);
        }
        assert!(listing.page_list.len() > 0);
        assert!(listing.posting_buffer.count() > 0);
        listing.commit(&mut cache);
        assert_eq!(listing.posting_buffer.count(), 0);
    }

    #[test]
    fn add_much() {
        let mut cache = new_cache("add_much");
        let mut listing = Listing::new();
        listing.add(&[Posting(DocId(0))], &mut cache);
        assert_eq!(listing.page_list.len(), 0);
        assert_eq!(listing.posting_buffer.count(), 1);
        for i in 0..10000 {
            listing.add(&[Posting(DocId(i))], &mut cache);
        }
        assert!(listing.page_list.len() > 0);
        assert!(listing.posting_buffer.count() > 0);
        listing.commit(&mut cache);
        assert_eq!(listing.posting_buffer.count(), 0);
    }

    #[test]
    fn multiple_listings() {
        let mut cache = new_cache("multiple_listings");
        let mut listings = (0..100).map(|_| Listing::new()).collect::<Vec<_>>();
        for i in 0..50000 {
            listings[i % 100].add(&[Posting(DocId(i as u64))], &mut cache);
        }
        for listing in listings.iter_mut() {
            assert!(listing.page_list.len() > 0);
            assert!(listing.posting_buffer.count() > 0);
            listing.commit(&mut cache);
        }
        for listing in listings {
            assert_eq!(listing.posting_buffer.count(), 0);
        }
    }

    #[test]
    fn biases() {
        let mut cache = new_cache("biases");
        let mut listing = Listing::new();

        listing.add(&[Posting(DocId(1))], &mut cache);
        assert_eq!(listing.block_start, Posting(DocId(0)));
        assert_eq!(listing.block_end, Posting(DocId(1)));
        listing.commit(&mut cache);
        assert_eq!(listing.block_start, Posting(DocId(1)));
        assert_eq!(listing.block_end, Posting(DocId(1)));
        listing.add(&[Posting(DocId(10))], &mut cache);
        assert_eq!(listing.block_end, Posting(DocId(10)));
        listing.commit(&mut cache);
        assert_eq!(listing.block_start, Posting(DocId(10)));
        assert_eq!(listing.block_biases,
                   vec![Posting(DocId(0)), Posting(DocId(1))]);
    }
}
<commit_msg>removed unneccessary pub<commit_after>use utils::ring_buffer::BiasedRingBuffer;
use utils::Baseable;

use compressor::{Compressor, NaiveCompressor};

use page_manager::{PageId, Block, BlockIter, BlockId, RamPageCache, BlockManager};

use index::posting::{Posting, DocId, PostingIterator};

pub type UsedCompressor = NaiveCompressor;

pub struct Listing {
    page_list: Vec<PageId>,
    block_biases: Vec<Posting>,
    block_counter: BlockId,
    block_start: Posting,
    block_end: Posting,
    posting_buffer: BiasedRingBuffer<Posting>,
}

impl Listing {
    pub fn new() -> Self {
        Listing {
            page_list: Vec::new(),
            block_biases: Vec::new(),
            block_counter: BlockId::first(),
            posting_buffer: BiasedRingBuffer::new(),
            block_start: Posting(DocId(0)),
            block_end: Posting(DocId(0)),
        }
    }

    pub fn add(&mut self, postings: &[Posting], page_cache: &mut RamPageCache) {
        for (i, posting) in postings.iter().enumerate() {
            self.block_end = *posting;
            self.posting_buffer.push_back(*posting);
            if i % 8 == 0 {
                self.compress_and_ship(page_cache, false);
            }
        }
        self.compress_and_ship(page_cache, false);
    }

    pub fn commit(&mut self, page_cache: &mut RamPageCache) {
        self.compress_and_ship(page_cache, true);
        page_cache.flush_page(*self.page_list.last().unwrap());
    }

    pub fn posting_iter<'a>(&self, cache: &'a RamPageCache) -> PostingIterator<'a> {
        let block_iter = BlockIter::new(cache, self.page_list.clone(), self.last_block_id());
        PostingIterator::new(block_iter, self.block_biases.clone())
    }

    fn last_block_id(&self) -> BlockId {
        let mut r = self.block_counter;
        r.dec();
        r
    }

    fn compress_and_ship(&mut self, page_cache: &mut RamPageCache, force: bool) {
        while let Some(block) = UsedCompressor::compress(&mut self.posting_buffer) {
            self.ship(page_cache, block);
        }
        if force && self.posting_buffer.count() > 0 {
            let block = UsedCompressor::force_compress(&mut self.posting_buffer);
            self.ship(page_cache, block);
        }
    }

    fn ship(&mut self, page_cache: &mut RamPageCache, block: Block) {
        // If the block is on a new page
        if self.block_counter == BlockId::first() {
            // Push it on a new page and store the page ide
            self.page_list.push(page_cache.store_block(block));
        } else {
            // Otherwise store it on an existing page
            page_cache.store_in_place(*self.page_list.last().unwrap(), self.block_counter, block)
        }
        // Save with what doc_id the block just stored block starts
        self.block_biases.push(self.block_start);
        // Count up the block
        self.block_counter.inc();
        if let Some(posting) = self.posting_buffer.peek_front() {
            self.block_start = *posting;
        } else {
            self.block_start = self.block_end;
        }
        self.posting_buffer.base_on(self.block_start);
    }
}

#[cfg(test)]
mod tests {

    use super::Listing;

    use test_utils::create_test_dir;

    use index::posting::{Posting, DocId};
    use page_manager::{FsPageManager, RamPageCache, BlockId};

    fn new_cache(name: &str) -> RamPageCache {
        let path = &create_test_dir(format!("listing/{}", name).as_str());
        let pmgr = FsPageManager::new(&path.join("pages.bin"));
        RamPageCache::new(pmgr)
    }

    #[test]
    fn basic_add() {
        let mut cache = new_cache("basic_add");
        let mut listing = Listing::new();
        listing.add(&[Posting(DocId(0))], &mut cache);
        assert_eq!(listing.page_list.len(), 0);
        assert_eq!(listing.posting_buffer.count(), 1);
    }

    #[test]
    fn commit() {
        let mut cache = new_cache("commit");
        let mut listing = Listing::new();
        listing.add(&[Posting(DocId(0))], &mut cache);
        assert_eq!(listing.page_list.len(), 0);
        assert_eq!(listing.posting_buffer.count(), 1);
        listing.commit(&mut cache);
        assert_eq!(listing.page_list.len(), 1);
        assert_eq!(listing.posting_buffer.count(), 0);
        assert_eq!(listing.last_block_id(), BlockId::first());
    }

    #[test]
    fn add() {
        let mut cache = new_cache("add");
        let mut listing = Listing::new();
        listing.add(&[Posting(DocId(0))], &mut cache);
        assert_eq!(listing.page_list.len(), 0);
        assert_eq!(listing.posting_buffer.count(), 1);
        for i in 0..100 {
            listing.add(&[Posting(DocId(i))], &mut cache);
        }
        assert!(listing.page_list.len() > 0);
        assert!(listing.posting_buffer.count() > 0);
        listing.commit(&mut cache);
        assert_eq!(listing.posting_buffer.count(), 0);
    }

    #[test]
    fn add_much() {
        let mut cache = new_cache("add_much");
        let mut listing = Listing::new();
        listing.add(&[Posting(DocId(0))], &mut cache);
        assert_eq!(listing.page_list.len(), 0);
        assert_eq!(listing.posting_buffer.count(), 1);
        for i in 0..10000 {
            listing.add(&[Posting(DocId(i))], &mut cache);
        }
        assert!(listing.page_list.len() > 0);
        assert!(listing.posting_buffer.count() > 0);
        listing.commit(&mut cache);
        assert_eq!(listing.posting_buffer.count(), 0);
    }

    #[test]
    fn multiple_listings() {
        let mut cache = new_cache("multiple_listings");
        let mut listings = (0..100).map(|_| Listing::new()).collect::<Vec<_>>();
        for i in 0..50000 {
            listings[i % 100].add(&[Posting(DocId(i as u64))], &mut cache);
        }
        for listing in listings.iter_mut() {
            assert!(listing.page_list.len() > 0);
            assert!(listing.posting_buffer.count() > 0);
            listing.commit(&mut cache);
        }
        for listing in listings {
            assert_eq!(listing.posting_buffer.count(), 0);
        }
    }

    #[test]
    fn biases() {
        let mut cache = new_cache("biases");
        let mut listing = Listing::new();

        listing.add(&[Posting(DocId(1))], &mut cache);
        assert_eq!(listing.block_start, Posting(DocId(0)));
        assert_eq!(listing.block_end, Posting(DocId(1)));
        listing.commit(&mut cache);
        assert_eq!(listing.block_start, Posting(DocId(1)));
        assert_eq!(listing.block_end, Posting(DocId(1)));
        listing.add(&[Posting(DocId(10))], &mut cache);
        assert_eq!(listing.block_end, Posting(DocId(10)));
        listing.commit(&mut cache);
        assert_eq!(listing.block_start, Posting(DocId(10)));
        assert_eq!(listing.block_biases,
                   vec![Posting(DocId(0)), Posting(DocId(1))]);
    }
}
<|endoftext|>


<commit_before>
use super::{DrmError, ModeError};

use super::devices;

use backend::graphics::GraphicsBackend;
use backend::graphics::egl::{EGLGraphicsBackend, EGLSurface, PixelFormat, SwapBuffersError};
use drm::buffer::Buffer;
use drm::control::{connector, crtc, framebuffer, Mode};
use drm::control::ResourceInfo;

use gbm::{BufferObject, BufferObjectFlags, Format as GbmFormat, Surface as GbmSurface, SurfaceBufferHandle};

use image::{ImageBuffer, Rgba};

use nix::c_void;

use std::cell::{Cell, RefCell};
use std::rc::Rc;

/// Backend based on a `DrmDevice` and a given crtc
pub struct DrmBackend(Rc<RefCell<DrmBackendInternal>>);

impl DrmBackend {
    pub(crate) fn new(drm: Rc<RefCell<DrmBackendInternal>>) -> DrmBackend {
        DrmBackend(drm)
    }
}

/*
    Dependency graph
    - drm
        - gbm
            - context
            - gbm_surface
                - egl_surface
                - gbm_buffers
            - cursor
*/

pub(crate) struct GbmTypes<'dev, 'context> {
    cursor: BufferObject<'dev, ()>,
    surface: Surface<'context>,
}

pub(crate) struct EGL<'gbm, 'context> {
    surface: EGLSurface<'context, 'gbm, GbmSurface<'context, framebuffer::Info>>,
    buffers: GbmBuffers<'gbm>,
}

pub(crate) struct GbmBuffers<'gbm> {
    front_buffer: Cell<SurfaceBufferHandle<'gbm, framebuffer::Info>>,
    next_buffer: Cell<Option<SurfaceBufferHandle<'gbm, framebuffer::Info>>>,
}

rental! {
    mod graphics {
        use drm::control::framebuffer;
        use gbm::Surface as GbmSurface;
        use std::rc::Rc;
        use super::devices::{Context, Context_Borrow};
        use super::GbmTypes;

        #[rental]
        pub(crate) struct Surface<'context> {
            gbm: Box<GbmSurface<'context, framebuffer::Info>>,
            egl: super::EGL<'gbm, 'context>,
        }

        #[rental]
        pub(crate) struct Graphics {
            #[subrental(arity = 3)]
            context: Rc<Context>,
            gbm: GbmTypes<'context_1, 'context_2>,
        }
    }
}
use self::graphics::{Graphics, Surface};

/// Id of a `DrmBackend` related to its `DrmDevice`.
///
/// Used to track which `DrmBackend` finished page-flipping
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct Id(usize);

impl Id {
    pub(crate) fn raw(&self) -> usize {
        self.0
    }
}

pub(crate) struct DrmBackendInternal {
    graphics: Graphics,
    crtc: crtc::Handle,
    mode: Mode,
    connectors: Vec<connector::Handle>,
    own_id: Id,
    logger: ::slog::Logger,
}

impl DrmBackendInternal {
    pub(crate) fn new<I, L>(context: Rc<devices::Context>, crtc: crtc::Handle, mode: Mode, connectors: I,
                            own_id: usize, logger: L)
                            -> Result<DrmBackendInternal, DrmError>
    where
        I: Into<Vec<connector::Handle>>,
        L: Into<Option<::slog::Logger>>,
    {
        let log = ::slog_or_stdlog(logger).new(o!("smithay_module" => "backend_drm"));
        info!(log, "Initializing a drm backend");

        let connectors = connectors.into();

        for connector in connectors.iter() {
            if !connector::Info::load_from_device(context.head().head(), *connector)?
                .modes()
                .contains(&mode)
            {
                return Err(DrmError::Mode(ModeError::ModeNotSuitable));
            }
        }

        let (w, h) = mode.size();

        info!(log, "Drm Backend initialized");

        Ok(DrmBackendInternal {
            graphics: Graphics::try_new(context, |context| {
                debug!(log, "GBM EGLContext initialized");

                Ok(GbmTypes {
                    cursor: {
                        // Create an unused cursor buffer (we don't want an Option here)
                        context.devices.gbm.create_buffer_object(
                            1,
                            1,
                            GbmFormat::ARGB8888,
                            &[BufferObjectFlags::Cursor, BufferObjectFlags::Write],
                        )?
                    },
                    surface: Surface::try_new(
                        Box::new(context.devices.gbm.create_surface(
                            w as u32,
                            h as u32,
                            GbmFormat::XRGB8888,
                            &[BufferObjectFlags::Scanout, BufferObjectFlags::Rendering],
                        )?),
                        |surface| {
                            let egl_surface = context.egl.create_surface(&surface)?;
                            unsafe { egl_surface.make_current()? };
                            egl_surface.swap_buffers()?;

                            let mut front_bo = surface.lock_front_buffer()?;
                            debug!(log, "FrontBuffer color format: {:?}", front_bo.format());
                            let fb = framebuffer::create(context.devices.drm, &*front_bo)?;
                            crtc::set(
                                context.devices.drm,
                                crtc,
                                fb.handle(),
                                &connectors,
                                (0, 0),
                                Some(mode),
                            )?;
                            front_bo.set_userdata(fb);

                            Ok(EGL {
                                surface: egl_surface,
                                buffers: GbmBuffers {
                                    front_buffer: Cell::new(front_bo),
                                    next_buffer: Cell::new(None),
                                },
                            })
                        },
                    ).map_err(DrmError::from)?,
                })
            })?,
            crtc,
            mode,
            connectors,
            own_id: Id(own_id),
            logger: log.clone(),
        })
    }

    pub(crate) fn unlock_buffer(&self) {
        self.graphics.rent(|gbm| {
            gbm.surface.rent(|egl| {
                let next_bo = egl.buffers.next_buffer.replace(None);

                if let Some(next_buffer) = next_bo {
                    egl.buffers.front_buffer.set(next_buffer);
                // drop and release the old buffer
                } else {
                    unreachable!();
                }
            })
        });
    }

    pub(crate) fn is_crtc(&self, crtc: crtc::Handle) -> bool {
        crtc == self.crtc
    }
}

impl DrmBackend {
    /// Add a connector to backend
    ///
    /// # Errors
    ///
    /// Errors if the new connector does not support the currently set `Mode`
    pub fn add_connector(&mut self, connector: connector::Handle) -> Result<(), ModeError> {
        let info =
            connector::Info::load_from_device(self.0.borrow().graphics.head().head().head(), connector)
                .map_err(|err| ModeError::FailedToLoad(err))?;
        let mut internal = self.0.borrow_mut();
        if info.modes().contains(&internal.mode) {
            internal.connectors.push(connector);
            Ok(())
        } else {
            Err(ModeError::ModeNotSuitable)
        }
    }

    /// Returns a copy of the currently set connectors
    pub fn used_connectors(&self) -> Vec<connector::Handle> {
        self.0.borrow().connectors.clone()
    }

    /// Removes a currently set connector
    pub fn remove_connector(&mut self, connector: connector::Handle) {
        self.0.borrow_mut().connectors.retain(|x| *x != connector);
    }

    /// Changes the currently set mode
    ///
    /// # Errors
    ///
    /// This will fail if not all set connectors support the new `Mode`.
    /// Several internal resources will need to be recreated to fit the new `Mode`.
    /// Other errors might occur.
    pub fn use_mode(&mut self, mode: Mode) -> Result<(), DrmError> {
        for connector in self.0.borrow().connectors.iter() {
            if !connector::Info::load_from_device(self.0.borrow().graphics.head().head().head(), *connector)?
                .modes()
                .contains(&mode)
            {
                return Err(DrmError::Mode(ModeError::ModeNotSuitable));
            }
        }

        let crtc = self.0.borrow().crtc;
        let mut internal = self.0.borrow_mut();
        let connectors = internal.connectors.clone();
        let logger = internal.logger.clone();

        let (w, h) = mode.size();

        internal
            .graphics
            .rent_all_mut(|graphics| -> Result<(), DrmError> {
                graphics.gbm.surface = Surface::try_new(
                    Box::new(graphics.context.devices.gbm.create_surface(
                        w as u32,
                        h as u32,
                        GbmFormat::XRGB8888,
                        &[BufferObjectFlags::Scanout, BufferObjectFlags::Rendering],
                    )?),
                    |surface| {
                        let egl_surface = graphics.context.egl.create_surface(&surface)?;
                        unsafe { egl_surface.make_current()? };
                        egl_surface.swap_buffers()?;

                        let mut front_bo = surface.lock_front_buffer()?;
                        debug!(logger, "FrontBuffer color format: {:?}", front_bo.format());
                        let fb = framebuffer::create(graphics.context.devices.drm, &*front_bo)?;
                        crtc::set(
                            graphics.context.devices.drm,
                            crtc,
                            fb.handle(),
                            &connectors,
                            (0, 0),
                            Some(mode),
                        )?;
                        front_bo.set_userdata(fb);

                        Ok(EGL {
                            surface: egl_surface,
                            buffers: GbmBuffers {
                                front_buffer: Cell::new(front_bo),
                                next_buffer: Cell::new(None),
                            },
                        })
                    },
                )?;

                Ok(())
            })?;

        internal.mode = mode;
        Ok(())
    }

    /// Checks of the `DrmBackend` is of the given `Id`
    ///
    /// Only produces valid results, if the `Id` is from the `DrmDevice`,
    /// that created this backend.
    pub fn is(&self, id: Id) -> bool {
        self.0.borrow().own_id == id
    }
}

impl GraphicsBackend for DrmBackend {
    type CursorFormat = ImageBuffer<Rgba<u8>, Vec<u8>>;
    type Error = DrmError;

    fn set_cursor_position(&self, x: u32, y: u32) -> Result<(), DrmError> {
        crtc::move_cursor(
            self.0.borrow().graphics.head().head().head(),
            self.0.borrow().crtc,
            (x as i32, y as i32),
        ).map_err(DrmError::from)
    }

    fn set_cursor_representation(&self, buffer: ImageBuffer<Rgba<u8>, Vec<u8>>, hotspot: (u32, u32))
                                 -> Result<(), DrmError> {
        let (w, h) = buffer.dimensions();
        self.0
            .borrow_mut()
            .graphics
            .rent_all_mut(|graphics| -> Result<(), DrmError> {
                graphics.gbm.cursor = {
                    let mut cursor = graphics.context.devices.gbm.create_buffer_object(
                        w,
                        h,
                        GbmFormat::ARGB8888,
                        &[BufferObjectFlags::Cursor, BufferObjectFlags::Write],
                    )?;
                    cursor.write(&*buffer.into_raw())?;
                    cursor
                };
                Ok(())
            })?;

        if crtc::set_cursor2(
            self.0.borrow().graphics.head().head().head(),
            self.0.borrow().crtc,
            self.0
                .borrow()
                .graphics
                .rent(|gbm| Buffer::handle(&gbm.cursor)),
            (w, h),
            (hotspot.0 as i32, hotspot.1 as i32),
        ).is_err()
        {
            crtc::set_cursor(
                self.0.borrow().graphics.head().head().head(),
                self.0.borrow().crtc,
                self.0
                    .borrow()
                    .graphics
                    .rent(|gbm| Buffer::handle(&gbm.cursor)),
                (w, h),
            ).map_err(DrmError::from)
        } else {
            Ok(())
        }
    }
}

impl EGLGraphicsBackend for DrmBackend {
    fn swap_buffers(&self) -> Result<(), SwapBuffersError> {
        self.0.borrow().graphics.rent_all(|graphics| {
            // We cannot call lock_front_buffer anymore without releasing the previous buffer, which will happen when the page flip is done
            if graphics.gbm.surface.rent(|egl| {
                let next = egl.buffers.next_buffer.take();
                let res = next.is_some();
                egl.buffers.next_buffer.set(next);
                res
            }) {
                return Err(SwapBuffersError::AlreadySwapped);
            }

            graphics.gbm.surface.rent(|egl| egl.surface.swap_buffers())?;

            graphics.gbm.surface.rent_all(|surface| {
                // supporting this would cause a lot of inconvinience and
                // would most likely result in a lot of flickering.
                // neither weston, wlc or wlroots bother with that as well.
                let mut next_bo = surface.gbm.lock_front_buffer().expect("Surface only has one front buffer. Not supported by smithay");

                let maybe_fb = next_bo.userdata().cloned();
                let fb = if let Some(info) = maybe_fb {
                    info
                } else {
                    let fb = framebuffer::create(graphics.context.devices.drm, &*next_bo).map_err(|_| SwapBuffersError::ContextLost)?;
                    next_bo.set_userdata(fb);
                    fb
                };
                surface.egl.buffers.next_buffer.set(Some(next_bo));

                trace!(self.0.borrow().logger, "Page flip queued");

                let id: Id = self.0.borrow().own_id;

                crtc::page_flip(graphics.context.devices.drm, self.0.borrow().crtc, fb.handle(), &[crtc::PageFlipFlags::PageFlipEvent], id).map_err(|_| SwapBuffersError::ContextLost)
            })
        })
    }

    unsafe fn get_proc_address(&self, symbol: &str) -> *const c_void {
        self.0
            .borrow()
            .graphics
            .head()
            .rent(|context| context.get_proc_address(symbol))
    }

    fn get_framebuffer_dimensions(&self) -> (u32, u32) {
        let (w, h) = self.0.borrow().mode.size();
        (w as u32, h as u32)
    }

    fn is_current(&self) -> bool {
        self.0
            .borrow()
            .graphics
            .head()
            .rent(|context| context.is_current())
    }

    unsafe fn make_current(&self) -> Result<(), SwapBuffersError> {
        self.0
            .borrow()
            .graphics
            .rent(|gbm| gbm.surface.rent(|egl| egl.surface.make_current()))
    }

    fn get_pixel_format(&self) -> PixelFormat {
        self.0
            .borrow()
            .graphics
            .head()
            .rent(|context| context.get_pixel_format())
    }
}
<commit_msg>drm: Add more inline comments explaining the drm code<commit_after>
use super::{DrmError, ModeError};

use super::devices;

use backend::graphics::GraphicsBackend;
use backend::graphics::egl::{EGLGraphicsBackend, EGLSurface, PixelFormat, SwapBuffersError};
use drm::buffer::Buffer;
use drm::control::{connector, crtc, framebuffer, Mode};
use drm::control::ResourceInfo;

use gbm::{BufferObject, BufferObjectFlags, Format as GbmFormat, Surface as GbmSurface, SurfaceBufferHandle};

use image::{ImageBuffer, Rgba};

use nix::c_void;

use std::cell::{Cell, RefCell};
use std::rc::Rc;

/// Backend based on a `DrmDevice` and a given crtc
pub struct DrmBackend(Rc<RefCell<DrmBackendInternal>>);

impl DrmBackend {
    pub(crate) fn new(drm: Rc<RefCell<DrmBackendInternal>>) -> DrmBackend {
        DrmBackend(drm)
    }
}

/*
    Dependency graph
    - drm
        - gbm
            - context
            - gbm_surface
                - egl_surface
                - gbm_buffers
            - cursor
*/

pub(crate) struct GbmTypes<'dev, 'context> {
    cursor: BufferObject<'dev, ()>,
    surface: Surface<'context>,
}

pub(crate) struct EGL<'gbm, 'context> {
    surface: EGLSurface<'context, 'gbm, GbmSurface<'context, framebuffer::Info>>,
    buffers: GbmBuffers<'gbm>,
}

pub(crate) struct GbmBuffers<'gbm> {
    front_buffer: Cell<SurfaceBufferHandle<'gbm, framebuffer::Info>>,
    next_buffer: Cell<Option<SurfaceBufferHandle<'gbm, framebuffer::Info>>>,
}

rental! {
    mod graphics {
        use drm::control::framebuffer;
        use gbm::Surface as GbmSurface;
        use std::rc::Rc;
        use super::devices::{Context, Context_Borrow};
        use super::GbmTypes;

        #[rental]
        pub(crate) struct Surface<'context> {
            gbm: Box<GbmSurface<'context, framebuffer::Info>>,
            egl: super::EGL<'gbm, 'context>,
        }

        #[rental]
        pub(crate) struct Graphics {
            #[subrental(arity = 3)]
            context: Rc<Context>,
            gbm: GbmTypes<'context_1, 'context_2>,
        }
    }
}
use self::graphics::{Graphics, Surface};

/// Id of a `DrmBackend` related to its `DrmDevice`.
///
/// Used to track which `DrmBackend` finished page-flipping
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct Id(usize);

impl Id {
    pub(crate) fn raw(&self) -> usize {
        self.0
    }
}

pub(crate) struct DrmBackendInternal {
    graphics: Graphics,
    crtc: crtc::Handle,
    mode: Mode,
    connectors: Vec<connector::Handle>,
    own_id: Id,
    logger: ::slog::Logger,
}

impl DrmBackendInternal {
    pub(crate) fn new<I, L>(context: Rc<devices::Context>, crtc: crtc::Handle, mode: Mode, connectors: I,
                            own_id: usize, logger: L)
                            -> Result<DrmBackendInternal, DrmError>
    where
        I: Into<Vec<connector::Handle>>,
        L: Into<Option<::slog::Logger>>,
    {
        let log = ::slog_or_stdlog(logger).new(o!("smithay_module" => "backend_drm"));
        info!(log, "Initializing a drm backend");

        let connectors = connectors.into();

        // check the connectors, if they suite the mode
        for connector in connectors.iter() {
            if !connector::Info::load_from_device(context.head().head(), *connector)?
                .modes()
                .contains(&mode)
            {
                return Err(DrmError::Mode(ModeError::ModeNotSuitable));
            }
        }

        let (w, h) = mode.size();

        info!(log, "Drm Backend initializing");

        Ok(DrmBackendInternal {
            graphics: Graphics::try_new(context, |context| {
                debug!(log, "GBM EGLContext initialized");

                Ok(GbmTypes {
                    cursor: {
                        // Create an unused cursor buffer (we don't want an Option here)
                        context.devices.gbm.create_buffer_object(
                            1,
                            1,
                            GbmFormat::ARGB8888,
                            &[BufferObjectFlags::Cursor, BufferObjectFlags::Write],
                        )?
                    },
                    surface: Surface::try_new(
                        // create a gbm surface
                        Box::new(context.devices.gbm.create_surface(
                            w as u32,
                            h as u32,
                            GbmFormat::XRGB8888,
                            &[BufferObjectFlags::Scanout, BufferObjectFlags::Rendering],
                        )?),
                        |surface| {
                            // create an egl surface from the gbm one
                            let egl_surface = context.egl.create_surface(&surface)?;

                            // set it to be able to use `crtc::set` once
                            unsafe { egl_surface.make_current()? };
                            egl_surface.swap_buffers()?;

                            // init the first screen
                            // (must be done before calling page_flip for the first time)
                            let mut front_bo = surface.lock_front_buffer()?;
                            debug!(log, "FrontBuffer color format: {:?}", front_bo.format());
                            // we need a framebuffer per front buffer
                            let fb = framebuffer::create(context.devices.drm, &*front_bo)?;
                            crtc::set(
                                context.devices.drm,
                                crtc,
                                fb.handle(),
                                &connectors,
                                (0, 0),
                                Some(mode),
                            )?;
                            front_bo.set_userdata(fb);

                            Ok(EGL {
                                surface: egl_surface,
                                buffers: GbmBuffers {
                                    front_buffer: Cell::new(front_bo),
                                    next_buffer: Cell::new(None),
                                },
                            })
                        },
                    ).map_err(DrmError::from)?,
                })
            })?,
            crtc,
            mode,
            connectors,
            own_id: Id(own_id),
            logger: log.clone(),
        })
    }

    pub(crate) fn unlock_buffer(&self) {
        // after the page swap is finished we need to release the rendered buffer.
        // this is called from the PageFlipHandler
        self.graphics.rent(|gbm| {
            gbm.surface.rent(|egl| {
                let next_bo = egl.buffers.next_buffer.replace(None);

                if let Some(next_buffer) = next_bo {
                    egl.buffers.front_buffer.set(next_buffer);
                // drop and release the old buffer
                } else {
                    unreachable!();
                }
            })
        });
    }

    pub(crate) fn is_crtc(&self, crtc: crtc::Handle) -> bool {
        crtc == self.crtc
    }
}

impl DrmBackend {
    /// Add a connector to backend
    ///
    /// # Errors
    ///
    /// Errors if the new connector does not support the currently set `Mode`
    pub fn add_connector(&mut self, connector: connector::Handle) -> Result<(), ModeError> {
        let info =
            connector::Info::load_from_device(self.0.borrow().graphics.head().head().head(), connector)
                .map_err(|err| ModeError::FailedToLoad(err))?;

        // check if the connector can handle the current mode
        let mut internal = self.0.borrow_mut();
        if info.modes().contains(&internal.mode) {
            internal.connectors.push(connector);
            Ok(())
        } else {
            Err(ModeError::ModeNotSuitable)
        }
    }

    /// Returns a copy of the currently set connectors
    pub fn used_connectors(&self) -> Vec<connector::Handle> {
        // thanks to the RefCell we can sadly not return a `&[connector::Handle]`
        self.0.borrow().connectors.clone()
    }

    /// Removes a currently set connector
    pub fn remove_connector(&mut self, connector: connector::Handle) {
        self.0.borrow_mut().connectors.retain(|x| *x != connector);
    }

    /// Changes the currently set mode
    ///
    /// # Errors
    ///
    /// This will fail if not all set connectors support the new `Mode`.
    /// Several internal resources will need to be recreated to fit the new `Mode`.
    /// Other errors might occur.
    pub fn use_mode(&mut self, mode: Mode) -> Result<(), DrmError> {
        // check the connectors
        for connector in self.0.borrow().connectors.iter() {
            if !connector::Info::load_from_device(self.0.borrow().graphics.head().head().head(), *connector)?
                .modes()
                .contains(&mode)
            {
                return Err(DrmError::Mode(ModeError::ModeNotSuitable));
            }
        }

        // borrow & clone stuff because rust cannot figure out the upcoming
        // closure otherwise.
        let crtc = self.0.borrow().crtc;
        let mut internal = self.0.borrow_mut();
        let connectors = internal.connectors.clone();
        let logger = internal.logger.clone();

        let (w, h) = mode.size();

        internal
            .graphics
            .rent_all_mut(|graphics| -> Result<(), DrmError> {
                // Recreate the surface and the related resources to match the new
                // resolution.
                graphics.gbm.surface = Surface::try_new(
                    Box::new(graphics.context.devices.gbm.create_surface(
                        w as u32,
                        h as u32,
                        GbmFormat::XRGB8888,
                        &[BufferObjectFlags::Scanout, BufferObjectFlags::Rendering],
                    )?),
                    |surface| {
                        let egl_surface = graphics.context.egl.create_surface(&surface)?;

                        // make it active for the first crtc::set
                        // (which is needed before the first page_flip)
                        unsafe { egl_surface.make_current()? };
                        egl_surface.swap_buffers()?;

                        let mut front_bo = surface.lock_front_buffer()?;
                        debug!(logger, "FrontBuffer color format: {:?}", front_bo.format());
                        // we need a framebuffer per front_buffer
                        let fb = framebuffer::create(graphics.context.devices.drm, &*front_bo)?;
                        // init the first screen
                        crtc::set(
                            graphics.context.devices.drm,
                            crtc,
                            fb.handle(),
                            &connectors,
                            (0, 0),
                            Some(mode),
                        )?;
                        front_bo.set_userdata(fb);

                        Ok(EGL {
                            surface: egl_surface,
                            buffers: GbmBuffers {
                                front_buffer: Cell::new(front_bo),
                                next_buffer: Cell::new(None),
                            },
                        })
                    },
                )?;

                Ok(())
            })?;

        internal.mode = mode;
        Ok(())
    }

    /// Checks of the `DrmBackend` is of the given `Id`
    ///
    /// Only produces valid results, if the `Id` is from the `DrmDevice`,
    /// that created this backend.
    pub fn is(&self, id: Id) -> bool {
        self.0.borrow().own_id == id
    }
}

impl GraphicsBackend for DrmBackend {
    type CursorFormat = ImageBuffer<Rgba<u8>, Vec<u8>>;
    type Error = DrmError;

    fn set_cursor_position(&self, x: u32, y: u32) -> Result<(), DrmError> {
        crtc::move_cursor(
            self.0.borrow().graphics.head().head().head(),
            self.0.borrow().crtc,
            (x as i32, y as i32),
        ).map_err(DrmError::from)
    }

    fn set_cursor_representation(&self, buffer: ImageBuffer<Rgba<u8>, Vec<u8>>, hotspot: (u32, u32))
                                 -> Result<(), DrmError> {
        let (w, h) = buffer.dimensions();
        /// import the cursor into a buffer we can render
        self.0
            .borrow_mut()
            .graphics
            .rent_all_mut(|graphics| -> Result<(), DrmError> {
                graphics.gbm.cursor = {
                    let mut cursor = graphics.context.devices.gbm.create_buffer_object(
                        w,
                        h,
                        GbmFormat::ARGB8888,
                        &[BufferObjectFlags::Cursor, BufferObjectFlags::Write],
                    )?;
                    cursor.write(&*buffer.into_raw())?;
                    cursor
                };
                Ok(())
            })?;

        // and set it
        if crtc::set_cursor2(
            self.0.borrow().graphics.head().head().head(),
            self.0.borrow().crtc,
            self.0
                .borrow()
                .graphics
                .rent(|gbm| Buffer::handle(&gbm.cursor)),
            (w, h),
            (hotspot.0 as i32, hotspot.1 as i32),
        ).is_err()
        {
            crtc::set_cursor(
                self.0.borrow().graphics.head().head().head(),
                self.0.borrow().crtc,
                self.0
                    .borrow()
                    .graphics
                    .rent(|gbm| Buffer::handle(&gbm.cursor)),
                (w, h),
            ).map_err(DrmError::from)
        } else {
            Ok(())
        }
    }
}

impl EGLGraphicsBackend for DrmBackend {
    fn swap_buffers(&self) -> Result<(), SwapBuffersError> {
        self.0.borrow().graphics.rent_all(|graphics| {
            // We cannot call lock_front_buffer anymore without releasing the previous buffer, which will happen when the page flip is done
            if graphics.gbm.surface.rent(|egl| {
                let next = egl.buffers.next_buffer.take();
                let res = next.is_some();
                egl.buffers.next_buffer.set(next);
                res
            }) {
                return Err(SwapBuffersError::AlreadySwapped);
            }

            /// flip normally
            graphics.gbm.surface.rent(|egl| egl.surface.swap_buffers())?;

            graphics.gbm.surface.rent_all(|surface| {
                // supporting this error would cause a lot of inconvinience and
                // would most likely result in a lot of flickering.
                // neither weston, wlc or wlroots bother with that as well.
                // so we just assume we got at least two buffers to do flipping
                let mut next_bo = surface.gbm.lock_front_buffer().expect("Surface only has one front buffer. Not supported by smithay");

                // create a framebuffer if the front buffer does not have one already
                // (they are reused by gbm)
                let maybe_fb = next_bo.userdata().cloned();
                let fb = if let Some(info) = maybe_fb {
                    info
                } else {
                    let fb = framebuffer::create(graphics.context.devices.drm, &*next_bo).map_err(|_| SwapBuffersError::ContextLost)?;
                    next_bo.set_userdata(fb);
                    fb
                };
                surface.egl.buffers.next_buffer.set(Some(next_bo));

                trace!(self.0.borrow().logger, "Page flip queued");

                let id: Id = self.0.borrow().own_id;

                // and flip
                crtc::page_flip(graphics.context.devices.drm, self.0.borrow().crtc, fb.handle(), &[crtc::PageFlipFlags::PageFlipEvent], id).map_err(|_| SwapBuffersError::ContextLost)
            })
        })
    }

    unsafe fn get_proc_address(&self, symbol: &str) -> *const c_void {
        self.0
            .borrow()
            .graphics
            .head()
            .rent(|context| context.get_proc_address(symbol))
    }

    fn get_framebuffer_dimensions(&self) -> (u32, u32) {
        let (w, h) = self.0.borrow().mode.size();
        (w as u32, h as u32)
    }

    fn is_current(&self) -> bool {
        self.0
            .borrow()
            .graphics
            .head()
            .rent(|context| context.is_current())
    }

    unsafe fn make_current(&self) -> Result<(), SwapBuffersError> {
        self.0
            .borrow()
            .graphics
            .rent(|gbm| gbm.surface.rent(|egl| egl.surface.make_current()))
    }

    fn get_pixel_format(&self) -> PixelFormat {
        self.0
            .borrow()
            .graphics
            .head()
            .rent(|context| context.get_pixel_format())
    }
}
<|endoftext|>