<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,
}

#[derive(Copy, Clone)]
pub enum Prompt {
    EpisodeCompleted,
    UpdateScore,
    SeriesCompleted,
    AlreadyCompleted,
    BeginRewatch,
    ResumePausedSeries,
    PauseSeries(Status),
}

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()
    }

    pub fn prepare_initial_state(&mut self) -> Result<(), SeriesError> {
        match self.season.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 !self.has_unwatched_episodes() {
                    self.prompt(Prompt::SeriesCompleted)?;
                }

                Ok(())
            }
            Status::PlanToWatch => self.set_list_entry_status(Status::Watching),
            Status::Completed => self.prompt(Prompt::AlreadyCompleted),
            Status::OnHold | Status::Dropped => self.prompt(Prompt::ResumePausedSeries),
        }
    }

    pub fn prompt(&mut self, prompt: Prompt) -> Result<(), SeriesError> {
        let state = &mut self.season.state;

        match prompt {
            Prompt::EpisodeCompleted => {
                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()?;
                self.prompt_next_episode_options()
            }
            Prompt::UpdateScore => {
                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(()),
                    Ok(input) => input,
                    Err(err) => {
                        eprintln!("failed to read score input: {}", err);
                        return Ok(());
                    }
                };

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

                Ok(())
            }
            Prompt::SeriesCompleted => {
                println!("[{}] completed!", state.info.title);
                self.set_list_entry_status(Status::Completed)?;
                self.prompt_series_completed_options()
            }
            Prompt::AlreadyCompleted => {
                println!("[{}] already completed", state.info.title);
                self.prompt_series_completed_options()
            }
            Prompt::BeginRewatch => {
                println!("[{}] starting rewatch", state.info.title);
                println!("do you want to reset the start and end dates? (Y/n)");

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

                self.set_list_entry_status(Status::Rewatching)
            }
            Prompt::ResumePausedSeries => {
                println!(
                    "[{}] was previously put on hold or dropped",
                    state.info.title
                );

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

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

                self.set_list_entry_status(Status::Watching)
            }
            Prompt::PauseSeries(status) => {
                self.set_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();
                }

                Ok(())
            }
        }
    }

    pub 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.prompt(Prompt::PauseSeries(status))?;
                Err(SeriesError::RequestExit)
            }
            "r" => {
                self.prompt(Prompt::UpdateScore)?;
                self.prompt_next_episode_options()
            }
            "x" => Err(SeriesError::RequestExit),
            _ => Ok(()),
        }
    }

    pub 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(Prompt::UpdateScore)?;
                self.prompt_series_completed_options()
            }
            "w" => {
                self.prompt(Prompt::BeginRewatch)?;
                self.play_all_episodes()
            }
            "d" => {
                self.dir.try_remove_dir();
                Err(SeriesError::RequestExit)
            }
            "x" => Err(SeriesError::RequestExit),
            _ => Ok(()),
        }
    }

    pub fn prompt_series_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:", self.season.state.info.title);
        println!(
            "\t[r] rate{}\n\t[dd] delete local files\n\t[d] drop\n\t[h] put on hold\n\t[x] exit",
            current_score_text
        );

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

        match input.as_str() {
            "r" => {
                self.prompt(Prompt::UpdateScore)?;
                self.prompt_series_options()
            }
            "dd" => {
                self.dir.try_remove_dir();
                Err(SeriesError::RequestExit)
            }
            "d" | "h" => {
                let status = if input == "d" {
                    Status::Dropped
                } else {
                    Status::OnHold
                };

                self.prompt(Prompt::PauseSeries(status))?;
                Err(SeriesError::RequestExit)
            }
            "x" => Err(SeriesError::RequestExit),
            _ => self.prompt_series_options(),
        }
    }

    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 set_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 => {
                if state.start_date.is_none() {
                    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());
                }
            }
            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()
    }

    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.get_episode(absolute_ep)?.clone();

        let status = process::open_with_default(path).map_err(SeriesError::FailedToOpenPlayer)?;

        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(());
            }
        }

        let state = &mut self.season.state;
        state.watched_episodes = episode.max(state.watched_episodes);

        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)?;

            if self.has_unwatched_episodes() {
                self.prompt(Prompt::EpisodeCompleted)?;
            } else {
                self.prompt(Prompt::SeriesCompleted)?;
                break;
            }
        }

        Ok(())
    }

    pub fn has_unwatched_episodes(&self) -> bool {
        let state = &self.season.state;

        match state.info.episodes {
            Some(total_eps) if total_eps > state.watched_episodes => true,
            _ => false,
        }
    }

    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)?;

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

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

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

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

    match index {
        0 => {
            println!("enter the name you want to search for:");
            let name = input::read_line()?;

            search_for_series_info(backend, &name, season)
        }
        _ => Ok(found.swap_remove(index - 1)),
    }
}

#[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.series_info.name.clone();
            return Ok(());
        }

        let info = search_for_series_info(&config.sync_service, &dir.series_info.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(())
    }
}
<commit_msg>[series] Read score input until a valid one is detected<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,
}

#[derive(Copy, Clone)]
pub enum Prompt {
    EpisodeCompleted,
    UpdateScore,
    SeriesCompleted,
    AlreadyCompleted,
    BeginRewatch,
    ResumePausedSeries,
    PauseSeries(Status),
}

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()
    }

    pub fn prepare_initial_state(&mut self) -> Result<(), SeriesError> {
        match self.season.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 !self.has_unwatched_episodes() {
                    self.prompt(Prompt::SeriesCompleted)?;
                }

                Ok(())
            }
            Status::PlanToWatch => self.set_list_entry_status(Status::Watching),
            Status::Completed => self.prompt(Prompt::AlreadyCompleted),
            Status::OnHold | Status::Dropped => self.prompt(Prompt::ResumePausedSeries),
        }
    }

    pub fn prompt(&mut self, prompt: Prompt) -> Result<(), SeriesError> {
        let state = &mut self.season.state;

        match prompt {
            Prompt::EpisodeCompleted => {
                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()?;
                self.prompt_next_episode_options()
            }
            Prompt::UpdateScore => {
                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
                );

                // Read & parse score input until we get a valid one
                loop {
                    let input = match input::read_line() {
                        Ok(ref input) if input.is_empty() => return Ok(()),
                        Ok(input) => input,
                        Err(err) => {
                            eprintln!("failed to read score input: {}", err);
                            return Ok(());
                        }
                    };

                    match self.config.sync_service.parse_score(&input) {
                        Ok(score) => {
                            state.score = Some(score);
                            break;
                        }
                        Err(err) => eprintln!("error: {}", err),
                    }
                }

                self.update_list_entry()
            }
            Prompt::SeriesCompleted => {
                println!("[{}] completed!", state.info.title);
                self.set_list_entry_status(Status::Completed)?;
                self.prompt_series_completed_options()
            }
            Prompt::AlreadyCompleted => {
                println!("[{}] already completed", state.info.title);
                self.prompt_series_completed_options()
            }
            Prompt::BeginRewatch => {
                println!("[{}] starting rewatch", state.info.title);
                println!("do you want to reset the start and end dates? (Y/n)");

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

                self.set_list_entry_status(Status::Rewatching)
            }
            Prompt::ResumePausedSeries => {
                println!(
                    "[{}] was previously put on hold or dropped",
                    state.info.title
                );

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

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

                self.set_list_entry_status(Status::Watching)
            }
            Prompt::PauseSeries(status) => {
                self.set_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();
                }

                Ok(())
            }
        }
    }

    pub 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.prompt(Prompt::PauseSeries(status))?;
                Err(SeriesError::RequestExit)
            }
            "r" => {
                self.prompt(Prompt::UpdateScore)?;
                self.prompt_next_episode_options()
            }
            "x" => Err(SeriesError::RequestExit),
            _ => Ok(()),
        }
    }

    pub 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(Prompt::UpdateScore)?;
                self.prompt_series_completed_options()
            }
            "w" => {
                self.prompt(Prompt::BeginRewatch)?;
                self.play_all_episodes()
            }
            "d" => {
                self.dir.try_remove_dir();
                Err(SeriesError::RequestExit)
            }
            "x" => Err(SeriesError::RequestExit),
            _ => Ok(()),
        }
    }

    pub fn prompt_series_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:", self.season.state.info.title);
        println!(
            "\t[r] rate{}\n\t[dd] delete local files\n\t[d] drop\n\t[h] put on hold\n\t[x] exit",
            current_score_text
        );

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

        match input.as_str() {
            "r" => {
                self.prompt(Prompt::UpdateScore)?;
                self.prompt_series_options()
            }
            "dd" => {
                self.dir.try_remove_dir();
                Err(SeriesError::RequestExit)
            }
            "d" | "h" => {
                let status = if input == "d" {
                    Status::Dropped
                } else {
                    Status::OnHold
                };

                self.prompt(Prompt::PauseSeries(status))?;
                Err(SeriesError::RequestExit)
            }
            "x" => Err(SeriesError::RequestExit),
            _ => self.prompt_series_options(),
        }
    }

    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 set_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 => {
                if state.start_date.is_none() {
                    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());
                }
            }
            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()
    }

    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.get_episode(absolute_ep)?.clone();

        let status = process::open_with_default(path).map_err(SeriesError::FailedToOpenPlayer)?;

        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(());
            }
        }

        let state = &mut self.season.state;
        state.watched_episodes = episode.max(state.watched_episodes);

        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)?;

            if self.has_unwatched_episodes() {
                self.prompt(Prompt::EpisodeCompleted)?;
            } else {
                self.prompt(Prompt::SeriesCompleted)?;
                break;
            }
        }

        Ok(())
    }

    pub fn has_unwatched_episodes(&self) -> bool {
        let state = &self.season.state;

        match state.info.episodes {
            Some(total_eps) if total_eps > state.watched_episodes => true,
            _ => false,
        }
    }

    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)?;

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

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

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

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

    match index {
        0 => {
            println!("enter the name you want to search for:");
            let name = input::read_line()?;

            search_for_series_info(backend, &name, season)
        }
        _ => Ok(found.swap_remove(index - 1)),
    }
}

