<commit_before>pub mod dir;

use self::dir::FolderData;
use backend::{AnimeEntry, AnimeInfo, Status, SyncBackend};
use chrono::Local;
use error::SeriesError;
use input::{self, Answer};
use process;
use std::borrow::Cow;

pub struct SeriesConfig<B>
where
    B: SyncBackend,
{
    pub offline_mode: bool,
    pub sync_service: B,
    pub season_num: usize,
}

pub struct Series<B>
where
    B: SyncBackend,
{
    config: SeriesConfig<B>,
    dir: FolderData,
    season: SeasonState,
    ep_offset: u32,
}

impl<B> Series<B>
where
    B: SyncBackend,
{
    pub fn init(config: SeriesConfig<B>, mut dir: FolderData) -> Result<Series<B>, SeriesError> {
        dir.populate_season_data(&config)?;
        let ep_offset = dir.calculate_season_offset(0..config.season_num);
        let season = dir.seasons()[config.season_num].clone();

        let series = Series {
            config,
            dir,
            season,
            ep_offset,
        };

        Ok(series)
    }

    pub fn sync_remote_states(&mut self) -> Result<(), SeriesError> {
        self.season
            .sync_data_from_remote(&self.config, &self.dir, self.config.season_num)?;

        self.save_updated_season_data()?;
        self.prepare_list_entry()?;

        Ok(())
    }

    fn prepare_list_entry(&mut self) -> Result<(), SeriesError> {
        let state = &self.season.state;

        match state.status {
            Status::Watching | Status::Rewatching => {
                // Handle potential edge-case where all episodes have already been watched
                // but the series is still set in a watching state
                if let Some(total_eps) = state.info.episodes {
                    if state.watched_episodes >= total_eps {
                        self.update_list_entry_status(Status::Completed)?;
                        self.prompt_series_completed_options()?;
                    }
                }

                Ok(())
            }
            Status::PlanToWatch => self.update_list_entry_status(Status::Watching),
            Status::Completed => {
                println!("[{}] already completed", state.info.title);
                self.prompt_series_completed_options()
            }
            Status::OnHold | Status::Dropped => self.prompt_to_watch_paused_series(),
        }
    }

    fn update_list_entry_status(&mut self, status: Status) -> Result<(), SeriesError> {
        let state = &mut self.season.state;

        match status {
            Status::Watching => {
                // A series that was on hold probably already has a starting date, and it would make
                // more sense to use that one instead of replacing it
                if state.status != Status::OnHold {
                    state.start_date = Some(Local::today().naive_local());
                }

                state.finish_date = None;
            }
            Status::Rewatching => {
                println!("[{}] starting rewatch", state.info.title);
                println!("do you want to reset the start and end dates of the series? (Y/n)");

                if input::read_yn(Answer::Yes)? {
                    state.start_date = Some(Local::today().naive_local());
                    state.finish_date = None;
                }

                state.watched_episodes = 0;
            }
            Status::Completed => {
                if state.finish_date.is_none() {
                    state.finish_date = Some(Local::today().naive_local());
                }

                println!("[{}] completed!", state.info.title);
            }
            Status::Dropped => {
                if state.finish_date.is_none() {
                    state.finish_date = Some(Local::today().naive_local());
                }
            }
            Status::OnHold | Status::PlanToWatch => (),
        }

        state.status = status;
        self.update_list_entry()?;

        Ok(())
    }

    pub fn update_list_entry(&mut self) -> Result<(), SeriesError> {
        self.season.needs_sync = self.config.offline_mode;
        self.save_updated_season_data()?;

        if self.config.offline_mode {
            return Ok(());
        }

        self.config
            .sync_service
            .update_list_entry(&self.season.state)?;

        Ok(())
    }

    pub fn play_episode(&mut self, episode: u32) -> Result<(), SeriesError> {
        let absolute_ep = self.ep_offset + episode;
        let path = self.dir.episodes.get_episode(absolute_ep)?.clone();

        let status = process::open_with_default(path).map_err(SeriesError::FailedToOpenPlayer)?;
        let state = &mut self.season.state;

        state.watched_episodes = episode.max(state.watched_episodes);

        if !status.success() {
            eprintln!("video player not exited normally");
            eprintln!("do you still want to count the episode as completed? (y/N)");

            if !input::read_yn(Answer::No)? {
                return Ok(());
            }
        }

        match state.info.episodes {
            Some(total_eps) if episode >= total_eps => {
                self.update_list_entry_status(Status::Completed)?;
                self.prompt_series_completed_options()?;

                return Err(SeriesError::RequestExit);
            }
            _ => self.prompt_episode_completed()?,
        }

        Ok(())
    }

    pub fn play_all_episodes(&mut self) -> Result<(), SeriesError> {
        loop {
            let next_episode = self.season.state.watched_episodes + 1;

            self.play_episode(next_episode)?;
            self.prompt_next_episode_options()?;
        }
    }

    fn prompt_episode_completed(&mut self) -> Result<(), SeriesError> {
        let state = &mut self.season.state;

        let total_episodes = state
            .info
            .episodes
            .map(|e| Cow::Owned(e.to_string()))
            .unwrap_or_else(|| Cow::Borrowed("?"));

        println!(
            "[{}] episode {}/{} completed",
            state.info.title, state.watched_episodes, total_episodes
        );

        self.update_list_entry()?;
        Ok(())
    }

    fn prompt_next_episode_options(&mut self) -> Result<(), SeriesError> {
        let current_score_text: Cow<str> = match self.format_entry_score() {
            Some(score) => Cow::Owned(format!(" [{}]", score)),
            None => Cow::Borrowed(""),
        };

        println!("series options:");
        println!("\t[d] drop\n\t[h] put on hold\n\t[r] rate{}\n\t[x] exit\n\t[n] watch next episode (default)", current_score_text);

        let input = input::read_line()?.to_lowercase();

        match input.as_str() {
            "d" | "h" => {
                let status = if input == "d" {
                    Status::Dropped
                } else {
                    Status::OnHold
                };

                self.update_list_entry_status(status)?;

                println!("do you want to remove the episodes on disk? (Y/n)");

                if input::read_yn(Answer::Yes)? {
                    self.dir.try_remove_dir();
                }

                Err(SeriesError::RequestExit)
            }
            "r" => {
                self.prompt_to_update_score();
                self.update_list_entry()?;

                self.prompt_next_episode_options()
            }
            "x" => Err(SeriesError::RequestExit),
            _ => Ok(()),
        }
    }

    fn prompt_series_completed_options(&mut self) -> Result<(), SeriesError> {
        let current_score_text: Cow<str> = match self.format_entry_score() {
            Some(score) => Cow::Owned(format!(" [{}]", score)),
            None => Cow::Borrowed(""),
        };

        println!("series options:");
        println!(
            "\t[r] rate{}\n\t[w] rewatch\n\t[d] delete local files\n\t[x] exit",
            current_score_text
        );

        let input = input::read_line()?.to_lowercase();

        match input.as_str() {
            "r" => {
                self.prompt_to_update_score();
                self.update_list_entry()?;

                self.prompt_series_completed_options()
            }
            "w" => {
                self.update_list_entry_status(Status::Rewatching)?;
                self.play_all_episodes()
            }
            "d" => {
                self.dir.try_remove_dir();
                Err(SeriesError::RequestExit)
            }
            "x" => Err(SeriesError::RequestExit),
            _ => Ok(()),
        }
    }

    fn prompt_to_update_score(&mut self) {
        let (min_score, max_score) = self.config.sync_service.formatted_score_range();

        println!(
            "enter your score between {} and {} (press return to skip):",
            min_score, max_score
        );

        // TODO: use read_range() with empty line bypassing
        let input = match input::read_line() {
            Ok(ref input) if input.is_empty() => return,
            Ok(input) => input,
            Err(err) => {
                eprintln!("failed to read score input: {}", err);
                return;
            }
        };

        match self.config.sync_service.parse_score(&input) {
            Ok(score) => self.season.state.score = Some(score),
            Err(err) => eprintln!("failed to parse score: {}", err),
        }
    }

    fn prompt_to_watch_paused_series(&mut self) -> Result<(), SeriesError> {
        println!(
            "[{}] was previously put on hold or dropped",
            self.season.state.info.title
        );

        println!("do you want to watch it from the beginning? (Y/n)");

        if input::read_yn(Answer::Yes)? {
            self.season.state.watched_episodes = 0;
        }

        self.update_list_entry_status(Status::Watching)?;
        Ok(())
    }

    fn format_entry_score(&self) -> Option<String> {
        let state = &self.season.state;

        match state.score {
            Some(score) => {
                let formatted_score = self.config.sync_service.format_score(score);

                match formatted_score {
                    Ok(score) => Some(score),
                    Err(err) => {
                        eprintln!("failed to read existing list entry score: {}", err);
                        None
                    }
                }
            }
            None => None,
        }
    }

    fn save_updated_season_data(&mut self) -> Result<(), SeriesError> {
        if self.config.season_num >= self.dir.seasons().len() {
            return Ok(());
        }

        self.dir.seasons_mut()[self.config.season_num] = self.season.clone();
        self.dir.save()
    }
}

pub fn search_for_series_info<B>(
    backend: &B,
    name: &str,
    season: usize,
) -> Result<AnimeInfo, SeriesError>
where
    B: SyncBackend,
{
    println!("[{}] searching on {}..", name, B::name());

    let mut found = backend.search_by_name(name)?;

    if !found.is_empty() {
        println!(
            "select season {} by entering the number next to its name:\n",
            1 + season
        );

        println!("0 [custom search]");

        for (i, series) in found.iter().enumerate() {
            println!("{} [{}]", 1 + i, series.title);
        }

        let index = input::read_range(0, found.len())?;

        if index == 0 {
            println!("enter the name you want to search for:");

            let name = input::read_line()?;
            search_for_series_info(backend, &name, season)
        } else {
            let info = found.swap_remove(index - 1);
            Ok(info)
        }
    } else {
        println!("no results found\nplease enter a custom search term:");

        let name = input::read_line()?;
        search_for_series_info(backend, &name, season)
    }
}

#[derive(Clone, Serialize, Deserialize)]
pub struct SeasonState {
    #[serde(flatten)]
    pub state: AnimeEntry,
    pub needs_info: bool,
    pub needs_sync: bool,
}

impl SeasonState {
    pub fn sync_info_from_remote<B>(
        &mut self,
        config: &SeriesConfig<B>,
        dir: &FolderData,
        season_num: usize,
    ) -> Result<(), SeriesError>
    where
        B: SyncBackend,
    {
        if !self.needs_info {
            return Ok(());
        }

        if config.offline_mode {
            // The series title is the only thing that can really be used when offline.
            // While things like the number of episode files present could be used as the
            // series episode count, it is common to only partially have a series downloaded
            self.state.info.title = dir.episodes.series_name.clone();
            return Ok(());
        }

        let info =
            search_for_series_info(&config.sync_service, &dir.episodes.series_name, season_num)?;

        self.state.info = info;
        // We only want to set this flag when online, since offline mode only provides
        // very basic information at best
        self.needs_info = false;
        Ok(())
    }

