<commit_before>#![cfg(feature = "glutin")]
/*!

Backend implementation for the glutin library

# Features

Only available if the 'glutin' feature is enabled.

*/
extern crate glutin;

use libc;

use DisplayBuild;
use Frame;
use GliumCreationError;

use context;
use backend;
use backend::Context;
use backend::Backend;
use version::Version;

use std::cell::{RefCell, Ref};
use std::rc::Rc;
use std::ops::Deref;

/// Facade implementation for glutin. Wraps both glium and glutin.
#[derive(Clone)]
pub struct GlutinFacade {
    // contains everything related to the current context and its state
    context: Rc<context::Context>,

    // contains the window
    backend: Rc<Option<RefCell<Rc<GlutinWindowBackend>>>>,
}

impl backend::Facade for GlutinFacade {
    fn get_context(&self) -> &Rc<Context> {
        &self.context
    }
}

/// Iterator for all the events received by the window.
pub struct PollEventsIter<'a> {
    window: Option<&'a RefCell<Rc<GlutinWindowBackend>>>,
}

impl<'a> Iterator for PollEventsIter<'a> {
    type Item = glutin::Event;

    fn next(&mut self) -> Option<glutin::Event> {
        if let Some(window) = self.window.as_ref() {
            window.borrow().poll_events().next()
        } else {
            None
        }
    }
}

/// Blocking iterator over all the events received by the window.
///
/// This iterator polls for events, until the window associated with its context
/// is closed.
pub struct WaitEventsIter<'a> {
    window: Option<&'a RefCell<Rc<GlutinWindowBackend>>>,
}

impl<'a> Iterator for WaitEventsIter<'a> {
    type Item = glutin::Event;

    fn next(&mut self) -> Option<glutin::Event> {
        if let Some(window) = self.window.as_ref() {
            window.borrow().wait_events().next()
        } else {
            None
        }
    }
}

/// Borrow of the glutin window.
pub struct WinRef<'a>(Ref<'a, Rc<GlutinWindowBackend>>);

impl<'a> Deref for WinRef<'a> {
    type Target = glutin::Window;

    fn deref(&self) -> &glutin::Window {
        self.0.get_window()
    }
}

impl GlutinFacade {
    /// Reads all events received by the window.
    ///
    /// This iterator polls for events and can be exhausted.
    pub fn poll_events(&self) -> PollEventsIter {
        PollEventsIter {
            window: self.backend.as_ref(),
        }
    }

    /// Reads all events received by the window.
    pub fn wait_events(&self) -> WaitEventsIter {
        WaitEventsIter {
            window: self.backend.as_ref(),
        }
    }

    /// Returns true if the window has been closed.
    pub fn is_closed(&self) -> bool {
        self.backend.as_ref().map(|b| b.borrow().is_closed()).unwrap_or(false)
    }

    /// Returns the underlying window, or `None` if glium uses a headless context.
    pub fn get_window(&self) -> Option<WinRef> {
        self.backend.as_ref().map(|w| WinRef(w.borrow()))
    }

    /// Returns the OpenGL version of the current context.
    // TODO: change Context so that this function derefs from it as well
    pub fn get_opengl_version(&self) -> Version {
        *self.context.get_version()
    }

    /// Start drawing on the backbuffer.
    ///
    /// This function returns a `Frame`, which can be used to draw on it. When the `Frame` is
    /// destroyed, the buffers are swapped.
    ///
    /// Note that destroying a `Frame` is immediate, even if vsync is enabled.
    pub fn draw(&self) -> Frame {
        Frame::new(self.context.clone(), self.get_framebuffer_dimensions())
    }
}

impl Deref for GlutinFacade {
    type Target = Context;

    fn deref(&self) -> &Context {
        &self.context
    }
}

impl DisplayBuild for glutin::WindowBuilder<'static> {
    type Facade = GlutinFacade;
    type Err = GliumCreationError<glutin::CreationError>;

    fn build_glium(self) -> Result<GlutinFacade, GliumCreationError<glutin::CreationError>> {
        let backend = Rc::new(try!(backend::glutin_backend::GlutinWindowBackend::new(self)));
        let context = try!(unsafe { context::Context::new(backend.clone(), true) });

        let display = GlutinFacade {
            context: context,
            backend: Rc::new(Some(RefCell::new(backend))),
        };

        Ok(display)
    }

    unsafe fn build_glium_unchecked(self) -> Result<GlutinFacade, GliumCreationError<glutin::CreationError>> {
        let backend = Rc::new(try!(backend::glutin_backend::GlutinWindowBackend::new(self)));
        let context = try!(context::Context::new(backend.clone(), false));

        let display = GlutinFacade {
            context: context,
            backend: Rc::new(Some(RefCell::new(backend))),
        };

        Ok(display)
    }

    fn rebuild_glium(self, display: &GlutinFacade) -> Result<(), GliumCreationError<glutin::CreationError>> {
        let mut existing_window = display.backend.as_ref()
                                         .expect("can't rebuild a headless display").borrow_mut();
        let new_backend = Rc::new(try!(existing_window.rebuild(self)));
        try!(unsafe { display.context.rebuild(new_backend.clone()) });
        *existing_window = new_backend;
        Ok(())
    }
}

impl DisplayBuild for glutin::HeadlessRendererBuilder {
    type Facade = GlutinFacade;
    type Err = GliumCreationError<glutin::CreationError>;

    fn build_glium(self) -> Result<GlutinFacade, GliumCreationError<glutin::CreationError>> {
        let backend = Rc::new(try!(backend::glutin_backend::GlutinHeadlessBackend::new(self)));
        let context = try!(unsafe { context::Context::new(backend.clone(), true) });

        let display = GlutinFacade {
            context: context,
            backend: Rc::new(None),
        };

        Ok(display)
    }

    unsafe fn build_glium_unchecked(self) -> Result<GlutinFacade, GliumCreationError<glutin::CreationError>> {
        let backend = Rc::new(try!(backend::glutin_backend::GlutinHeadlessBackend::new(self)));
        let context = try!(context::Context::new(backend.clone(), true));

        let display = GlutinFacade {
            context: context,
            backend: Rc::new(None),
        };

        Ok(display)
    }

    fn rebuild_glium(self, _: &GlutinFacade) -> Result<(), GliumCreationError<glutin::CreationError>> {
        unimplemented!()
    }
}

/// An implementation of the `Backend` trait for a glutin window.
pub struct GlutinWindowBackend {
    window: glutin::Window,
}

unsafe impl Backend for GlutinWindowBackend {
    fn swap_buffers(&self) {
        self.window.swap_buffers();
    }

    unsafe fn get_proc_address(&self, symbol: &str) -> *const libc::c_void {
        self.window.get_proc_address(symbol)
    }

    fn get_framebuffer_dimensions(&self) -> (u32, u32) {
        let (width, height) = self.window.get_inner_size().unwrap_or((800, 600));      // TODO: 800x600 ?
        let scale = self.window.hidpi_factor();
        ((width as f32 * scale) as u32, (height as f32 * scale) as u32)
    }

    fn is_current(&self) -> bool {
        self.window.is_current()
    }

    unsafe fn make_current(&self) {
        self.window.make_current();
    }
}

#[allow(missing_docs)]
impl GlutinWindowBackend {
    /// Builds a new backend from the builder.
    pub fn new(builder: glutin::WindowBuilder)
               -> Result<GlutinWindowBackend, GliumCreationError<glutin::CreationError>>
    {
        let window = try!(builder.build());

        Ok(GlutinWindowBackend {
            window: window,
        })
    }

    pub fn get_window(&self) -> &glutin::Window {
        &self.window
    }

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

    pub fn poll_events(&self) -> glutin::PollEventsIterator {
        self.window.poll_events()
    }

    pub fn wait_events(&self) -> glutin::WaitEventsIterator {
        self.window.wait_events()
    }

    pub fn rebuild(&self, builder: glutin::WindowBuilder)
                   -> Result<GlutinWindowBackend, GliumCreationError<glutin::CreationError>>
    {
        let window = try!(builder.with_shared_lists(&self.window).build());

        Ok(GlutinWindowBackend {
            window: window,
        })
    }
}

/// An implementation of the `Backend` trait for a glutin headless context.
pub struct GlutinHeadlessBackend {
    context: glutin::HeadlessContext,
}

unsafe impl Backend for GlutinHeadlessBackend {
    fn swap_buffers(&self) {
    }

    unsafe fn get_proc_address(&self, symbol: &str) -> *const libc::c_void {
        self.context.get_proc_address(symbol)
    }

    fn get_framebuffer_dimensions(&self) -> (u32, u32) {
        (800, 600)      // FIXME: these are random
    }

    fn is_current(&self) -> bool {
        self.context.is_current()
    }

    unsafe fn make_current(&self) {
        self.context.make_current();
    }
}

impl GlutinHeadlessBackend {
    /// Builds a new backend from the builder.
    pub fn new(builder: glutin::HeadlessRendererBuilder)
               -> Result<GlutinHeadlessBackend, GliumCreationError<glutin::CreationError>>
    {
        let context = try!(builder.build());

        Ok(GlutinHeadlessBackend {
            context: context,
        })
    }
}
<commit_msg>Fix code for Rust nightly<commit_after>#![cfg(feature = "glutin")]
/*!

Backend implementation for the glutin library

# Features

Only available if the 'glutin' feature is enabled.

*/
extern crate glutin;

use libc;

use DisplayBuild;
use Frame;
use GliumCreationError;

use context;
use backend;
use backend::Context;
use backend::Backend;
use version::Version;

use std::cell::{RefCell, Ref};
use std::rc::Rc;
use std::ops::Deref;

/// Facade implementation for glutin. Wraps both glium and glutin.
#[derive(Clone)]
pub struct GlutinFacade {
    // contains everything related to the current context and its state
    context: Rc<context::Context>,

    // contains the window
    backend: Rc<Option<RefCell<Rc<GlutinWindowBackend>>>>,
}

impl backend::Facade for GlutinFacade {
    fn get_context(&self) -> &Rc<Context> {
        &self.context
    }
}

/// Iterator for all the events received by the window.
pub struct PollEventsIter<'a> {
    window: Option<&'a RefCell<Rc<GlutinWindowBackend>>>,
}

impl<'a> Iterator for PollEventsIter<'a> {
    type Item = glutin::Event;

    fn next(&mut self) -> Option<glutin::Event> {
        if let Some(window) = self.window.as_ref() {
            window.borrow().poll_events().next()
        } else {
            None
        }
    }
}

/// Blocking iterator over all the events received by the window.
///
/// This iterator polls for events, until the window associated with its context
/// is closed.
pub struct WaitEventsIter<'a> {
    window: Option<&'a RefCell<Rc<GlutinWindowBackend>>>,
}

impl<'a> Iterator for WaitEventsIter<'a> {
    type Item = glutin::Event;

    fn next(&mut self) -> Option<glutin::Event> {
        if let Some(window) = self.window.as_ref() {
            window.borrow().wait_events().next()
        } else {
            None
        }
    }
}

/// Borrow of the glutin window.
pub struct WinRef<'a>(Ref<'a, Rc<GlutinWindowBackend>>);

impl<'a> Deref for WinRef<'a> {
    type Target = glutin::Window;

    fn deref(&self) -> &glutin::Window {
        self.0.get_window()
    }
}

impl GlutinFacade {
    /// Reads all events received by the window.
    ///
    /// This iterator polls for events and can be exhausted.
    pub fn poll_events(&self) -> PollEventsIter {
        PollEventsIter {
            window: Option::as_ref(&self.backend),
        }
    }

    /// Reads all events received by the window.
    pub fn wait_events(&self) -> WaitEventsIter {
        WaitEventsIter {
            window: Option::as_ref(&self.backend),
        }
    }

    /// Returns true if the window has been closed.
    pub fn is_closed(&self) -> bool {
        Option::as_ref(&self.backend).map(|b| b.borrow().is_closed()).unwrap_or(false)
    }

    /// Returns the underlying window, or `None` if glium uses a headless context.
    pub fn get_window(&self) -> Option<WinRef> {
        Option::as_ref(&self.backend).map(|w| WinRef(w.borrow()))
    }