#[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.series_info.name.clone();
            return Ok(());
        }

        let info = search_for_series_info(&config.sync_service, &dir.series_info.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 bifrost::rpc;
use bifrost::conshash::ConsistentHashing;
use bifrost::raft;
use bifrost::raft::client::RaftClient;
use bifrost::raft::state_machine::{master as sm_master};
use bifrost::raft::state_machine::callback::client::SubscriptionService;
use bifrost::membership::server::Membership;
use bifrost::membership::member::MemberService;
use bifrost::membership::client::ObserverClient;
use bifrost::conshash::weights::Weights;
use bifrost::tcp::{STANDALONE_ADDRESS_STRING, STANDALONE_SERVER_ID};
use ram::chunk::Chunks;
use ram::schema::Schemas;
use ram::schema::{sm as schema_sm};
use ram::types::Id;
use std::rc::Rc;
use std::sync::Arc;
use std::io;

pub mod cell_rpc;
pub mod transactions;

#[derive(Debug)]
pub enum ServerError {
    CannotJoinCluster,
    CannotJoinClusterGroup(sm_master::ExecError),
    CannotInitMemberTable,
    CannotSetServerWeight,
    CannotInitConsistentHashTable,
    CannotLoadMetaClient,
    MetaServerCannotBeStandalone,
}

pub struct ServerOptions {
    pub chunk_count: usize,
    pub memory_size: usize,
    pub standalone: bool,
    pub is_meta: bool,
    pub meta_members: Vec<String>,
    pub address: String,
    pub backup_storage: Option<String>,
    pub meta_storage: Option<String>,
    pub group_name: String,
}

pub struct ServerMeta {
    pub schemas: Arc<Schemas>
}

pub struct NebServer {
    pub chunks: Arc<Chunks>,
    pub meta: Arc<ServerMeta>,
    pub rpc: Arc<rpc::Server>,
    pub consh: Option<Arc<ConsistentHashing>>,
    pub member_pool: rpc::ClientPool,
    pub txn_peer: transactions::Peer,
    pub server_id: u64
}

impl NebServer {
    fn load_meta_server(opt: &ServerOptions, rpc_server: &Arc<rpc::Server>) -> Result<(), ServerError> {
        let server_addr = &opt.address;
        let raft_service = raft::RaftService::new(raft::Options {
            storage: match opt.meta_storage {
                Some(ref path) => raft::Storage::DISK(path.clone()),
                None => raft::Storage::MEMORY,
            },
            address: server_addr.clone(),
            service_id: raft::DEFAULT_SERVICE_ID,
        });
        rpc_server.register_service(
            raft::DEFAULT_SERVICE_ID,
            &raft_service
        );
        raft_service.register_state_machine(Box::new(schema_sm::SchemasSM::new(&raft_service)));
        raft::RaftService::start(&raft_service);
        match raft_service.join(&opt.meta_members) {
            Err(sm_master::ExecError::CannotConstructClient) => {
                info!("Cannot join meta cluster, will bootstrap one.");
                raft_service.bootstrap();
            },
            Ok(Ok(())) => {
                info!("Joined meta cluster, number of members: {}", raft_service.num_members());
            },
            e => {
                error!("Cannot join into cluster: {:?}", e);
                return Err(ServerError::CannotJoinCluster)
            }
        }
        Membership::new(rpc_server, &raft_service);
        Weights::new(&raft_service);
        return Ok(());
    }
    fn load_subscription(rpc_server: &Arc<rpc::Server>, raft_client: &RaftClient) {
        let subs_service = SubscriptionService::initialize(rpc_server);
        raft_client.set_subscription(&subs_service);
    }
    fn join_group(opt: &ServerOptions, raft_client: &Arc<RaftClient>) -> Result<(), ServerError> {
        let member_service = MemberService::new(&opt.address, raft_client);
        match member_service.join_group(&opt.group_name) {
            Ok(_) => {Ok(())},
            Err(e) => {
                error!("Cannot join cluster group");
                Err(ServerError::CannotJoinClusterGroup(e))
            }
        }
    }
    fn init_conshash(opt: &ServerOptions, raft_client: &Arc<RaftClient>)
        -> Result<Arc<ConsistentHashing>, ServerError> {
        match ConsistentHashing::new(&opt.group_name, raft_client) {
            Ok(ch) => {
                ch.set_weight(&opt.address, opt.memory_size as u64);
                if !ch.init_table().is_ok() {
                    error!("Cannot initialize member table");
                    return Err(ServerError::CannotInitMemberTable);
                }
                return Ok(ch);
            },
            _ => {
                error!("Cannot initialize consistent hash table");
                return Err(ServerError::CannotInitConsistentHashTable);
            }
        }
    }
    fn load_cluster_clients
    (opt: &ServerOptions, schemas: &mut Arc<Schemas>, rpc_server: &Arc<rpc::Server>)
    -> Result<Arc<ConsistentHashing>, ServerError>{
        let raft_client = RaftClient::new(&opt.meta_members, raft::DEFAULT_SERVICE_ID);
        match raft_client {
            Ok(raft_client) => {
                *schemas = Schemas::new(Some(&raft_client));
                NebServer::load_subscription(rpc_server, &raft_client);
                NebServer::join_group(opt, &raft_client)?;
                Ok(NebServer::init_conshash(opt, &raft_client)?)
            },
            Err(e) => {
                error!("Cannot load meta client: {:?}", e);
                Err(ServerError::CannotLoadMetaClient)
            }
        }
    }
    pub fn new(opt: ServerOptions) -> Result<Arc<NebServer>, ServerError> {
        let server_addr = if opt.standalone {&STANDALONE_ADDRESS_STRING} else {&opt.address};
        let rpc_server = rpc::Server::new(server_addr);
        let mut conshasing = None;
        let mut schemas = Schemas::new(None);
        rpc::Server::listen_and_resume(&rpc_server);
        if !opt.standalone {
            if opt.is_meta {
                NebServer::load_meta_server(&opt, &rpc_server)?;
            }
            conshasing = Some(NebServer::load_cluster_clients(&opt, &mut schemas, &rpc_server)?);
        } else if opt.is_meta {
            error!("Meta server cannot be standalone");
            return Err(ServerError::MetaServerCannotBeStandalone)
        }
        let meta_rc = Arc::new(ServerMeta {
            schemas: schemas.clone()
        });
        let chunks = Chunks::new(
            opt.chunk_count,
            opt.memory_size,
            meta_rc.clone(),
            opt.backup_storage.clone(),
        );
        let server = Arc::new(NebServer {
            chunks: chunks,
            meta: meta_rc,
            rpc: rpc_server.clone(),
            consh: conshasing,
            member_pool: rpc::ClientPool::new(),
            txn_peer: transactions::Peer::new(server_addr),
            server_id: rpc_server.server_id
        });
        rpc_server.register_service(
            cell_rpc::DEFAULT_SERVICE_ID,
            &cell_rpc::NebRPCService::new(&server)
        );
        rpc_server.register_service(
           transactions::manager::DEFAULT_SERVICE_ID,
           &transactions::manager::TransactionManager::new(&server)
        );
        rpc_server.register_service(
            transactions::data_site::DEFAULT_SERVICE_ID,
            &transactions::data_site::DataManager::new(&server)
        );
        Ok(server)
    }
    pub fn get_server_id_by_id(&self, id: &Id) -> Option<u64> {
        match self.consh {
            Some(ref consh) => {
                if let Some(server_id) = consh.get_server_id(id.higher) {
                    Some(server_id)
                } else {
                    None
                }
            },
            _ => Some(STANDALONE_SERVER_ID.clone())
        }
    }
    pub fn get_member_by_server_id(&self, server_id: u64) -> io::Result<Arc<rpc::RPCClient>> {
        match self.consh {
            Some(ref consh) => {
                if let Some(ref server_name) = consh.to_server_name(Some(server_id)) {
                    self.member_pool.get(server_name)
                } else {
                    Err(io::Error::new(io::ErrorKind::Other, "Cannot find server id in consistent hash table"))
                }
            },
            _ => self.member_pool.get(&STANDALONE_ADDRESS_STRING)
        }
    }
}<commit_msg>update server parameter policy with new bifrost with shortcut<commit_after>use bifrost::rpc;
use bifrost::conshash::ConsistentHashing;
use bifrost::raft;
use bifrost::raft::client::RaftClient;
use bifrost::raft::state_machine::{master as sm_master};
use bifrost::raft::state_machine::callback::client::SubscriptionService;
use bifrost::membership::server::Membership;
use bifrost::membership::member::MemberService;
use bifrost::membership::client::ObserverClient;
use bifrost::conshash::weights::Weights;
use bifrost::tcp::{STANDALONE_ADDRESS_STRING, STANDALONE_SERVER_ID};
use ram::chunk::Chunks;
use ram::schema::Schemas;
use ram::schema::{sm as schema_sm};
use ram::types::Id;
use std::rc::Rc;
use std::sync::Arc;
use std::io;

pub mod cell_rpc;
pub mod transactions;

#[derive(Debug)]
pub enum ServerError {
    CannotJoinCluster,
    CannotJoinClusterGroup(sm_master::ExecError),
    CannotInitMemberTable,
    CannotSetServerWeight,
    CannotInitConsistentHashTable,
    CannotLoadMetaClient,
    StandaloneMustAlsoBeMetaServer,
}

pub struct ServerOptions {
    pub chunk_count: usize,
    pub memory_size: usize,
    pub standalone: bool,
    pub is_meta: bool,
    pub meta_members: Vec<String>,
    pub address: String,
    pub backup_storage: Option<String>,
    pub meta_storage: Option<String>,
    pub group_name: String,
}

pub struct ServerMeta {
    pub schemas: Arc<Schemas>
}

pub struct NebServer {
    pub chunks: Arc<Chunks>,
    pub meta: Arc<ServerMeta>,
    pub rpc: Arc<rpc::Server>,
    pub consh: Option<Arc<ConsistentHashing>>,
    pub member_pool: rpc::ClientPool,
    pub txn_peer: transactions::Peer,
    pub server_id: u64
}

impl NebServer {
    fn load_meta_server(opt: &ServerOptions, rpc_server: &Arc<rpc::Server>) -> Result<(), ServerError> {
        let server_addr = &opt.address;
        let raft_service = raft::RaftService::new(raft::Options {
            storage: match opt.meta_storage {
                Some(ref path) => raft::Storage::DISK(path.clone()),
                None => raft::Storage::MEMORY,
            },
            address: server_addr.clone(),
            service_id: raft::DEFAULT_SERVICE_ID,
        });
        rpc_server.register_service(
            raft::DEFAULT_SERVICE_ID,
            &raft_service
        );
        raft_service.register_state_machine(Box::new(schema_sm::SchemasSM::new(&raft_service)));
        raft::RaftService::start(&raft_service);
        match raft_service.join(&opt.meta_members) {
            Err(sm_master::ExecError::CannotConstructClient) => {
                info!("Cannot join meta cluster, will bootstrap one.");
                raft_service.bootstrap();
            },
            Ok(Ok(())) => {
                info!("Joined meta cluster, number of members: {}", raft_service.num_members());
            },
            e => {
                error!("Cannot join into cluster: {:?}", e);
                return Err(ServerError::CannotJoinCluster)
            }
        }
        Membership::new(rpc_server, &raft_service);
        Weights::new(&raft_service);
        return Ok(());
    }
    fn load_subscription(rpc_server: &Arc<rpc::Server>, raft_client: &RaftClient) {
        let subs_service = SubscriptionService::initialize(rpc_server);
        raft_client.set_subscription(&subs_service);
    }
    fn join_group(opt: &ServerOptions, raft_client: &Arc<RaftClient>) -> Result<(), ServerError> {
        let member_service = MemberService::new(&opt.address, raft_client);
        match member_service.join_group(&opt.group_name) {
            Ok(_) => {Ok(())},
            Err(e) => {
                error!("Cannot join cluster group");
                Err(ServerError::CannotJoinClusterGroup(e))
            }
        }
    }
    fn init_conshash(opt: &ServerOptions, raft_client: &Arc<RaftClient>)
        -> Result<Arc<ConsistentHashing>, ServerError> {
        match ConsistentHashing::new(&opt.group_name, raft_client) {
            Ok(ch) => {
                ch.set_weight(&opt.address, opt.memory_size as u64);
                if !ch.init_table().is_ok() {
                    error!("Cannot initialize member table");
                    return Err(ServerError::CannotInitMemberTable);
                }
                return Ok(ch);
            },
            _ => {
                error!("Cannot initialize consistent hash table");
                return Err(ServerError::CannotInitConsistentHashTable);
            }
        }
    }
    fn load_cluster_clients
    (opt: &ServerOptions, schemas: &mut Arc<Schemas>, rpc_server: &Arc<rpc::Server>)
    -> Result<Arc<ConsistentHashing>, ServerError>{
        let raft_client = RaftClient::new(&opt.meta_members, raft::DEFAULT_SERVICE_ID);
        match raft_client {
            Ok(raft_client) => {
                *schemas = Schemas::new(Some(&raft_client));
                NebServer::load_subscription(rpc_server, &raft_client);
                NebServer::join_group(opt, &raft_client)?;
                Ok(NebServer::init_conshash(opt, &raft_client)?)
            },
            Err(e) => {
                error!("Cannot load meta client: {:?}", e);
                Err(ServerError::CannotLoadMetaClient)
            }
        }
    }
    pub fn new(opt: ServerOptions) -> Result<Arc<NebServer>, ServerError> {
        let server_addr = if opt.standalone {&STANDALONE_ADDRESS_STRING} else {&opt.address};
        let rpc_server = rpc::Server::new(server_addr);
        let mut conshasing = None;
        let mut schemas = Schemas::new(None);
        rpc::Server::listen_and_resume(&rpc_server);
        if opt.is_meta {
            NebServer::load_meta_server(&opt, &rpc_server)?;
        }
        if !opt.standalone {
            conshasing = Some(NebServer::load_cluster_clients(&opt, &mut schemas, &rpc_server)?);
        }
        if !opt.is_meta && !opt.standalone {
            return Err(ServerError::StandaloneMustAlsoBeMetaServer)
        }
        let meta_rc = Arc::new(ServerMeta {
            schemas: schemas
        });
        let chunks = Chunks::new(
            opt.chunk_count,
            opt.memory_size,
            meta_rc.clone(),
            opt.backup_storage.clone(),
        );
        let server = Arc::new(NebServer {
            chunks: chunks,
            meta: meta_rc,
            rpc: rpc_server.clone(),
            consh: conshasing,
            member_pool: rpc::ClientPool::new(),
            txn_peer: transactions::Peer::new(server_addr),
            server_id: rpc_server.server_id
        });
        rpc_server.register_service(
            cell_rpc::DEFAULT_SERVICE_ID,
            &cell_rpc::NebRPCService::new(&server)
        );
        rpc_server.register_service(
           transactions::manager::DEFAULT_SERVICE_ID,
           &transactions::manager::TransactionManager::new(&server)
        );
        rpc_server.register_service(
            transactions::data_site::DEFAULT_SERVICE_ID,
            &transactions::data_site::DataManager::new(&server)
        );
        Ok(server)
    }
    pub fn get_server_id_by_id(&self, id: &Id) -> Option<u64> {
        match self.consh {
            Some(ref consh) => {
                if let Some(server_id) = consh.get_server_id(id.higher) {
                    Some(server_id)
                } else {
                    None
                }
            },
            _ => Some(STANDALONE_SERVER_ID.clone())
        }
    }
    pub fn get_member_by_server_id(&self, server_id: u64) -> io::Result<Arc<rpc::RPCClient>> {
        match self.consh {
            Some(ref consh) => {
                if let Some(ref server_name) = consh.to_server_name(Some(server_id)) {
                    self.member_pool.get(server_name)
                } else {
                    Err(io::Error::new(io::ErrorKind::Other, "Cannot find server id in consistent hash table"))
                }
            },
            _ => self.member_pool.get(&STANDALONE_ADDRESS_STRING)
        }
    }
}<|endoftext|>


<commit_before>//
// sprocketnes/util.rs
//
// Author: Patrick Walton
//

use std::io::File;
use std::libc::{c_int, c_void, time_t};
use std::ptr::null;

//
// A tiny custom serialization infrastructure, used for savestates.
//
// TODO: Use the standard library's ToBytes and add a FromBytes -- or don't; this is such a small
// amount of code it barely seems worth it.
//

pub trait Save {
    fn save(&mut self, fd: &mut File);
    fn load(&mut self, fd: &mut File);
}

impl Save for u8 {
    fn save(&mut self, fd: &mut File) { fd.write([ *self ]) }
    fn load(&mut self, fd: &mut File) { let mut buf = [ 0 ]; fd.read(buf); *self = buf[0]; }
}

impl Save for u16 {
    fn save(&mut self, fd: &mut File) { fd.write([ *self as u8, (*self >> 8) as u8 ]) }
    fn load(&mut self, fd: &mut File) {
        let mut buf = [ 0, 0 ];
        fd.read(buf);
        *self = (buf[0] as u16) | ((buf[1] as u16) << 8);
    }
}

impl Save for u64 {
    fn save(&mut self, fd: &mut File) {
        let mut buf = [ 0, ..8 ];
        for i in range(0, 8) {
            buf[i] = ((*self) >> (i * 8)) as u8;
        }
        fd.write(buf);
    }
    fn load(&mut self, fd: &mut File) {
        let mut buf = [ 0, ..8 ];
        fd.read(buf);
        *self = 0;
        for i in range(0, 8) {
            *self = *self | (buf[i] as u64 << (i * 8));
        }
    }
}

impl<'a> Save for &'a mut [u8] {
    fn save(&mut self, fd: &mut File) {
        fd.write(*self);
    }
    fn load(&mut self, fd: &mut File) {
        fd.read(*self);
    }
}

impl Save for bool {
    fn save(&mut self, fd: &mut File) { fd.write([ if *self { 0 } else { 1 } ]) }
    fn load(&mut self, fd: &mut File) {
        let mut val: [u8, ..1] = [ 0 ];
        fd.read(val);
        *self = val[0] != 0
    }
}

// A convenience macro to save and load entire structs.
macro_rules! save_struct(
    ($name:ident { $($field:ident),* }) => (
        impl Save for $name {
            fn save(&mut self, fd: &mut File) {
                $(self.$field.save(fd);)*
            }
            fn load(&mut self, fd: &mut File) {
                $(self.$field.load(fd);)*
            }
        }
    )
)

macro_rules! save_enum(
    ($name:ident { $val_0:ident, $val_1:ident }) => (
        impl Save for $name {
            fn save(&mut self, fd: &mut File) {
                let mut val: u8 = match *self { $val_0 => 0, $val_1 => 1 };
                val.save(fd)
            }
            fn load(&mut self, fd: &mut File) {
                let mut val: u8 = 0;
                val.load(fd);
                *self = if val == 0 { $val_0 } else { $val_1 };
            }
        }
    )
)

//
// Random number generation
//

pub struct Xorshift {
    x: u32,
    y: u32,
    z: u32,
    w: u32,
}

impl Xorshift {
    pub fn new() -> Xorshift {
        Xorshift { x: 123456789, y: 362436069, z: 521288629, w: 88675123 }
    }

    pub fn next(&mut self) -> u32 {
        let t = self.x ^ (self.x << 11);
        self.x = self.y; self.y = self.z; self.z = self.w;
        self.w = self.w ^ (self.w >> 19) ^ (t ^ (t >> 8));
        self.w
    }
}

//
// Simple assertions
//

#[cfg(debug)]
pub fn debug_assert(cond: bool, msg: &str) {
    if !cond {
        println!(msg);
    }
}

#[cfg(not(debug))]
pub fn debug_assert(_: bool, _: &str) {}

#[cfg(debug)]
pub fn debug_print(msg: &str) {
    println!(msg);
}

#[cfg(not(debug))]
pub fn debug_print(_: &str) {}

//
// Bindings for `gettimeofday(2)`
//

struct timeval {
    tv_sec: time_t,
    tv_usec: u32,
}

extern {
    fn gettimeofday(tp: *mut timeval, tzp: *c_void) -> c_int;
}

pub fn current_time_millis() -> u64 {
    unsafe {
        let mut tv = timeval { tv_sec: 0, tv_usec: 0 };
        gettimeofday(&mut tv, null());
        (tv.tv_sec as u64) * 1000 + (tv.tv_usec as u64) / 1000
    }
}

<commit_msg>Don't return.<commit_after>//
// sprocketnes/util.rs
//
// Author: Patrick Walton
//

use std::io::File;
use std::libc::{c_int, c_void, time_t};
use std::ptr::null;

//
// A tiny custom serialization infrastructure, used for savestates.
//
// TODO: Use the standard library's ToBytes and add a FromBytes -- or don't; this is such a small
// amount of code it barely seems worth it.
//

pub trait Save {
    fn save(&mut self, fd: &mut File);
    fn load(&mut self, fd: &mut File);
}

impl Save for u8 {
    fn save(&mut self, fd: &mut File) { fd.write([ *self ]); }
    fn load(&mut self, fd: &mut File) { let mut buf = [ 0 ]; fd.read(buf); *self = buf[0]; }
}

impl Save for u16 {
    fn save(&mut self, fd: &mut File) { fd.write([ *self as u8, (*self >> 8) as u8 ]); }
    fn load(&mut self, fd: &mut File) {
        let mut buf = [ 0, 0 ];
        fd.read(buf);
        *self = (buf[0] as u16) | ((buf[1] as u16) << 8);
    }
}

impl Save for u64 {
    fn save(&mut self, fd: &mut File) {
        let mut buf = [ 0, ..8 ];
        for i in range(0, 8) {
            buf[i] = ((*self) >> (i * 8)) as u8;
        }
        fd.write(buf);
    }
    fn load(&mut self, fd: &mut File) {
        let mut buf = [ 0, ..8 ];
        fd.read(buf);
        *self = 0;
        for i in range(0, 8) {
            *self = *self | (buf[i] as u64 << (i * 8));
        }
    }
}

impl<'a> Save for &'a mut [u8] {
    fn save(&mut self, fd: &mut File) {
        fd.write(*self);
    }
    fn load(&mut self, fd: &mut File) {
        fd.read(*self);
    }
}

impl Save for bool {
    fn save(&mut self, fd: &mut File) { fd.write([ if *self { 0 } else { 1 } ]); }
    fn load(&mut self, fd: &mut File) {
        let mut val: [u8, ..1] = [ 0 ];
        fd.read(val);
        *self = val[0] != 0
    }
}

// A convenience macro to save and load entire structs.
macro_rules! save_struct(
    ($name:ident { $($field:ident),* }) => (
        impl Save for $name {
            fn save(&mut self, fd: &mut File) {
                $(self.$field.save(fd);)*
            }
            fn load(&mut self, fd: &mut File) {
                $(self.$field.load(fd);)*
            }
        }
    )
)

macro_rules! save_enum(
    ($name:ident { $val_0:ident, $val_1:ident }) => (
        impl Save for $name {
            fn save(&mut self, fd: &mut File) {
                let mut val: u8 = match *self { $val_0 => 0, $val_1 => 1 };
                val.save(fd)
            }
            fn load(&mut self, fd: &mut File) {
                let mut val: u8 = 0;
                val.load(fd);
                *self = if val == 0 { $val_0 } else { $val_1 };
            }
        }
    )
)

//
// Random number generation
//

pub struct Xorshift {
    x: u32,
    y: u32,
    z: u32,
    w: u32,
}

impl Xorshift {
    pub fn new() -> Xorshift {
        Xorshift { x: 123456789, y: 362436069, z: 521288629, w: 88675123 }
    }

    pub fn next(&mut self) -> u32 {
        let t = self.x ^ (self.x << 11);
        self.x = self.y; self.y = self.z; self.z = self.w;
        self.w = self.w ^ (self.w >> 19) ^ (t ^ (t >> 8));
        self.w
    }
}

//
// Simple assertions
//

#[cfg(debug)]
pub fn debug_assert(cond: bool, msg: &str) {
    if !cond {
        println!(msg);
    }
}

#[cfg(not(debug))]
pub fn debug_assert(_: bool, _: &str) {}

#[cfg(debug)]
pub fn debug_print(msg: &str) {
    println!(msg);
}

#[cfg(not(debug))]
pub fn debug_print(_: &str) {}

//
// Bindings for `gettimeofday(2)`
//

struct timeval {
    tv_sec: time_t,
    tv_usec: u32,
}

extern {
    fn gettimeofday(tp: *mut timeval, tzp: *c_void) -> c_int;
}

pub fn current_time_millis() -> u64 {
    unsafe {
        let mut tv = timeval { tv_sec: 0, tv_usec: 0 };
        gettimeofday(&mut tv, null());
        (tv.tv_sec as u64) * 1000 + (tv.tv_usec as u64) / 1000
    }
}

<|endoftext|>


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

#![allow(unused_assignments)]

struct Wrap<'a> { w: &'a mut u32 }