    pub fn sync_data_from_remote<B>(
        &mut self,
        config: &SeriesConfig<B>,
        dir: &FolderData,
        season_num: usize,
    ) -> Result<(), SeriesError>
    where
        B: SyncBackend,
    {
        self.sync_info_from_remote(config, dir, season_num)?;

        // Sync data from the backend when not offline
        if !config.offline_mode {
            let entry = config
                .sync_service
                .get_list_entry(self.state.info.clone())?;

            if let Some(entry) = entry {
                // If we don't have new data to report, we should sync the data from the backend to keep up with
                // any changes made outside of the program
                if !self.needs_sync {
                    self.state = entry;
                }
            }
        }

        Ok(())
    }
}
<commit_msg>[series] Simplify offline mode check in SeasonState::sync_data_from_remote<commit_after>pub mod dir;

use self::dir::FolderData;
use backend::{AnimeEntry, AnimeInfo, Status, SyncBackend};
use chrono::Local;
use error::SeriesError;
use input::{self, Answer};
use process;
use std::borrow::Cow;

pub struct SeriesConfig<B>
where
    B: SyncBackend,
{
    pub offline_mode: bool,
    pub sync_service: B,
    pub season_num: usize,
}

pub struct Series<B>
where
    B: SyncBackend,
{
    config: SeriesConfig<B>,
    dir: FolderData,
    season: SeasonState,
    ep_offset: u32,
}

impl<B> Series<B>
where
    B: SyncBackend,
{
    pub fn init(config: SeriesConfig<B>, mut dir: FolderData) -> Result<Series<B>, SeriesError> {
        dir.populate_season_data(&config)?;
        let ep_offset = dir.calculate_season_offset(0..config.season_num);
        let season = dir.seasons()[config.season_num].clone();

        let series = Series {
            config,
            dir,
            season,
            ep_offset,
        };

        Ok(series)
    }

    pub fn sync_remote_states(&mut self) -> Result<(), SeriesError> {
        self.season
            .sync_data_from_remote(&self.config, &self.dir, self.config.season_num)?;

        self.save_updated_season_data()?;
        self.prepare_list_entry()?;

        Ok(())
    }

    fn prepare_list_entry(&mut self) -> Result<(), SeriesError> {
        let state = &self.season.state;

        match state.status {
            Status::Watching | Status::Rewatching => {
                // Handle potential edge-case where all episodes have already been watched
                // but the series is still set in a watching state
                if let Some(total_eps) = state.info.episodes {
                    if state.watched_episodes >= total_eps {
                        self.update_list_entry_status(Status::Completed)?;
                        self.prompt_series_completed_options()?;
                    }
                }

                Ok(())
            }
            Status::PlanToWatch => self.update_list_entry_status(Status::Watching),
            Status::Completed => {
                println!("[{}] already completed", state.info.title);
                self.prompt_series_completed_options()
            }
            Status::OnHold | Status::Dropped => self.prompt_to_watch_paused_series(),
        }
    }

    fn update_list_entry_status(&mut self, status: Status) -> Result<(), SeriesError> {
        let state = &mut self.season.state;

        match status {
            Status::Watching => {
                // A series that was on hold probably already has a starting date, and it would make
                // more sense to use that one instead of replacing it
                if state.status != Status::OnHold {
                    state.start_date = Some(Local::today().naive_local());
                }

                state.finish_date = None;
            }
            Status::Rewatching => {
                println!("[{}] starting rewatch", state.info.title);
                println!("do you want to reset the start and end dates of the series? (Y/n)");

                if input::read_yn(Answer::Yes)? {
                    state.start_date = Some(Local::today().naive_local());
                    state.finish_date = None;
                }

                state.watched_episodes = 0;
            }
            Status::Completed => {
                if state.finish_date.is_none() {
                    state.finish_date = Some(Local::today().naive_local());
                }

                println!("[{}] completed!", state.info.title);
            }
            Status::Dropped => {
                if state.finish_date.is_none() {
                    state.finish_date = Some(Local::today().naive_local());
                }
            }
            Status::OnHold | Status::PlanToWatch => (),
        }

        state.status = status;
        self.update_list_entry()?;

        Ok(())
    }

    pub fn update_list_entry(&mut self) -> Result<(), SeriesError> {
        self.season.needs_sync = self.config.offline_mode;
        self.save_updated_season_data()?;

        if self.config.offline_mode {
            return Ok(());
        }

        self.config
            .sync_service
            .update_list_entry(&self.season.state)?;

        Ok(())
    }

    pub fn play_episode(&mut self, episode: u32) -> Result<(), SeriesError> {
        let absolute_ep = self.ep_offset + episode;
        let path = self.dir.episodes.get_episode(absolute_ep)?.clone();

        let status = process::open_with_default(path).map_err(SeriesError::FailedToOpenPlayer)?;
        let state = &mut self.season.state;

        state.watched_episodes = episode.max(state.watched_episodes);

        if !status.success() {
            eprintln!("video player not exited normally");
            eprintln!("do you still want to count the episode as completed? (y/N)");

            if !input::read_yn(Answer::No)? {
                return Ok(());
            }
        }

        match state.info.episodes {
            Some(total_eps) if episode >= total_eps => {
                self.update_list_entry_status(Status::Completed)?;
                self.prompt_series_completed_options()?;

                return Err(SeriesError::RequestExit);
            }
            _ => self.prompt_episode_completed()?,
        }

        Ok(())
    }

    pub fn play_all_episodes(&mut self) -> Result<(), SeriesError> {
        loop {
            let next_episode = self.season.state.watched_episodes + 1;

            self.play_episode(next_episode)?;
            self.prompt_next_episode_options()?;
        }
    }

    fn prompt_episode_completed(&mut self) -> Result<(), SeriesError> {
        let state = &mut self.season.state;

        let total_episodes = state
            .info
            .episodes
            .map(|e| Cow::Owned(e.to_string()))
            .unwrap_or_else(|| Cow::Borrowed("?"));

        println!(
            "[{}] episode {}/{} completed",
            state.info.title, state.watched_episodes, total_episodes
        );

        self.update_list_entry()?;
        Ok(())
    }

    fn prompt_next_episode_options(&mut self) -> Result<(), SeriesError> {
        let current_score_text: Cow<str> = match self.format_entry_score() {
            Some(score) => Cow::Owned(format!(" [{}]", score)),
            None => Cow::Borrowed(""),
        };

        println!("series options:");
        println!("\t[d] drop\n\t[h] put on hold\n\t[r] rate{}\n\t[x] exit\n\t[n] watch next episode (default)", current_score_text);

        let input = input::read_line()?.to_lowercase();

        match input.as_str() {
            "d" | "h" => {
                let status = if input == "d" {
                    Status::Dropped
                } else {
                    Status::OnHold
                };

                self.update_list_entry_status(status)?;

                println!("do you want to remove the episodes on disk? (Y/n)");

                if input::read_yn(Answer::Yes)? {
                    self.dir.try_remove_dir();
                }

                Err(SeriesError::RequestExit)
            }
            "r" => {
                self.prompt_to_update_score();
                self.update_list_entry()?;

                self.prompt_next_episode_options()
            }
            "x" => Err(SeriesError::RequestExit),
            _ => Ok(()),
        }
    }

    fn prompt_series_completed_options(&mut self) -> Result<(), SeriesError> {
        let current_score_text: Cow<str> = match self.format_entry_score() {
            Some(score) => Cow::Owned(format!(" [{}]", score)),
            None => Cow::Borrowed(""),
        };

        println!("series options:");
        println!(
            "\t[r] rate{}\n\t[w] rewatch\n\t[d] delete local files\n\t[x] exit",
            current_score_text
        );

        let input = input::read_line()?.to_lowercase();

        match input.as_str() {
            "r" => {
                self.prompt_to_update_score();
                self.update_list_entry()?;

                self.prompt_series_completed_options()
            }
            "w" => {
                self.update_list_entry_status(Status::Rewatching)?;
                self.play_all_episodes()
            }
            "d" => {
                self.dir.try_remove_dir();
                Err(SeriesError::RequestExit)
            }
            "x" => Err(SeriesError::RequestExit),
            _ => Ok(()),
        }
    }

    fn prompt_to_update_score(&mut self) {
        let (min_score, max_score) = self.config.sync_service.formatted_score_range();

        println!(
            "enter your score between {} and {} (press return to skip):",
            min_score, max_score
        );

        // TODO: use read_range() with empty line bypassing
        let input = match input::read_line() {
            Ok(ref input) if input.is_empty() => return,
            Ok(input) => input,
            Err(err) => {
                eprintln!("failed to read score input: {}", err);
                return;
            }
        };

        match self.config.sync_service.parse_score(&input) {
            Ok(score) => self.season.state.score = Some(score),
            Err(err) => eprintln!("failed to parse score: {}", err),
        }
    }

    fn prompt_to_watch_paused_series(&mut self) -> Result<(), SeriesError> {
        println!(
            "[{}] was previously put on hold or dropped",
            self.season.state.info.title
        );

        println!("do you want to watch it from the beginning? (Y/n)");

        if input::read_yn(Answer::Yes)? {
            self.season.state.watched_episodes = 0;
        }

        self.update_list_entry_status(Status::Watching)?;
        Ok(())
    }

    fn format_entry_score(&self) -> Option<String> {
        let state = &self.season.state;

        match state.score {
            Some(score) => {
                let formatted_score = self.config.sync_service.format_score(score);

                match formatted_score {
                    Ok(score) => Some(score),
                    Err(err) => {
                        eprintln!("failed to read existing list entry score: {}", err);
                        None
                    }
                }
            }
            None => None,
        }
    }

    fn save_updated_season_data(&mut self) -> Result<(), SeriesError> {
        if self.config.season_num >= self.dir.seasons().len() {
            return Ok(());
        }

        self.dir.seasons_mut()[self.config.season_num] = self.season.clone();
        self.dir.save()
    }
}

pub fn search_for_series_info<B>(
    backend: &B,
    name: &str,
    season: usize,
) -> Result<AnimeInfo, SeriesError>
where
    B: SyncBackend,
{
    println!("[{}] searching on {}..", name, B::name());

    let mut found = backend.search_by_name(name)?;

    if !found.is_empty() {
        println!(
            "select season {} by entering the number next to its name:\n",
            1 + season
        );

        println!("0 [custom search]");

        for (i, series) in found.iter().enumerate() {
            println!("{} [{}]", 1 + i, series.title);
        }

        let index = input::read_range(0, found.len())?;

        if index == 0 {
            println!("enter the name you want to search for:");

            let name = input::read_line()?;
            search_for_series_info(backend, &name, season)
        } else {
            let info = found.swap_remove(index - 1);
            Ok(info)
        }
    } else {
        println!("no results found\nplease enter a custom search term:");

        let name = input::read_line()?;
        search_for_series_info(backend, &name, season)
    }
}

#[derive(Clone, Serialize, Deserialize)]
pub struct SeasonState {
    #[serde(flatten)]
    pub state: AnimeEntry,
    pub needs_info: bool,
    pub needs_sync: bool,
}

impl SeasonState {
    pub fn sync_info_from_remote<B>(
        &mut self,
        config: &SeriesConfig<B>,
        dir: &FolderData,
        season_num: usize,
    ) -> Result<(), SeriesError>
    where
        B: SyncBackend,
    {
        if !self.needs_info {
            return Ok(());
        }

        if config.offline_mode {
            // The series title is the only thing that can really be used when offline.
            // While things like the number of episode files present could be used as the
            // series episode count, it is common to only partially have a series downloaded
            self.state.info.title = dir.episodes.series_name.clone();
            return Ok(());
        }

        let info =
            search_for_series_info(&config.sync_service, &dir.episodes.series_name, season_num)?;

        self.state.info = info;
        // We only want to set this flag when online, since offline mode only provides
        // very basic information at best
        self.needs_info = false;
        Ok(())
    }