    /// Returns the OpenGL version of the current context.
    // TODO: change Context so that this function derefs from it as well
    pub fn get_opengl_version(&self) -> Version {
        *self.context.get_version()
    }

    /// Start drawing on the backbuffer.
    ///
    /// This function returns a `Frame`, which can be used to draw on it. When the `Frame` is
    /// destroyed, the buffers are swapped.
    ///
    /// Note that destroying a `Frame` is immediate, even if vsync is enabled.
    pub fn draw(&self) -> Frame {
        Frame::new(self.context.clone(), self.get_framebuffer_dimensions())
    }
}

impl Deref for GlutinFacade {
    type Target = Context;

    fn deref(&self) -> &Context {
        &self.context
    }
}

impl DisplayBuild for glutin::WindowBuilder<'static> {
    type Facade = GlutinFacade;
    type Err = GliumCreationError<glutin::CreationError>;

    fn build_glium(self) -> Result<GlutinFacade, GliumCreationError<glutin::CreationError>> {
        let backend = Rc::new(try!(backend::glutin_backend::GlutinWindowBackend::new(self)));
        let context = try!(unsafe { context::Context::new(backend.clone(), true) });

        let display = GlutinFacade {
            context: context,
            backend: Rc::new(Some(RefCell::new(backend))),
        };

        Ok(display)
    }

    unsafe fn build_glium_unchecked(self) -> Result<GlutinFacade, GliumCreationError<glutin::CreationError>> {
        let backend = Rc::new(try!(backend::glutin_backend::GlutinWindowBackend::new(self)));
        let context = try!(context::Context::new(backend.clone(), false));

        let display = GlutinFacade {
            context: context,
            backend: Rc::new(Some(RefCell::new(backend))),
        };

        Ok(display)
    }

    fn rebuild_glium(self, display: &GlutinFacade) -> Result<(), GliumCreationError<glutin::CreationError>> {
        let mut existing_window = Option::as_ref(&display.backend)
                                         .expect("can't rebuild a headless display").borrow_mut();
        let new_backend = Rc::new(try!(existing_window.rebuild(self)));
        try!(unsafe { display.context.rebuild(new_backend.clone()) });
        *existing_window = new_backend;
        Ok(())
    }
}

impl DisplayBuild for glutin::HeadlessRendererBuilder {
    type Facade = GlutinFacade;
    type Err = GliumCreationError<glutin::CreationError>;

    fn build_glium(self) -> Result<GlutinFacade, GliumCreationError<glutin::CreationError>> {
        let backend = Rc::new(try!(backend::glutin_backend::GlutinHeadlessBackend::new(self)));
        let context = try!(unsafe { context::Context::new(backend.clone(), true) });

        let display = GlutinFacade {
            context: context,
            backend: Rc::new(None),
        };

        Ok(display)
    }

    unsafe fn build_glium_unchecked(self) -> Result<GlutinFacade, GliumCreationError<glutin::CreationError>> {
        let backend = Rc::new(try!(backend::glutin_backend::GlutinHeadlessBackend::new(self)));
        let context = try!(context::Context::new(backend.clone(), true));

        let display = GlutinFacade {
            context: context,
            backend: Rc::new(None),
        };

        Ok(display)
    }

    fn rebuild_glium(self, _: &GlutinFacade) -> Result<(), GliumCreationError<glutin::CreationError>> {
        unimplemented!()
    }
}

/// An implementation of the `Backend` trait for a glutin window.
pub struct GlutinWindowBackend {
    window: glutin::Window,
}

unsafe impl Backend for GlutinWindowBackend {
    fn swap_buffers(&self) {
        self.window.swap_buffers();
    }

    unsafe fn get_proc_address(&self, symbol: &str) -> *const libc::c_void {
        self.window.get_proc_address(symbol)
    }

    fn get_framebuffer_dimensions(&self) -> (u32, u32) {
        let (width, height) = self.window.get_inner_size().unwrap_or((800, 600));      // TODO: 800x600 ?
        let scale = self.window.hidpi_factor();
        ((width as f32 * scale) as u32, (height as f32 * scale) as u32)
    }

    fn is_current(&self) -> bool {
        self.window.is_current()
    }

    unsafe fn make_current(&self) {
        self.window.make_current();
    }
}

#[allow(missing_docs)]
impl GlutinWindowBackend {
    /// Builds a new backend from the builder.
    pub fn new(builder: glutin::WindowBuilder)
               -> Result<GlutinWindowBackend, GliumCreationError<glutin::CreationError>>
    {
        let window = try!(builder.build());

        Ok(GlutinWindowBackend {
            window: window,
        })
    }

    pub fn get_window(&self) -> &glutin::Window {
        &self.window
    }

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

    pub fn poll_events(&self) -> glutin::PollEventsIterator {
        self.window.poll_events()
    }

    pub fn wait_events(&self) -> glutin::WaitEventsIterator {
        self.window.wait_events()
    }

    pub fn rebuild(&self, builder: glutin::WindowBuilder)
                   -> Result<GlutinWindowBackend, GliumCreationError<glutin::CreationError>>
    {
        let window = try!(builder.with_shared_lists(&self.window).build());

        Ok(GlutinWindowBackend {
            window: window,
        })
    }
}

/// An implementation of the `Backend` trait for a glutin headless context.
pub struct GlutinHeadlessBackend {
    context: glutin::HeadlessContext,
}

unsafe impl Backend for GlutinHeadlessBackend {
    fn swap_buffers(&self) {
    }

    unsafe fn get_proc_address(&self, symbol: &str) -> *const libc::c_void {
        self.context.get_proc_address(symbol)
    }

    fn get_framebuffer_dimensions(&self) -> (u32, u32) {
        (800, 600)      // FIXME: these are random
    }

    fn is_current(&self) -> bool {
        self.context.is_current()
    }

    unsafe fn make_current(&self) {
        self.context.make_current();
    }
}

impl GlutinHeadlessBackend {
    /// Builds a new backend from the builder.
    pub fn new(builder: glutin::HeadlessRendererBuilder)
               -> Result<GlutinHeadlessBackend, GliumCreationError<glutin::CreationError>>
    {
        let context = try!(builder.build());

        Ok(GlutinHeadlessBackend {
            context: context,
        })
    }
}
<|endoftext|>


<commit_before>use crate::custom::credential::NewAwsCredsForStsCreds;
use crate::{AssumeRoleWithWebIdentityRequest, Sts, StsClient};
use rusoto_core::credential::{
    AwsCredentials, CredentialsError, ProvideAwsCredentials, Secret, Variable,
};
use rusoto_core::request::HttpClient;
use rusoto_core::{Client, Region};

use async_trait::async_trait;

const AWS_WEB_IDENTITY_TOKEN_FILE: &str = "AWS_WEB_IDENTITY_TOKEN_FILE";

const AWS_ROLE_ARN: &str = "AWS_ROLE_ARN";

const AWS_ROLE_SESSION_NAME: &str = "AWS_ROLE_SESSION_NAME";

/// WebIdentityProvider using OpenID Connect bearer token to retrieve AWS IAM credentials.
///
/// See https://docs.aws.amazon.com/STS/latest/APIReference/API_AssumeRoleWithWebIdentity.html for
/// more details.
#[derive(Debug, Clone)]
pub struct WebIdentityProvider {
    /// The OAuth 2.0 access token or OpenID Connect ID token that is provided by the identity provider.
    /// Your application must get this token by authenticating the user who is using your application
    /// with a web identity provider before the application makes an AssumeRoleWithWebIdentity call.
    pub web_identity_token: Variable<Secret, CredentialsError>,
    /// The Amazon Resource Name (ARN) of the role that the caller is assuming.
    pub role_arn: Variable<String, CredentialsError>,
    /// An identifier for the assumed role session. Typically, you pass the name or identifier that is
    /// associated with the user who is using your application. That way, the temporary security credentials
    /// that your application will use are associated with that user. This session name is included as part
    /// of the ARN and assumed role ID in the AssumedRoleUser response element.
    pub role_session_name: Variable<String, CredentialsError>,
}

impl WebIdentityProvider {
    /// Create new WebIdentityProvider by explicitly passing its configuration.
    pub fn new<A, B, C>(web_identity_token: A, role_arn: B, role_session_name: Option<C>) -> Self
    where
        A: Into<Variable<Secret, CredentialsError>>,
        B: Into<Variable<String, CredentialsError>>,
        C: Into<Variable<String, CredentialsError>>,
    {
        Self {
            web_identity_token: web_identity_token.into(),
            role_arn: role_arn.into(),
            role_session_name: role_session_name
                .map(|v| v.into())
                .unwrap_or_else(|| Variable::with_value(Self::create_session_name())),
        }
    }

    /// Creat a WebIdentityProvider from the following environment variables:
    ///
    /// - `AWS_WEB_IDENTITY_TOKEN_FILE` path to the web identity token file.
    /// - `AWS_ROLE_ARN` ARN of the role to assume.
    /// - `AWS_ROLE_SESSION_NAME` (optional) name applied to the assume-role session.
    ///
    /// See https://docs.aws.amazon.com/eks/latest/userguide/iam-roles-for-service-accounts-technical-overview.html
    /// for more information about how IAM Roles for Kubernetes Service Accounts works.
    pub fn from_k8s_env() -> Self {
        Self::_from_k8s_env(
            Variable::from_env_var(AWS_WEB_IDENTITY_TOKEN_FILE),
            Variable::from_env_var(AWS_ROLE_ARN),
            Some(Variable::from_env_var(AWS_ROLE_SESSION_NAME)),
        )
    }

    /// Used by unit testing
    pub(crate) fn _from_k8s_env(
        token_file: Variable<String, CredentialsError>,
        role: Variable<String, CredentialsError>,
        session_name: Option<Variable<String, CredentialsError>>,
    ) -> Self {
        Self::new(
            Variable::dynamic(move || Variable::from_text_file(token_file.resolve()?).resolve()),
            role,
            session_name,
        )
    }

    #[cfg(test)]
    pub(crate) fn load_token(&self) -> Result<Secret, CredentialsError> {
        self.web_identity_token.resolve()
    }

    fn create_session_name() -> String {
        // TODO can we do better here?
        // - Pod service account, Pod name and Pod namespace
        // - EC2 Instance ID if available
        // - IP address if available
        // - ...
        // Having some information in the session name that identifies the client would enable
        // better correlation analysis in CloudTrail.
        "WebIdentitySession".to_string()
    }
}