fn foo() {
    let mut x = 22u64;
    let wrapper = Wrap { w: &mut x };
    x += 1;  //~ ERROR cannot assign to `x`
    *wrapper.w += 1;
}

fn main() { }<commit_msg>Fix failing test<commit_after>// Copyright 2017 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.
//revisions: ast mir
//[mir] compile-flags: -Z emit-end-regions -Z borrowck-mir -Z nll

#![allow(unused_assignments)]

struct Wrap<'a> { w: &'a mut u32 }

fn foo() {
    let mut x = 22;
    let wrapper = Wrap { w: &mut x };
    //~^ ERROR cannot assign to `x` because it is borrowed (Mir) [E0506]
    //~^^ ERROR cannot use `x` because it was mutably borrowed (Mir) [E0503]
    x += 1; //[ast]~ ERROR cannot assign to `x` because it is borrowed [E0506]
    //[mir]~^ ERROR cannot assign to `x` because it is borrowed (Ast) [E0506]
    //[mir]~^^ ERROR cannot assign to `x` because it is borrowed (Mir) [E0506]
    //[mir]~^^^ ERROR cannot use `x` because it was mutably borrowed (Mir) [E0503]
    *wrapper.w += 1;
}

fn main() { }<|endoftext|>


<commit_before>#![allow(unsafe_code)]