    pub fn sync_data_from_remote<B>(
        &mut self,
        config: &SeriesConfig<B>,
        dir: &FolderData,
        season_num: usize,
    ) -> Result<(), SeriesError>
    where
        B: SyncBackend,
    {
        self.sync_info_from_remote(config, dir, season_num)?;

        // There's no information we can sync when offline
        if config.offline_mode {
            return Ok(());
        }

        let entry = config
            .sync_service
            .get_list_entry(self.state.info.clone())?;

        if let Some(entry) = entry {
            // If we don't have new data to report, we should sync the data from the backend to keep up with
            // any changes made outside of the program
            if !self.needs_sync {
                self.state = entry;
            }
        }

        Ok(())
    }
}
<|endoftext|>


<commit_before>use super::Blob;
use failure::Error;
use futures::Future;
use postgres::Connection;
use rusoto_core::region::Region;
use rusoto_credential::DefaultCredentialsProvider;
use rusoto_s3::{GetObjectRequest, PutObjectRequest, S3Client, S3};
use std::convert::TryInto;
use std::io::Read;
use time::Timespec;
use log::{error, warn};

pub(crate) struct S3Backend<'a> {
    client: S3Client,
    bucket: &'a str,
}

impl<'a> S3Backend<'a> {
    pub(crate) fn new(client: S3Client, bucket: &'a str) -> Self {
        Self { client, bucket }
    }

    pub(super) fn get(&self, path: &str) -> Result<Blob, Error> {
        let res = self
            .client
            .get_object(GetObjectRequest {
                bucket: self.bucket.to_string(),
                key: path.into(),
                ..Default::default()
            })
            .sync()?;

        let mut b = res.body.unwrap().into_blocking_read();
        let mut content = Vec::with_capacity(
            res.content_length
                .and_then(|l| l.try_into().ok())
                .unwrap_or(0),
        );
        b.read_to_end(&mut content).unwrap();

        let date_updated = parse_timespec(&res.last_modified.unwrap())?;

        Ok(Blob {
            path: path.into(),
            mime: res.content_type.unwrap(),
            date_updated,
            content,
        })
    }

    pub(super) fn store_batch(&self, batch: &[Blob]) -> Result<(), Error> {
        let mut rt = tokio::runtime::Runtime::new().unwrap();
        let mut attempts = 0;

        loop {
            let mut futures = Vec::with_capacity(batch.len());
            for blob in batch {
                futures.push(
                    self.client
                        .put_object(PutObjectRequest {
                            bucket: self.bucket.to_string(),
                            key: blob.path.clone(),
                            body: Some(blob.content.clone().into()),
                            content_type: Some(blob.mime.clone()),
                            ..Default::default()
                        })
                        .inspect(|_| {
                            crate::web::metrics::UPLOADED_FILES_TOTAL.inc_by(1);
                        }),
                );
            }
            attempts += 1;

            match rt.block_on(::futures::future::join_all(futures)) {
                // this batch was successful, start another batch if there are still more files
                Ok(_) => break,
                Err(err) => {
                    error!("failed to upload to s3: {:?}", err);
                    // if a futures error occurs, retry the batch
                    if attempts > 2 {
                        panic!("failed to upload 3 times, exiting");
                    }
                }
            }
        }
        Ok(())
    }
}

// public for testing
pub(crate) const TIME_FMT: &str = "%a, %d %b %Y %H:%M:%S %Z";

fn parse_timespec(raw: &str) -> Result<Timespec, Error> {
    Ok(time::strptime(raw, TIME_FMT)?.to_timespec())
}

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

    fn assert_s3_404(env: &TestEnvironment, path: &'static str) {
        use rusoto_core::RusotoError;
        use rusoto_s3::GetObjectError;

        let s3 = env.s3().not_found(path);
        let backend = S3Backend::new(s3.client, s3.bucket);
        let err = backend.get(path).unwrap_err();
        let status = match err
            .downcast_ref::<RusotoError<GetObjectError>>()
            .expect("wanted GetObject")
        {
            RusotoError::Unknown(http) => http.status,
            _ => panic!("wrong error"),
        };
        assert_eq!(status, 404);
    }

    #[test]
    fn test_parse_timespec() {
        crate::test::wrapper(|_| {
            // Test valid conversions
            assert_eq!(
                parse_timespec("Thu, 1 Jan 1970 00:00:00 GMT")?,
                Timespec::new(0, 0)
            );
            assert_eq!(
                parse_timespec("Mon, 16 Apr 2018 04:33:50 GMT")?,
                Timespec::new(1523853230, 0)
            );

            // Test invalid conversion
            assert!(parse_timespec("foo").is_err());

            Ok(())
        })
    }

    #[test]
    fn test_path_get() {
        crate::test::wrapper(|env| {
            let blob = Blob {
                path: "dir/foo.txt".into(),
                mime: "text/plain".into(),
                date_updated: Timespec::new(42, 0),
                content: "Hello world!".into(),
            };

            // Add a test file to the database
            let s3 = env.s3().upload(blob.clone());

            let backend = S3Backend::new(s3.client, &s3.bucket);

            // Test that the proper file was returned
            assert_eq!(blob, backend.get("dir/foo.txt")?);

            // Test that other files are not returned
            assert_s3_404(&env, "dir/bar.txt");
            assert_s3_404(&env, "foo.txt");

            Ok(())
        });
    }
}

pub(crate) static S3_BUCKET_NAME: &str = "rust-docs-rs";

pub(crate) fn s3_client() -> Option<S3Client> {
    // If AWS keys aren't configured, then presume we should use the DB exclusively
    // for file storage.
    if std::env::var_os("AWS_ACCESS_KEY_ID").is_none() && std::env::var_os("FORCE_S3").is_none() {
        return None;
    }
    let creds = match DefaultCredentialsProvider::new() {
        Ok(creds) => creds,
        Err(err) => {
            warn!("failed to retrieve AWS credentials: {}", err);
            return None;
        }
    };
    Some(S3Client::new_with(
        rusoto_core::request::HttpClient::new().unwrap(),
        creds,
        std::env::var("S3_ENDPOINT")
            .ok()
            .map(|e| Region::Custom {
                name: std::env::var("S3_REGION").unwrap_or_else(|_| "us-west-1".to_owned()),
                endpoint: e,
            })
            .unwrap_or(Region::UsWest1),
    ))
}

pub fn move_to_s3(conn: &Connection, n: usize) -> Result<usize, Error> {
    let trans = conn.transaction()?;
    let client = s3_client().expect("configured s3");

    let rows = trans.query(
        &format!(
            "SELECT path, mime, content FROM files WHERE content != E'in-s3' LIMIT {}",
            n
        ),
        &[],
    )?;
    let count = rows.len();

    let mut rt = ::tokio::runtime::Runtime::new().unwrap();
    let mut futures = Vec::with_capacity(rows.len());
    for row in &rows {
        let path: String = row.get(0);
        let mime: String = row.get(1);
        let content: Vec<u8> = row.get(2);
        let path_1 = path.clone();
        futures.push(
            client
                .put_object(PutObjectRequest {
                    bucket: S3_BUCKET_NAME.into(),
                    key: path.clone(),
                    body: Some(content.into()),
                    content_type: Some(mime),
                    ..Default::default()
                })
                .map(move |_| path_1)
                .map_err(move |e| panic!("failed to upload to {}: {:?}", path, e)),
        );
    }

    match rt.block_on(::futures::future::join_all(futures)) {
        Ok(paths) => {
            let statement = trans.prepare("DELETE FROM files WHERE path = $1").unwrap();
            for path in paths {
                statement.execute(&[&path]).unwrap();
            }
        }
        Err(e) => {
            panic!("results err: {:?}", e);
        }
    }

    trans.commit()?;

    Ok(count)
}
<commit_msg>Use 2018 idioms<commit_after>use super::Blob;
use failure::Error;
use futures::Future;
use postgres::Connection;
use rusoto_core::region::Region;
use rusoto_credential::DefaultCredentialsProvider;
use rusoto_s3::{GetObjectRequest, PutObjectRequest, S3Client, S3};
use std::convert::TryInto;
use std::io::Read;
use time::Timespec;
use log::{error, warn};

pub(crate) struct S3Backend<'a> {
    client: S3Client,
    bucket: &'a str,
}

impl<'a> S3Backend<'a> {
    pub(crate) fn new(client: S3Client, bucket: &'a str) -> Self {
        Self { client, bucket }
    }

    pub(super) fn get(&self, path: &str) -> Result<Blob, Error> {
        let res = self
            .client
            .get_object(GetObjectRequest {
                bucket: self.bucket.to_string(),
                key: path.into(),
                ..Default::default()
            })
            .sync()?;

        let mut b = res.body.unwrap().into_blocking_read();
        let mut content = Vec::with_capacity(
            res.content_length
                .and_then(|l| l.try_into().ok())
                .unwrap_or(0),
        );
        b.read_to_end(&mut content).unwrap();

        let date_updated = parse_timespec(&res.last_modified.unwrap())?;

        Ok(Blob {
            path: path.into(),
            mime: res.content_type.unwrap(),
            date_updated,
            content,
        })
    }

    pub(super) fn store_batch(&self, batch: &[Blob]) -> Result<(), Error> {
        let mut rt = tokio::runtime::Runtime::new().unwrap();
        let mut attempts = 0;

        loop {
            let mut futures = Vec::with_capacity(batch.len());
            for blob in batch {
                futures.push(
                    self.client
                        .put_object(PutObjectRequest {
                            bucket: self.bucket.to_string(),
                            key: blob.path.clone(),
                            body: Some(blob.content.clone().into()),
                            content_type: Some(blob.mime.clone()),
                            ..Default::default()
                        })
                        .inspect(|_| {
                            crate::web::metrics::UPLOADED_FILES_TOTAL.inc_by(1);
                        }),
                );
            }
            attempts += 1;

            match rt.block_on(::futures::future::join_all(futures)) {
                // this batch was successful, start another batch if there are still more files
                Ok(_) => break,
                Err(err) => {
                    error!("failed to upload to s3: {:?}", err);
                    // if a futures error occurs, retry the batch
                    if attempts > 2 {
                        panic!("failed to upload 3 times, exiting");
                    }
                }
            }
        }
        Ok(())
    }
}

// public for testing
pub(crate) const TIME_FMT: &str = "%a, %d %b %Y %H:%M:%S %Z";

fn parse_timespec(raw: &str) -> Result<Timespec, Error> {
    Ok(time::strptime(raw, TIME_FMT)?.to_timespec())
}

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

    fn assert_s3_404(env: &TestEnvironment, path: &'static str) {
        use rusoto_core::RusotoError;
        use rusoto_s3::GetObjectError;

        let s3 = env.s3().not_found(path);
        let backend = S3Backend::new(s3.client, s3.bucket);
        let err = backend.get(path).unwrap_err();
        let status = match err
            .downcast_ref::<RusotoError<GetObjectError>>()
            .expect("wanted GetObject")
        {
            RusotoError::Unknown(http) => http.status,
            _ => panic!("wrong error"),
        };
        assert_eq!(status, 404);
    }

    #[test]
    fn test_parse_timespec() {
        crate::test::wrapper(|_| {
            // Test valid conversions
            assert_eq!(
                parse_timespec("Thu, 1 Jan 1970 00:00:00 GMT")?,
                Timespec::new(0, 0)
            );
            assert_eq!(
                parse_timespec("Mon, 16 Apr 2018 04:33:50 GMT")?,
                Timespec::new(1523853230, 0)
            );

            // Test invalid conversion
            assert!(parse_timespec("foo").is_err());

            Ok(())
        })
    }

    #[test]
    fn test_path_get() {
        crate::test::wrapper(|env| {
            let blob = Blob {
                path: "dir/foo.txt".into(),
                mime: "text/plain".into(),
                date_updated: Timespec::new(42, 0),
                content: "Hello world!".into(),
            };

            // Add a test file to the database
            let s3 = env.s3().upload(blob.clone());

            let backend = S3Backend::new(s3.client, &s3.bucket);

            // Test that the proper file was returned
            assert_eq!(blob, backend.get("dir/foo.txt")?);

            // Test that other files are not returned
            assert_s3_404(&env, "dir/bar.txt");
            assert_s3_404(&env, "foo.txt");

            Ok(())
        });
    }
}