#[async_trait]
impl ProvideAwsCredentials for WebIdentityProvider {
    async fn credentials(&self) -> Result<AwsCredentials, CredentialsError> {
        let http_client = match HttpClient::new() {
            Ok(c) => c,
            Err(e) => return Err(CredentialsError::new(e)),
        };
        let client = Client::new_not_signing(http_client);
        let sts = StsClient::new_with_client(client, Region::default());
        let mut req = AssumeRoleWithWebIdentityRequest::default();

        req.role_arn = self.role_arn.resolve()?;
        req.web_identity_token = self.web_identity_token.resolve()?.to_string();
        req.role_session_name = self.role_session_name.resolve()?;

        let assume_role = sts.assume_role_with_web_identity(req).await;
        match assume_role {
            Err(e) => Err(CredentialsError::new(e)),
            Ok(role) => match role.credentials {
                None => Err(CredentialsError::new(format!(
                    "No credentials found in AssumeRoleWithWebIdentityResponse: {:?}",
                    role
                ))),
                Some(c) => AwsCredentials::new_for_credentials(c),
            },
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::io::Write;
    use tempfile::NamedTempFile;

    #[test]
    fn api_ergonomy() {
        WebIdentityProvider::new(Secret::from("".to_string()), "", Some("".to_string()));
    }

    #[test]
    fn from_k8s_env() -> Result<(), CredentialsError> {
        const TOKEN_VALUE: &str = "secret";
        const ROLE_ARN: &str = "role";
        const SESSION_NAME: &str = "session";
        let mut file = NamedTempFile::new()?;
        // We use writeln to add an extra newline at the end of the token, which should be
        // removed by Variable::from_text_file.
        writeln!(file, "{}", TOKEN_VALUE)?;
        let p = WebIdentityProvider::_from_k8s_env(
            Variable::with_value(file.path().to_string_lossy().to_string()),
            Variable::with_value(ROLE_ARN.to_string()),
            Some(Variable::with_value(SESSION_NAME.to_string())),
        );
        let token = p.load_token()?;
        assert_eq!(token.as_ref(), TOKEN_VALUE);
        Ok(())
    }
}
<commit_msg>Fix serializing jwt token when sending request<commit_after>use crate::custom::credential::NewAwsCredsForStsCreds;
use crate::{AssumeRoleWithWebIdentityRequest, Sts, StsClient};
use rusoto_core::credential::{
    AwsCredentials, CredentialsError, ProvideAwsCredentials, Secret, Variable,
};
use rusoto_core::request::HttpClient;
use rusoto_core::{Client, Region};

use async_trait::async_trait;

const AWS_WEB_IDENTITY_TOKEN_FILE: &str = "AWS_WEB_IDENTITY_TOKEN_FILE";

const AWS_ROLE_ARN: &str = "AWS_ROLE_ARN";

const AWS_ROLE_SESSION_NAME: &str = "AWS_ROLE_SESSION_NAME";

/// WebIdentityProvider using OpenID Connect bearer token to retrieve AWS IAM credentials.
///
/// See https://docs.aws.amazon.com/STS/latest/APIReference/API_AssumeRoleWithWebIdentity.html for
/// more details.
#[derive(Debug, Clone)]
pub struct WebIdentityProvider {
    /// The OAuth 2.0 access token or OpenID Connect ID token that is provided by the identity provider.
    /// Your application must get this token by authenticating the user who is using your application
    /// with a web identity provider before the application makes an AssumeRoleWithWebIdentity call.
    pub web_identity_token: Variable<Secret, CredentialsError>,
    /// The Amazon Resource Name (ARN) of the role that the caller is assuming.
    pub role_arn: Variable<String, CredentialsError>,
    /// An identifier for the assumed role session. Typically, you pass the name or identifier that is
    /// associated with the user who is using your application. That way, the temporary security credentials
    /// that your application will use are associated with that user. This session name is included as part
    /// of the ARN and assumed role ID in the AssumedRoleUser response element.
    pub role_session_name: Variable<String, CredentialsError>,
}

impl WebIdentityProvider {
    /// Create new WebIdentityProvider by explicitly passing its configuration.
    pub fn new<A, B, C>(web_identity_token: A, role_arn: B, role_session_name: Option<C>) -> Self
    where
        A: Into<Variable<Secret, CredentialsError>>,
        B: Into<Variable<String, CredentialsError>>,
        C: Into<Variable<String, CredentialsError>>,
    {
        Self {
            web_identity_token: web_identity_token.into(),
            role_arn: role_arn.into(),
            role_session_name: role_session_name
                .map(|v| v.into())
                .unwrap_or_else(|| Variable::with_value(Self::create_session_name())),
        }
    }

    /// Creat a WebIdentityProvider from the following environment variables:
    ///
    /// - `AWS_WEB_IDENTITY_TOKEN_FILE` path to the web identity token file.
    /// - `AWS_ROLE_ARN` ARN of the role to assume.
    /// - `AWS_ROLE_SESSION_NAME` (optional) name applied to the assume-role session.
    ///
    /// See https://docs.aws.amazon.com/eks/latest/userguide/iam-roles-for-service-accounts-technical-overview.html
    /// for more information about how IAM Roles for Kubernetes Service Accounts works.
    pub fn from_k8s_env() -> Self {
        Self::_from_k8s_env(
            Variable::from_env_var(AWS_WEB_IDENTITY_TOKEN_FILE),
            Variable::from_env_var(AWS_ROLE_ARN),
            Some(Variable::from_env_var(AWS_ROLE_SESSION_NAME)),
        )
    }

    /// Used by unit testing
    pub(crate) fn _from_k8s_env(
        token_file: Variable<String, CredentialsError>,
        role: Variable<String, CredentialsError>,
        session_name: Option<Variable<String, CredentialsError>>,
    ) -> Self {
        Self::new(
            Variable::dynamic(move || Variable::from_text_file(token_file.resolve()?).resolve()),
            role,
            session_name,
        )
    }

    #[cfg(test)]
    pub(crate) fn load_token(&self) -> Result<Secret, CredentialsError> {
        self.web_identity_token.resolve()
    }

    fn create_session_name() -> String {
        // TODO can we do better here?
        // - Pod service account, Pod name and Pod namespace
        // - EC2 Instance ID if available
        // - IP address if available
        // - ...
        // Having some information in the session name that identifies the client would enable
        // better correlation analysis in CloudTrail.
        "WebIdentitySession".to_string()
    }
}

#[async_trait]
impl ProvideAwsCredentials for WebIdentityProvider {
    async fn credentials(&self) -> Result<AwsCredentials, CredentialsError> {
        let http_client = match HttpClient::new() {
            Ok(c) => c,
            Err(e) => return Err(CredentialsError::new(e)),
        };
        let client = Client::new_not_signing(http_client);
        let sts = StsClient::new_with_client(client, Region::default());
        let mut req = AssumeRoleWithWebIdentityRequest::default();

        req.role_arn = self.role_arn.resolve()?;
        req.web_identity_token = self.web_identity_token.resolve()?.as_ref().to_string();
        req.role_session_name = self.role_session_name.resolve()?;

        let assume_role = sts.assume_role_with_web_identity(req).await;
        match assume_role {
            Err(e) => Err(CredentialsError::new(e)),
            Ok(role) => match role.credentials {
                None => Err(CredentialsError::new(format!(
                    "No credentials found in AssumeRoleWithWebIdentityResponse: {:?}",
                    role
                ))),
                Some(c) => AwsCredentials::new_for_credentials(c),
            },
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::io::Write;
    use tempfile::NamedTempFile;

    #[test]
    fn api_ergonomy() {
        WebIdentityProvider::new(Secret::from("".to_string()), "", Some("".to_string()));
    }

    #[test]
    fn from_k8s_env() -> Result<(), CredentialsError> {
        const TOKEN_VALUE: &str = "secret";
        const ROLE_ARN: &str = "role";
        const SESSION_NAME: &str = "session";
        let mut file = NamedTempFile::new()?;
        // We use writeln to add an extra newline at the end of the token, which should be
        // removed by Variable::from_text_file.
        writeln!(file, "{}", TOKEN_VALUE)?;
        let p = WebIdentityProvider::_from_k8s_env(
            Variable::with_value(file.path().to_string_lossy().to_string()),
            Variable::with_value(ROLE_ARN.to_string()),
            Some(Variable::with_value(SESSION_NAME.to_string())),
        );
        let token = p.load_token()?;
        assert_eq!(token.as_ref(), TOKEN_VALUE);
        Ok(())
    }
}
<|endoftext|>


<commit_before><commit_msg>Impl expr parsing.<commit_after><|endoftext|>


<commit_before><commit_msg>Fix imag-notes for new error interface<commit_after><|endoftext|>


<commit_before><commit_msg>Switch to using LOAD_OP_LOAD for both colour and depth in in renderpasses<commit_after><|endoftext|>


<commit_before><commit_msg>First tranche of rustdoc comments for renderervk.rs<commit_after><|endoftext|>


<commit_before><commit_msg>Ensure validation logic for m.key.verification.start is run when deserializing the event, not just the content.<commit_after><|endoftext|>


<commit_before><commit_msg>Add autobackup in print<commit_after><|endoftext|>


<commit_before><commit_msg>Use a macro for parser errors in expr.<commit_after><|endoftext|>


<commit_before>use super::{LzfResult, LzfError};
use std::{cmp, mem};

const HLOG    : usize = 16;
const HSIZE   : u32 = 1 << HLOG;
const MAX_OFF : usize = 1 << 13;
const MAX_REF : usize = ((1 << 8) + (1 << 3));
const MAX_LIT : i32 = (1 << 5);

fn first(p: &[u8], off: usize) -> u32 {
    ((p[off] as u32) << 8) | p[off+1] as u32
}

fn next(v: u32, p: &[u8], off: usize) -> u32 {
    (v << 8) | p[off+2] as u32
}

fn idx(h: u32) -> usize {
    let h = h as u64;
    (
        // 8 = 3*8-HLOG, but HLOG is constant at 16
        (h.wrapping_shr(8).wrapping_sub(h*5))
        & (HSIZE-1) as u64
    ) as usize
}

fn not(i: i32) -> i32 {
    if i == 0 {
        1
    } else {
        0
    }
}

/// Compress the given data, if possible.
/// The return value will be set to the error if compression fails.
///
/// The buffer is always set to the same size as the input buffer.
/// If that is not enough to hold the lzf-compressed data,
/// an error will be returned.
///
/// Example:
///
/// ```rust
/// let data = "foobar";
/// let compressed = lzf::compress(data.as_bytes());
/// ```
pub fn compress(data: &[u8]) -> LzfResult<Vec<u8>> {
    let in_len = data.len();
    let out_buf_len = in_len;
    let mut out = Vec::with_capacity(out_buf_len);
    unsafe { out.set_len(out_buf_len) };

    let mut out_len: i32 = 1; /* start run by default */

    /* This goes against all of Rust's statically verifiable guarantees,
     * but for the below use-case accessing uninitialized memory is ok,
     * as we have other checks to make sure the read memory is not used.
     *
     * The otherwise happening memset slows down the code by a factor of 20-30
     */
    let mut htab: [usize; 1 << HLOG] = unsafe { mem::uninitialized() };

    let mut current_offset = 0;

    if in_len == 0 {
        return Err(LzfError::DataCorrupted);
    }

    let mut lit: i32 = 0;

    let mut hval: u32;
    let mut ref_offset;

    hval = first(data, current_offset);

    while current_offset < in_len - 2 {
        hval = next(hval, data, current_offset);
        let hslot_idx = idx(hval);

        ref_offset = htab[hslot_idx];
        htab[hslot_idx] = current_offset;

        let off = current_offset.wrapping_sub(ref_offset).wrapping_sub(1);
        if off < MAX_OFF && current_offset + 4 < in_len && ref_offset > 0 &&
           data[ref_offset] == data[current_offset] &&
           data[ref_offset+1] == data[current_offset+1] &&
           data[ref_offset+2] == data[current_offset+2] {

            let mut len = 2;
            let maxlen = cmp::min(in_len - current_offset - len, MAX_REF);

            /* stop run */
            out[(out_len - lit - 1) as usize] = (lit as u8).wrapping_sub(1);
            out_len -= not(lit); /* undo run if length is zero */

            if out_len as i32 + 3 + 1 >= out_buf_len as i32 {
                return Err(LzfError::NoCompressionPossible);
            }

            loop {
                len += 1;
                while len < maxlen && data[ref_offset+len] == data[current_offset+len] {
                    len += 1;
                }
                break;
            }

            len -= 2; /* len is now #octets - 1 */
            current_offset += 1;

            if len < 7 {
                out[out_len as usize] = (off >> 8) as u8 + (len << 5) as u8;
                out_len += 1;
            } else {
                out[out_len as usize] = (off >> 8) as u8 + (7 << 5);
                out[out_len as usize + 1] = (len as u8).wrapping_sub(7);
                out_len += 2;
            }

            out[out_len as usize] = off as u8;
            out_len += 2; /* start run */
            lit = 0;

            /* we add here, because we later substract from the total length */
            current_offset += len - 1;

            if current_offset >= in_len {
                break;
            }

            hval = first(data, current_offset);

            hval = next(hval, data, current_offset);
            htab[idx(hval)] = current_offset;
            current_offset += 1;

            hval = next(hval, data, current_offset);
            htab[idx(hval)] = current_offset;
            current_offset += 1;
        } else {
            /* one more literal byte we must copy */
            if out_len >= out_buf_len as i32 {
                return Err(LzfError::NoCompressionPossible);
            }

            lit += 1;
            out[out_len as usize] = data[current_offset];
            out_len += 1;
            current_offset += 1;

            if lit == MAX_LIT {
                /* stop run */
                out[(out_len - lit - 1) as usize] = (lit as u8).wrapping_sub(1);
                lit = 0;
                out_len += 1; /* start run */
            }
        }
    }

    /* at most 3 bytes can be missing here */
    if out_len + 3 > out_buf_len as i32 {
        return Err(LzfError::NoCompressionPossible);
    }

    while current_offset < in_len {
        lit += 1;
        out[out_len as usize] = data[current_offset];
        out_len += 1;
        current_offset += 1;

        if lit == MAX_LIT {
            /* stop run */
            out[(out_len - lit - 1) as usize] = (lit as u8).wrapping_sub(1);
            lit = 0;
            out_len += 1; /* start run */
        }
    }

    /* end run */
    out[(out_len - lit - 1) as usize] = (lit as u8).wrapping_sub(1);
    out_len -= not(lit); /* undo run if length is zero */

    unsafe { out.set_len(out_len as usize) };

    Ok(out)
}

#[test]
fn test_compress_skips_short() {
    match compress("foo".as_bytes()) {
        Ok(_) => panic!("Compression did _something_, which is wrong for 'foo'"),
        Err(err) => assert_eq!(LzfError::NoCompressionPossible, err),
    }
}

#[test]
fn test_compress_lorem() {
    let lorem = "Lorem ipsum dolor sit amet, consetetur sadipscing elitr, sed diam nonumy eirmod \
                 tempor invidunt ut labore et dolore magna aliquyam erat, sed diam voluptua. At \
                 vero eos et accusam et justo duo dolores et ea rebum. Stet clita kasd gubergren, \
                 no sea takimata sanctus est Lorem ipsum dolor sit amet. Lorem ipsum dolor sit \
                 amet, consetetur sadipscing elitr, sed diam nonumy eirmod tempor invidunt ut \
                 labore et dolore magna aliquyam erat, sed diam voluptua.";

    match compress(lorem.as_bytes()) {
        Ok(compressed) => {
            assert_eq!(272, compressed.len())
        }
        Err(err) => panic!("Compression failed with error {:?}", err),
    }
}

#[test]
fn test_compress_decompress_lorem_round() {
    use super::decompress;

    let lorem = "Lorem ipsum dolor sit amet, consetetur sadipscing elitr, sed diam nonumy eirmod \
                 tempor invidunt ut labore et dolore magna aliquyam erat, sed diam voluptua. At \
                 vero eos et accusam et justo duo dolores et ea rebum. Stet clita kasd gubergren, \
                 no sea takimata sanctus est Lorem ipsum dolor sit amet. Lorem ipsum dolor sit \
                 amet, consetetur sadipscing elitr, sed diam nonumy eirmod tempor invidunt ut \
                 labore et dolore magna aliquyam erat, sed diam voluptua.";

    let compressed = match compress(lorem.as_bytes()) {
        Ok(c) => c,
        Err(err) => panic!("Compression failed with error {:?}", err),
    };

    match decompress(&compressed, lorem.len()) {
        Ok(decompressed) => {
            assert_eq!(lorem.len(), decompressed.len());
            assert_eq!(lorem.as_bytes(), &decompressed[..]);
        }
        Err(err) => panic!("Decompression failed with error {:?}", err),
    };
}

#[test]
fn test_alice_wonderland_both() {
    let alice = "\r\n\r\n\r\n\r\n                ALICE'S ADVENTURES IN WONDERLAND\r\n";

    let compressed = match compress(alice.as_bytes()) {
        Ok(c) => c,
        Err(err) => panic!("Compression failed with error {:?}", err),
    };

    let c_compressed = match super::compress(alice.as_bytes()) {
        Ok(c) => c,
        Err(err) => panic!("Compression failed with error {:?}", err),
    };

    assert_eq!(&compressed[..], &c_compressed[..]);
}

#[test]
fn quickcheck_found_bug() {
    let inp = vec![0, 0, 0, 0, 1, 0, 0, 2, 0, 0, 3, 0, 0, 4, 0, 1, 1, 0, 1, 2, 0, 1, 3, 0, 1, 4, 0, 0, 5, 0, 0, 6, 0, 0, 7, 0, 0, 8, 0, 0, 9, 0, 0, 10, 0, 0, 11, 0, 1, 5, 0, 1, 6, 0, 1, 7, 0, 1, 8, 0, 1, 9, 0, 1, 10, 0, 0];

    assert_eq!(LzfError::NoCompressionPossible, compress(&inp).unwrap_err());
}
<commit_msg>fix: Avoid compression if there is not enough data<commit_after>use super::{LzfResult, LzfError};
use std::{cmp, mem};

const HLOG    : usize = 16;
const HSIZE   : u32 = 1 << HLOG;
const MAX_OFF : usize = 1 << 13;
const MAX_REF : usize = ((1 << 8) + (1 << 3));
const MAX_LIT : i32 = (1 << 5);

fn first(p: &[u8], off: usize) -> u32 {
    ((p[off] as u32) << 8) | p[off+1] as u32
}

fn next(v: u32, p: &[u8], off: usize) -> u32 {
    (v << 8) | p[off+2] as u32
}

fn idx(h: u32) -> usize {
    let h = h as u64;
    (
        // 8 = 3*8-HLOG, but HLOG is constant at 16
        (h.wrapping_shr(8).wrapping_sub(h*5))
        & (HSIZE-1) as u64
    ) as usize
}

fn not(i: i32) -> i32 {
    if i == 0 {
        1
    } else {
        0
    }
}

/// Compress the given data, if possible.
/// The return value will be set to the error if compression fails.
///
/// The buffer is always set to the same size as the input buffer.
/// If that is not enough to hold the lzf-compressed data,
/// an error will be returned.
///
/// Example:
///
/// ```rust
/// let data = "foobar";
/// let compressed = lzf::compress(data.as_bytes());
/// ```
pub fn compress(data: &[u8]) -> LzfResult<Vec<u8>> {
    let in_len = data.len();
    let out_buf_len = in_len;
    let mut out = Vec::with_capacity(out_buf_len);
    unsafe { out.set_len(out_buf_len) };

    let mut out_len: i32 = 1; /* start run by default */

    /* This goes against all of Rust's statically verifiable guarantees,
     * but for the below use-case accessing uninitialized memory is ok,
     * as we have other checks to make sure the read memory is not used.
     *
     * The otherwise happening memset slows down the code by a factor of 20-30
     */
    let mut htab: [usize; 1 << HLOG] = unsafe { mem::uninitialized() };

    let mut current_offset = 0;

    if in_len < 2 {
        return Err(LzfError::NoCompressionPossible);
    }

    let mut lit: i32 = 0;

    let mut hval: u32;
    let mut ref_offset;

    hval = first(data, current_offset);

    while current_offset < in_len - 2 {
        hval = next(hval, data, current_offset);
        let hslot_idx = idx(hval);

        ref_offset = htab[hslot_idx];
        htab[hslot_idx] = current_offset;

        let off = current_offset.wrapping_sub(ref_offset).wrapping_sub(1);
        if off < MAX_OFF && current_offset + 4 < in_len && ref_offset > 0 &&
           data[ref_offset] == data[current_offset] &&
           data[ref_offset+1] == data[current_offset+1] &&
           data[ref_offset+2] == data[current_offset+2] {

            let mut len = 2;
            let maxlen = cmp::min(in_len - current_offset - len, MAX_REF);

            /* stop run */
            out[(out_len - lit - 1) as usize] = (lit as u8).wrapping_sub(1);
            out_len -= not(lit); /* undo run if length is zero */

            if out_len as i32 + 3 + 1 >= out_buf_len as i32 {
                return Err(LzfError::NoCompressionPossible);
            }

            loop {
                len += 1;
                while len < maxlen && data[ref_offset+len] == data[current_offset+len] {
                    len += 1;
                }
                break;
            }

            len -= 2; /* len is now #octets - 1 */
            current_offset += 1;

            if len < 7 {
                out[out_len as usize] = (off >> 8) as u8 + (len << 5) as u8;
                out_len += 1;
            } else {
                out[out_len as usize] = (off >> 8) as u8 + (7 << 5);
                out[out_len as usize + 1] = (len as u8).wrapping_sub(7);
                out_len += 2;
            }

            out[out_len as usize] = off as u8;
            out_len += 2; /* start run */
            lit = 0;

            /* we add here, because we later substract from the total length */
            current_offset += len - 1;

            if current_offset >= in_len {
                break;
            }

            hval = first(data, current_offset);

            hval = next(hval, data, current_offset);
            htab[idx(hval)] = current_offset;
            current_offset += 1;

            hval = next(hval, data, current_offset);
            htab[idx(hval)] = current_offset;
            current_offset += 1;
        } else {
            /* one more literal byte we must copy */
            if out_len >= out_buf_len as i32 {
                return Err(LzfError::NoCompressionPossible);
            }

            lit += 1;
            out[out_len as usize] = data[current_offset];
            out_len += 1;
            current_offset += 1;

            if lit == MAX_LIT {
                /* stop run */
                out[(out_len - lit - 1) as usize] = (lit as u8).wrapping_sub(1);
                lit = 0;
                out_len += 1; /* start run */
            }
        }
    }

    /* at most 3 bytes can be missing here */
    if out_len + 3 > out_buf_len as i32 {
        return Err(LzfError::NoCompressionPossible);
    }

    while current_offset < in_len {
        lit += 1;
        out[out_len as usize] = data[current_offset];
        out_len += 1;
        current_offset += 1;

        if lit == MAX_LIT {
            /* stop run */
            out[(out_len - lit - 1) as usize] = (lit as u8).wrapping_sub(1);
            lit = 0;
            out_len += 1; /* start run */
        }
    }

    /* end run */
    out[(out_len - lit - 1) as usize] = (lit as u8).wrapping_sub(1);
    out_len -= not(lit); /* undo run if length is zero */

    unsafe { out.set_len(out_len as usize) };

    Ok(out)
}

#[test]
fn test_compress_skips_short() {
    match compress("foo".as_bytes()) {
        Ok(_) => panic!("Compression did _something_, which is wrong for 'foo'"),
        Err(err) => assert_eq!(LzfError::NoCompressionPossible, err),
    }
}

#[test]
fn test_compress_lorem() {
    let lorem = "Lorem ipsum dolor sit amet, consetetur sadipscing elitr, sed diam nonumy eirmod \
                 tempor invidunt ut labore et dolore magna aliquyam erat, sed diam voluptua. At \
                 vero eos et accusam et justo duo dolores et ea rebum. Stet clita kasd gubergren, \
                 no sea takimata sanctus est Lorem ipsum dolor sit amet. Lorem ipsum dolor sit \
                 amet, consetetur sadipscing elitr, sed diam nonumy eirmod tempor invidunt ut \
                 labore et dolore magna aliquyam erat, sed diam voluptua.";

    match compress(lorem.as_bytes()) {
        Ok(compressed) => {
            assert_eq!(272, compressed.len())
        }
        Err(err) => panic!("Compression failed with error {:?}", err),
    }
}

#[test]
fn test_compress_decompress_lorem_round() {
    use super::decompress;

    let lorem = "Lorem ipsum dolor sit amet, consetetur sadipscing elitr, sed diam nonumy eirmod \
                 tempor invidunt ut labore et dolore magna aliquyam erat, sed diam voluptua. At \
                 vero eos et accusam et justo duo dolores et ea rebum. Stet clita kasd gubergren, \
                 no sea takimata sanctus est Lorem ipsum dolor sit amet. Lorem ipsum dolor sit \
                 amet, consetetur sadipscing elitr, sed diam nonumy eirmod tempor invidunt ut \
                 labore et dolore magna aliquyam erat, sed diam voluptua.";

    let compressed = match compress(lorem.as_bytes()) {
        Ok(c) => c,
        Err(err) => panic!("Compression failed with error {:?}", err),
    };

    match decompress(&compressed, lorem.len()) {
        Ok(decompressed) => {
            assert_eq!(lorem.len(), decompressed.len());
            assert_eq!(lorem.as_bytes(), &decompressed[..]);
        }
        Err(err) => panic!("Decompression failed with error {:?}", err),
    };
}

#[test]
fn test_alice_wonderland_both() {
    let alice = "\r\n\r\n\r\n\r\n                ALICE'S ADVENTURES IN WONDERLAND\r\n";

    let compressed = match compress(alice.as_bytes()) {
        Ok(c) => c,
        Err(err) => panic!("Compression failed with error {:?}", err),
    };

    let c_compressed = match super::compress(alice.as_bytes()) {
        Ok(c) => c,
        Err(err) => panic!("Compression failed with error {:?}", err),
    };

    assert_eq!(&compressed[..], &c_compressed[..]);
}

#[test]
fn quickcheck_found_bug() {
    let inp = vec![0, 0, 0, 0, 1, 0, 0, 2, 0, 0, 3, 0, 0, 4, 0, 1, 1, 0, 1, 2, 0, 1, 3, 0, 1, 4, 0, 0, 5, 0, 0, 6, 0, 0, 7, 0, 0, 8, 0, 0, 9, 0, 0, 10, 0, 0, 11, 0, 1, 5, 0, 1, 6, 0, 1, 7, 0, 1, 8, 0, 1, 9, 0, 1, 10, 0, 0];

    assert_eq!(LzfError::NoCompressionPossible, compress(&inp).unwrap_err());
}

#[test]
fn quickcheck_found_bug2() {
    let inp = vec![0];

    assert_eq!(LzfError::NoCompressionPossible, compress(&inp).unwrap_err());
}
<|endoftext|>


<commit_before>// Copyright 2015-2016 Intecture Developers. See the COPYRIGHT file at the
// top-level directory of this distribution and at
// https://intecture.io/COPYRIGHT.
//
// Licensed under the Mozilla Public License 2.0 <LICENSE or
// https://www.tldrlegal.com/l/mpl-2.0>. This file may not be copied,
// modified, or distributed except according to those terms.

//! Parser for Intecture data files.
//!
//! # Examples
//!
//! ```no_run
//! ```

mod condition;
// pub mod ffi;

use error::{Error, Result};
use serde_json::{self, Value};
use std::collections::BTreeMap;
use std::fs;
use std::path::Path;

/// Parser for Intecture data files.
pub struct DataParser;

impl DataParser {
    /// Open a new file and recursively parse its contents.
    ///
    /// # Examples
    ///
    /// ```no_run
    /// ```
    pub fn open<P: AsRef<Path>>(path: P) -> Result<Value> {
        let file = try!(DataFile::new(path));
        Ok(try!(file.merge(Value::Null)))
    }
}

struct DataFile {
    v: Value,
}

impl DataFile {
    fn new<P: AsRef<Path>>(path: P) -> Result<DataFile> {
        let mut fh = try!(fs::File::open(path.as_ref()));
        let data: Value = try!(serde_json::from_reader(&mut fh));

        if !data.is_object() {
            Err(Error::Generic("Value is not an object".into()))
        } else {
            Ok(DataFile {
                v: data,
            })
        }
    }

    fn dependencies(&self) -> Result<Vec<DataFile>> {
        let mut deps = Vec::new();

        if let Some(inc) = self.v.find("_include") {
            if !inc.is_array() {
                return Err(Error::Generic("Value of `_include` is not an array".into()));
            }

            // Loop in reverse order to get lowest importance first
            for i in inc.as_array().unwrap().iter().rev() {
                deps.push(try!(DataFile::new(i.as_str().unwrap())));
            }
        }

        Ok(deps)
    }

    fn merge(self, mut last_value: Value) -> Result<Value> {
        for dep in try!(self.dependencies()) {
            last_value = try!(dep.merge(last_value));
        }

        Ok(try!(Self::merge_values(self.v, last_value)))
    }

    fn merge_values(into: Value, mut from: Value) -> Result<Value> {
        match into {
            Value::Null |
            Value::Bool(_) |
            Value::I64(_) |
            Value::U64(_) |
            Value::F64(_) |
            Value::String(_) => Ok(into),
            Value::Array(mut a) => {
                if from.is_array() {
                    a.append(from.as_array_mut().unwrap());
                } else {
                    a.push(from);
                }

                Ok(Value::Array(a))
            },
            Value::Object(o) => {
                let mut new: BTreeMap<String, Value> = BTreeMap::new();

                for (mut key, mut value) in o {
                    if key.ends_with("?") || key.ends_with("?!") {
                        if key.pop().unwrap() == '!' {
                            key.pop();
                            key.push('!');
                        }

                        value = try!(Self::query_value(&from, value)).unwrap_or(Value::Null);
                    }

                    if key.ends_with("!") {
                        key.pop();
                    }
                    else if let Some(o1) = from.find(&key) {
                        value = try!(Self::merge_values(value, o1.clone()));
                    }

                    new.insert(key, value);
                }

                // Insert any missing values
                if let Some(o1) = from.as_object() {
                    for (key, value) in o1 {
                        if !new.contains_key(key) {
                            new.insert(key.clone(), value.clone());
                        }
                    }
                }

                Ok(Value::Object(new))
            }
        }
    }

    fn query_value(data: &Value, value: Value) -> Result<Option<Value>> {
        match value {
            Value::Array(a) => {
                for opt in a {
                    if let Some(v) = try!(Self::query_value(data, opt)) {
                        return Ok(Some(v));
                    }
                }
            },
            Value::Object(mut o) => {
                if let Some(v) = o.remove("_") {
                    if let Some(q) = o.get("?") {
                        match *q {
                            Value::String(ref s) if !try!(condition::eval(data, s)) => return Ok(None),
                            _ => return Err(Error::Generic("Query must be string".into())),
                        };
                    }

                    return Ok(Some(v));
                }
            },
            _ => return Ok(Some(value)),
        }

        Ok(None)
    }
}

#[cfg(test)]
mod tests {
    use serde_json::{self, Value};
    use std::fs::File;
    use std::io::Write;
    use std::path::PathBuf;
    use super::*;
    use tempdir::TempDir;

    #[test]
    fn test_parser() {
        let tempdir = TempDir::new("parser_test").unwrap();
        let mut path = tempdir.path().to_owned();
        let expected_value = create_data(&mut path);

        path.push("top.json");
        let value = DataParser::open(&path).unwrap();

        assert_eq!(value, expected_value);
    }

    fn create_data(path: &mut PathBuf) -> Value {
        path.push("top.json");
        let mut fh = File::create(&path).unwrap();
        path.pop();
        fh.write_all(format!("{{
            \"a\": 1,
            \"payload\": {{
                \"b\": [ 1, 2 ],
                \"c?\": [
                    {{
                        \"_\": [ 6, 7 ],
                        \"?\": \"/variable = false\"
                    }},
                    {{
                        \"_\": [ 8, 9 ]
                    }}
                ],
                \"d\": [ 123 ]
            }},
            \"variable\": {{
                \"one!\": true,
                \"two\": false
            }},
            \"_include\": [
                \"{0}/middle.json\",
                \"{0}/bottom.json\"
            ]
        }}", &path.to_str().unwrap()).as_bytes()).unwrap();

        path.push("middle.json");
        let mut fh = File::create(&path).unwrap();
        path.pop();
        fh.write_all(b"{
            \"payload\": {
                \"b!\": [ 3, 4 ],
                \"d\": [ 987 ]
            },
            \"variable\": false,
            \"d\": 4
        }").unwrap();

        path.push("bottom.json");
        let mut fh = File::create(&path).unwrap();
        path.pop();
        fh.write_all(b"{
            \"moo\": \"cow\",
            \"payload\": {
                \"b\": [ 5 ],
                \"d\": [ 999 ]
            }
        }").unwrap();

        serde_json::from_str(&format!("{{
            \"_include\": [
                \"{0}/middle.json\",
                \"{0}/bottom.json\"
            ],
            \"a\": 1,
            \"d\": 4,
            \"moo\": \"cow\",
            \"payload\": {{
                \"b\": [ 1, 2, 3, 4 ],
                \"c\": [ 6, 7 ],
                \"d\": [ 123, 987, 999 ]
            }},
            \"variable\": {{
                \"one\": true,
                \"two\": false
            }}
        }}", &path.to_str().unwrap())).unwrap()
    }
}
<commit_msg>Fixed bugs with sharing data with query parser<commit_after>// Copyright 2015-2016 Intecture Developers. See the COPYRIGHT file at the
// top-level directory of this distribution and at
// https://intecture.io/COPYRIGHT.
//
// Licensed under the Mozilla Public License 2.0 <LICENSE or
// https://www.tldrlegal.com/l/mpl-2.0>. This file may not be copied,
// modified, or distributed except according to those terms.