#[cfg(any(target_os = "linux", target_os = "macos"))]
use std::io;
#[cfg(any(target_os = "linux"))]
use {std::fs::File, std::io::Read};

/// See the Kernel's documentation for more information about this subsystem,
/// found at:  [Documentation/cgroup-v1/memory.txt](https://www.kernel.org/doc/Documentation/cgroup-v1/memory.txt)
///
/// If there's no memory limit specified on the container this may return
/// 0x7FFFFFFFFFFFF000 (2^63-1 rounded down to 4k which is a common page size).
/// So we know we are not running in a memory restricted environment.
#[cfg(target_os = "linux")]
fn get_cgroup_memory_limit() -> io::Result<u64> {
    File::open("/sys/fs/cgroup/memory/memory.limit_in_bytes")
        .and_then(read_u64_from)
}

#[cfg(target_os = "linux")]
fn read_u64_from(mut file: File) -> io::Result<u64> {
    let mut s = String::new();
    file.read_to_string(&mut s).and_then(|_| {
        s.trim()
            .parse()
            .map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))
    })
}

/// Returns the maximum size of total available memory of the process, in bytes.
/// If this limit is exceeded, the malloc() and mmap() functions shall fail with
/// errno set to [ENOMEM].
#[cfg(any(target_os = "linux", target_os = "macos"))]
fn get_rlimit_as() -> io::Result<libc::rlimit> {
    let mut limit = std::mem::MaybeUninit::<libc::rlimit>::uninit();

    let ret = unsafe { libc::getrlimit(libc::RLIMIT_AS, limit.as_mut_ptr()) };

    if ret == 0 {
        Ok(unsafe { limit.assume_init() })
    } else {
        Err(io::Error::last_os_error())
    }
}