pub(crate) static S3_BUCKET_NAME: &str = "rust-docs-rs";

pub(crate) fn s3_client() -> Option<S3Client> {
    // If AWS keys aren't configured, then presume we should use the DB exclusively
    // for file storage.
    if std::env::var_os("AWS_ACCESS_KEY_ID").is_none() && std::env::var_os("FORCE_S3").is_none() {
        return None;
    }
    let creds = match DefaultCredentialsProvider::new() {
        Ok(creds) => creds,
        Err(err) => {
            warn!("failed to retrieve AWS credentials: {}", err);
            return None;
        }
    };
    Some(S3Client::new_with(
        rusoto_core::request::HttpClient::new().unwrap(),
        creds,
        std::env::var("S3_ENDPOINT")
            .ok()
            .map(|e| Region::Custom {
                name: std::env::var("S3_REGION").unwrap_or_else(|_| "us-west-1".to_owned()),
                endpoint: e,
            })
            .unwrap_or(Region::UsWest1),
    ))
}

pub fn move_to_s3(conn: &Connection, n: usize) -> Result<usize, Error> {
    let trans = conn.transaction()?;
    let client = s3_client().expect("configured s3");

    let rows = trans.query(
        &format!(
            "SELECT path, mime, content FROM files WHERE content != E'in-s3' LIMIT {}",
            n
        ),
        &[],
    )?;
    let count = rows.len();

    let mut rt = tokio::runtime::Runtime::new().unwrap();
    let mut futures = Vec::with_capacity(rows.len());
    for row in &rows {
        let path: String = row.get(0);
        let mime: String = row.get(1);
        let content: Vec<u8> = row.get(2);
        let path_1 = path.clone();
        futures.push(
            client
                .put_object(PutObjectRequest {
                    bucket: S3_BUCKET_NAME.into(),
                    key: path.clone(),
                    body: Some(content.into()),
                    content_type: Some(mime),
                    ..Default::default()
                })
                .map(move |_| path_1)
                .map_err(move |e| panic!("failed to upload to {}: {:?}", path, e)),
        );
    }

    match rt.block_on(::futures::future::join_all(futures)) {
        Ok(paths) => {
            let statement = trans.prepare("DELETE FROM files WHERE path = $1").unwrap();
            for path in paths {
                statement.execute(&[&path]).unwrap();
            }
        }
        Err(e) => {
            panic!("results err: {:?}", e);
        }
    }

    trans.commit()?;

    Ok(count)
}
<|endoftext|>


<commit_before>// (c) 2016 Joost Yervante Damad <joost@productize.be>

use std::fmt;
use std::str::FromStr;
use std::path::PathBuf;
use std::collections::HashMap;

// get from parent
use ERes;
use err;
use read_file;
use parse_split_quote_aware;
use schematic;

#[derive(Debug)]
pub struct SymbolLib {
    pub symbols:Vec<Symbol>,
}

// DEF name reference unused text_offset draw_pinnumber draw_pinname unit_count units_locked option_flag
#[derive(Debug)]
pub struct Symbol {
    pub name:String,
    pub reference:String,
    pub text_offset:f64,
    pub draw_pinnumber: bool,
    pub draw_pinname: bool,
    pub unit_count: i64,
    pub unit_locked: bool,
    pub is_power: bool,
    pub fields:Vec<Field>,
    pub draw:Vec<String>, // TODO parse draw
}
// F n “text” posx posy dimension orientation visibility hjustify vjustify/italic/bold “name”
// F0 "#PWR" 0 0 30 H I C CNN

#[derive(Debug)]
pub struct Field {
    pub i:i64,
    pub value:String,
    pub x:f64,
    pub y:f64,
    pub dimention:i64,
    pub orientation:schematic::Orientation,
    pub visibility:bool,
    pub hjustify:schematic::Justify,
    pub vjustify:schematic::Justify,
    pub italic:bool,
    pub bold:bool,
    pub name:String,
}

impl SymbolLib {
    fn new() -> SymbolLib {
        SymbolLib {
            symbols:vec![]
        }
    }
}

impl Symbol {
    fn new(name:String, reference:String) -> Symbol {
        Symbol {
            name:name,
            reference:reference,
            text_offset:0.0,
            draw_pinnumber:false,
            draw_pinname:false,
            unit_count:1,
            unit_locked:false,
            is_power:false,
            fields:vec![],
            draw:vec![],
        }
    }
}

macro_rules! assume_line {
    ($s:expr, $exp:expr) => (
        if $s.eof() {
            return err("end of file reached")
        }
        if $s.here() != $exp {
            return Err(format!("expected '{}', got '{}'", $exp, $s.here()))
        }
        $s.i += 1;
    )
}

#[derive(Debug)]
struct ParseState {
    i:usize,
    v:Vec<String>,
}

impl ParseState {
    fn new(v2:Vec<&str>) -> ParseState {
        ParseState {
            i:0,
            v:v2.iter().map(|x| String::from(*x)).collect(),
        }
    }

    fn here(&self) -> String {
        let s =  (self.v[self.i]).clone();
        //println!("{}", s);
        s
    }

    fn next(&mut self) {
        self.i += 1;
    }

    fn eof(&self) -> bool {
        self.i >= self.v.len()
    }
}

fn char_at(s:&String, p:usize) -> char {
    let v:Vec<char> = s.chars().collect();
    v[..][p]
}

fn assume_string(e:&'static str, s:&String) -> ERes<()> {
    if String::from(e) != *s {
        return Err(format!("expecting: {}, actually: {}", e, s))
    }
    return Ok(())
}

fn i64_from_string(p:&ParseState, s:&String) -> ERes<i64> {
    match i64::from_str(&s[..]) {
        Ok(i) => Ok(i),
        _ => Err(format!("int parse error in {}; line: {}", s, p.here()))
    }
}

fn f64_from_string(p:&ParseState, s:&String) -> ERes<f64> {
    match f64::from_str(&s[..]) {
        Ok(i) => Ok(i),
        _ => Err(format!("float parse error in {}; line: {}", s, p.here()))
    }
}

fn bool_from_string(s:&String, t:&'static str, f:&'static str) -> ERes<bool> {
    if &s[..] == t {
        return Ok(true)
    }
    if &s[..] == f {
        return Ok(false)
    }
    Err(format!("unknown boolean {}, expected {} or {}", s, t, f))
}

fn bool_from<T: PartialEq + fmt::Display>(i:T, t:T, f:T) -> ERes<bool> {
    if i == t {
        return Ok(true)
    }
    if i == f {
        return Ok(false)
    }
    Err(format!("unknown boolean {}, expected {} or {}", i, t, f))
}

fn parse_symbol(p:&mut ParseState) -> ERes<Symbol> {
    assume_line!(p,"#");
    p.next(); // skip line like # name
    assume_line!(p,"#");
    let s = p.here();
    let v = &parse_split_quote_aware(&s);
    if v.len() != 10 {
        return Err(format!("unexpected elements in {}", s))
    }
    assume_string("DEF", &v[0]);
    let mut s = Symbol::new(v[1].clone(), v[2].clone());
    s.text_offset = try!(f64_from_string(p, &v[4]));
    s.draw_pinnumber = try!(bool_from_string(&v[5], "Y", "N"));
    s.draw_pinname = try!(bool_from_string(&v[6], "Y", "N"));
    s.unit_count = try!(i64_from_string(p, &v[7]));
    s.unit_locked = try!(bool_from_string(&v[8], "L", "F"));
    s.is_power = try!(bool_from_string(&v[9], "P", "N"));
    // TODO fields
    Ok(s)
}

fn parse(filename:Option<PathBuf>, s: &str) -> ERes<SymbolLib> {
    let mut lib = SymbolLib::new();
    let v:Vec<&str> = s.lines().collect();
    let p = &mut ParseState::new(v);
    assume_line!(p, "EESchema-LIBRARY Version 2.3");
    assume_line!(p, "#encoding utf-8");
    while !p.eof() {
        let s = try!(parse_symbol(p));
        lib.symbols.push(s)
    }
    Ok(lib)
}

pub fn parse_str(s:&str) -> ERes<SymbolLib> {
    parse(None, s)
}

pub fn parse_file(filename:&PathBuf) -> ERes<SymbolLib> {
    let name = filename.to_str().unwrap();
    let s = try!(read_file(name));
    parse(Some(filename.clone()), &s[..])
}
<commit_msg>fix compilation<commit_after>// (c) 2016 Joost Yervante Damad <joost@productize.be>

use std::fmt;
use std::str::FromStr;
use std::path::PathBuf;
use std::collections::HashMap;

// get from parent
use ERes;
use err;
use read_file;
use parse_split_quote_aware;
use schematic;

#[derive(Debug)]
pub struct SymbolLib {
    pub symbols:Vec<Symbol>,
}

// DEF name reference unused text_offset draw_pinnumber draw_pinname unit_count units_locked option_flag
#[derive(Debug)]
pub struct Symbol {
    pub name:String,
    pub reference:String,
    pub text_offset:f64,
    pub draw_pinnumber: bool,
    pub draw_pinname: bool,
    pub unit_count: i64,
    pub unit_locked: bool,
    pub is_power: bool,
    pub fields:Vec<Field>,
    pub draw:Vec<String>, // TODO parse draw
}
// F n “text” posx posy dimension orientation visibility hjustify vjustify/italic/bold “name”
// F0 "#PWR" 0 0 30 H I C CNN

#[derive(Debug)]
pub struct Field {
    pub i:i64,
    pub value:String,
    pub x:f64,
    pub y:f64,
    pub dimension:i64,
    pub orientation:schematic::Orientation,
    pub visible:bool,
    pub hjustify:schematic::Justify,
    pub vjustify:schematic::Justify,
    pub italic:bool,
    pub bold:bool,
    pub name:String,
}

impl SymbolLib {
    fn new() -> SymbolLib {
        SymbolLib {
            symbols:vec![]
        }
    }
}

impl fmt::Display for SymbolLib {
    fn fmt(&self, f: &mut fmt::Formatter) -> Result<(), fmt::Error> {
        try!(writeln!(f, "EESchema-LIBRARY Version 2.3"));
        try!(writeln!(f, "#encoding utf-8"));
        try!(writeln!(f, "#"));
        for i in &self.symbols {
            try!(write!(f, "{}", i))
        };
        writeln!(f, "#End Library")
    }
}

impl Symbol {
    fn new(name:String, reference:String) -> Symbol {
        Symbol {
            name:name,
            reference:reference,
            text_offset:0.0,
            draw_pinnumber:false,
            draw_pinname:false,
            unit_count:1,
            unit_locked:false,
            is_power:false,
            fields:vec![],
            draw:vec![],
        }
    }
}