//! Parser for Intecture data files.
//!
//! # Examples
//!
//! ```no_run
//! ```

mod condition;
// pub mod ffi;

use error::{Error, Result};
use serde_json::{self, Value};
use std::collections::BTreeMap;
use std::fs;
use std::path::Path;

/// Parser for Intecture data files.
pub struct DataParser;

impl DataParser {
    /// Open a new file and recursively parse its contents.
    ///
    /// # Examples
    ///
    /// ```no_run
    /// ```
    pub fn open<P: AsRef<Path>>(path: P) -> Result<Value> {
        let file = try!(DataFile::new(path));
        Ok(try!(file.merge(Value::Null)))
    }
}

struct DataFile {
    v: Value,
}

impl DataFile {
    fn new<P: AsRef<Path>>(path: P) -> Result<DataFile> {
        let mut fh = try!(fs::File::open(path.as_ref()));
        let data: Value = try!(serde_json::from_reader(&mut fh));

        if !data.is_object() {
            Err(Error::Generic("Value is not an object".into()))
        } else {
            Ok(DataFile {
                v: data,
            })
        }
    }

    fn dependencies(&self) -> Result<Vec<DataFile>> {
        let mut deps = Vec::new();

        if let Some(inc) = self.v.find("_include") {
            if !inc.is_array() {
                return Err(Error::Generic("Value of `_include` is not an array".into()));
            }

            // Loop in reverse order to get lowest importance first
            for i in inc.as_array().unwrap().iter().rev() {
                deps.push(try!(DataFile::new(i.as_str().unwrap())));
            }
        }

        Ok(deps)
    }