#[cfg(any(target_os = "linux", target_os = "macos"))]
pub fn get_available_memory() -> io::Result<u64> {
    let pages = unsafe { libc::sysconf(libc::_SC_PHYS_PAGES) };
    if pages == -1 {
        return Err(io::Error::last_os_error());
    }

    let page_size = unsafe { libc::sysconf(libc::_SC_PAGE_SIZE) };
    if page_size == -1 {
        return Err(io::Error::last_os_error());
    }

    Ok((pages as u64) * (page_size as u64))
}

pub fn get_memory_limit() -> u64 {
    // Maximum addressable memory space limit in u64
    static MAX_USIZE: u64 = usize::max_value() as u64;

    let mut max: u64 = 0;

    #[cfg(target_os = "linux")]
    {
        if let Ok(mem) = get_cgroup_memory_limit() {
            max = mem;
        }

        // If there's no memory limit specified on the container this
        // actually returns 0x7FFFFFFFFFFFF000 (2^63-1 rounded down to
        // 4k which is a common page size). So we know we are not
        // running in a memory restricted environment.
        // src: https://github.com/dotnet/coreclr/blob/master/src/pal/src/misc/cgroup.cpp#L385-L428
        if max > 0x7FFF_FFFF_0000_0000 {
            return 0;
        }
    }

    #[cfg(any(target_os = "linux", target_os = "macos"))]
    {
        if let Ok(rlim) = get_rlimit_as() {
            let rlim_cur = rlim.rlim_cur;
            if rlim_cur < max || max == 0 {
                max = rlim_cur;
            }
        }

        if let Ok(available) = get_available_memory() {
            if available < max || max == 0 {
                max = available;
            }
        }
    }

    if max > MAX_USIZE {
        // It is observed in practice when the memory is unrestricted, Linux
        // control group returns a physical limit that is bigger than
        // the address space
        max = MAX_USIZE;
    }

    max
}
<commit_msg>return rlimit_cur u64 type to be comparable with max: u64 on non 64 bit platforms<commit_after>#![allow(unsafe_code)]

#[cfg(any(target_os = "linux", target_os = "macos"))]
use std::io;
#[cfg(any(target_os = "linux"))]
use {std::fs::File, std::io::Read};

/// See the Kernel's documentation for more information about this subsystem,
/// found at:  [Documentation/cgroup-v1/memory.txt](https://www.kernel.org/doc/Documentation/cgroup-v1/memory.txt)
///
/// If there's no memory limit specified on the container this may return
/// 0x7FFFFFFFFFFFF000 (2^63-1 rounded down to 4k which is a common page size).
/// So we know we are not running in a memory restricted environment.
#[cfg(target_os = "linux")]
fn get_cgroup_memory_limit() -> io::Result<u64> {
    File::open("/sys/fs/cgroup/memory/memory.limit_in_bytes")
        .and_then(read_u64_from)
}

#[cfg(target_os = "linux")]
fn read_u64_from(mut file: File) -> io::Result<u64> {
    let mut s = String::new();
    file.read_to_string(&mut s).and_then(|_| {
        s.trim()
            .parse()
            .map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))
    })
}

/// Returns the maximum size of total available memory of the process, in bytes.
/// If this limit is exceeded, the malloc() and mmap() functions shall fail with
/// errno set to [ENOMEM].
#[cfg(any(target_os = "linux", target_os = "macos"))]
fn get_rlimit_as() -> io::Result<libc::rlimit> {
    let mut limit = std::mem::MaybeUninit::<libc::rlimit>::uninit();

    let ret = unsafe { libc::getrlimit(libc::RLIMIT_AS, limit.as_mut_ptr()) };

    if ret == 0 {
        Ok(unsafe { limit.assume_init() })
    } else {
        Err(io::Error::last_os_error())
    }
}

#[cfg(any(target_os = "linux", target_os = "macos"))]
pub fn get_available_memory() -> io::Result<u64> {
    let pages = unsafe { libc::sysconf(libc::_SC_PHYS_PAGES) };
    if pages == -1 {
        return Err(io::Error::last_os_error());
    }

    let page_size = unsafe { libc::sysconf(libc::_SC_PAGE_SIZE) };
    if page_size == -1 {
        return Err(io::Error::last_os_error());
    }

    Ok((pages as u64) * (page_size as u64))
}

pub fn get_memory_limit() -> u64 {
    // Maximum addressable memory space limit in u64
    static MAX_USIZE: u64 = usize::max_value() as u64;

    let mut max: u64 = 0;

    #[cfg(target_os = "linux")]
    {
        if let Ok(mem) = get_cgroup_memory_limit() {
            max = mem;
        }

        // If there's no memory limit specified on the container this
        // actually returns 0x7FFFFFFFFFFFF000 (2^63-1 rounded down to
        // 4k which is a common page size). So we know we are not
        // running in a memory restricted environment.
        // src: https://github.com/dotnet/coreclr/blob/master/src/pal/src/misc/cgroup.cpp#L385-L428
        if max > 0x7FFF_FFFF_0000_0000 {
            return 0;
        }
    }

    #[cfg(any(target_os = "linux", target_os = "macos"))]
    {
        if let Ok(rlim) = get_rlimit_as() {
            let rlim_cur = Into::<u64>::into(rlim.rlim_cur);
            if rlim_cur < max || max == 0 {
                max = rlim_cur;
            }
        }

        if let Ok(available) = get_available_memory() {
            if available < max || max == 0 {
                max = available;
            }
        }
    }

    if max > MAX_USIZE {
        // It is observed in practice when the memory is unrestricted, Linux
        // control group returns a physical limit that is bigger than
        // the address space
        max = MAX_USIZE;
    }

    max
}
<|endoftext|>


<commit_before>use http;
use super::server;

#[test]
pub fn test_proxy() {
    let srv = server!(
        recv!(b"POST http://www.google.com HTTP/1.1\r\n\
                Host: www.google.com\r\n\
                Accept: */*\r\n\
                Proxy-Connection: Keep-Alive\r\n\
                Content-Type: application/octet-stream\r\n\
                Content-Length: 11\r\n\
                \r\n\
                Foo Bar Baz"),
        send!(b"HTTP/1.1 200 OK\r\n\
                Content-Length: 5\r\n\r\n\
                Hello\r\n\r\n"));

    let res = http::handle()
        .proxy(server::url("/"))
        .post("http://www.google.com", "Foo Bar Baz")
        .exec().unwrap();

    srv.assert();

    assert!(res.get_code() == 200);
    assert!(res.get_body() == "Hello".as_bytes());
}
<commit_msg>Fix failing proxy test<commit_after>use http;
use super::server;

#[test]
pub fn test_proxy() {
    let srv = server!(
        recv!(b"POST http://www.google.com/ HTTP/1.1\r\n\
                Host: www.google.com\r\n\
                Accept: */*\r\n\
                Proxy-Connection: Keep-Alive\r\n\
                Content-Type: application/octet-stream\r\n\
                Content-Length: 11\r\n\
                \r\n\
                Foo Bar Baz"),
        send!(b"HTTP/1.1 200 OK\r\n\
                Content-Length: 5\r\n\r\n\
                Hello\r\n\r\n"));

    let res = http::handle()
        .proxy(server::url("/"))
        .post("http://www.google.com", "Foo Bar Baz")
        .exec().unwrap();

    srv.assert();

    assert!(res.get_code() == 200);
    assert!(res.get_body() == "Hello".as_bytes());
}
<|endoftext|>


<commit_before>use tls::*;
use tls_alert::TlsAlertSeverity;

/// Error types for the state machine
pub enum StateChangeError {
    InvalidTransition,
    ParseError,
}

/// TLS machine possible states
#[derive(Copy,Clone,Debug,PartialEq)]
pub enum TlsState {
    None,
    ClientHello,
    AskResumeSession,
    ResumeSession,
    ServerHello,
    Certificate,
    CertificateSt,
    ServerKeyExchange,
    ServerHelloDone,
    ClientKeyExchange,
    ClientChangeCipherSpec,

    CRCertRequest,
    CRHelloDone,
    CRCert,
    CRClientKeyExchange,
    CRCertVerify,

    NoCertSKE,
    NoCertHelloDone,
    NoCertCKE,

    PskHelloDone,
    PskCKE,

    SessionEncrypted,

    Alert,

    Invalid,
}