impl fmt::Display for Symbol {
    fn fmt(&self, f: &mut fmt::Formatter) -> Result<(), fmt::Error> {
        try!(writeln!(f, "#"));
        try!(writeln!(f, "# {}", self.name));
        try!(writeln!(f, "#"));
        try!(writeln!(f, "DEF {} {} 0 {} {} {} {} {} {}",
                     self.name, self.reference, self.text_offset,
                     if self.draw_pinnumber { "Y" } else { "N" },
                     if self.draw_pinname { "Y" } else { "N" },
                     self.unit_count,
                     if self.unit_locked { "L" } else { "F" },
                     if self.is_power { "P" } else { "N" },
                     ));
        for field in &self.fields {
            try!(writeln!(f, "{}", field))
        };
        for draw in &self.draw {
            try!(writeln!(f, "{}", draw))
        }
        writeln!(f, "ENDDEF")
    }
}

impl Field {
    fn new() -> Field {
        Field {
            i:0,
            value:String::from(""),
            x:0.0,
            y:0.0,
            dimension:0,
            orientation:schematic::Orientation::Horizontal,
            visible:false,
            hjustify:schematic::Justify::Center,
            vjustify:schematic::Justify::Center,
            italic:false,
            bold:false,
            name:String::from(""),
        }
    }
}

impl fmt::Display for Field {
    fn fmt(&self, f: &mut fmt::Formatter) -> Result<(), fmt::Error> {
        try!(write!(f, "F{} {} {} {} {} {} {} {} {}{}{}",
               self.i, self.value, self.x, self.y, self.dimension,
               self.orientation,
               if self.visible { "V" } else { "I" },
               self.hjustify, self.vjustify,
               if self.italic { "I" } else { "N" },
               if self.bold { "I" } else { "N" },
                    ));
        if self.i > 3 {
            try!(write!(f, " \"{}\"", self.name))
        };
        Ok(())
    }
}

macro_rules! assume_line {
    ($s:expr, $exp:expr) => (
        if $s.eof() {
            return err("end of file reached")
        }
        if $s.here() != $exp {
            return Err(format!("expected '{}', got '{}'", $exp, $s.here()))
        }
        $s.i += 1;
    )
}

#[derive(Debug)]
struct ParseState {
    i:usize,
    v:Vec<String>,
}

impl ParseState {
    fn new(v2:Vec<&str>) -> ParseState {
        ParseState {
            i:0,
            v:v2.iter().map(|x| String::from(*x)).collect(),
        }
    }

    fn here(&self) -> String {
        let s =  (self.v[self.i]).clone();
        //println!("{}", s);
        s
    }

    fn next(&mut self) {
        self.i += 1;
    }

    fn eof(&self) -> bool {
        self.i >= self.v.len()
    }
}

fn char_at(s:&String, p:usize) -> char {
    let v:Vec<char> = s.chars().collect();
    v[..][p]
}

fn assume_string(e:&'static str, s:&String) -> ERes<()> {
    if String::from(e) != *s {
        return Err(format!("expecting: {}, actually: {}", e, s))
    }
    return Ok(())
}

fn i64_from_string(p:&ParseState, s:&String) -> ERes<i64> {
    match i64::from_str(&s[..]) {
        Ok(i) => Ok(i),
        _ => Err(format!("int parse error in {}; line: {}", s, p.here()))
    }
}

fn f64_from_string(p:&ParseState, s:&String) -> ERes<f64> {
    match f64::from_str(&s[..]) {
        Ok(i) => Ok(i),
        _ => Err(format!("float parse error in {}; line: {}", s, p.here()))
    }
}

fn bool_from_string(s:&String, t:&'static str, f:&'static str) -> ERes<bool> {
    if &s[..] == t {
        return Ok(true)
    }
    if &s[..] == f {
        return Ok(false)
    }
    Err(format!("unknown boolean {}, expected {} or {}", s, t, f))
}

fn bool_from<T: PartialEq + fmt::Display>(i:T, t:T, f:T) -> ERes<bool> {
    if i == t {
        return Ok(true)
    }
    if i == f {
        return Ok(false)
    }
    Err(format!("unknown boolean {}, expected {} or {}", i, t, f))
}

fn parse_symbol(p:&mut ParseState) -> ERes<Symbol> {
    assume_line!(p,"#");
    p.next(); // skip line like # name
    assume_line!(p,"#");
    let s = p.here();
    let v = &parse_split_quote_aware(&s);
    if v.len() != 10 {
        return Err(format!("unexpected elements in {}", s))
    }
    assume_string("DEF", &v[0]);
    let mut s = Symbol::new(v[1].clone(), v[2].clone());
    s.text_offset = try!(f64_from_string(p, &v[4]));
    s.draw_pinnumber = try!(bool_from_string(&v[5], "Y", "N"));
    s.draw_pinname = try!(bool_from_string(&v[6], "Y", "N"));
    s.unit_count = try!(i64_from_string(p, &v[7]));
    s.unit_locked = try!(bool_from_string(&v[8], "L", "F"));
    s.is_power = try!(bool_from_string(&v[9], "P", "N"));
    // TODO fields
    Ok(s)
}

fn parse(filename:Option<PathBuf>, s: &str) -> ERes<SymbolLib> {
    let mut lib = SymbolLib::new();
    let v:Vec<&str> = s.lines().collect();
    let p = &mut ParseState::new(v);
    assume_line!(p, "EESchema-LIBRARY Version 2.3");
    assume_line!(p, "#encoding utf-8");
    while !p.eof() {
        let s = try!(parse_symbol(p));
        lib.symbols.push(s)
    }
    Ok(lib)
}

pub fn parse_str(s:&str) -> ERes<SymbolLib> {
    parse(None, s)
}

pub fn parse_file(filename:&PathBuf) -> ERes<SymbolLib> {
    let name = filename.to_str().unwrap();
    let s = try!(read_file(name));
    parse(Some(filename.clone()), &s[..])
}
<|endoftext|>


<commit_before>use crate::{
    error::{ParameterError, ParameterErrorKind, UnsupportedError, UnsupportedErrorKind},
    ColorType, ImageError, ImageFormat, ImageResult,
};
use byteorder::{LittleEndian, ReadBytesExt, WriteBytesExt};
use std::convert::TryFrom;
use std::io::{Read, Write};

pub(crate) const ALPHA_BIT_MASK: u8 = 0b1111;
pub(crate) const SCREEN_ORIGIN_BIT_MASK: u8 = 0b10_0000;

pub(crate) enum ImageType {
    NoImageData = 0,
    /// Uncompressed images.
    RawColorMap = 1,
    RawTrueColor = 2,
    RawGrayScale = 3,
    /// Run length encoded images.
    RunColorMap = 9,
    RunTrueColor = 10,
    RunGrayScale = 11,
    Unknown,
}

impl ImageType {
    /// Create a new image type from a u8.
    pub(crate) fn new(img_type: u8) -> ImageType {
        match img_type {
            0 => ImageType::NoImageData,

            1 => ImageType::RawColorMap,
            2 => ImageType::RawTrueColor,
            3 => ImageType::RawGrayScale,

            9 => ImageType::RunColorMap,
            10 => ImageType::RunTrueColor,
            11 => ImageType::RunGrayScale,

            _ => ImageType::Unknown,
        }
    }

    /// Check if the image format uses colors as opposed to gray scale.
    pub(crate) fn is_color(&self) -> bool {
        match *self {
            ImageType::RawColorMap
            | ImageType::RawTrueColor
            | ImageType::RunTrueColor
            | ImageType::RunColorMap => true,
            _ => false,
        }
    }

    /// Does the image use a color map.
    pub(crate) fn is_color_mapped(&self) -> bool {
        match *self {
            ImageType::RawColorMap | ImageType::RunColorMap => true,
            _ => false,
        }
    }

    /// Is the image run length encoded.
    pub(crate) fn is_encoded(&self) -> bool {
        match *self {
            ImageType::RunColorMap | ImageType::RunTrueColor | ImageType::RunGrayScale => true,
            _ => false,
        }
    }
}

/// Header used by TGA image files.
#[derive(Debug, Default)]
pub(crate) struct Header {
    pub(crate) id_length: u8,      // length of ID string
    pub(crate) map_type: u8,       // color map type
    pub(crate) image_type: u8,     // image type code
    pub(crate) map_origin: u16,    // starting index of map
    pub(crate) map_length: u16,    // length of map
    pub(crate) map_entry_size: u8, // size of map entries in bits
    pub(crate) x_origin: u16,      // x-origin of image
    pub(crate) y_origin: u16,      // y-origin of image
    pub(crate) image_width: u16,   // width of image
    pub(crate) image_height: u16,  // height of image
    pub(crate) pixel_depth: u8,    // bits per pixel
    pub(crate) image_desc: u8,     // image descriptor
}

impl Header {
    /// Load the header with values from pixel information.
    pub(crate) fn from_pixel_info(
        color_type: ColorType,
        width: u32,
        height: u32,
    ) -> ImageResult<Header> {
        let mut header = Self::default();

        if width > 0 && height > 0 {
            header.image_width = u16::try_from(width).map_err(|_| {
                ImageError::Parameter(ParameterError::from_kind(
                    ParameterErrorKind::DimensionMismatch,
                ))
            })?;

            header.image_height = u16::try_from(height).map_err(|_| {
                ImageError::Parameter(ParameterError::from_kind(
                    ParameterErrorKind::DimensionMismatch,
                ))
            })?;

            let (num_alpha_bits, other_channel_bits, image_type) = match color_type {
                ColorType::Rgba8 | ColorType::Bgra8 => (8, 24, ImageType::RawTrueColor),
                ColorType::Rgb8 | ColorType::Bgr8 => (0, 24, ImageType::RawTrueColor),
                ColorType::La8 => (8, 8, ImageType::RawGrayScale),
                ColorType::L8 => (0, 8, ImageType::RawGrayScale),
                _ => {
                    return Err(ImageError::Unsupported(
                        UnsupportedError::from_format_and_kind(
                            ImageFormat::Tga.into(),
                            UnsupportedErrorKind::Color(color_type.into()),
                        ),
                    ))
                }
            };

            header.image_type = image_type as u8;
            header.pixel_depth = num_alpha_bits + other_channel_bits;
            header.image_desc = num_alpha_bits & ALPHA_BIT_MASK;
            header.image_desc |= SCREEN_ORIGIN_BIT_MASK; // Upper left origin.
        }

        Ok(header)
    }

    /// Load the header with values from the reader.
    pub(crate) fn from_reader(r: &mut dyn Read) -> ImageResult<Header> {
        Ok(Header {
            id_length: r.read_u8()?,
            map_type: r.read_u8()?,
            image_type: r.read_u8()?,
            map_origin: r.read_u16::<LittleEndian>()?,
            map_length: r.read_u16::<LittleEndian>()?,
            map_entry_size: r.read_u8()?,
            x_origin: r.read_u16::<LittleEndian>()?,
            y_origin: r.read_u16::<LittleEndian>()?,
            image_width: r.read_u16::<LittleEndian>()?,
            image_height: r.read_u16::<LittleEndian>()?,
            pixel_depth: r.read_u8()?,
            image_desc: r.read_u8()?,
        })
    }