    fn merge(self, mut last_value: Value) -> Result<Value> {
        for dep in try!(self.dependencies()) {
            last_value = try!(dep.merge(last_value));
        }

        let lv_clone = last_value.clone();
        Ok(try!(Self::merge_values(self.v, last_value, &lv_clone)))
    }

    fn merge_values(into: Value, mut from: Value, parent_from: &Value) -> Result<Value> {
        match into {
            Value::Null |
            Value::Bool(_) |
            Value::I64(_) |
            Value::U64(_) |
            Value::F64(_) |
            Value::String(_) => Ok(into),
            Value::Array(mut a) => {
                if from.is_array() {
                    a.append(from.as_array_mut().unwrap());
                } else {
                    a.push(from);
                }

                Ok(Value::Array(a))
            },
            Value::Object(o) => {
                let mut new: BTreeMap<String, Value> = BTreeMap::new();

                for (mut key, mut value) in o {
                    if key.ends_with("?") || key.ends_with("?!") {
                        if key.pop().unwrap() == '!' {
                            key.pop();
                            key.push('!');
                        }

                        value = try!(Self::query_value(&parent_from, value)).unwrap_or(Value::Null);
                    }

                    if key.ends_with("!") {
                        key.pop();
                    }
                    else if let Some(o1) = from.find(&key) {
                        value = try!(Self::merge_values(value, o1.clone(), &parent_from));
                    }

                    new.insert(key, value);
                }

                // Insert any missing values
                if let Some(o1) = from.as_object() {
                    for (key, value) in o1 {
                        if !new.contains_key(key) {
                            new.insert(key.clone(), value.clone());
                        }
                    }
                }

                Ok(Value::Object(new))
            }
        }
    }