fn tls_state_transition_handshake(state: TlsState, msg: &TlsMessageHandshake) -> Result<TlsState,StateChangeError> {
    match (state,msg) {
        (TlsState::None,             &TlsMessageHandshake::ClientHello(ref msg)) => {
            match msg.session_id {
                Some(_) => Ok(TlsState::AskResumeSession),
                _       => Ok(TlsState::ClientHello)
            }
        },
        // Server certificate
        (TlsState::ClientHello,      &TlsMessageHandshake::ServerHello(_))       => Ok(TlsState::ServerHello),
        (TlsState::ServerHello,      &TlsMessageHandshake::Certificate(_))       => Ok(TlsState::Certificate),
        // Server certificate, no client certificate requested
        (TlsState::Certificate,      &TlsMessageHandshake::ServerKeyExchange(_)) => Ok(TlsState::ServerKeyExchange),
        (TlsState::Certificate,      &TlsMessageHandshake::CertificateStatus(_)) => Ok(TlsState::CertificateSt),
        (TlsState::CertificateSt,    &TlsMessageHandshake::ServerKeyExchange(_)) => Ok(TlsState::ServerKeyExchange),
        (TlsState::ServerKeyExchange,&TlsMessageHandshake::ServerDone(_))        => Ok(TlsState::ServerHelloDone),
        (TlsState::ServerHelloDone  ,&TlsMessageHandshake::ClientKeyExchange(_)) => Ok(TlsState::ClientKeyExchange),
        // Server certificate, client certificate requested
        (TlsState::Certificate,      &TlsMessageHandshake::CertificateRequest(_))=> Ok(TlsState::CRCertRequest),
        (TlsState::ServerKeyExchange,&TlsMessageHandshake::CertificateRequest(_))=> Ok(TlsState::CRCertRequest),
        (TlsState::CRCertRequest,    &TlsMessageHandshake::ServerDone(_))        => Ok(TlsState::CRHelloDone),
        (TlsState::CRHelloDone,      &TlsMessageHandshake::Certificate(_))       => Ok(TlsState::CRCert),
        (TlsState::CRCert,           &TlsMessageHandshake::ClientKeyExchange(_)) => Ok(TlsState::CRClientKeyExchange),
        (TlsState::CRClientKeyExchange, &TlsMessageHandshake::CertificateVerify(_)) => Ok(TlsState::CRCertVerify),
        // Server has no certificate (but accepts anonymous)
        (TlsState::ServerHello,      &TlsMessageHandshake::ServerKeyExchange(_)) => Ok(TlsState::NoCertSKE),
        (TlsState::NoCertSKE,        &TlsMessageHandshake::ServerDone(_))        => Ok(TlsState::NoCertHelloDone),
        (TlsState::NoCertHelloDone,  &TlsMessageHandshake::ClientKeyExchange(_)) => Ok(TlsState::NoCertCKE),
        // PSK
        (TlsState::Certificate,      &TlsMessageHandshake::ServerDone(_))        => Ok(TlsState::PskHelloDone),
        (TlsState::PskHelloDone,     &TlsMessageHandshake::ClientKeyExchange(_)) => Ok(TlsState::PskCKE),
        // Resuming session
        (TlsState::AskResumeSession, &TlsMessageHandshake::ServerHello(_))       => Ok(TlsState::ResumeSession),
        // Resume session failed
        (TlsState::ResumeSession,    &TlsMessageHandshake::Certificate(_))       => Ok(TlsState::Certificate),
        // TLS 1.3 Draft 18 1-RTT
        // Re-use the ClientChangeCipherSpec state to indicate the next message will be encrypted
        (TlsState::ClientHello,      &TlsMessageHandshake::ServerHelloV13Draft18(_))    => Ok(TlsState::ClientChangeCipherSpec),
        // Hello requests must be accepted at any time (except start), but ignored [RFC5246] 7.4.1.1
        (TlsState::None,             &TlsMessageHandshake::HelloRequest)         => Err(StateChangeError::InvalidTransition),
        (s,                          &TlsMessageHandshake::HelloRequest)         => Ok(s),
        // All other transitions are considered invalid
        _ => Err(StateChangeError::InvalidTransition),
    }
}

/// Update the TLS state machine, doing one transition
///
/// Given the previous state and the parsed message, return the new state or a state machine error.
///
/// This state machine only implements the TLS handshake.
///
/// Some transitions only check the new message type, while some others must match the content
/// (for example, to check if the client asked to resume a session).
///
/// If the previous state is `Invalid`, the state machine will not return an error, but keep the
/// same `Invalid` state. This is used to raise error only once if the state machine keeps being
/// updated by new messages.
pub fn tls_state_transition(state: TlsState, msg: &TlsMessage) -> Result<TlsState,StateChangeError> {
    match (state,msg) {
        (TlsState::Invalid,_) => Ok(TlsState::Invalid),
        (_,&TlsMessage::Handshake(ref m)) => tls_state_transition_handshake(state,m),
        // Server certificate
        (TlsState::ClientKeyExchange,     &TlsMessage::ChangeCipherSpec) => Ok(TlsState::ClientChangeCipherSpec),
        (TlsState::ClientChangeCipherSpec,&TlsMessage::ChangeCipherSpec) => Ok(TlsState::SessionEncrypted),
        // Server certificate, client certificate requested
        (TlsState::CRClientKeyExchange,   &TlsMessage::ChangeCipherSpec) => Ok(TlsState::ClientChangeCipherSpec),
        (TlsState::CRCertVerify,          &TlsMessage::ChangeCipherSpec) => Ok(TlsState::ClientChangeCipherSpec),
        // No server certificate
        (TlsState::NoCertCKE,             &TlsMessage::ChangeCipherSpec) => Ok(TlsState::ClientChangeCipherSpec),
        // PSK
        (TlsState::PskCKE,                &TlsMessage::ChangeCipherSpec) => Ok(TlsState::ClientChangeCipherSpec),
        // Resume session
        (TlsState::ResumeSession,         &TlsMessage::ChangeCipherSpec) => Ok(TlsState::ClientChangeCipherSpec),
        // non-fatal alerts
        (s,                               &TlsMessage::Alert(ref a)) => {
            if a.severity == TlsAlertSeverity::Warning { Ok(s) } else { Err(StateChangeError::InvalidTransition) }
        },
        (_,_) => Err(StateChangeError::InvalidTransition),
    }
}
<commit_msg>Add new state to indicate connection is closed (after fatal alert)<commit_after>use tls::*;
use tls_alert::TlsAlertSeverity;

/// Error types for the state machine
pub enum StateChangeError {
    InvalidTransition,
    ParseError,
}

/// TLS machine possible states
#[derive(Copy,Clone,Debug,PartialEq)]
pub enum TlsState {
    None,
    ClientHello,
    AskResumeSession,
    ResumeSession,
    ServerHello,
    Certificate,
    CertificateSt,
    ServerKeyExchange,
    ServerHelloDone,
    ClientKeyExchange,
    ClientChangeCipherSpec,

    CRCertRequest,
    CRHelloDone,
    CRCert,
    CRClientKeyExchange,
    CRCertVerify,

    NoCertSKE,
    NoCertHelloDone,
    NoCertCKE,

    PskHelloDone,
    PskCKE,

    SessionEncrypted,

    Alert,

    Finished,

    Invalid,
}

fn tls_state_transition_handshake(state: TlsState, msg: &TlsMessageHandshake) -> Result<TlsState,StateChangeError> {
    match (state,msg) {
        (TlsState::None,             &TlsMessageHandshake::ClientHello(ref msg)) => {
            match msg.session_id {
                Some(_) => Ok(TlsState::AskResumeSession),
                _       => Ok(TlsState::ClientHello)
            }
        },
        // Server certificate
        (TlsState::ClientHello,      &TlsMessageHandshake::ServerHello(_))       => Ok(TlsState::ServerHello),
        (TlsState::ServerHello,      &TlsMessageHandshake::Certificate(_))       => Ok(TlsState::Certificate),
        // Server certificate, no client certificate requested
        (TlsState::Certificate,      &TlsMessageHandshake::ServerKeyExchange(_)) => Ok(TlsState::ServerKeyExchange),
        (TlsState::Certificate,      &TlsMessageHandshake::CertificateStatus(_)) => Ok(TlsState::CertificateSt),
        (TlsState::CertificateSt,    &TlsMessageHandshake::ServerKeyExchange(_)) => Ok(TlsState::ServerKeyExchange),
        (TlsState::ServerKeyExchange,&TlsMessageHandshake::ServerDone(_))        => Ok(TlsState::ServerHelloDone),
        (TlsState::ServerHelloDone  ,&TlsMessageHandshake::ClientKeyExchange(_)) => Ok(TlsState::ClientKeyExchange),
        // Server certificate, client certificate requested
        (TlsState::Certificate,      &TlsMessageHandshake::CertificateRequest(_))=> Ok(TlsState::CRCertRequest),
        (TlsState::ServerKeyExchange,&TlsMessageHandshake::CertificateRequest(_))=> Ok(TlsState::CRCertRequest),
        (TlsState::CRCertRequest,    &TlsMessageHandshake::ServerDone(_))        => Ok(TlsState::CRHelloDone),
        (TlsState::CRHelloDone,      &TlsMessageHandshake::Certificate(_))       => Ok(TlsState::CRCert),
        (TlsState::CRCert,           &TlsMessageHandshake::ClientKeyExchange(_)) => Ok(TlsState::CRClientKeyExchange),
        (TlsState::CRClientKeyExchange, &TlsMessageHandshake::CertificateVerify(_)) => Ok(TlsState::CRCertVerify),
        // Server has no certificate (but accepts anonymous)
        (TlsState::ServerHello,      &TlsMessageHandshake::ServerKeyExchange(_)) => Ok(TlsState::NoCertSKE),
        (TlsState::NoCertSKE,        &TlsMessageHandshake::ServerDone(_))        => Ok(TlsState::NoCertHelloDone),
        (TlsState::NoCertHelloDone,  &TlsMessageHandshake::ClientKeyExchange(_)) => Ok(TlsState::NoCertCKE),
        // PSK
        (TlsState::Certificate,      &TlsMessageHandshake::ServerDone(_))        => Ok(TlsState::PskHelloDone),
        (TlsState::PskHelloDone,     &TlsMessageHandshake::ClientKeyExchange(_)) => Ok(TlsState::PskCKE),
        // Resuming session
        (TlsState::AskResumeSession, &TlsMessageHandshake::ServerHello(_))       => Ok(TlsState::ResumeSession),
        // Resume session failed
        (TlsState::ResumeSession,    &TlsMessageHandshake::Certificate(_))       => Ok(TlsState::Certificate),
        // TLS 1.3 Draft 18 1-RTT
        // Re-use the ClientChangeCipherSpec state to indicate the next message will be encrypted
        (TlsState::ClientHello,      &TlsMessageHandshake::ServerHelloV13Draft18(_))    => Ok(TlsState::ClientChangeCipherSpec),
        // Hello requests must be accepted at any time (except start), but ignored [RFC5246] 7.4.1.1
        (TlsState::None,             &TlsMessageHandshake::HelloRequest)         => Err(StateChangeError::InvalidTransition),
        (s,                          &TlsMessageHandshake::HelloRequest)         => Ok(s),
        // All other transitions are considered invalid
        _ => Err(StateChangeError::InvalidTransition),
    }
}