    /// Write out the header values.
    pub(crate) fn write_to(&self, w: &mut dyn Write) -> ImageResult<()> {
        w.write_u8(self.id_length)?;
        w.write_u8(self.map_type)?;
        w.write_u8(self.image_type)?;
        w.write_u16::<LittleEndian>(self.map_origin)?;
        w.write_u16::<LittleEndian>(self.map_length)?;
        w.write_u8(self.map_entry_size)?;
        w.write_u16::<LittleEndian>(self.x_origin)?;
        w.write_u16::<LittleEndian>(self.y_origin)?;
        w.write_u16::<LittleEndian>(self.image_width)?;
        w.write_u16::<LittleEndian>(self.image_height)?;
        w.write_u8(self.pixel_depth)?;
        w.write_u8(self.image_desc)?;
        Ok(())
    }
}
<commit_msg>Switched TGA header methods to using Self.<commit_after>use crate::{
    error::{ParameterError, ParameterErrorKind, UnsupportedError, UnsupportedErrorKind},
    ColorType, ImageError, ImageFormat, ImageResult,
};
use byteorder::{LittleEndian, ReadBytesExt, WriteBytesExt};
use std::convert::TryFrom;
use std::io::{Read, Write};

pub(crate) const ALPHA_BIT_MASK: u8 = 0b1111;
pub(crate) const SCREEN_ORIGIN_BIT_MASK: u8 = 0b10_0000;

pub(crate) enum ImageType {
    NoImageData = 0,
    /// Uncompressed images.
    RawColorMap = 1,
    RawTrueColor = 2,
    RawGrayScale = 3,
    /// Run length encoded images.
    RunColorMap = 9,
    RunTrueColor = 10,
    RunGrayScale = 11,
    Unknown,
}

impl ImageType {
    /// Create a new image type from a u8.
    pub(crate) fn new(img_type: u8) -> ImageType {
        match img_type {
            0 => ImageType::NoImageData,

            1 => ImageType::RawColorMap,
            2 => ImageType::RawTrueColor,
            3 => ImageType::RawGrayScale,

            9 => ImageType::RunColorMap,
            10 => ImageType::RunTrueColor,
            11 => ImageType::RunGrayScale,

            _ => ImageType::Unknown,
        }
    }

    /// Check if the image format uses colors as opposed to gray scale.
    pub(crate) fn is_color(&self) -> bool {
        match *self {
            ImageType::RawColorMap
            | ImageType::RawTrueColor
            | ImageType::RunTrueColor
            | ImageType::RunColorMap => true,
            _ => false,
        }
    }

    /// Does the image use a color map.
    pub(crate) fn is_color_mapped(&self) -> bool {
        match *self {
            ImageType::RawColorMap | ImageType::RunColorMap => true,
            _ => false,
        }
    }

    /// Is the image run length encoded.
    pub(crate) fn is_encoded(&self) -> bool {
        match *self {
            ImageType::RunColorMap | ImageType::RunTrueColor | ImageType::RunGrayScale => true,
            _ => false,
        }
    }
}

/// Header used by TGA image files.
#[derive(Debug, Default)]
pub(crate) struct Header {
    pub(crate) id_length: u8,      // length of ID string
    pub(crate) map_type: u8,       // color map type
    pub(crate) image_type: u8,     // image type code
    pub(crate) map_origin: u16,    // starting index of map
    pub(crate) map_length: u16,    // length of map
    pub(crate) map_entry_size: u8, // size of map entries in bits
    pub(crate) x_origin: u16,      // x-origin of image
    pub(crate) y_origin: u16,      // y-origin of image
    pub(crate) image_width: u16,   // width of image
    pub(crate) image_height: u16,  // height of image
    pub(crate) pixel_depth: u8,    // bits per pixel
    pub(crate) image_desc: u8,     // image descriptor
}

impl Header {
    /// Load the header with values from pixel information.
    pub(crate) fn from_pixel_info(
        color_type: ColorType,
        width: u32,
        height: u32,
    ) -> ImageResult<Self> {
        let mut header = Self::default();

        if width > 0 && height > 0 {
            header.image_width = u16::try_from(width).map_err(|_| {
                ImageError::Parameter(ParameterError::from_kind(
                    ParameterErrorKind::DimensionMismatch,
                ))
            })?;

            header.image_height = u16::try_from(height).map_err(|_| {
                ImageError::Parameter(ParameterError::from_kind(
                    ParameterErrorKind::DimensionMismatch,
                ))
            })?;

            let (num_alpha_bits, other_channel_bits, image_type) = match color_type {
                ColorType::Rgba8 | ColorType::Bgra8 => (8, 24, ImageType::RawTrueColor),
                ColorType::Rgb8 | ColorType::Bgr8 => (0, 24, ImageType::RawTrueColor),
                ColorType::La8 => (8, 8, ImageType::RawGrayScale),
                ColorType::L8 => (0, 8, ImageType::RawGrayScale),
                _ => {
                    return Err(ImageError::Unsupported(
                        UnsupportedError::from_format_and_kind(
                            ImageFormat::Tga.into(),
                            UnsupportedErrorKind::Color(color_type.into()),
                        ),
                    ))
                }
            };

            header.image_type = image_type as u8;
            header.pixel_depth = num_alpha_bits + other_channel_bits;
            header.image_desc = num_alpha_bits & ALPHA_BIT_MASK;
            header.image_desc |= SCREEN_ORIGIN_BIT_MASK; // Upper left origin.
        }

        Ok(header)
    }

    /// Load the header with values from the reader.
    pub(crate) fn from_reader(r: &mut dyn Read) -> ImageResult<Self> {
        Ok(Self {
            id_length: r.read_u8()?,
            map_type: r.read_u8()?,
            image_type: r.read_u8()?,
            map_origin: r.read_u16::<LittleEndian>()?,
            map_length: r.read_u16::<LittleEndian>()?,
            map_entry_size: r.read_u8()?,
            x_origin: r.read_u16::<LittleEndian>()?,
            y_origin: r.read_u16::<LittleEndian>()?,
            image_width: r.read_u16::<LittleEndian>()?,
            image_height: r.read_u16::<LittleEndian>()?,
            pixel_depth: r.read_u8()?,
            image_desc: r.read_u8()?,
        })
    }

    /// Write out the header values.
    pub(crate) fn write_to(&self, w: &mut dyn Write) -> ImageResult<()> {
        w.write_u8(self.id_length)?;
        w.write_u8(self.map_type)?;
        w.write_u8(self.image_type)?;
        w.write_u16::<LittleEndian>(self.map_origin)?;
        w.write_u16::<LittleEndian>(self.map_length)?;
        w.write_u8(self.map_entry_size)?;
        w.write_u16::<LittleEndian>(self.x_origin)?;
        w.write_u16::<LittleEndian>(self.y_origin)?;
        w.write_u16::<LittleEndian>(self.image_width)?;
        w.write_u16::<LittleEndian>(self.image_height)?;
        w.write_u8(self.pixel_depth)?;
        w.write_u8(self.image_desc)?;
        Ok(())
    }
}
<|endoftext|>


<commit_before>use std::ops::{Index, IndexMut};

/// A 2 dimensional matrix in row-major order backed by a contiguous slice.
#[derive(Debug)]
pub struct Vec2D<T> {
    data: Box<[T]>,
    cols: usize,
}

impl<T> Vec2D<T> {
    /// Initialize a grid of size (`rows`, `cols`) with the given data element.
    pub fn init(data: T, size: (usize, usize)) -> Vec2D<T>
    where
        T: Clone,
    {
        let (rows, cols) = size;
        let len = rows
            .checked_mul(cols)
            .unwrap_or_else(|| panic!("{} rows by {} cols exceeds usize::MAX", rows, cols));
        Vec2D {
            data: vec![data; len].into_boxed_slice(),
            cols,
        }
    }

    /// Fills the grid with elements by cloning `value`.
    pub fn fill(&mut self, value: T)
    where
        T: Clone,
    {
        self.data.fill(value)
    }
}

impl<T> Index<usize> for Vec2D<T> {
    type Output = [T];

    #[inline]
    fn index(&self, row: usize) -> &Self::Output {
        let start_row = row
            .checked_mul(self.cols)
            .unwrap_or_else(|| panic!("{} row by {} cols exceeds usize::MAX", row, self.cols));
        &self.data[start_row..start_row + self.cols]
    }
}

impl<T> IndexMut<usize> for Vec2D<T> {
    #[inline]
    fn index_mut(&mut self, row: usize) -> &mut Self::Output {
        let start_row = row
            .checked_mul(self.cols)
            .unwrap_or_else(|| panic!("{} row by {} cols exceeds usize::MAX", row, self.cols));
        &mut self.data[start_row..start_row + self.cols]
    }
}

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

    #[test]
    fn init() {
        let vec2d = Vec2D::init(1, (2, 3));
        assert_eq!(vec2d[0], [1, 1, 1]);
        assert_eq!(vec2d[1], [1, 1, 1]);
    }

    #[test]
    #[should_panic]
    fn init_overflow() {
        Vec2D::init(1, (usize::MAX, usize::MAX));
    }

    #[test]
    fn fill() {
        let mut vec2d = Vec2D::init(0, (2, 3));
        vec2d.fill(7);
        assert_eq!(vec2d[0], [7, 7, 7]);
        assert_eq!(vec2d[1], [7, 7, 7]);
    }

    #[test]
    fn index() {
        let vec2d = Vec2D {
            data: vec![0, 1, 2, 3, 4, 5, 6, 7].into_boxed_slice(),
            cols: 2,
        };
        assert_eq!(vec2d[0], [0, 1]);
        assert_eq!(vec2d[1], [2, 3]);
        assert_eq!(vec2d[2], [4, 5]);
        assert_eq!(vec2d[3], [6, 7]);
    }

    #[test]
    fn index_mut() {
        let mut vec2d = Vec2D {
            data: vec![0, 1, 2, 3, 4, 5, 6, 7].into_boxed_slice(),
            cols: 2,
        };

        vec2d[1][1] = 9;
        assert_eq!(vec2d[0], [0, 1]);
        // 1,1 should be 9.
        assert_eq!(vec2d[1], [2, 9]);
        assert_eq!(vec2d[2], [4, 5]);
        assert_eq!(vec2d[3], [6, 7]);
    }

    #[test]
    #[should_panic]
    fn index_out_of_bounds() {
        let vec2d = Vec2D::init(1, (2, 3));
        let _x = vec2d[2][3];
    }

    #[test]
    #[should_panic]
    fn index_out_of_bounds_vec_edge() {
        let vec2d = Vec2D::init(1, (2, 3));
        let _x = vec2d[1][3];
    }

    #[test]
    #[should_panic]
    fn column_out_of_bounds() {
        let vec2d = Vec2D::init(1, (2, 3));
        let _x = vec2d[0][3];
    }

    #[test]
    #[should_panic]
    fn row_out_of_bounds() {
        let vec2d = Vec2D::init(1, (2, 3));
        let _x = vec2d[2][0];
    }

    #[test]
    #[should_panic]
    fn index_mul_overflow() {
        // Matrix with 4 columns.
        let matrix = Vec2D::init(0, (3, 4));
        // 2^{usize.numbits() - 2}.
        let row = (usize::MAX / 4) + 1;
        // Returns the same as matrix[0] if overflow is not caught.
        let _ = matrix[row];
    }

    #[test]
    #[should_panic]
    fn index_add_overflow() {
        // Matrix with 5 columns.
        let matrix = Vec2D::init(0, (3, 5));
        // Somehow, as long as numbits(usize) is a multiple of 4, then 5 divides usize::MAX.
        // This is clear in hexadecimal: usize::MAX is 0xFFF...F and usize::MAX / 5 is 0x333...3.
        let row = usize::MAX / 5;
        // This will therefore try to index data[usize::MAX..4].
        let _ = matrix[row];
    }

    #[test]
    #[should_panic]
    fn indexmut_out_of_bounds() {
        let mut vec2d = Vec2D::init(1, (2, 3));
        vec2d[2][3] = 0;
    }

    #[test]
    #[should_panic]
    fn indexmut_out_of_bounds_vec_edge() {
        let mut vec2d = Vec2D::init(1, (2, 3));
        vec2d[1][3] = 0;
    }

    #[test]
    #[should_panic]
    fn mut_column_out_of_bounds() {
        let mut vec2d = Vec2D::init(1, (2, 3));
        vec2d[0][3] = 0;
    }

    #[test]
    #[should_panic]
    fn mut_row_out_of_bounds() {
        let mut vec2d = Vec2D::init(1, (2, 3));
        vec2d[2][0] = 0;
    }

    #[test]
    #[should_panic]
    fn index_mut_mul_overflow() {
        // Matrix with 4 columns.
        let mut matrix = Vec2D::init(0, (3, 4));
        // 2^{usize.numbits() - 2}.
        let row = (usize::MAX / 4) + 1;
        // Returns the same as matrix[0] if overflow is not caught.
        matrix[row][0] = 9;
    }

    #[test]
    #[should_panic]
    fn index_mut_add_overflow() {
        // Matrix with 5 columns.
        let mut matrix = Vec2D::init(0, (3, 5));
        // Somehow, as long as numbits(usize) is a multiple of 4, then 5 divides usize::MAX.
        // This is clear in hexadecimal: usize::MAX is 0xFFF...F and usize::MAX / 5 is 0x333...3.
        let row = usize::MAX / 5;
        // This will therefore try to index data[usize::MAX..4].
        matrix[row][0] = 9;
    }
}
<commit_msg>Unify test names in src/util/vec2d.rs.<commit_after>use std::ops::{Index, IndexMut};