    fn query_value(data: &Value, value: Value) -> Result<Option<Value>> {
        match value {
            Value::Array(a) => {
                for opt in a {
                    if let Some(v) = try!(Self::query_value(data, opt)) {
                        return Ok(Some(v));
                    }
                }
            },
            Value::Object(mut o) => {
                if let Some(v) = o.remove("_") {
                    if let Some(q) = o.get("?") {
                        match *q {
                            Value::String(ref s) => {
                                if try!(condition::eval(data, s)) {
                                    return Ok(Some(v));
                                }
                            },
                            _ => return Err(Error::Generic("Query must be string".into())),
                        };
                    }

                    return Ok(Some(v));
                }
            },
            _ => return Ok(Some(value)),
        }

        Ok(None)
    }
}

#[cfg(test)]
mod tests {
    use serde_json::{self, Value};
    use std::fs::File;
    use std::io::Write;
    use std::path::PathBuf;
    use super::*;
    use tempdir::TempDir;

    #[test]
    fn test_parser() {
        let tempdir = TempDir::new("parser_test").unwrap();
        let mut path = tempdir.path().to_owned();
        let expected_value = create_data(&mut path);

        path.push("top.json");
        let value = DataParser::open(&path).unwrap();

        assert_eq!(value, expected_value);
    }

    fn create_data(path: &mut PathBuf) -> Value {
        path.push("top.json");
        let mut fh = File::create(&path).unwrap();
        path.pop();
        fh.write_all(format!("{{
            \"a\": 1,
            \"payload\": {{
                \"b\": [ 1, 2 ],
                \"c?\": [
                    {{
                        \"_\": [ 6, 7 ],
                        \"?\": \"/variable = false\"
                    }},
                    {{
                        \"_\": [ 8, 9 ]
                    }}
                ],
                \"d\": [ 123 ]
            }},
            \"variable\": {{
                \"one!\": true,
                \"two\": false
            }},
            \"_include\": [
                \"{0}/middle.json\",
                \"{0}/bottom.json\"
            ]
        }}", &path.to_str().unwrap()).as_bytes()).unwrap();

        path.push("middle.json");
        let mut fh = File::create(&path).unwrap();
        path.pop();
        fh.write_all(b"{
            \"payload\": {
                \"b!\": [ 3, 4 ],
                \"d\": [ 987 ]
            },
            \"variable\": false,
            \"d\": 4
        }").unwrap();

        path.push("bottom.json");
        let mut fh = File::create(&path).unwrap();
        path.pop();
        fh.write_all(b"{
            \"moo\": \"cow\",
            \"payload\": {
                \"b\": [ 5 ],
                \"d\": [ 999 ]
            }
        }").unwrap();

        serde_json::from_str(&format!("{{
            \"_include\": [
                \"{0}/middle.json\",
                \"{0}/bottom.json\"
            ],
            \"a\": 1,
            \"d\": 4,
            \"moo\": \"cow\",
            \"payload\": {{
                \"b\": [ 1, 2, 3, 4 ],
                \"c\": [ 6, 7 ],
                \"d\": [ 123, 987, 999 ]
            }},
            \"variable\": {{
                \"one\": true,
                \"two\": false
            }}
        }}", &path.to_str().unwrap())).unwrap()
    }
}
<|endoftext|>


<commit_before><commit_msg>Add download.rs<commit_after>use std::ffi::{OsStr, OsString};
use std::io::Write;
use std::path::{Path, PathBuf};

use atomicwrites::{AtomicFile, OverwriteBehavior};
use rayon::iter::{IntoParallelIterator, IntoParallelRefIterator, ParallelIterator};
use s3::bucket::Bucket;
use s3::creds::Credentials;
//use s3::S3Error;

fn s3_download(bucket: &Bucket, remote_path: &str, local_path: &Path) -> Result<(), anyhow::Error> {
    if !local_path.exists() {
        let mut i = 0;
        let contents;
        loop {
            match bucket.get_object_blocking(&remote_path) {
                Ok((c, code)) => {
                    contents = c;
                    if code != 200 {
                        println!("{}", remote_path);
                        return Ok(());
                    }
                    break;
                }
                Err(e) => {
                    if i < 20 {
                        i += 1;
                        continue;
                    }
                    return Err(e.into());
                }
            }
        }

        AtomicFile::new(local_path, OverwriteBehavior::AllowOverwrite)
            .write(|f| f.write_all(&contents))?;
    }
    Ok(())
}

fn make_vrt(directory: &Path, extension: &OsStr) -> Result<(), anyhow::Error> {
    let files: Vec<OsString> = std::fs::read_dir(directory)?
        .filter_map(Result::ok)
        .filter(|f| f.path().extension() == Some(extension))
        .map(|f| f.file_name())
        .collect();

    let mut args = vec![OsString::from("merged.vrt")];
    args.extend(files);

    let output = std::process::Command::new("gdalbuildvrt")
        .current_dir(directory)
        .args(args)
        .output()
        .expect("Failed to run gdalbuildvrt. Is gdal installed?");

    println!("stdout: {}", String::from_utf8_lossy(&output.stdout));
    println!("stderr: {}", String::from_utf8_lossy(&output.stderr));

    Ok(())
}

// pub fn download_nasadem(path: &Path) -> Result<(), anyhow::Error> {
//     let directory = path.join("nasadem");
//     std::fs::create_dir_all(&directory)?;

//     // TODO: Actually do the download step

//     let files = std::fs::read_dir(directory)?.filter_map(Result::ok).filter(|f|f.path().extension() == ).collect();

//     Ok(())
// }

// Download global watermask from Copernicus dataset.
//
//  | Pixel | Value Meaning |
//  |-------|---------------|
//  |   0   | No water      |
//  |   1   | Ocean         |
//  |   2   | Lake          |
//  |   3   | River         |
//  |-------|---------------|
pub fn download_copernicus_wbm(path: &Path) -> Result<(), anyhow::Error> {
    let directory = path.join("copernicus-wbm");

    std::fs::create_dir_all(&directory)?;

    let bucket =
        Bucket::new("copernicus-dem-30m", "eu-central-1".parse()?, Credentials::anonymous()?)?;
    let bucket_fallback =
        Bucket::new("copernicus-dem-90m", "eu-central-1".parse()?, Credentials::anonymous()?)?;

    let (tile_list, code) = bucket.get_object_blocking("tileList.txt")?;
    assert_eq!(200, code);
    let (missing, code) = bucket.get_object_blocking("blacklist.txt")?;
    assert_eq!(200, code);

    let tile_list = String::from_utf8(tile_list)?;
    let tile_list: Vec<_> = tile_list.split_ascii_whitespace().collect();
    tile_list.par_iter().try_for_each(|name| -> Result<(), anyhow::Error> {
        let filename = format!("{}WBM.tif", &name[..name.len() - 3]);
        let local_path = directory.join(&filename);
        let remote_path = format!("{}/AUXFILES/{}", name, filename);
        s3_download(&bucket, &remote_path, &local_path)
    })?;

    let missing = String::from_utf8(missing)?;
    let missing: Vec<_> = missing.split_ascii_whitespace().collect();
    missing.par_iter().try_for_each(|name| -> Result<(), anyhow::Error> {
        let name = name.replace("DSM_10", "DSM_COG_30").replace(".tif", "");
        let filename = format!("{}WBM.tif", &name[..name.len() - 3]);
        let local_path = directory.join(&filename);
        let remote_path = format!("{}/AUXFILES/{}", name, filename);
        s3_download(&bucket_fallback, &remote_path, &local_path)
    })?;

    make_vrt(&directory, OsStr::new("tif"))?;

    Ok(())
}

// Download heights from Copernicus dataset.
//
// See https://registry.opendata.aws/copernicus-dem/
pub fn download_copernicus_hgt(path: &Path) -> Result<(), anyhow::Error> {
    let directory = path.join("copernicus-hgt");

    std::fs::create_dir_all(&directory)?;

    let bucket =
        Bucket::new("copernicus-dem-30m", "eu-central-1".parse()?, Credentials::anonymous()?)?;
    let bucket_fallback =
        Bucket::new("copernicus-dem-90m", "eu-central-1".parse()?, Credentials::anonymous()?)?;

    let (tile_list, code) = bucket.get_object_blocking("tileList.txt")?;
    assert_eq!(200, code);
    let (missing, code) = bucket.get_object_blocking("blacklist.txt")?;
    assert_eq!(200, code);

    let tile_list = String::from_utf8(tile_list)?;
    let tile_list: Vec<_> = tile_list.split_ascii_whitespace().collect();
    tile_list.into_par_iter().try_for_each(|name| -> Result<(), anyhow::Error> {
        let filename = format!("{}DEM.tif", &name[..name.len() - 3]);
        let local_path = directory.join(&filename);
        let remote_path = format!("{}/{}", name, filename);
        s3_download(&bucket, &remote_path, &local_path)
    })?;

    let missing = String::from_utf8(missing)?;
    let missing: Vec<_> = missing.split_ascii_whitespace().collect();
    missing.into_par_iter().try_for_each(|name| -> Result<(), anyhow::Error> {
        let name = name.replace("DSM_10", "DSM_COG_30").replace(".tif", "");
        let filename = format!("{}DEM.tif", &name[..name.len() - 3]);
        let local_path = directory.join(&filename);
        let remote_path = format!("{}/{}", name, filename);
        s3_download(&bucket_fallback, &remote_path, &local_path)
    })?;

    make_vrt(&directory, OsStr::new("tif"))?;

    Ok(())
}
<|endoftext|>


<commit_before>// Copyright 2014 Johannes Köster, Christopher Schröder.
// Licensed under the MIT license (http://opensource.org/licenses/MIT)
// This file may not be copied, modified, or distributed
// except according to those terms.

//! FASTA format reading and writing.
//!
//! # Example
//!
//! ```
//! use std::io;
//! use bio::io::fasta;
//! let reader = fasta::Reader::new(io::stdin());
//! ```

use std::io;
use std::io::prelude::*;
use std::ascii::AsciiExt;
use std::collections;
use std::fs;
use std::path::Path;
use std::convert::AsRef;
use std::cmp::min;

use csv;

use utils::{TextSlice, Text};

/// A FASTA reader.
pub struct Reader<R: io::Read> {
    reader: io::BufReader<R>,
    line: String,
}

impl Reader<fs::File> {
    /// Read FASTA from given file path.
    pub fn from_file<P: AsRef<Path>>(path: P) -> io::Result<Self> {
        fs::File::open(path).map(Reader::new)
    }
}

impl<R: io::Read> Reader<R> {
    /// Create a new Fasta reader given an instance of `io::Read`.
    pub fn new(reader: R) -> Self {
        Reader {
            reader: io::BufReader::new(reader),
            line: String::new(),
        }
    }

    /// Read next FASTA record into the given `Record`.
    pub fn read(&mut self, record: &mut Record) -> io::Result<()> {
        record.clear();
        if self.line.is_empty() {
            try!(self.reader.read_line(&mut self.line));
            if self.line.is_empty() {
                return Ok(());
            }
        }

        if !self.line.starts_with('>') {
            return Err(io::Error::new(io::ErrorKind::Other, "Expected > at record start."));
        }
        record.header.push_str(&self.line);
        loop {
            self.line.clear();
            try!(self.reader.read_line(&mut self.line));
            if self.line.is_empty() || self.line.starts_with('>') {
                break;
            }
            record.seq.push_str(&self.line.trim_right());
        }

        Ok(())
    }

    /// Return an iterator over the records of this FastQ file.
    pub fn records(self) -> Records<R> {
        Records { reader: self }
    }
}

/// A FASTA index as created by SAMtools (.fai).
pub struct Index {
    inner: collections::HashMap<String, IndexRecord>,
    seqs: Vec<String>,
}

impl Index {
    /// Open a FASTA index from a given `io::Read` instance.
    pub fn new<R: io::Read>(fai: R) -> csv::Result<Self> {
        let mut inner = collections::HashMap::new();
        let mut seqs = vec![];
        let mut fai_reader = csv::Reader::from_reader(fai).delimiter(b'\t').has_headers(false);
        for row in fai_reader.decode() {
            let (name, record): (String, IndexRecord) = try!(row);
            seqs.push(name.clone());
            inner.insert(name, record);
        }
        Ok(Index {
            inner: inner,
            seqs: seqs,
        })
    }