/// Update the TLS state machine, doing one transition
///
/// Given the previous state and the parsed message, return the new state or a state machine error.
///
/// This state machine only implements the TLS handshake.
///
/// Some transitions only check the new message type, while some others must match the content
/// (for example, to check if the client asked to resume a session).
///
/// If the previous state is `Invalid`, the state machine will not return an error, but keep the
/// same `Invalid` state. This is used to raise error only once if the state machine keeps being
/// updated by new messages.
pub fn tls_state_transition(state: TlsState, msg: &TlsMessage) -> Result<TlsState,StateChangeError> {
    match (state,msg) {
        (TlsState::Invalid,_) => Ok(TlsState::Invalid),
        (TlsState::Finished,_) => Ok(TlsState::Invalid),
        (_,&TlsMessage::Handshake(ref m)) => tls_state_transition_handshake(state,m),
        // Server certificate
        (TlsState::ClientKeyExchange,     &TlsMessage::ChangeCipherSpec) => Ok(TlsState::ClientChangeCipherSpec),
        (TlsState::ClientChangeCipherSpec,&TlsMessage::ChangeCipherSpec) => Ok(TlsState::SessionEncrypted),
        // Server certificate, client certificate requested
        (TlsState::CRClientKeyExchange,   &TlsMessage::ChangeCipherSpec) => Ok(TlsState::ClientChangeCipherSpec),
        (TlsState::CRCertVerify,          &TlsMessage::ChangeCipherSpec) => Ok(TlsState::ClientChangeCipherSpec),
        // No server certificate
        (TlsState::NoCertCKE,             &TlsMessage::ChangeCipherSpec) => Ok(TlsState::ClientChangeCipherSpec),
        // PSK
        (TlsState::PskCKE,                &TlsMessage::ChangeCipherSpec) => Ok(TlsState::ClientChangeCipherSpec),
        // Resume session
        (TlsState::ResumeSession,         &TlsMessage::ChangeCipherSpec) => Ok(TlsState::ClientChangeCipherSpec),
        // non-fatal alerts
        (s,                               &TlsMessage::Alert(ref a)) => {
            if a.severity == TlsAlertSeverity::Warning { Ok(s) } else { Ok(TlsState::Finished) }
        },
        (_,_) => Err(StateChangeError::InvalidTransition),
    }
}
<|endoftext|>


<commit_before>pub fn remove_mean_offset(signal: &[f32]) -> Vec<f32> {
    let mean = signal.iter().sum::<f32>()/signal.len() as f32;
    signal.iter().map(|x| x - mean).collect()
}

pub fn correlation(signal: &[f32]) -> Vec<f32> {
    (0..signal.len()).map(|offset| {
        signal.iter().take(signal.len() - offset)
            .zip(signal.iter().skip(offset))
            .map(|(sig_i, sig_j)| sig_i * sig_j)
            .sum()
    }).collect()
}

pub fn find_fundamental_frequency_correlation(signal: &[f32], sample_rate: f32) -> Option<f32> {
    let normalized_signal = remove_mean_offset(signal);

    if normalized_signal.iter().all(|&x| x.abs() < 0.05) {
        // silence
        return None;
    }

    let correlation = correlation(&normalized_signal);

    let first_peak_end = match correlation.iter().position(|&c| c < 0.0) {
        Some(p) => p,
        None => {
            // musical signals will drop below 0 at some point
            return None
        }
    };

    let peak = correlation.iter()
        .enumerate()
        .skip(first_peak_end)
        .fold((first_peak_end, 0.0), |(xi, xmag), (yi, &ymag)| if ymag > xmag { (yi, ymag) } else { (xi, xmag) });

    let (peak_index, _) = peak;

    let refined_peak_index = refine_fundamentals(&correlation, peak_index as f32 - 0.5, peak_index as f32 + 0.5);

    if is_noise(&correlation, refined_peak_index) {
        None
    }
    else {
        Some(sample_rate / refined_peak_index)
    }
}

fn refine_fundamentals(correlation: &[f32], low_bound: f32, high_bound: f32) -> f32 {
    let data_points = 2 * correlation.len() / high_bound.ceil() as usize;
    let range = high_bound - low_bound;
    let midpoint = (low_bound + high_bound) / 2.0;

    if (range * data_points as f32) < 1.0 {
        midpoint
    }
    else {
        let low_guess = score_guess(correlation, low_bound, data_points);
        let high_guess = score_guess(correlation, high_bound, data_points);

        if high_guess > low_guess {
            refine_fundamentals(correlation, midpoint, high_bound)
        }
        else {
            refine_fundamentals(correlation, low_bound, midpoint)
        }
    }
}

fn is_noise(correlation: &[f32], fundamental: f32) -> bool {
    let value_at_point = interpolate(correlation, fundamental);
    let score_data_points = 2 * correlation.len() / fundamental.ceil() as usize;
    let score = score_guess(correlation, fundamental, score_data_points);

    value_at_point > 2.0*score
}

fn score_guess(correlation: &[f32], period: f32, data_points: usize) -> f32 {
    (1..data_points).map(|i| {
        let expected_sign = if i % 2 == 0 { 1.0 } else { -1.0 };
        let x = i as f32 * period / 2.0;
        let weight = 0.5 * i as f32;
        expected_sign * weight * interpolate(correlation, x)
    }).sum()
}

fn interpolate(correlation: &[f32], x: f32) -> f32 {
    if x < 0.0 {
        correlation[0]
    }
    else if x >= correlation.len() as f32 {
        correlation[correlation.len()-1]
    }
    else {
        let x0 = x.floor();
        let y0 = correlation[x0 as usize];
        let x1 = x.ceil();
        let y1 = correlation[x1 as usize];

        if x0 as usize == x1 as usize {
            y0
        }
        else {
            (y0*(x1-x) + y1*(x-x0)) / (x1-x0)
        }
    }
}

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

    const SAMPLE_RATE: f32 = 44100.0;
    const FRAMES: usize = 512;

    fn frequency_resolution() -> f32 {
        SAMPLE_RATE / 2.0 / FRAMES as f32
    }

    fn sin_arg(f: f32, t: f32, phase: f32) -> f32 {
        2.0 as f32 * PI * f * t + phase
    }

    fn sample_sinusoud(amplitude: f32, frequency: f32, phase: f32) -> Vec<f32> {
        (0..FRAMES)
            .map(|x| {
                let t = x as f32 / SAMPLE_RATE;
                sin_arg(frequency, t, phase).sin() * amplitude
            }).collect()
    }

    #[test]
    fn correlation_on_sine_wave() {
        let frequency = 440.0 as f32; //concert A

        let samples = sample_sinusoud(1.0, frequency, 0.0);
        let fundamental = find_fundamental_frequency_correlation(&samples, SAMPLE_RATE).expect("Find fundamental returned None");
        assert!((fundamental-frequency).abs() < frequency_resolution(), "expected={}, actual={}", frequency, fundamental);
    }

    #[test]
    fn correlation_on_two_sine_waves() {
        //Unfortunately, real signals won't be this neat
        let samples1a = sample_sinusoud(2.0, 440.0, 0.0);
        let samples2a = sample_sinusoud(1.0, 880.0, 0.0);
        let expected_fundamental = 440.0;

        let samples: Vec<f32> = samples1a.iter().zip(samples2a.iter())
            .map(|(a, b)| a+b)
            .collect();

        let fundamental = find_fundamental_frequency_correlation(&samples, SAMPLE_RATE).expect("Find fundamental returned None");

        assert!((fundamental-expected_fundamental).abs() < frequency_resolution(), "expected_fundamental={}, actual={}", expected_fundamental, fundamental);
    }

    #[test]
    fn interpolate_half_way() {
        assert_eq!(0.5, interpolate(&vec!(0.0, 1.0), 0.5))
    }
}