/// A 2 dimensional matrix in row-major order backed by a contiguous slice.
#[derive(Debug)]
pub struct Vec2D<T> {
    data: Box<[T]>,
    cols: usize,
}

impl<T> Vec2D<T> {
    /// Initialize a grid of size (`rows`, `cols`) with the given data element.
    pub fn init(data: T, size: (usize, usize)) -> Vec2D<T>
    where
        T: Clone,
    {
        let (rows, cols) = size;
        let len = rows
            .checked_mul(cols)
            .unwrap_or_else(|| panic!("{} rows by {} cols exceeds usize::MAX", rows, cols));
        Vec2D {
            data: vec![data; len].into_boxed_slice(),
            cols,
        }
    }

    /// Fills the grid with elements by cloning `value`.
    pub fn fill(&mut self, value: T)
    where
        T: Clone,
    {
        self.data.fill(value)
    }
}

impl<T> Index<usize> for Vec2D<T> {
    type Output = [T];

    #[inline]
    fn index(&self, row: usize) -> &Self::Output {
        let start_row = row
            .checked_mul(self.cols)
            .unwrap_or_else(|| panic!("{} row by {} cols exceeds usize::MAX", row, self.cols));
        &self.data[start_row..start_row + self.cols]
    }
}

impl<T> IndexMut<usize> for Vec2D<T> {
    #[inline]
    fn index_mut(&mut self, row: usize) -> &mut Self::Output {
        let start_row = row
            .checked_mul(self.cols)
            .unwrap_or_else(|| panic!("{} row by {} cols exceeds usize::MAX", row, self.cols));
        &mut self.data[start_row..start_row + self.cols]
    }
}

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

    #[test]
    fn init() {
        let vec2d = Vec2D::init(1, (2, 3));
        assert_eq!(vec2d[0], [1, 1, 1]);
        assert_eq!(vec2d[1], [1, 1, 1]);
    }

    #[test]
    #[should_panic]
    fn init_overflow() {
        Vec2D::init(1, (usize::MAX, usize::MAX));
    }

    #[test]
    fn fill() {
        let mut vec2d = Vec2D::init(0, (2, 3));
        vec2d.fill(7);
        assert_eq!(vec2d[0], [7, 7, 7]);
        assert_eq!(vec2d[1], [7, 7, 7]);
    }

    #[test]
    fn index() {
        let vec2d = Vec2D {
            data: vec![0, 1, 2, 3, 4, 5, 6, 7].into_boxed_slice(),
            cols: 2,
        };
        assert_eq!(vec2d[0], [0, 1]);
        assert_eq!(vec2d[1], [2, 3]);
        assert_eq!(vec2d[2], [4, 5]);
        assert_eq!(vec2d[3], [6, 7]);
    }

    #[test]
    fn indexmut() {
        let mut vec2d = Vec2D {
            data: vec![0, 1, 2, 3, 4, 5, 6, 7].into_boxed_slice(),
            cols: 2,
        };

        vec2d[1][1] = 9;
        assert_eq!(vec2d[0], [0, 1]);
        // (1, 1) should be 9.
        assert_eq!(vec2d[1], [2, 9]);
        assert_eq!(vec2d[2], [4, 5]);
        assert_eq!(vec2d[3], [6, 7]);
    }

    #[test]
    #[should_panic]
    fn index_out_of_bounds() {
        let vec2d = Vec2D::init(1, (2, 3));
        let _x = vec2d[2][3];
    }

    #[test]
    #[should_panic]
    fn index_out_of_bounds_vec_edge() {
        let vec2d = Vec2D::init(1, (2, 3));
        let _x = vec2d[1][3];
    }

    #[test]
    #[should_panic]
    fn index_column_out_of_bounds() {
        let vec2d = Vec2D::init(1, (2, 3));
        let _x = vec2d[0][3];
    }

    #[test]
    #[should_panic]
    fn index_row_out_of_bounds() {
        let vec2d = Vec2D::init(1, (2, 3));
        let _x = vec2d[2][0];
    }

    #[test]
    #[should_panic]
    fn index_mul_overflow() {
        // Matrix with 4 columns.
        let matrix = Vec2D::init(0, (3, 4));
        // 2^{usize.numbits() - 2}.
        let row = (usize::MAX / 4) + 1;
        // Returns the same as matrix[0] if overflow is not caught.
        let _ = matrix[row];
    }

    #[test]
    #[should_panic]
    fn index_add_overflow() {
        // Matrix with 5 columns.
        let matrix = Vec2D::init(0, (3, 5));
        // Somehow, as long as numbits(usize) is a multiple of 4, then 5 divides usize::MAX.
        // This is clear in hexadecimal: usize::MAX is 0xFFF...F and usize::MAX / 5 is 0x333...3.
        let row = usize::MAX / 5;
        // This will therefore try to index data[usize::MAX..4].
        let _ = matrix[row];
    }

    #[test]
    #[should_panic]
    fn indexmut_out_of_bounds() {
        let mut vec2d = Vec2D::init(1, (2, 3));
        vec2d[2][3] = 0;
    }

    #[test]
    #[should_panic]
    fn indexmut_out_of_bounds_vec_edge() {
        let mut vec2d = Vec2D::init(1, (2, 3));
        vec2d[1][3] = 0;
    }

    #[test]
    #[should_panic]
    fn indexmut_column_out_of_bounds() {
        let mut vec2d = Vec2D::init(1, (2, 3));
        vec2d[0][3] = 0;
    }

    #[test]
    #[should_panic]
    fn indexmut_row_out_of_bounds() {
        let mut vec2d = Vec2D::init(1, (2, 3));
        vec2d[2][0] = 0;
    }

    #[test]
    #[should_panic]
    fn indexmut_mul_overflow() {
        // Matrix with 4 columns.
        let mut matrix = Vec2D::init(0, (3, 4));
        // 2^{usize.numbits() - 2}.
        let row = (usize::MAX / 4) + 1;
        // Returns the same as matrix[0] if overflow is not caught.
        matrix[row][0] = 9;
    }

    #[test]
    #[should_panic]
    fn indexmut_add_overflow() {
        // Matrix with 5 columns.
        let mut matrix = Vec2D::init(0, (3, 5));
        // Somehow, as long as numbits(usize) is a multiple of 4, then 5 divides usize::MAX.
        // This is clear in hexadecimal: usize::MAX is 0xFFF...F and usize::MAX / 5 is 0x333...3.
        let row = usize::MAX / 5;
        // This will therefore try to index data[usize::MAX..4].
        matrix[row][0] = 9;
    }
}
<|endoftext|>


<commit_before>use crate::crypto::SigningAlgorithm;
use crate::utils::{
    base64url,
    pem::{self, ParsedKeyPair},
    SecureRandom,
};
use err_derive::Error;
use ring::{
    digest,
    io::Positive,
    signature::{self, Ed25519KeyPair, KeyPair, RsaKeyPair},
};
use serde_json::{json, Value as JsonValue};
use std::ffi::OsString;
use std::process::{Command, Stdio};

#[derive(Debug, Error)]
pub enum SignError {
    #[error(display = "unsupported signing algorithm {}", _0)]
    UnsupportedAlgorithm(SigningAlgorithm),
    #[error(display = "unspecified signing error")]
    Unspecified,
}

impl From<ring::error::Unspecified> for SignError {
    fn from(_: ring::error::Unspecified) -> Self {
        Self::Unspecified
    }
}

/// A named key pair, for use in JWS signing.
pub struct NamedKeyPair<T: KeyPairExt> {
    pub kid: String,
    pub key_pair: T,
}

impl<T: KeyPairExt> NamedKeyPair<T> {
    /// Create a JSON Web Signature (JWS) for the given JSON structure.
    pub fn sign_jws(&self, payload: &JsonValue, rng: &SecureRandom) -> Result<String, SignError> {
        self.key_pair.sign_jws(&self.kid, payload, rng)
    }

    /// Return JSON represenation of the public key for use in JWK key sets.
    pub fn public_jwk(&self) -> JsonValue {
        self.key_pair.public_jwk(&self.kid)
    }
}

impl<T: KeyPairExt> From<T> for NamedKeyPair<T> {
    fn from(key_pair: T) -> Self {
        let kid = key_pair.generate_kid();
        Self { kid, key_pair }
    }
}

/// Additional `KeyPair` methods we implement for key pair types we support.
pub trait KeyPairExt {
    /// Generate an ID for the key by hashing the public components.
    ///
    /// Note that this hash is not a standard format, but that's okay, because it's only used as
    /// a simple identifier in JWKs.
    fn generate_kid(&self) -> String;

    /// Get the signing algorithm for this key type.
    fn signing_alg(&self) -> SigningAlgorithm;

    /// Create a JSON Web Signature (JWS) for the given JSON structure.
    fn sign_jws(
        &self,
        kid: &str,
        payload: &JsonValue,
        rng: &SecureRandom,
    ) -> Result<String, SignError>;

    /// Return JSON represenation of the public key for use in JWK key sets.
    fn public_jwk(&self, kid: &str) -> JsonValue;
}

impl KeyPairExt for Ed25519KeyPair {
    fn generate_kid(&self) -> String {
        let mut ctx = digest::Context::new(&digest::SHA256);
        ctx.update(b"ed25519.");
        ctx.update(self.public_key().as_ref());
        base64url::encode(&ctx.finish())
    }

    fn signing_alg(&self) -> SigningAlgorithm {
        SigningAlgorithm::EdDsa
    }