    /// Open a FASTA index from a given file path.
    pub fn from_file<P: AsRef<Path>>(path: &P) -> csv::Result<Self> {
        match fs::File::open(path) {
            Ok(fai) => Self::new(fai),
            Err(e) => Err(csv::Error::Io(e)),
        }
    }

    /// Open a FASTA index given the corresponding FASTA file path (e.g. for ref.fasta we expect ref.fasta.fai).
    pub fn with_fasta_file<P: AsRef<Path>>(fasta_path: &P) -> csv::Result<Self> {
        let mut ext = fasta_path.as_ref().extension().unwrap().to_str().unwrap().to_owned();
        ext.push_str(".fai");
        let fai_path = fasta_path.as_ref().with_extension(ext);

        Self::from_file(&fai_path)
    }

    /// Return a vector of sequences described in the index.
    pub fn sequences(&self) -> Vec<Sequence> {
        self.seqs
            .iter()
            .map(|name| {
                Sequence {
                    name: name.clone(),
                    len: self.inner.get(name).unwrap().len,
                }
            })
            .collect()
    }
}

/// A FASTA reader with an index as created by SAMtools (.fai).
pub struct IndexedReader<R: io::Read + io::Seek> {
    reader: io::BufReader<R>,
    pub index: Index,
}

impl IndexedReader<fs::File> {
    /// Read from a given file path. This assumes the index ref.fasta.fai to be present for FASTA ref.fasta.
    pub fn from_file<P: AsRef<Path>>(path: &P) -> csv::Result<Self> {
        let index = try!(Index::with_fasta_file(path));

        match fs::File::open(path) {
            Ok(fasta) => Ok(IndexedReader::with_index(fasta, index)),
            Err(e) => Err(csv::Error::Io(e)),
        }
    }
}

impl<R: io::Read + io::Seek> IndexedReader<R> {
    /// Read from a FASTA and its index, both given as `io::Read`. FASTA has to be `io::Seek` in addition.
    pub fn new<I: io::Read>(fasta: R, fai: I) -> csv::Result<Self> {
        let index = try!(Index::new(fai));
        Ok(IndexedReader {
            reader: io::BufReader::new(fasta),
            index: index,
        })
    }

    /// Read from a FASTA and its index, the first given as `io::Read`, the second given as index object.
    pub fn with_index(fasta: R, index: Index) -> Self {
        IndexedReader {
            reader: io::BufReader::new(fasta),
            index: index,
        }
    }

    /// For a given seqname, read the whole sequence into the given vector.
    pub fn read_all(&mut self, seqname: &str, seq: &mut Text) -> io::Result<()> {
        match self.index.inner.get(seqname) {
            Some(&idx) => self.read(seqname, 0, idx.len, seq),
            None => Err(io::Error::new(io::ErrorKind::Other, "Unknown sequence name.")),
        }
    }

    /// Read the given interval of the given seqname into the given vector (stop position is exclusive).
    pub fn read(&mut self,
                seqname: &str,
                start: u64,
                stop: u64,
                seq: &mut Text)
                -> io::Result<()> {
        match self.index.inner.get(seqname) {
            Some(idx) => {
                seq.clear();

                let length = stop - start as u64;
                let mut buf = vec![0u8; idx.line_bases as usize];

                loop {
                    let current_start = start + seq.len() as u64;
                    let line_start = current_start / idx.line_bases * idx.line_bytes;
                    let line_offset = current_start % idx.line_bases;
                    let offset = idx.offset + line_start + line_offset;
                    try!(self.reader.seek(io::SeekFrom::Start(offset)));

                    let left_to_read = length - seq.len() as u64;
                    let left_in_line = min(left_to_read, idx.line_bases - line_offset) as usize;

                    try!(self.reader.read(&mut buf[..left_in_line]));
                    seq.extend_from_slice(&buf[..left_in_line]);

                    if seq.len() as u64 == length
                    {
                        break;
                    }
                }

                Ok(())
            }
            None => Err(io::Error::new(io::ErrorKind::Other, "Unknown sequence name.")),
        }
    }
}

/// Record of a FASTA index.
#[derive(RustcDecodable, Debug, Copy, Clone)]
struct IndexRecord {
    len: u64,
    offset: u64,
    line_bases: u64,
    line_bytes: u64,
}

/// A sequence record returned by the FASTA index.
pub struct Sequence {
    pub name: String,
    pub len: u64,
}

/// A Fasta writer.
pub struct Writer<W: io::Write> {
    writer: io::BufWriter<W>,
}

impl Writer<fs::File> {
    /// Write to the given file path.
    pub fn to_file<P: AsRef<Path>>(path: P) -> io::Result<Self> {
        fs::File::create(path).map(Writer::new)
    }
}

impl<W: io::Write> Writer<W> {
    /// Create a new Fasta writer.
    pub fn new(writer: W) -> Self {
        Writer { writer: io::BufWriter::new(writer) }
    }

    /// Directly write a Fasta record.
    pub fn write_record(&mut self, record: Record) -> io::Result<()> {
        self.write(record.id().unwrap_or(""), record.desc(), record.seq())
    }

    /// Write a Fasta record with given id, optional description and sequence.
    pub fn write(&mut self, id: &str, desc: Option<&str>, seq: TextSlice) -> io::Result<()> {
        try!(self.writer.write(b">"));
        try!(self.writer.write(id.as_bytes()));
        if desc.is_some() {
            try!(self.writer.write(b" "));
            try!(self.writer.write(desc.unwrap().as_bytes()));
        }
        try!(self.writer.write(b"\n"));
        try!(self.writer.write(seq));
        try!(self.writer.write(b"\n"));

        Ok(())
    }

    /// Flush the writer, ensuring that everything is written.
    pub fn flush(&mut self) -> io::Result<()> {
        self.writer.flush()
    }
}

/// A FASTA record.
pub struct Record {
    header: String,
    seq: String,
}

impl Record {
    /// Create a new instance.
    pub fn new() -> Self {
        Record {
            header: String::new(),
            seq: String::new(),
        }
    }

    /// Check if record is empty.
    pub fn is_empty(&self) -> bool {
        self.header.is_empty() && self.seq.is_empty()
    }

    /// Check validity of Fasta record.
    pub fn check(&self) -> Result<(), &str> {
        if self.id().is_none() {
            return Err("Expecting id for FastQ record.");
        }
        if !self.seq.is_ascii() {
            return Err("Non-ascii character found in sequence.");
        }

        Ok(())
    }

    /// Return the id of the record.
    pub fn id(&self) -> Option<&str> {
        self.header[1..].trim_right().splitn(2, ' ').next()
    }

    /// Return descriptions if present.
    pub fn desc(&self) -> Option<&str> {
        self.header[1..].trim_right().splitn(2, ' ').skip(1).next()
    }

    /// Return the sequence of the record.
    pub fn seq(&self) -> TextSlice {
        self.seq.as_bytes()
    }

    /// Clear the record.
    fn clear(&mut self) {
        self.header.clear();
        self.seq.clear();
    }
}

/// An iterator over the records of a Fasta file.
pub struct Records<R: io::Read> {
    reader: Reader<R>,
}

impl<R: io::Read> Iterator for Records<R> {
    type Item = io::Result<Record>;

    fn next(&mut self) -> Option<io::Result<Record>> {
        let mut record = Record::new();
        match self.reader.read(&mut record) {
            Ok(()) if record.is_empty() => None,
            Ok(()) => Some(Ok(record)),
            Err(err) => Some(Err(err)),
        }
    }
}

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

    const FASTA_FILE: &'static [u8] = b">id desc
ACCGTAGGCTGA
CCGTAGGCTGAA
CGTAGGCTGAAA
GTAGGCTGAAAA
CCCC
>id2
ATTGTTGTTTTA
ATTGTTGTTTTA
ATTGTTGTTTTA
GGGG
";
    const FAI_FILE: &'static [u8] = b"id\t52\t9\t12\t13
id2\t40\t71\t12\t13
";
    const WRITE_FASTA_FILE: &'static [u8] = b">id desc
ACCGTAGGCTGA
>id2
ATTGTTGTTTTA
";

    #[test]
    fn test_reader() {
        let reader = Reader::new(FASTA_FILE);
        let ids = [Some("id"), Some("id2")];
        let descs = [Some("desc"), None];
        let seqs: [&[u8]; 2] = [b"ACCGTAGGCTGACCGTAGGCTGAACGTAGGCTGAAAGTAGGCTGAAAACCCC",
                                b"ATTGTTGTTTTAATTGTTGTTTTAATTGTTGTTTTAGGGG"];

        for (i, r) in reader.records().enumerate() {
            let record = r.ok().expect("Error reading record");
            assert_eq!(record.check(), Ok(()));
            assert_eq!(record.id(), ids[i]);
            assert_eq!(record.desc(), descs[i]);
            assert_eq!(record.seq(), seqs[i]);
        }

        // let record = records.ok().nth(1).unwrap();
    }

    #[test]
    fn test_indexed_reader() {
        let mut reader = IndexedReader::new(io::Cursor::new(FASTA_FILE), FAI_FILE)
                             .ok()
                             .expect("Error reading index");
        let mut seq = Vec::new();

        // Test reading various substrings of the sequence
        reader.read("id", 1, 5, &mut seq).ok().expect("Error reading sequence.");
        assert_eq!(seq, b"CCGT");

        reader.read("id", 1, 31, &mut seq).ok().expect("Error reading sequence.");
        assert_eq!(seq, b"CCGTAGGCTGACCGTAGGCTGAACGTAGGC");

        reader.read("id", 13, 23, &mut seq).ok().expect("Error reading sequence.");
        assert_eq!(seq, b"CGTAGGCTGA");

        reader.read("id", 36, 52, &mut seq).ok().expect("Error reading sequence.");
        assert_eq!(seq, b"GTAGGCTGAAAACCCC");

        reader.read("id2", 12, 40, &mut seq).ok().expect("Error reading sequence.");
        assert_eq!(seq, b"ATTGTTGTTTTAATTGTTGTTTTAGGGG");
    }

    #[test]
    fn test_writer() {
        let mut writer = Writer::new(Vec::new());
        writer.write("id", Some("desc"), b"ACCGTAGGCTGA").ok().expect("Expected successful write");
        writer.write("id2", None, b"ATTGTTGTTTTA").ok().expect("Expected successful write");
        writer.flush().ok().expect("Expected successful write");
        assert_eq!(writer.writer.get_ref(), &WRITE_FASTA_FILE);
    }
}
<commit_msg>read can read less than the requested number of bytes. use read_exact instead. truncate the stop position to the size of the fasta entry<commit_after>// Copyright 2014 Johannes Köster, Christopher Schröder.
// Licensed under the MIT license (http://opensource.org/licenses/MIT)
// This file may not be copied, modified, or distributed
// except according to those terms.

//! FASTA format reading and writing.
//!
//! # Example
//!
//! ```
//! use std::io;
//! use bio::io::fasta;
//! let reader = fasta::Reader::new(io::stdin());
//! ```

use std::io;
use std::io::prelude::*;
use std::ascii::AsciiExt;
use std::collections;
use std::fs;
use std::path::Path;
use std::convert::AsRef;
use std::cmp::min;

use csv;

use utils::{TextSlice, Text};

/// A FASTA reader.
pub struct Reader<R: io::Read> {
    reader: io::BufReader<R>,
    line: String,
}

impl Reader<fs::File> {
    /// Read FASTA from given file path.
    pub fn from_file<P: AsRef<Path>>(path: P) -> io::Result<Self> {
        fs::File::open(path).map(Reader::new)
    }
}

impl<R: io::Read> Reader<R> {
    /// Create a new Fasta reader given an instance of `io::Read`.
    pub fn new(reader: R) -> Self {
        Reader {
            reader: io::BufReader::new(reader),
            line: String::new(),
        }
    }