pub fn hz_to_midi_number(hz: f32) -> f32 {
    69.0 + 12.0 * (hz / 440.0).log2()
}

pub fn hz_to_cents_error(hz: f32) -> f32 {
    let midi_number = hz_to_midi_number(hz);
    let cents = (midi_number * 100.0).round() % 100.0;
    if cents >= 50.0 {
        cents - 100.0
    }
    else {
        cents
    }
}

pub fn hz_to_pitch(hz: f32) -> String {
    let pitch_names = [
        "C",
        "C♯",
        "D",
        "E♭",
        "E",
        "F",
        "F♯",
        "G",
        "G♯",
        "A",
        "B♭",
        "B"
    ];

    let midi_number = hz_to_midi_number(hz);
    //midi_number of 0 is C-1.

    let rounded_pitch = midi_number.round() as i32;
    let name = pitch_names[rounded_pitch as usize % pitch_names.len()];
    let octave = rounded_pitch / pitch_names.len() as i32 - 1; //0 is C-1
    if octave < 0 {
        return "< C 1".to_string();
    }

    format!("{: <2}{}", name, octave)
}

#[test]
fn a4_is_correct() {
    assert_eq!(hz_to_pitch(440.0), "A 4");
}

#[test]
fn a2_is_correct() {
    assert_eq!(hz_to_pitch(110.0), "A 2");
}

#[test]
fn c4_is_correct() {
    assert_eq!(hz_to_pitch(261.63), "C 4");
}

#[test]
fn f5_is_correct() {
    assert_eq!(hz_to_pitch(698.46), "F 5");
}

pub fn align_to_rising_edge(samples: &[f32]) -> Vec<f32> {
    remove_mean_offset(samples)
        .iter()
        .skip_while(|x| !x.is_sign_negative())
        .skip_while(|x| x.is_sign_negative())
        .cloned()
        .collect()
}
<commit_msg>Removed unnecessary rounding of error<commit_after>
pub fn remove_mean_offset(signal: &[f32]) -> Vec<f32> {
    let mean = signal.iter().sum::<f32>()/signal.len() as f32;
    signal.iter().map(|x| x - mean).collect()
}

pub fn correlation(signal: &[f32]) -> Vec<f32> {
    (0..signal.len()).map(|offset| {
        signal.iter().take(signal.len() - offset)
            .zip(signal.iter().skip(offset))
            .map(|(sig_i, sig_j)| sig_i * sig_j)
            .sum()
    }).collect()
}

pub fn find_fundamental_frequency_correlation(signal: &[f32], sample_rate: f32) -> Option<f32> {
    let normalized_signal = remove_mean_offset(signal);

    if normalized_signal.iter().all(|&x| x.abs() < 0.05) {
        // silence
        return None;
    }

    let correlation = correlation(&normalized_signal);

    let first_peak_end = match correlation.iter().position(|&c| c < 0.0) {
        Some(p) => p,
        None => {
            // musical signals will drop below 0 at some point
            return None
        }
    };

    let peak = correlation.iter()
        .enumerate()
        .skip(first_peak_end)
        .fold((first_peak_end, 0.0), |(xi, xmag), (yi, &ymag)| if ymag > xmag { (yi, ymag) } else { (xi, xmag) });

    let (peak_index, _) = peak;

    let refined_peak_index = refine_fundamentals(&correlation, peak_index as f32 - 0.5, peak_index as f32 + 0.5);

    if is_noise(&correlation, refined_peak_index) {
        None
    }
    else {
        Some(sample_rate / refined_peak_index)
    }
}

fn refine_fundamentals(correlation: &[f32], low_bound: f32, high_bound: f32) -> f32 {
    let data_points = 2 * correlation.len() / high_bound.ceil() as usize;
    let range = high_bound - low_bound;
    let midpoint = (low_bound + high_bound) / 2.0;

    if (range * data_points as f32) < 1.0 {
        midpoint
    }
    else {
        let low_guess = score_guess(correlation, low_bound, data_points);
        let high_guess = score_guess(correlation, high_bound, data_points);

        if high_guess > low_guess {
            refine_fundamentals(correlation, midpoint, high_bound)
        }
        else {
            refine_fundamentals(correlation, low_bound, midpoint)
        }
    }
}

fn is_noise(correlation: &[f32], fundamental: f32) -> bool {
    let value_at_point = interpolate(correlation, fundamental);
    let score_data_points = 2 * correlation.len() / fundamental.ceil() as usize;
    let score = score_guess(correlation, fundamental, score_data_points);

    value_at_point > 2.0*score
}

fn score_guess(correlation: &[f32], period: f32, data_points: usize) -> f32 {
    (1..data_points).map(|i| {
        let expected_sign = if i % 2 == 0 { 1.0 } else { -1.0 };
        let x = i as f32 * period / 2.0;
        let weight = 0.5 * i as f32;
        expected_sign * weight * interpolate(correlation, x)
    }).sum()
}

fn interpolate(correlation: &[f32], x: f32) -> f32 {
    if x < 0.0 {
        correlation[0]
    }
    else if x >= correlation.len() as f32 {
        correlation[correlation.len()-1]
    }
    else {
        let x0 = x.floor();
        let y0 = correlation[x0 as usize];
        let x1 = x.ceil();
        let y1 = correlation[x1 as usize];

        if x0 as usize == x1 as usize {
            y0
        }
        else {
            (y0*(x1-x) + y1*(x-x0)) / (x1-x0)
        }
    }
}

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

    const SAMPLE_RATE: f32 = 44100.0;
    const FRAMES: usize = 512;

    fn frequency_resolution() -> f32 {
        SAMPLE_RATE / 2.0 / FRAMES as f32
    }

    fn sin_arg(f: f32, t: f32, phase: f32) -> f32 {
        2.0 as f32 * PI * f * t + phase
    }

    fn sample_sinusoud(amplitude: f32, frequency: f32, phase: f32) -> Vec<f32> {
        (0..FRAMES)
            .map(|x| {
                let t = x as f32 / SAMPLE_RATE;
                sin_arg(frequency, t, phase).sin() * amplitude
            }).collect()
    }

    #[test]
    fn correlation_on_sine_wave() {
        let frequency = 440.0 as f32; //concert A

        let samples = sample_sinusoud(1.0, frequency, 0.0);
        let fundamental = find_fundamental_frequency_correlation(&samples, SAMPLE_RATE).expect("Find fundamental returned None");
        assert!((fundamental-frequency).abs() < frequency_resolution(), "expected={}, actual={}", frequency, fundamental);
    }

    #[test]
    fn correlation_on_two_sine_waves() {
        //Unfortunately, real signals won't be this neat
        let samples1a = sample_sinusoud(2.0, 440.0, 0.0);
        let samples2a = sample_sinusoud(1.0, 880.0, 0.0);
        let expected_fundamental = 440.0;

        let samples: Vec<f32> = samples1a.iter().zip(samples2a.iter())
            .map(|(a, b)| a+b)
            .collect();

        let fundamental = find_fundamental_frequency_correlation(&samples, SAMPLE_RATE).expect("Find fundamental returned None");

        assert!((fundamental-expected_fundamental).abs() < frequency_resolution(), "expected_fundamental={}, actual={}", expected_fundamental, fundamental);
    }

    #[test]
    fn interpolate_half_way() {
        assert_eq!(0.5, interpolate(&vec!(0.0, 1.0), 0.5))
    }
}

pub fn hz_to_midi_number(hz: f32) -> f32 {
    69.0 + 12.0 * (hz / 440.0).log2()
}

pub fn hz_to_cents_error(hz: f32) -> f32 {
    let midi_number = hz_to_midi_number(hz);
    let cents = (midi_number % 1.0) * 100.0;
    if cents >= 50.0 {
        cents - 100.0
    }
    else {
        cents
    }
}

pub fn hz_to_pitch(hz: f32) -> String {
    let pitch_names = [
        "C",
        "C♯",
        "D",
        "E♭",
        "E",
        "F",
        "F♯",
        "G",
        "G♯",
        "A",
        "B♭",
        "B"
    ];

    let midi_number = hz_to_midi_number(hz);
    //midi_number of 0 is C-1.

    let rounded_pitch = midi_number.round() as i32;
    let name = pitch_names[rounded_pitch as usize % pitch_names.len()];
    let octave = rounded_pitch / pitch_names.len() as i32 - 1; //0 is C-1
    if octave < 0 {
        return "< C 1".to_string();
    }

    format!("{: <2}{}", name, octave)
}

#[test]
fn a4_is_correct() {
    assert_eq!(hz_to_pitch(440.0), "A 4");
}

#[test]
fn a2_is_correct() {
    assert_eq!(hz_to_pitch(110.0), "A 2");
}

#[test]
fn c4_is_correct() {
    assert_eq!(hz_to_pitch(261.63), "C 4");
}

#[test]
fn f5_is_correct() {
    assert_eq!(hz_to_pitch(698.46), "F 5");
}

pub fn align_to_rising_edge(samples: &[f32]) -> Vec<f32> {
    remove_mean_offset(samples)
        .iter()
        .skip_while(|x| !x.is_sign_negative())
        .skip_while(|x| x.is_sign_negative())
        .cloned()
        .collect()
}
<|endoftext|>