    fn sign_jws(
        &self,
        kid: &str,
        payload: &JsonValue,
        _rng: &SecureRandom,
    ) -> Result<String, SignError> {
        let header = json!({ "kid": kid, "alg": "EdDSA" }).to_string();
        let mut data = String::new();
        data.push_str(&base64url::encode(&header));
        data.push('.');
        data.push_str(&base64url::encode(&payload.to_string()));
        let sig = self.sign(data.as_bytes());
        data.push('.');
        data.push_str(&base64url::encode(&sig));
        Ok(data)
    }

    fn public_jwk(&self, kid: &str) -> JsonValue {
        let public = self.public_key();
        json!({
            "kty": "OKP",
            "alg": "EdDSA",
            "crv": "Ed25519",
            "use": "sig",
            "kid": &kid,
            "x": base64url::encode(&public),
        })
    }
}

impl KeyPairExt for RsaKeyPair {
    fn generate_kid(&self) -> String {
        let public = self.public_key();
        let (n, e) = (public.modulus(), public.exponent());
        let mut ctx = digest::Context::new(&digest::SHA256);
        ctx.update(b"rsa.");
        ctx.update(e.big_endian_without_leading_zero());
        ctx.update(b".");
        ctx.update(n.big_endian_without_leading_zero());
        base64url::encode(&ctx.finish())
    }

    fn signing_alg(&self) -> SigningAlgorithm {
        SigningAlgorithm::Rs256
    }

    fn sign_jws(
        &self,
        kid: &str,
        payload: &JsonValue,
        rng: &SecureRandom,
    ) -> Result<String, SignError> {
        let header = json!({ "kid": kid, "alg": "RS256" }).to_string();
        let mut data = String::new();
        data.push_str(&base64url::encode(&header));
        data.push('.');
        data.push_str(&base64url::encode(&payload.to_string()));
        let mut sig = vec![0; self.public_modulus_len()];
        self.sign(
            &signature::RSA_PKCS1_SHA256,
            &rng.generator,
            data.as_bytes(),
            &mut sig,
        )?;
        data.push('.');
        data.push_str(&base64url::encode(&sig));
        Ok(data)
    }

    fn public_jwk(&self, kid: &str) -> JsonValue {
        fn json_big_num(v: Positive) -> String {
            base64url::encode(v.big_endian_without_leading_zero())
        }

        let public = self.public_key();
        let (n, e) = (public.modulus(), public.exponent());
        json!({
            "kty": "RSA",
            "alg": "RS256",
            "use": "sig",
            "kid": &kid,
            "n": json_big_num(n),
            "e": json_big_num(e),
        })
    }
}

/// Trait for key pair types we can generate.
pub trait GeneratedKeyPair: KeyPairExt + Sized {
    /// Configuration required for generating a key pair.
    type Config;

    /// Generate a new key pair.
    ///
    /// If this fails, we panic, because it may happen at an arbitrary moment at run-time.
    fn generate(config: Self::Config) -> String;

    /// Convert a ParsedKeyPair, if it is of the correct type.
    fn from_parsed(parsed: ParsedKeyPair) -> Option<Self>;
}

impl GeneratedKeyPair for Ed25519KeyPair {
    type Config = SecureRandom;

    fn generate(config: Self::Config) -> String {
        let doc =
            Self::generate_pkcs8(&config.generator).expect("could not generate Ed25519 key pair");
        pem::from_der(doc.as_ref())
    }

    fn from_parsed(parsed: ParsedKeyPair) -> Option<Self> {
        match parsed {
            ParsedKeyPair::Ed25519(inner) => Some(inner),
            _ => None,
        }
    }
}

impl GeneratedKeyPair for RsaKeyPair {
    type Config = Vec<String>;

    fn generate(config: Vec<String>) -> String {
        let mut args: Vec<OsString> = config.iter().map(|arg| arg.into()).collect();
        let program = args.remove(0);
        let output = Command::new(program)
            .args(args)
            .stdin(Stdio::null())
            .stdout(Stdio::piped())
            .output()
            .expect("Failed to run command to generate RSA key");
        if !output.status.success() {
            panic!(
                "Command to generate RSA key failed with status {}",
                output.status
            );
        }
        String::from_utf8(output.stdout).expect("Generated RSA is not UTF-8")
    }

    fn from_parsed(parsed: ParsedKeyPair) -> Option<Self> {
        match parsed {
            ParsedKeyPair::Rsa(inner) => Some(inner),
            _ => None,
        }
    }
}
<commit_msg>Make RSA key IDs backwards compatible<commit_after>use crate::crypto::SigningAlgorithm;
use crate::utils::{
    base64url,
    pem::{self, ParsedKeyPair},
    SecureRandom,
};
use err_derive::Error;
use ring::{
    digest,
    io::Positive,
    signature::{self, Ed25519KeyPair, KeyPair, RsaKeyPair},
};
use serde_json::{json, Value as JsonValue};
use std::ffi::OsString;
use std::process::{Command, Stdio};

#[derive(Debug, Error)]
pub enum SignError {
    #[error(display = "unsupported signing algorithm {}", _0)]
    UnsupportedAlgorithm(SigningAlgorithm),
    #[error(display = "unspecified signing error")]
    Unspecified,
}

impl From<ring::error::Unspecified> for SignError {
    fn from(_: ring::error::Unspecified) -> Self {
        Self::Unspecified
    }
}

/// A named key pair, for use in JWS signing.
pub struct NamedKeyPair<T: KeyPairExt> {
    pub kid: String,
    pub key_pair: T,
}

impl<T: KeyPairExt> NamedKeyPair<T> {
    /// Create a JSON Web Signature (JWS) for the given JSON structure.
    pub fn sign_jws(&self, payload: &JsonValue, rng: &SecureRandom) -> Result<String, SignError> {
        self.key_pair.sign_jws(&self.kid, payload, rng)
    }

    /// Return JSON represenation of the public key for use in JWK key sets.
    pub fn public_jwk(&self) -> JsonValue {
        self.key_pair.public_jwk(&self.kid)
    }
}

impl<T: KeyPairExt> From<T> for NamedKeyPair<T> {
    fn from(key_pair: T) -> Self {
        let kid = key_pair.generate_kid();
        Self { kid, key_pair }
    }
}

/// Additional `KeyPair` methods we implement for key pair types we support.
pub trait KeyPairExt {
    /// Generate an ID for the key by hashing the public components.
    ///
    /// Note that this hash is not a standard format, but that's okay, because it's only used as
    /// a simple identifier in JWKs.
    fn generate_kid(&self) -> String;

    /// Get the signing algorithm for this key type.
    fn signing_alg(&self) -> SigningAlgorithm;

    /// Create a JSON Web Signature (JWS) for the given JSON structure.
    fn sign_jws(
        &self,
        kid: &str,
        payload: &JsonValue,
        rng: &SecureRandom,
    ) -> Result<String, SignError>;

    /// Return JSON represenation of the public key for use in JWK key sets.
    fn public_jwk(&self, kid: &str) -> JsonValue;
}

impl KeyPairExt for Ed25519KeyPair {
    fn generate_kid(&self) -> String {
        let mut ctx = digest::Context::new(&digest::SHA256);
        ctx.update(b"ed25519.");
        ctx.update(self.public_key().as_ref());
        base64url::encode(&ctx.finish())
    }

    fn signing_alg(&self) -> SigningAlgorithm {
        SigningAlgorithm::EdDsa
    }

    fn sign_jws(
        &self,
        kid: &str,
        payload: &JsonValue,
        _rng: &SecureRandom,
    ) -> Result<String, SignError> {
        let header = json!({ "kid": kid, "alg": "EdDSA" }).to_string();
        let mut data = String::new();
        data.push_str(&base64url::encode(&header));
        data.push('.');
        data.push_str(&base64url::encode(&payload.to_string()));
        let sig = self.sign(data.as_bytes());
        data.push('.');
        data.push_str(&base64url::encode(&sig));
        Ok(data)
    }

    fn public_jwk(&self, kid: &str) -> JsonValue {
        let public = self.public_key();
        json!({
            "kty": "OKP",
            "alg": "EdDSA",
            "crv": "Ed25519",
            "use": "sig",
            "kid": &kid,
            "x": base64url::encode(&public),
        })
    }
}

impl KeyPairExt for RsaKeyPair {
    fn generate_kid(&self) -> String {
        let public = self.public_key();
        let (n, e) = (public.modulus(), public.exponent());
        let mut ctx = digest::Context::new(&digest::SHA256);
        ctx.update(e.big_endian_without_leading_zero());
        ctx.update(b".");
        ctx.update(n.big_endian_without_leading_zero());
        base64url::encode(&ctx.finish())
    }

    fn signing_alg(&self) -> SigningAlgorithm {
        SigningAlgorithm::Rs256
    }

    fn sign_jws(
        &self,
        kid: &str,
        payload: &JsonValue,
        rng: &SecureRandom,
    ) -> Result<String, SignError> {
        let header = json!({ "kid": kid, "alg": "RS256" }).to_string();
        let mut data = String::new();
        data.push_str(&base64url::encode(&header));
        data.push('.');
        data.push_str(&base64url::encode(&payload.to_string()));
        let mut sig = vec![0; self.public_modulus_len()];
        self.sign(
            &signature::RSA_PKCS1_SHA256,
            &rng.generator,
            data.as_bytes(),
            &mut sig,
        )?;
        data.push('.');
        data.push_str(&base64url::encode(&sig));
        Ok(data)
    }

    fn public_jwk(&self, kid: &str) -> JsonValue {
        fn json_big_num(v: Positive) -> String {
            base64url::encode(v.big_endian_without_leading_zero())
        }

        let public = self.public_key();
        let (n, e) = (public.modulus(), public.exponent());
        json!({
            "kty": "RSA",
            "alg": "RS256",
            "use": "sig",
            "kid": &kid,
            "n": json_big_num(n),
            "e": json_big_num(e),
        })
    }
}

/// Trait for key pair types we can generate.
pub trait GeneratedKeyPair: KeyPairExt + Sized {
    /// Configuration required for generating a key pair.
    type Config;

    /// Generate a new key pair.
    ///
    /// If this fails, we panic, because it may happen at an arbitrary moment at run-time.
    fn generate(config: Self::Config) -> String;

    /// Convert a ParsedKeyPair, if it is of the correct type.
    fn from_parsed(parsed: ParsedKeyPair) -> Option<Self>;
}

impl GeneratedKeyPair for Ed25519KeyPair {
    type Config = SecureRandom;

    fn generate(config: Self::Config) -> String {
        let doc =
            Self::generate_pkcs8(&config.generator).expect("could not generate Ed25519 key pair");
        pem::from_der(doc.as_ref())
    }

    fn from_parsed(parsed: ParsedKeyPair) -> Option<Self> {
        match parsed {
            ParsedKeyPair::Ed25519(inner) => Some(inner),
            _ => None,
        }
    }
}

impl GeneratedKeyPair for RsaKeyPair {
    type Config = Vec<String>;

    fn generate(config: Vec<String>) -> String {
        let mut args: Vec<OsString> = config.iter().map(|arg| arg.into()).collect();
        let program = args.remove(0);
        let output = Command::new(program)
            .args(args)
            .stdin(Stdio::null())
            .stdout(Stdio::piped())
            .output()
            .expect("Failed to run command to generate RSA key");
        if !output.status.success() {
            panic!(
                "Command to generate RSA key failed with status {}",
                output.status
            );
        }
        String::from_utf8(output.stdout).expect("Generated RSA is not UTF-8")
    }

    fn from_parsed(parsed: ParsedKeyPair) -> Option<Self> {
        match parsed {
            ParsedKeyPair::Rsa(inner) => Some(inner),
            _ => None,
        }
    }
}
<|endoftext|>