    /// Read next FASTA record into the given `Record`.
    pub fn read(&mut self, record: &mut Record) -> io::Result<()> {
        record.clear();
        if self.line.is_empty() {
            try!(self.reader.read_line(&mut self.line));
            if self.line.is_empty() {
                return Ok(());
            }
        }

        if !self.line.starts_with('>') {
            return Err(io::Error::new(io::ErrorKind::Other, "Expected > at record start."));
        }
        record.header.push_str(&self.line);
        loop {
            self.line.clear();
            try!(self.reader.read_line(&mut self.line));
            if self.line.is_empty() || self.line.starts_with('>') {
                break;
            }
            record.seq.push_str(&self.line.trim_right());
        }

        Ok(())
    }

    /// Return an iterator over the records of this FastQ file.
    pub fn records(self) -> Records<R> {
        Records { reader: self }
    }
}

/// A FASTA index as created by SAMtools (.fai).
pub struct Index {
    inner: collections::HashMap<String, IndexRecord>,
    seqs: Vec<String>,
}

impl Index {
    /// Open a FASTA index from a given `io::Read` instance.
    pub fn new<R: io::Read>(fai: R) -> csv::Result<Self> {
        let mut inner = collections::HashMap::new();
        let mut seqs = vec![];
        let mut fai_reader = csv::Reader::from_reader(fai).delimiter(b'\t').has_headers(false);
        for row in fai_reader.decode() {
            let (name, record): (String, IndexRecord) = try!(row);
            seqs.push(name.clone());
            inner.insert(name, record);
        }
        Ok(Index {
            inner: inner,
            seqs: seqs,
        })
    }

    /// Open a FASTA index from a given file path.
    pub fn from_file<P: AsRef<Path>>(path: &P) -> csv::Result<Self> {
        match fs::File::open(path) {
            Ok(fai) => Self::new(fai),
            Err(e) => Err(csv::Error::Io(e)),
        }
    }

    /// Open a FASTA index given the corresponding FASTA file path (e.g. for ref.fasta we expect ref.fasta.fai).
    pub fn with_fasta_file<P: AsRef<Path>>(fasta_path: &P) -> csv::Result<Self> {
        let mut ext = fasta_path.as_ref().extension().unwrap().to_str().unwrap().to_owned();
        ext.push_str(".fai");
        let fai_path = fasta_path.as_ref().with_extension(ext);

        Self::from_file(&fai_path)
    }

    /// Return a vector of sequences described in the index.
    pub fn sequences(&self) -> Vec<Sequence> {
        self.seqs
            .iter()
            .map(|name| {
                Sequence {
                    name: name.clone(),
                    len: self.inner.get(name).unwrap().len,
                }
            })
            .collect()
    }
}

/// A FASTA reader with an index as created by SAMtools (.fai).
pub struct IndexedReader<R: io::Read + io::Seek> {
    reader: io::BufReader<R>,
    pub index: Index,
}

impl IndexedReader<fs::File> {
    /// Read from a given file path. This assumes the index ref.fasta.fai to be present for FASTA ref.fasta.
    pub fn from_file<P: AsRef<Path>>(path: &P) -> csv::Result<Self> {
        let index = try!(Index::with_fasta_file(path));

        match fs::File::open(path) {
            Ok(fasta) => Ok(IndexedReader::with_index(fasta, index)),
            Err(e) => Err(csv::Error::Io(e)),
        }
    }
}

impl<R: io::Read + io::Seek> IndexedReader<R> {
    /// Read from a FASTA and its index, both given as `io::Read`. FASTA has to be `io::Seek` in addition.
    pub fn new<I: io::Read>(fasta: R, fai: I) -> csv::Result<Self> {
        let index = try!(Index::new(fai));
        Ok(IndexedReader {
            reader: io::BufReader::new(fasta),
            index: index,
        })
    }

    /// Read from a FASTA and its index, the first given as `io::Read`, the second given as index object.
    pub fn with_index(fasta: R, index: Index) -> Self {
        IndexedReader {
            reader: io::BufReader::new(fasta),
            index: index,
        }
    }

    /// For a given seqname, read the whole sequence into the given vector.
    pub fn read_all(&mut self, seqname: &str, seq: &mut Text) -> io::Result<()> {
        match self.index.inner.get(seqname) {
            Some(&idx) => self.read(seqname, 0, idx.len, seq),
            None => Err(io::Error::new(io::ErrorKind::Other, "Unknown sequence name.")),
        }
    }

    /// Read the given interval of the given seqname into the given vector (stop position is exclusive).
    pub fn read(&mut self,
                seqname: &str,
                start: u64,
                stop: u64,
                seq: &mut Text)
                -> io::Result<()> {
        match self.index.inner.get(seqname) {
            Some(idx) => {
                seq.clear();

                let stop = min(stop, idx.len);
                let length = stop - start as u64;
                let mut buf = vec![0u8; idx.line_bases as usize];

                loop {
                    let current_start = start + seq.len() as u64;
                    let line_start = current_start / idx.line_bases * idx.line_bytes;
                    let line_offset = current_start % idx.line_bases;
                    let offset = idx.offset + line_start + line_offset;

                    let left_to_read = length - seq.len() as u64;
                    let left_in_line = min(left_to_read, idx.line_bases - line_offset) as usize;

                    try!(self.reader.seek(io::SeekFrom::Start(offset)));
                    try!(self.reader.read_exact(&mut buf[..left_in_line]));

                    seq.extend_from_slice(&buf[..left_in_line]);

                    if seq.len() as u64 == length
                    {
                        break;
                    }
                }

                Ok(())
            }
            None => Err(io::Error::new(io::ErrorKind::Other, "Unknown sequence name.")),
        }
    }
}

/// Record of a FASTA index.
#[derive(RustcDecodable, Debug, Copy, Clone)]
struct IndexRecord {
    len: u64,
    offset: u64,
    line_bases: u64,
    line_bytes: u64,
}

/// A sequence record returned by the FASTA index.
pub struct Sequence {
    pub name: String,
    pub len: u64,
}

/// A Fasta writer.
pub struct Writer<W: io::Write> {
    writer: io::BufWriter<W>,
}

impl Writer<fs::File> {
    /// Write to the given file path.
    pub fn to_file<P: AsRef<Path>>(path: P) -> io::Result<Self> {
        fs::File::create(path).map(Writer::new)
    }
}

impl<W: io::Write> Writer<W> {
    /// Create a new Fasta writer.
    pub fn new(writer: W) -> Self {
        Writer { writer: io::BufWriter::new(writer) }
    }

    /// Directly write a Fasta record.
    pub fn write_record(&mut self, record: Record) -> io::Result<()> {
        self.write(record.id().unwrap_or(""), record.desc(), record.seq())
    }

    /// Write a Fasta record with given id, optional description and sequence.
    pub fn write(&mut self, id: &str, desc: Option<&str>, seq: TextSlice) -> io::Result<()> {
        try!(self.writer.write(b">"));
        try!(self.writer.write(id.as_bytes()));
        if desc.is_some() {
            try!(self.writer.write(b" "));
            try!(self.writer.write(desc.unwrap().as_bytes()));
        }
        try!(self.writer.write(b"\n"));
        try!(self.writer.write(seq));
        try!(self.writer.write(b"\n"));

        Ok(())
    }

    /// Flush the writer, ensuring that everything is written.
    pub fn flush(&mut self) -> io::Result<()> {
        self.writer.flush()
    }
}

/// A FASTA record.
pub struct Record {
    header: String,
    seq: String,
}

impl Record {
    /// Create a new instance.
    pub fn new() -> Self {
        Record {
            header: String::new(),
            seq: String::new(),
        }
    }

    /// Check if record is empty.
    pub fn is_empty(&self) -> bool {
        self.header.is_empty() && self.seq.is_empty()
    }

    /// Check validity of Fasta record.
    pub fn check(&self) -> Result<(), &str> {
        if self.id().is_none() {
            return Err("Expecting id for FastQ record.");
        }
        if !self.seq.is_ascii() {
            return Err("Non-ascii character found in sequence.");
        }

        Ok(())
    }

    /// Return the id of the record.
    pub fn id(&self) -> Option<&str> {
        self.header[1..].trim_right().splitn(2, ' ').next()
    }

    /// Return descriptions if present.
    pub fn desc(&self) -> Option<&str> {
        self.header[1..].trim_right().splitn(2, ' ').skip(1).next()
    }

    /// Return the sequence of the record.
    pub fn seq(&self) -> TextSlice {
        self.seq.as_bytes()
    }

    /// Clear the record.
    fn clear(&mut self) {
        self.header.clear();
        self.seq.clear();
    }
}

/// An iterator over the records of a Fasta file.
pub struct Records<R: io::Read> {
    reader: Reader<R>,
}

impl<R: io::Read> Iterator for Records<R> {
    type Item = io::Result<Record>;

    fn next(&mut self) -> Option<io::Result<Record>> {
        let mut record = Record::new();
        match self.reader.read(&mut record) {
            Ok(()) if record.is_empty() => None,
            Ok(()) => Some(Ok(record)),
            Err(err) => Some(Err(err)),
        }
    }
}

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

    const FASTA_FILE: &'static [u8] = b">id desc
ACCGTAGGCTGA
CCGTAGGCTGAA
CGTAGGCTGAAA
GTAGGCTGAAAA
CCCC
>id2
ATTGTTGTTTTA
ATTGTTGTTTTA
ATTGTTGTTTTA
GGGG
";
    const FAI_FILE: &'static [u8] = b"id\t52\t9\t12\t13
id2\t40\t71\t12\t13
";
    const WRITE_FASTA_FILE: &'static [u8] = b">id desc
ACCGTAGGCTGA
>id2
ATTGTTGTTTTA
";

    #[test]
    fn test_reader() {
        let reader = Reader::new(FASTA_FILE);
        let ids = [Some("id"), Some("id2")];
        let descs = [Some("desc"), None];
        let seqs: [&[u8]; 2] = [b"ACCGTAGGCTGACCGTAGGCTGAACGTAGGCTGAAAGTAGGCTGAAAACCCC",
                                b"ATTGTTGTTTTAATTGTTGTTTTAATTGTTGTTTTAGGGG"];

        for (i, r) in reader.records().enumerate() {
            let record = r.ok().expect("Error reading record");
            assert_eq!(record.check(), Ok(()));
            assert_eq!(record.id(), ids[i]);
            assert_eq!(record.desc(), descs[i]);
            assert_eq!(record.seq(), seqs[i]);
        }

        // let record = records.ok().nth(1).unwrap();
    }

    #[test]
    fn test_indexed_reader() {
        let mut reader = IndexedReader::new(io::Cursor::new(FASTA_FILE), FAI_FILE)
                             .ok()
                             .expect("Error reading index");
        let mut seq = Vec::new();

        // Test reading various substrings of the sequence
        reader.read("id", 1, 5, &mut seq).ok().expect("Error reading sequence.");
        assert_eq!(seq, b"CCGT");

        reader.read("id", 1, 31, &mut seq).ok().expect("Error reading sequence.");
        assert_eq!(seq, b"CCGTAGGCTGACCGTAGGCTGAACGTAGGC");

        reader.read("id", 13, 23, &mut seq).ok().expect("Error reading sequence.");
        assert_eq!(seq, b"CGTAGGCTGA");

        reader.read("id", 36, 52, &mut seq).ok().expect("Error reading sequence.");
        assert_eq!(seq, b"GTAGGCTGAAAACCCC");

        reader.read("id2", 12, 40, &mut seq).ok().expect("Error reading sequence.");
        assert_eq!(seq, b"ATTGTTGTTTTAATTGTTGTTTTAGGGG");
    }

    #[test]
    fn test_writer() {
        let mut writer = Writer::new(Vec::new());
        writer.write("id", Some("desc"), b"ACCGTAGGCTGA").ok().expect("Expected successful write");
        writer.write("id2", None, b"ATTGTTGTTTTA").ok().expect("Expected successful write");
        writer.flush().ok().expect("Expected successful write");
        assert_eq!(writer.writer.get_ref(), &WRITE_FASTA_FILE);
    }
}
<|endoftext|>