<commit_before>use super::{SeasonState, SeriesConfig};
use backend::{AnimeEntry, AnimeInfo, SyncBackend};
use error::SeriesError;
use input::{self, Answer};
use regex::Regex;
use std::borrow::Cow;
use std::collections::HashMap;
use std::fs;
use std::ops::Range;
use std::path::{Path, PathBuf};
use toml;

pub struct FolderData {
    pub series_info: SeriesInfo,
    pub savefile: SaveData,
    pub subseries: String,
    pub path: PathBuf,
}

impl FolderData {
    const DEFAULT_SUBSERIES_NAME: &'static str = "main";

    pub fn load_dir(path: &Path, subseries: Option<String>) -> Result<FolderData, SeriesError> {
        let subseries = subseries.unwrap_or_else(|| FolderData::DEFAULT_SUBSERIES_NAME.into());

        let mut savefile = SaveData::from_dir(path)?;
        let series_info = FolderData::load_series_info(path, &mut savefile, &subseries)?;

        Ok(FolderData {
            series_info,
            savefile,
            subseries,
            path: PathBuf::from(path),
        })
    }

    fn load_series_info(
        path: &Path,
        savefile: &mut SaveData,
        subseries: &str,
    ) -> Result<SeriesInfo, SeriesError> {
        let mut ep_data = parse_episode_files_until_valid(path, &mut savefile.episode_matcher)?;
        let subseries_data = savefile.get_mut_subseries_entry(subseries);

        if let Some(info) = SeriesInfo::select_from_subseries(&mut ep_data, &subseries_data) {
            return Ok(info);
        }

        let info = prompt_select_series_info(ep_data)?;
        subseries_data.files_title = Some(info.name.clone());

        Ok(info)
    }

    pub fn populate_season_data<B>(&mut self, config: &SeriesConfig<B>) -> Result<(), SeriesError>
    where
        B: SyncBackend,
    {
        let num_seasons = self.seasons().len();

        if num_seasons > config.season_num {
            return Ok(());
        }

        for cur_season in num_seasons..=config.season_num {
            let info = AnimeInfo::default();
            let entry = AnimeEntry::new(info);

            let mut season = SeasonState {
                state: entry,
                needs_info: true,
                needs_sync: config.offline_mode,
            };

            season.sync_info_from_remote(config, &self, cur_season)?;

            let subseries = self
                .savefile
                .get_mut_subseries_entry(self.subseries.clone());

            subseries.season_states.push(season);
        }

        Ok(())
    }

    pub fn calculate_season_offset(&self, mut range: Range<usize>) -> u32 {
        let num_seasons = self.seasons().len();
        range.start = num_seasons.min(range.start);
        range.end = num_seasons.min(range.end);

        let mut offset = 0;
        let seasons = self.seasons();

        for i in range {
            let season = &seasons[i];

            match season.state.info.episodes {
                Some(eps) => offset += eps,
                None => return offset,
            }
        }

        offset
    }

    pub fn try_remove_dir(&self) {
        let path = self.path.to_string_lossy();

        println!("WARNING: {} will be deleted", path);
        println!("is this ok? (y/N)");

        match input::read_yn(Answer::No) {
            Ok(true) => match fs::remove_dir_all(&self.path) {
                Ok(_) => (),
                Err(err) => {
                    eprintln!("failed to remove directory: {}", err);
                }
            },
            Ok(false) => (),
            Err(err) => {
                eprintln!("failed to read input: {}", err);
            }
        }
    }

    pub fn get_episode(&self, episode: u32) -> Result<&PathBuf, SeriesError> {
        self.series_info.get_episode(episode)
    }

    pub fn save(&self) -> Result<(), SeriesError> {
        self.savefile.write_to_file()
    }

    pub fn seasons(&self) -> &Vec<SeasonState> {
        &self.savefile.subseries[&self.subseries].season_states
    }

    pub fn seasons_mut(&mut self) -> &mut Vec<SeasonState> {
        &mut self
            .savefile
            .get_mut_subseries_entry(self.subseries.clone())
            .season_states
    }
}

pub type SubSeriesName = String;

#[derive(Serialize, Deserialize)]
pub struct SaveData {
    pub episode_matcher: Option<String>,
    #[serde(flatten)]
    pub subseries: HashMap<SubSeriesName, SubSeriesData>,
    #[serde(skip)]
    pub path: PathBuf,
}

impl SaveData {
    const DATA_FILE_NAME: &'static str = ".anup";

    pub fn new(path: PathBuf) -> SaveData {
        SaveData {
            episode_matcher: None,
            subseries: HashMap::new(),
            path,
        }
    }

    pub fn get_mut_subseries_entry<S>(&mut self, name: S) -> &mut SubSeriesData
    where
        S: Into<String>,
    {
        self.subseries
            .entry(name.into())
            .or_insert_with(SubSeriesData::new)
    }

    pub fn from_dir(path: &Path) -> Result<SaveData, SeriesError> {
        let path = PathBuf::from(path).join(SaveData::DATA_FILE_NAME);

        if !path.exists() {
            return Ok(SaveData::new(path));
        }

        let file_contents = fs::read_to_string(&path)?;

        let mut save_data: SaveData = toml::from_str(&file_contents)?;
        save_data.path = path;

        Ok(save_data)
    }

    pub fn write_to_file(&self) -> Result<(), SeriesError> {
        let toml = toml::to_string(self)?;
        fs::write(&self.path, toml)?;

        Ok(())
    }
}

#[derive(Serialize, Deserialize)]
pub struct SubSeriesData {
    pub files_title: Option<String>,
    pub season_states: Vec<SeasonState>,
}

impl SubSeriesData {
    pub fn new() -> SubSeriesData {
        SubSeriesData {
            files_title: None,
            season_states: Vec::new(),
        }
    }
}

pub struct SeriesInfo {
    pub name: String,
    pub episodes: HashMap<u32, PathBuf>,
}

impl SeriesInfo {
    pub fn get_episode(&self, episode: u32) -> Result<&PathBuf, SeriesError> {
        self.episodes
            .get(&episode)
            .ok_or_else(|| SeriesError::EpisodeNotFound(episode))
    }

    pub fn select_from_subseries(
        ep_data: &mut SeriesEpisodes,
        subseries: &SubSeriesData,
    ) -> Option<SeriesInfo> {
        if let Some(name) = &subseries.files_title {
            let entry = ep_data.remove_entry(name);

            if let Some((name, episodes)) = entry {
                return Some(SeriesInfo { name, episodes });
            }
        }

        None
    }
}

struct EpisodeFile {
    series_name: String,
    episode_num: u32,
}

impl EpisodeFile {
    fn parse(path: &Path, matcher: &Regex) -> Result<EpisodeFile, SeriesError> {
        // Replace certain characters with spaces since they can prevent proper series
        // identification or prevent it from being found on a sync backend
        let filename = path
            .file_name()
            .and_then(|path| path.to_str())
            .ok_or(SeriesError::UnableToGetFilename)?;

        let caps = matcher
            .captures(&filename)
            .ok_or(SeriesError::EpisodeRegexCaptureFailed)?;

        let series_name = caps
            .name("name")
            .ok_or_else(|| SeriesError::UnknownRegexCapture("name".into()))
            .map(|name| EpisodeFile::cleanup_title(name.as_str()))?;

        let episode = match caps.name("episode") {
            Some(capture) => capture
                .as_str()
                .parse()
                .map_err(SeriesError::EpisodeNumParseFailed)?,
            None => 1,
        };

        Ok(EpisodeFile {
            series_name,
            episode_num: episode,
        })
    }

    fn cleanup_title(title: &str) -> String {
        let mut clean_title = String::from(title);
        let bytes = unsafe { clean_title.as_bytes_mut() };

        for byte in bytes {
            match byte {
                b'.' | b'_' => *byte = b' ',
                _ => (),
            }
        }

        clean_title
    }
}

pub type EpisodePaths = HashMap<u32, PathBuf>;
pub type SeriesEpisodes = HashMap<String, EpisodePaths>;

pub fn prompt_select_series_info(info: SeriesEpisodes) -> Result<SeriesInfo, SeriesError> {
    if info.is_empty() {
        return Err(SeriesError::NoSeriesFound);
    }

    let mut info = info
        .into_iter()
        .map(|(name, eps)| SeriesInfo {
            name,
            episodes: eps,
        })
        .collect::<Vec<_>>();

    if info.len() == 1 {
        return Ok(info.swap_remove(0));
    }

    println!("multiple series found in directory");
    println!("please enter the number next to the episode files you want to use:");

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

    let index = input::read_range(1, info.len())? - 1;
    let series = info.swap_remove(index);

    Ok(series)
}

// This default pattern will match episodes in several common formats, such as:
// [Group] Series Name - 01.mkv
// [Group]_Series_Name_-_01.mkv
// [Group].Series.Name.-.01.mkv
// [Group] Series Name - 01 [tag 1][tag 2].mkv
// [Group]_Series_Name_-_01_[tag1][tag2].mkv
// [Group].Series.Name.-.01.[tag1][tag2].mkv
// Series Name - 01.mkv
// Series_Name_-_01.mkv
// Series.Name.-.01.mkv
const EP_FORMAT_REGEX: &str =
    r"(?:\[.+?\](?:_+|\.+|\s*))?(?P<name>.+?)(?:\s*|_*|\.*)-(?:\s*|_*|\.*)(?P<episode>\d+)";

fn format_episode_parser_regex<'a, S>(pattern: Option<S>) -> Result<Cow<'a, Regex>, SeriesError>
where
    S: AsRef<str>,
{
    lazy_static! {
        static ref EP_FORMAT: Regex = Regex::new(EP_FORMAT_REGEX).unwrap();
    }

    match pattern {
        Some(pattern) => {
            let pattern = pattern
                .as_ref()
                .replace("{name}", "(?P<name>.+?)")
                .replace("{episode}", r"(?P<episode>\d+)");

            let regex = Regex::new(&pattern)?;
            Ok(Cow::Owned(regex))
        }
        None => Ok(Cow::Borrowed(&*EP_FORMAT)),
    }
}

pub fn parse_episode_files<S>(
    path: &Path,
    pattern: Option<S>,
) -> Result<SeriesEpisodes, SeriesError>
where
    S: AsRef<str>,
{
    if !path.is_dir() {
        return Err(SeriesError::NotADirectory(path.to_string_lossy().into()));
    }

    let pattern = format_episode_parser_regex(pattern)?;
    let mut data = HashMap::with_capacity(1);

    for entry in fs::read_dir(path).map_err(SeriesError::Io)? {
        let entry = entry.map_err(SeriesError::Io)?.path();

        if !entry.is_file() {
            continue;
        }

        let episode = match EpisodeFile::parse(&entry, pattern.as_ref()) {
            Ok(episode) => episode,
            Err(SeriesError::EpisodeRegexCaptureFailed) => continue,
            Err(err) => return Err(err),
        };

        let series = data
            .entry(episode.series_name)
            .or_insert_with(|| HashMap::with_capacity(1));

        series.insert(episode.episode_num, entry);
    }

    if data.is_empty() {
        return Err(SeriesError::NoSeriesFound);
    }

    Ok(data)
}

pub fn parse_episode_files_until_valid<S>(
    path: &Path,
    pattern: &mut Option<S>,
) -> Result<SeriesEpisodes, SeriesError>
where
    S: AsRef<str> + From<String>,
{
    loop {
        match parse_episode_files(path, pattern.as_ref()) {
            Ok(data) => break Ok(data),
            Err(SeriesError::NoSeriesFound) => {
                eprintln!("error: no series found");
                println!("you will now be prompted to enter a custom regex pattern");
                println!("when entering the pattern, please mark the series name and episode number with {{name}} and {{episode}}, respectively");
                println!("note: if you're trying to play a single file (like a movie) and it doesn't have an episode number, you can omit the {{episode}} marker");
                println!("example:");
                println!("  filename: [SubGroup] Series Name - Ep01 [1080p].mkv");
                println!(r"  pattern: \[.+?\] {{name}} - Ep{{episode}} \[1080p\].mkv");
                println!("please enter your custom pattern:");

                *pattern = Some(input::read_line()?.into());
            }
            Err(err @ SeriesError::Regex(_)) | Err(err @ SeriesError::UnknownRegexCapture(_)) => {
                eprintln!("error parsing regex pattern: {}", err);
                println!("please try again:");

                *pattern = Some(input::read_line()?.into());
            }
            Err(err) => return Err(err),
        }
    }
}
<commit_msg>[series] Remove duplicate code from parse_episode_files_until_valid()<commit_after>use super::{SeasonState, SeriesConfig};
use backend::{AnimeEntry, AnimeInfo, SyncBackend};
use error::SeriesError;
use input::{self, Answer};
use regex::Regex;
use std::borrow::Cow;
use std::collections::HashMap;
use std::fs;
use std::ops::Range;
use std::path::{Path, PathBuf};
use toml;

pub struct FolderData {
    pub series_info: SeriesInfo,
    pub savefile: SaveData,
    pub subseries: String,
    pub path: PathBuf,
}

impl FolderData {
    const DEFAULT_SUBSERIES_NAME: &'static str = "main";

    pub fn load_dir(path: &Path, subseries: Option<String>) -> Result<FolderData, SeriesError> {
        let subseries = subseries.unwrap_or_else(|| FolderData::DEFAULT_SUBSERIES_NAME.into());

        let mut savefile = SaveData::from_dir(path)?;
        let series_info = FolderData::load_series_info(path, &mut savefile, &subseries)?;

        Ok(FolderData {
            series_info,
            savefile,
            subseries,
            path: PathBuf::from(path),
        })
    }

    fn load_series_info(
        path: &Path,
        savefile: &mut SaveData,
        subseries: &str,
    ) -> Result<SeriesInfo, SeriesError> {
        let mut ep_data = parse_episode_files_until_valid(path, &mut savefile.episode_matcher)?;
        let subseries_data = savefile.get_mut_subseries_entry(subseries);

        if let Some(info) = SeriesInfo::select_from_subseries(&mut ep_data, &subseries_data) {
            return Ok(info);
        }

        let info = prompt_select_series_info(ep_data)?;
        subseries_data.files_title = Some(info.name.clone());

        Ok(info)
    }

    pub fn populate_season_data<B>(&mut self, config: &SeriesConfig<B>) -> Result<(), SeriesError>
    where
        B: SyncBackend,
    {
        let num_seasons = self.seasons().len();

        if num_seasons > config.season_num {
            return Ok(());
        }

        for cur_season in num_seasons..=config.season_num {
            let info = AnimeInfo::default();
            let entry = AnimeEntry::new(info);

            let mut season = SeasonState {
                state: entry,
                needs_info: true,
                needs_sync: config.offline_mode,
            };

            season.sync_info_from_remote(config, &self, cur_season)?;

            let subseries = self
                .savefile
                .get_mut_subseries_entry(self.subseries.clone());

            subseries.season_states.push(season);
        }

        Ok(())
    }

    pub fn calculate_season_offset(&self, mut range: Range<usize>) -> u32 {
        let num_seasons = self.seasons().len();
        range.start = num_seasons.min(range.start);
        range.end = num_seasons.min(range.end);

        let mut offset = 0;
        let seasons = self.seasons();

        for i in range {
            let season = &seasons[i];

            match season.state.info.episodes {
                Some(eps) => offset += eps,
                None => return offset,
            }
        }

        offset
    }

    pub fn try_remove_dir(&self) {
        let path = self.path.to_string_lossy();

        println!("WARNING: {} will be deleted", path);
        println!("is this ok? (y/N)");

        match input::read_yn(Answer::No) {
            Ok(true) => match fs::remove_dir_all(&self.path) {
                Ok(_) => (),
                Err(err) => {
                    eprintln!("failed to remove directory: {}", err);
                }
            },
            Ok(false) => (),
            Err(err) => {
                eprintln!("failed to read input: {}", err);
            }
        }
    }

    pub fn get_episode(&self, episode: u32) -> Result<&PathBuf, SeriesError> {
        self.series_info.get_episode(episode)
    }

    pub fn save(&self) -> Result<(), SeriesError> {
        self.savefile.write_to_file()
    }

    pub fn seasons(&self) -> &Vec<SeasonState> {
        &self.savefile.subseries[&self.subseries].season_states
    }

    pub fn seasons_mut(&mut self) -> &mut Vec<SeasonState> {
        &mut self
            .savefile
            .get_mut_subseries_entry(self.subseries.clone())
            .season_states
    }
}

pub type SubSeriesName = String;

#[derive(Serialize, Deserialize)]
pub struct SaveData {
    pub episode_matcher: Option<String>,
    #[serde(flatten)]
    pub subseries: HashMap<SubSeriesName, SubSeriesData>,
    #[serde(skip)]
    pub path: PathBuf,
}

impl SaveData {
    const DATA_FILE_NAME: &'static str = ".anup";

    pub fn new(path: PathBuf) -> SaveData {
        SaveData {
            episode_matcher: None,
            subseries: HashMap::new(),
            path,
        }
    }

    pub fn get_mut_subseries_entry<S>(&mut self, name: S) -> &mut SubSeriesData
    where
        S: Into<String>,
    {
        self.subseries
            .entry(name.into())
            .or_insert_with(SubSeriesData::new)
    }

    pub fn from_dir(path: &Path) -> Result<SaveData, SeriesError> {
        let path = PathBuf::from(path).join(SaveData::DATA_FILE_NAME);

        if !path.exists() {
            return Ok(SaveData::new(path));
        }

        let file_contents = fs::read_to_string(&path)?;

        let mut save_data: SaveData = toml::from_str(&file_contents)?;
        save_data.path = path;

        Ok(save_data)
    }

    pub fn write_to_file(&self) -> Result<(), SeriesError> {
        let toml = toml::to_string(self)?;
        fs::write(&self.path, toml)?;

        Ok(())
    }
}

#[derive(Serialize, Deserialize)]
pub struct SubSeriesData {
    pub files_title: Option<String>,
    pub season_states: Vec<SeasonState>,
}

impl SubSeriesData {
    pub fn new() -> SubSeriesData {
        SubSeriesData {
            files_title: None,
            season_states: Vec::new(),
        }
    }
}

pub struct SeriesInfo {
    pub name: String,
    pub episodes: HashMap<u32, PathBuf>,
}

impl SeriesInfo {
    pub fn get_episode(&self, episode: u32) -> Result<&PathBuf, SeriesError> {
        self.episodes
            .get(&episode)
            .ok_or_else(|| SeriesError::EpisodeNotFound(episode))
    }

    pub fn select_from_subseries(
        ep_data: &mut SeriesEpisodes,
        subseries: &SubSeriesData,
    ) -> Option<SeriesInfo> {
        if let Some(name) = &subseries.files_title {
            let entry = ep_data.remove_entry(name);

            if let Some((name, episodes)) = entry {
                return Some(SeriesInfo { name, episodes });
            }
        }

        None
    }
}

struct EpisodeFile {
    series_name: String,
    episode_num: u32,
}

impl EpisodeFile {
    fn parse(path: &Path, matcher: &Regex) -> Result<EpisodeFile, SeriesError> {
        // Replace certain characters with spaces since they can prevent proper series
        // identification or prevent it from being found on a sync backend
        let filename = path
            .file_name()
            .and_then(|path| path.to_str())
            .ok_or(SeriesError::UnableToGetFilename)?;

        let caps = matcher
            .captures(&filename)
            .ok_or(SeriesError::EpisodeRegexCaptureFailed)?;

        let series_name = caps
            .name("name")
            .ok_or_else(|| SeriesError::UnknownRegexCapture("name".into()))
            .map(|name| EpisodeFile::cleanup_title(name.as_str()))?;

        let episode = match caps.name("episode") {
            Some(capture) => capture
                .as_str()
                .parse()
                .map_err(SeriesError::EpisodeNumParseFailed)?,
            None => 1,
        };

        Ok(EpisodeFile {
            series_name,
            episode_num: episode,
        })
    }

    fn cleanup_title(title: &str) -> String {
        let mut clean_title = String::from(title);
        let bytes = unsafe { clean_title.as_bytes_mut() };

        for byte in bytes {
            match byte {
                b'.' | b'_' => *byte = b' ',
                _ => (),
            }
        }

        clean_title
    }
}

pub type EpisodePaths = HashMap<u32, PathBuf>;
pub type SeriesEpisodes = HashMap<String, EpisodePaths>;

pub fn prompt_select_series_info(info: SeriesEpisodes) -> Result<SeriesInfo, SeriesError> {
    if info.is_empty() {
        return Err(SeriesError::NoSeriesFound);
    }

    let mut info = info
        .into_iter()
        .map(|(name, eps)| SeriesInfo {
            name,
            episodes: eps,
        })
        .collect::<Vec<_>>();

    if info.len() == 1 {
        return Ok(info.swap_remove(0));
    }

    println!("multiple series found in directory");
    println!("please enter the number next to the episode files you want to use:");

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

    let index = input::read_range(1, info.len())? - 1;
    let series = info.swap_remove(index);

    Ok(series)
}

// This default pattern will match episodes in several common formats, such as:
// [Group] Series Name - 01.mkv
// [Group]_Series_Name_-_01.mkv
// [Group].Series.Name.-.01.mkv
// [Group] Series Name - 01 [tag 1][tag 2].mkv
// [Group]_Series_Name_-_01_[tag1][tag2].mkv
// [Group].Series.Name.-.01.[tag1][tag2].mkv
// Series Name - 01.mkv
// Series_Name_-_01.mkv
// Series.Name.-.01.mkv
const EP_FORMAT_REGEX: &str =
    r"(?:\[.+?\](?:_+|\.+|\s*))?(?P<name>.+?)(?:\s*|_*|\.*)-(?:\s*|_*|\.*)(?P<episode>\d+)";

fn format_episode_parser_regex<'a, S>(pattern: Option<S>) -> Result<Cow<'a, Regex>, SeriesError>
where
    S: AsRef<str>,
{
    lazy_static! {
        static ref EP_FORMAT: Regex = Regex::new(EP_FORMAT_REGEX).unwrap();
    }

    match pattern {
        Some(pattern) => {
            let pattern = pattern
                .as_ref()
                .replace("{name}", "(?P<name>.+?)")
                .replace("{episode}", r"(?P<episode>\d+)");

            let regex = Regex::new(&pattern)?;
            Ok(Cow::Owned(regex))
        }
        None => Ok(Cow::Borrowed(&*EP_FORMAT)),
    }
}

pub fn parse_episode_files<S>(
    path: &Path,
    pattern: Option<S>,
) -> Result<SeriesEpisodes, SeriesError>
where
    S: AsRef<str>,
{
    if !path.is_dir() {
        return Err(SeriesError::NotADirectory(path.to_string_lossy().into()));
    }

    let pattern = format_episode_parser_regex(pattern)?;
    let mut data = HashMap::with_capacity(1);

    for entry in fs::read_dir(path).map_err(SeriesError::Io)? {
        let entry = entry.map_err(SeriesError::Io)?.path();

        if !entry.is_file() {
            continue;
        }

        let episode = match EpisodeFile::parse(&entry, pattern.as_ref()) {
            Ok(episode) => episode,
            Err(SeriesError::EpisodeRegexCaptureFailed) => continue,
            Err(err) => return Err(err),
        };

        let series = data
            .entry(episode.series_name)
            .or_insert_with(|| HashMap::with_capacity(1));

        series.insert(episode.episode_num, entry);
    }

    if data.is_empty() {
        return Err(SeriesError::NoSeriesFound);
    }

    Ok(data)
}

pub fn parse_episode_files_until_valid<S>(
    path: &Path,
    pattern: &mut Option<S>,
) -> Result<SeriesEpisodes, SeriesError>
where
    S: AsRef<str> + From<String>,
{
    loop {
        match parse_episode_files(path, pattern.as_ref()) {
            Ok(data) => return Ok(data),
            Err(SeriesError::NoSeriesFound) => {
                eprintln!("error: no series found");
                println!("you will now be prompted to enter a custom regex pattern");
                println!("when entering the pattern, please mark the series name and episode number with {{name}} and {{episode}}, respectively");
                println!("note: if you're trying to play a single file (like a movie) and it doesn't have an episode number, you can omit the {{episode}} marker");
                println!("example:");
                println!("  filename: [SubGroup] Series Name - Ep01 [1080p].mkv");
                println!(r"  pattern: \[.+?\] {{name}} - Ep{{episode}} \[1080p\].mkv");
                println!("please enter your custom pattern:");
            }
            Err(err @ SeriesError::Regex(_)) | Err(err @ SeriesError::UnknownRegexCapture(_)) => {
                eprintln!("error parsing regex pattern: {}", err);
                println!("please try again:");
            }
            Err(err) => return Err(err),
        }

        let line = input::read_line()?;
        *pattern = Some(line.into());
    }
}
<|endoftext|>


<commit_before>mod endpoint;
mod model;

use crate::database::settings::{PRIVACY_ADMIN_ONLY, PRIVACY_EVERYONE, PRIVACY_STAFF_ONLY};
use crate::discord::util::is_staff;
use crate::server::endpoint::ban::{get_ban, get_bans, update_ban};
use crate::server::endpoint::captcha::{get_captcha_page, submit_captcha};
use crate::server::endpoint::hardban::{get_hardban, get_hardbans, update_hardban};
use crate::server::endpoint::kick::{get_kick, get_kicks, update_kick};
use crate::server::endpoint::login::login;
use crate::server::endpoint::mute::{get_mute, get_mutes, update_mute};
use crate::server::endpoint::self_user::get_self;
use crate::server::endpoint::settings::{get_setting, reset_setting, update_setting};
use crate::server::endpoint::softban::{get_softban, get_softbans, update_softban};
use crate::server::endpoint::warn::{get_warn, get_warns, update_warn};
use crate::service::guild::{CachedMember, GuildService};
use crate::service::invalid_uuid::InvalidUUIDService;
use crate::service::setting::SettingService;
use crate::util::now;
use crate::Config;
use actix_cors::Cors;
use actix_web::http::header::{HeaderName, CONTENT_TYPE};
use actix_web::{get, web, App, HttpRequest, HttpResponse, HttpServer, Responder};
use jsonwebtoken::errors::ErrorKind;
use jsonwebtoken::{decode, encode, Algorithm, DecodingKey, EncodingKey, Header, Validation};
use lazy_static::lazy_static;
use serde::{Deserialize, Serialize};
use serenity::model::id::{GuildId, UserId};
use serenity::model::Permissions;
use std::error::Error;
use std::net::Ipv4Addr;
use std::sync::Arc;
use tracing::error;
use typemap_rev::TypeMap;
use uuid::Uuid;

lazy_static! {
    static ref VALIDATION: Validation = Validation {
        algorithms: vec![Algorithm::HS512],
        ..Validation::default()
    };
    static ref ENCODING_HEADER: Header = Header::new(Algorithm::HS512);
}

#[derive(Deserialize, Serialize)]
struct JwtClaims {
    #[serde(rename = "userId")]
    user_id: String,
    uuid: Uuid,
    exp: u64,
}

fn verify_token(secret: &str, token: &str) -> Option<JwtClaims> {
    match decode::<JwtClaims>(
        token,
        &DecodingKey::from_secret(secret.as_bytes()),
        &VALIDATION,
    ) {
        Ok(data) => Some(data.claims),
        Err(error) => match error.kind() {
            ErrorKind::InvalidSignature | ErrorKind::ExpiredSignature => None,
            _ => {
                error!("failed to decode jwt token {:?}", error.kind());
                None
            }
        },
    }
}

fn generate_token(secret: &str, user_id: UserId) -> Option<String> {
    let user_id = user_id.to_string();
    let expires_at = now() + 7 * 24 * 60 * 60;
    let uuid = Uuid::new_v4();
    let claims = JwtClaims {
        user_id,
        uuid,
        exp: expires_at,
    };

    match encode(
        &ENCODING_HEADER,
        &claims,
        &EncodingKey::from_secret(secret.as_bytes()),
    ) {
        Ok(token) => Some(token),
        Err(err) => {
            error!("failed to create jwt token {:?}", err);
            None
        }
    }
}

async fn is_authenticated(
    config: &Config,
    services: &TypeMap,
    req: &web::HttpRequest,
) -> Option<UserId> {
    let token = req
        .headers()
        .get(HeaderName::from_static("token")) // optimize into constant?
        .map(|value| value.to_str())
        .map(|result| result.ok())
        .flatten()?;

    let claims = verify_token(&config.server_secret, token)?;

    let invalid_uuid_service = services.get::<InvalidUUIDService>()?;

    if invalid_uuid_service.is_uuid_invalid(claims.uuid).await {
        None
    } else {
        match claims.user_id.parse::<u64>() {
            Ok(id) => Some(UserId(id)),
            Err(_) => {
                // secret leaked or we have a problem with token generation
                error!("received a token with valid signature with invalid user id");
                None
            }
        }
    }
}

pub async fn check_authentication(
    config: &Config,
    services: &TypeMap,
    req: &web::HttpRequest,
) -> Result<UserId, HttpResponse> {
    is_authenticated(config, services, req)
        .await
        .ok_or_else(|| HttpResponse::Unauthorized().json("Invalid token!"))
}

pub enum PrivateEndpointKind {
    ModLog,
    Settings,
}

async fn is_authorized_to_view_private_endpoint(
    services: &TypeMap,
    guild_id: GuildId,
    permissions: Permissions,
    endpoint_kind: PrivateEndpointKind,
) -> bool {
    let setting_service = if let Some(service) = services.get::<SettingService>() {
        service
    } else {
        return false;
    };

    let setting = setting_service.get_setting(guild_id).await;
    let privacy_setting = match endpoint_kind {
        PrivateEndpointKind::Settings => setting.privacy_settings,
        PrivateEndpointKind::ModLog => setting.privacy_mod_log,
    };

    if privacy_setting == PRIVACY_EVERYONE {
        true
    } else if privacy_setting == PRIVACY_STAFF_ONLY {
        is_staff(permissions)
    } else if privacy_setting == PRIVACY_ADMIN_ONLY {
        permissions.administrator()
    } else {
        false
    }
}

pub async fn check_authorization_to_view_private_endpoint(
    services: &TypeMap,
    guild_id: GuildId,
    permissions: Permissions,
    endpoint_kind: PrivateEndpointKind,
) -> Option<HttpResponse> {
    if is_authorized_to_view_private_endpoint(services, guild_id, permissions, endpoint_kind).await
    {
        None
    } else {
        Some(
            HttpResponse::Forbidden()
                .json("Server settings prevent you from viewing private information!"),
        )
    }
}

pub async fn is_guild_endpoint(
    services: &TypeMap,
    guild_id: GuildId,
    user_id: UserId,
) -> Option<(Arc<CachedMember>, Permissions)> {
    let guild_service = services.get::<GuildService>()?;

    let member = guild_service.get_member(guild_id, user_id).await.ok()?;

    let permissions = if let Ok(permissions) = guild_service
        .get_permissions(user_id, &member.roles, guild_id)
        .await
    {
        permissions
    } else {
        return None;
    };

    Some((member, permissions))
}

pub async fn check_guild_endpoint(
    services: &TypeMap,
    guild_id: GuildId,
    user_id: UserId,
) -> Result<(Arc<CachedMember>, Permissions), HttpResponse> {
    is_guild_endpoint(services, guild_id, user_id)
        .await
        .ok_or_else(|| {
            HttpResponse::Forbidden().json("This server either doesn't exist or you aren't in it!")
        })
}

pub async fn apply_private_endpoint_fetch_checks(
    config: &Config,
    services: &TypeMap,
    req: &HttpRequest,
    guild_id: GuildId,
    endpoint_kind: PrivateEndpointKind,
) -> Result<(), HttpResponse> {
    let user_id = match check_authentication(config, services, req).await {
        Ok(user_id) => user_id,
        Err(response) => return Err(response),
    };

    let (_, permissions) = match check_guild_endpoint(services, guild_id, user_id).await {
        Ok((member, permissions)) => (member, permissions),
        Err(response) => return Err(response),
    };

    if let Some(response) =
        check_authorization_to_view_private_endpoint(services, guild_id, permissions, endpoint_kind)
            .await
    {
        return Err(response);
    }

    Ok(())
}

pub async fn apply_mod_log_update_checks(
    config: &Config,
    services: &TypeMap,
    req: &HttpRequest,
    guild_id: GuildId,
    required_permission: Permissions,
) -> Result<(), HttpResponse> {
    let user_id = match check_authentication(config, services, req).await {
        Ok(user_id) => user_id,
        Err(response) => return Err(response),
    };

    let (_, permissions) = match check_guild_endpoint(services, guild_id, user_id).await {
        Ok((member, permissions)) => (member, permissions),
        Err(response) => return Err(response),
    };

    if permissions.contains(Permissions::ADMINISTRATOR) || permissions.contains(required_permission)
    {
        Ok(())
    } else {
        Err(HttpResponse::Forbidden()
            .json("You don't have permissions to change moderator log history!"))
    }
}

pub async fn apply_setting_update_checks(
    config: &Config,
    services: &TypeMap,
    req: &HttpRequest,
    guild_id: GuildId,
) -> Result<(), HttpResponse> {
    let user_id = match check_authentication(config, services, req).await {
        Ok(user_id) => user_id,
        Err(response) => return Err(response),
    };

    let (_, permissions) = match check_guild_endpoint(services, guild_id, user_id).await {
        Ok((member, permissions)) => (member, permissions),
        Err(response) => return Err(response),
    };

    if permissions.contains(Permissions::ADMINISTRATOR) {
        Ok(())
    } else {
        Err(HttpResponse::Forbidden().json("You don't have permissions to change server settings!"))
    }
}

#[get("/")]
async fn index() -> impl Responder {
    "Welcome to Safety Jim API."
}

pub async fn run_server(config: Arc<Config>, services: Arc<TypeMap>) -> Result<(), Box<dyn Error>> {
    HttpServer::new(move || {
        App::new()
            .app_data(web::Data::from(config.clone()))
            .app_data(web::Data::from(services.clone()))
            .service(index)
            .service(get_self)
            .service(get_bans)
            .service(get_ban)
            .service(update_ban)
            .service(get_captcha_page)
            .service(submit_captcha)
            .service(get_hardbans)
            .service(get_hardban)
            .service(update_hardban)
            .service(get_kicks)
            .service(get_kick)
            .service(update_kick)
            .service(login)
            .service(get_mutes)
            .service(get_mute)
            .service(update_mute)
            .service(get_setting)
            .service(update_setting)
            .service(reset_setting)
            .service(get_softbans)
            .service(get_softban)
            .service(update_softban)
            .service(get_warns)
            .service(get_warn)
            .service(update_warn)
            .wrap(
                Cors::default()
                    .allowed_origin(&config.cors_origin)
                    // needed for captcha endpoint, actix-cors doesn't actually check if request is for CORS
                    .allowed_origin(&config.self_url)
                    .allowed_methods(vec!["GET", "POST", "DELETE"])
                    .allowed_headers(vec![CONTENT_TYPE, HeaderName::from_static("token")])
                    .max_age(None),
            )
    })
    .bind((Ipv4Addr::LOCALHOST, config.server_port))?
    .run()
    .await?;

    Ok(())
}
<commit_msg>Fix lifetime issues<commit_after>mod endpoint;
mod model;

use crate::database::settings::{PRIVACY_ADMIN_ONLY, PRIVACY_EVERYONE, PRIVACY_STAFF_ONLY};
use crate::discord::util::is_staff;
use crate::server::endpoint::ban::{get_ban, get_bans, update_ban};
use crate::server::endpoint::captcha::{get_captcha_page, submit_captcha};
use crate::server::endpoint::hardban::{get_hardban, get_hardbans, update_hardban};
use crate::server::endpoint::kick::{get_kick, get_kicks, update_kick};
use crate::server::endpoint::login::login;
use crate::server::endpoint::mute::{get_mute, get_mutes, update_mute};
use crate::server::endpoint::self_user::get_self;
use crate::server::endpoint::settings::{get_setting, reset_setting, update_setting};
use crate::server::endpoint::softban::{get_softban, get_softbans, update_softban};
use crate::server::endpoint::warn::{get_warn, get_warns, update_warn};
use crate::service::guild::{CachedMember, GuildService};
use crate::service::invalid_uuid::InvalidUUIDService;
use crate::service::setting::SettingService;
use crate::util::now;
use crate::Config;
use actix_cors::Cors;
use actix_web::http::header::{HeaderName, CONTENT_TYPE};
use actix_web::{get, web, App, HttpRequest, HttpResponse, HttpServer, Responder};
use jsonwebtoken::errors::ErrorKind;
use jsonwebtoken::{decode, encode, Algorithm, DecodingKey, EncodingKey, Header, Validation};
use lazy_static::lazy_static;
use serde::{Deserialize, Serialize};
use serenity::model::id::{GuildId, UserId};
use serenity::model::Permissions;
use std::error::Error;
use std::net::Ipv4Addr;
use std::sync::Arc;
use tracing::error;
use typemap_rev::TypeMap;
use uuid::Uuid;

lazy_static! {
    static ref VALIDATION: Validation = Validation {
        algorithms: vec![Algorithm::HS512],
        ..Validation::default()
    };
    static ref ENCODING_HEADER: Header = Header::new(Algorithm::HS512);
}

#[derive(Deserialize, Serialize)]
struct JwtClaims {
    #[serde(rename = "userId")]
    user_id: String,
    uuid: Uuid,
    exp: u64,
}

fn verify_token(secret: &str, token: &str) -> Option<JwtClaims> {
    match decode::<JwtClaims>(
        token,
        &DecodingKey::from_secret(secret.as_bytes()),
        &VALIDATION,
    ) {
        Ok(data) => Some(data.claims),
        Err(error) => match error.kind() {
            ErrorKind::InvalidSignature | ErrorKind::ExpiredSignature => None,
            _ => {
                error!("failed to decode jwt token {:?}", error.kind());
                None
            }
        },
    }
}

fn generate_token(secret: &str, user_id: UserId) -> Option<String> {
    let user_id = user_id.to_string();
    let expires_at = now() + 7 * 24 * 60 * 60;
    let uuid = Uuid::new_v4();
    let claims = JwtClaims {
        user_id,
        uuid,
        exp: expires_at,
    };

    match encode(
        &ENCODING_HEADER,
        &claims,
        &EncodingKey::from_secret(secret.as_bytes()),
    ) {
        Ok(token) => Some(token),
        Err(err) => {
            error!("failed to create jwt token {:?}", err);
            None
        }
    }
}

async fn is_authenticated(
    config: &Config,
    services: &TypeMap,
    req: &web::HttpRequest,
) -> Option<UserId> {
    let token = req
        .headers()
        .get(HeaderName::from_static("token")) // optimize into constant?
        .map(|value| value.to_str())
        .map(|result| result.ok())
        .flatten()?;

    let claims = verify_token(&config.server_secret, token)?;

    let invalid_uuid_service = services.get::<InvalidUUIDService>()?;

    if invalid_uuid_service.is_uuid_invalid(claims.uuid).await {
        None
    } else {
        match claims.user_id.parse::<u64>() {
            Ok(id) => Some(UserId(id)),
            Err(_) => {
                // secret leaked or we have a problem with token generation
                error!("received a token with valid signature with invalid user id");
                None
            }
        }
    }
}

pub async fn check_authentication(
    config: &Config,
    services: &TypeMap,
    req: &web::HttpRequest,
) -> Result<UserId, HttpResponse> {
    is_authenticated(config, services, req)
        .await
        .ok_or_else(|| HttpResponse::Unauthorized().json("Invalid token!"))
}

pub enum PrivateEndpointKind {
    ModLog,
    Settings,
}

async fn is_authorized_to_view_private_endpoint(
    services: &TypeMap,
    guild_id: GuildId,
    permissions: Permissions,
    endpoint_kind: PrivateEndpointKind,
) -> bool {
    let setting_service = if let Some(service) = services.get::<SettingService>() {
        service
    } else {
        return false;
    };

    let setting = setting_service.get_setting(guild_id).await;
    let privacy_setting = match endpoint_kind {
        PrivateEndpointKind::Settings => setting.privacy_settings,
        PrivateEndpointKind::ModLog => setting.privacy_mod_log,
    };

    if privacy_setting == PRIVACY_EVERYONE {
        true
    } else if privacy_setting == PRIVACY_STAFF_ONLY {
        is_staff(permissions)
    } else if privacy_setting == PRIVACY_ADMIN_ONLY {
        permissions.administrator()
    } else {
        false
    }
}

pub async fn check_authorization_to_view_private_endpoint(
    services: &TypeMap,
    guild_id: GuildId,
    permissions: Permissions,
    endpoint_kind: PrivateEndpointKind,
) -> Option<HttpResponse> {
    if is_authorized_to_view_private_endpoint(services, guild_id, permissions, endpoint_kind).await
    {
        None
    } else {
        Some(
            HttpResponse::Forbidden()
                .json("Server settings prevent you from viewing private information!"),
        )
    }
}

pub async fn is_guild_endpoint(
    services: &TypeMap,
    guild_id: GuildId,
    user_id: UserId,
) -> Option<(Arc<CachedMember>, Permissions)> {
    let guild_service = services.get::<GuildService>()?;

    let member = guild_service.get_member(guild_id, user_id).await.ok()?;

    let permissions = if let Ok(permissions) = guild_service
        .get_permissions(user_id, &member.roles, guild_id)
        .await
    {
        permissions
    } else {
        return None;
    };

    Some((member, permissions))
}

pub async fn check_guild_endpoint(
    services: &TypeMap,
    guild_id: GuildId,
    user_id: UserId,
) -> Result<(Arc<CachedMember>, Permissions), HttpResponse> {
    is_guild_endpoint(services, guild_id, user_id)
        .await
        .ok_or_else(|| {
            HttpResponse::Forbidden().json("This server either doesn't exist or you aren't in it!")
        })
}

pub async fn apply_private_endpoint_fetch_checks(
    config: &Config,
    services: &TypeMap,
    req: &HttpRequest,
    guild_id: GuildId,
    endpoint_kind: PrivateEndpointKind,
) -> Result<(), HttpResponse> {
    let user_id = match check_authentication(config, services, req).await {
        Ok(user_id) => user_id,
        Err(response) => return Err(response),
    };

    let (_, permissions) = match check_guild_endpoint(services, guild_id, user_id).await {
        Ok((member, permissions)) => (member, permissions),
        Err(response) => return Err(response),
    };

    if let Some(response) =
        check_authorization_to_view_private_endpoint(services, guild_id, permissions, endpoint_kind)
            .await
    {
        return Err(response);
    }

    Ok(())
}

pub async fn apply_mod_log_update_checks(
    config: &Config,
    services: &TypeMap,
    req: &HttpRequest,
    guild_id: GuildId,
    required_permission: Permissions,
) -> Result<(), HttpResponse> {
    let user_id = match check_authentication(config, services, req).await {
        Ok(user_id) => user_id,
        Err(response) => return Err(response),
    };

    let (_, permissions) = match check_guild_endpoint(services, guild_id, user_id).await {
        Ok((member, permissions)) => (member, permissions),
        Err(response) => return Err(response),
    };

    if permissions.contains(Permissions::ADMINISTRATOR) || permissions.contains(required_permission)
    {
        Ok(())
    } else {
        Err(HttpResponse::Forbidden()
            .json("You don't have permissions to change moderator log history!"))
    }
}

pub async fn apply_setting_update_checks(
    config: &Config,
    services: &TypeMap,
    req: &HttpRequest,
    guild_id: GuildId,
) -> Result<(), HttpResponse> {
    let user_id = match check_authentication(config, services, req).await {
        Ok(user_id) => user_id,
        Err(response) => return Err(response),
    };

    let (_, permissions) = match check_guild_endpoint(services, guild_id, user_id).await {
        Ok((member, permissions)) => (member, permissions),
        Err(response) => return Err(response),
    };

    if permissions.contains(Permissions::ADMINISTRATOR) {
        Ok(())
    } else {
        Err(HttpResponse::Forbidden().json("You don't have permissions to change server settings!"))
    }
}

#[get("/")]
async fn index() -> impl Responder {
    "Welcome to Safety Jim API."
}

pub async fn run_server(config: Arc<Config>, services: Arc<TypeMap>) -> Result<(), Box<dyn Error>> {
    let port = config.server_port;
    HttpServer::new(move || {
        App::new()
            .app_data(web::Data::from(config.clone()))
            .app_data(web::Data::from(services.clone()))
            .service(index)
            .service(get_self)
            .service(get_bans)
            .service(get_ban)
            .service(update_ban)
            .service(get_captcha_page)
            .service(submit_captcha)
            .service(get_hardbans)
            .service(get_hardban)
            .service(update_hardban)
            .service(get_kicks)
            .service(get_kick)
            .service(update_kick)
            .service(login)
            .service(get_mutes)
            .service(get_mute)
            .service(update_mute)
            .service(get_setting)
            .service(update_setting)
            .service(reset_setting)
            .service(get_softbans)
            .service(get_softban)
            .service(update_softban)
            .service(get_warns)
            .service(get_warn)
            .service(update_warn)
            .wrap(
                Cors::default()
                    .allowed_origin(&config.cors_origin)
                    // needed for captcha endpoint, actix-cors doesn't actually check if request is for CORS
                    .allowed_origin(&config.self_url)
                    .allowed_methods(vec!["GET", "POST", "DELETE"])
                    .allowed_headers(vec![CONTENT_TYPE, HeaderName::from_static("token")])
                    .max_age(None),
            )
    })
    .bind((Ipv4Addr::LOCALHOST, port))?
    .run()
    .await?;

    Ok(())
}
<|endoftext|>


<commit_before>use std::io::{self, Read};
use std::iter::Iterator;
use std::path::{Component, Path, PathBuf};
use std::slice;

use flate2::read::GzDecoder;
use tar;
use time::Timespec;

use backend::Backend;
use collections::{CollectionsStatus, SignatureFile};

pub struct BackupFiles {
    chains: Vec<Chain>,
}

pub struct Snapshot<'a> {
    index: u8,
    chain: &'a Chain,
}

#[derive(Debug)]
pub struct File<'a> {
    pub path: &'a Path,
    pub last_modified: Timespec,
}

/// Iterator over a list of backup snapshots.
pub struct Snapshots<'a> {
    chain_iter: slice::Iter<'a, Chain>,
    chain: Option<&'a Chain>,
    snapshot_id: u8,
}

pub struct SnapshotFiles<'a> {
    index: u8,
    iter: slice::Iter<'a, PathSnapshots>,
}

enum DiffType {
    Signature,
    Snapshot,
    Deleted,
}

/// Store separately informations about the signatures and informations about the paths in the
/// signatures. This allows to reuse informations between snapshots and avoid duplicating them.
struct Chain {
    timestamps: Vec<Timespec>,
    files: Vec<PathSnapshots>,
}

struct PathSnapshots {
    // the directory or file path
    path: PathBuf,
    // all the snapshots for this path
    snapshots: Vec<PathSnapshot>,
}

struct PathSnapshot {
    // info are None if the snapshot has deleted this path
    info: Option<PathInfo>,
    // the index of the snapshot in the chain
    index: u8,
}

/// Informations about a path inside a snapshot.
struct PathInfo {
    mtime: Timespec,
}

impl BackupFiles {
    pub fn new<B: Backend>(backend: &B) -> io::Result<BackupFiles> {
        let collection = {
            let filenames = try!(backend.get_file_names());
            CollectionsStatus::from_filenames(&filenames)
        };
        let mut chains: Vec<Chain> = Vec::new();
        let coll_chains = collection.signature_chains();
        for coll_chain in coll_chains {
            // translate collections::SignatureChain into a Chain
            let mut chain = Chain {
                timestamps: Vec::new(),
                files: Vec::new(),
            };
            // add to the chain the full signature and all the incremental signatures
            // if an error occurs in the full signature exit
            let file = try!(backend.open_file(coll_chain.fullsig.file_name.as_ref()));
            try!(add_sigfile_to_chain(&mut chain, file, &coll_chain.fullsig));
            for inc in &coll_chain.inclist {
                // TODO: if an error occurs here, do not exit with an error, instead
                // break the iteration and store the error inside the chain
                let file = try!(backend.open_file(inc.file_name.as_ref()));
                try!(add_sigfile_to_chain(&mut chain, file, &inc));
            }
            chains.push(chain);
        }
        Ok(BackupFiles { chains: chains })
    }

    pub fn snapshots(&self) -> Snapshots {
        let mut iter = self.chains.iter();
        let first_chain = iter.next();
        Snapshots {
            chain_iter: iter,
            chain: first_chain,
            snapshot_id: 0,
        }
    }
}

impl<'a> Iterator for Snapshots<'a> {
    type Item = Snapshot<'a>;

    fn next(&mut self) -> Option<Snapshot<'a>> {
        loop {
            if let Some(chain) = self.chain {
                if let Some(_) = chain.timestamps.get(self.snapshot_id as usize) {
                    let result = Some(Snapshot {
                        index: self.snapshot_id,
                        chain: chain,
                    });
                    self.snapshot_id += 1;
                    return result;
                } else {
                    // this chain is completed
                    // go to next chain
                    self.chain = self.chain_iter.next();
                    self.snapshot_id = 0;
                }
            } else {
                // no other chains are present
                return None;
            }
        }
    }
}

impl<'a> Snapshot<'a> {
    pub fn time(&self) -> Timespec {
        self.chain.timestamps[self.index as usize]
    }

    pub fn files(&self) -> SnapshotFiles<'a> {
        SnapshotFiles {
            index: self.index,
            iter: self.chain.files.iter(),
        }
    }
}

impl<'a> Iterator for SnapshotFiles<'a> {
    type Item = File<'a>;

    fn next(&mut self) -> Option<File<'a>> {
        while let Some(path_snapshots) = self.iter.next() {
            if let Some(s) = path_snapshots.snapshots.iter().rev().find(|s| s.index <= self.index) {
                // now we have a path info present in this snapshot
                // if it is not deleted return it
                if let Some(ref info) = s.info {
                    return Some(File {
                        path: path_snapshots.path.as_ref(),
                        last_modified: info.mtime,
                    });
                }
            }
        }
        None
    }
}

fn add_sigfile_to_chain<R: Read>(chain: &mut Chain,
                                 file: R,
                                 sigfile: &SignatureFile)
                                 -> io::Result<()> {
    let result = {
        let snapshot_id = chain.files.len() as u8;
        if sigfile.compressed {
            let gz_decoder = try!(GzDecoder::new(file));
            add_sigtar_to_snapshots(&mut chain.files, tar::Archive::new(gz_decoder), snapshot_id)
        } else {
            add_sigtar_to_snapshots(&mut chain.files, tar::Archive::new(file), snapshot_id)
        }
    };
    if result.is_ok() {
        // add to the list of signatures only if everything is ok
        // we do not need to cleanup the chain if someting went wrong, because if the
        // list of signatures is not updated, the change is not observable
        chain.timestamps.push(sigfile.time);
    }
    result
}

fn add_sigtar_to_snapshots<R: Read>(snapshots: &mut Vec<PathSnapshots>,
                                    mut tar: tar::Archive<R>,
                                    snapshot_id: u8)
                                    -> io::Result<()> {
    let mut new_files: Vec<PathSnapshots> = Vec::new();
    {
        let mut old_snapshots = snapshots.iter_mut();
        for tarfile in try!(tar.files_mut()) {
            // we can ignore paths with errors
            // the only problem here is that we miss some change in the chain, but it is
            // better than abort the whole signature
            let tarfile = unwrap_or_continue!(tarfile);
            let header = tarfile.header();
            let path = unwrap_or_continue!(header.path());
            let (difftype, path) = unwrap_opt_or_continue!(parse_snapshot_path(&path));
            let info = match difftype {
                DiffType::Signature | DiffType::Snapshot => {
                    let time = Timespec::new(header.mtime().unwrap_or(0) as i64, 0);
                    Some(PathInfo { mtime: time })
                }
                _ => None,
            };
            let new_snapshot = PathSnapshot {
                info: info,
                index: snapshot_id,
            };
            // find the current path in the old snapshots
            // note: they are ordered
            let position = {
                let mut position: Option<&mut PathSnapshots> = None;
                while let Some(path_snapshots) = old_snapshots.next() {
                    if path_snapshots.path.as_path() < path {
                        continue;
                    }
                    if path_snapshots.path.as_path() == path {
                        // this path is already present in old snapshots: update them
                        position = Some(path_snapshots);
                    } else {
                        // we've already reached the first item next to the current path
                        // so, the path is not present in old snapshots
                        position = None;
                    }
                    break;
                }
                position
            };
            if let Some(path_snapshots) = position {
                path_snapshots.snapshots.push(new_snapshot);
            } else {
                // the path is not present in the old snapshots: add to new list
                new_files.push(PathSnapshots {
                    path: path.to_path_buf(),
                    snapshots: vec![new_snapshot],
                });
            }
        }
    }
    // merge the new files with old snapshots
    if !new_files.is_empty() {
        // TODO: Performance hurt here: we have two sorted arrays to merge,
        // better to use this algorithm: http://stackoverflow.com/a/4553321/1667955
        snapshots.extend(new_files.into_iter());
        snapshots.sort_by(|a, b| a.path.cmp(&b.path));
    }
    Ok(())
}

fn parse_snapshot_path(path: &Path) -> Option<(DiffType, &Path)> {
    // split the path in (first directory, the remaining path)
    // the first is the type, the remaining is the real path
    let mut pcomps = path.components();
    let pfirst = try_opt!(pcomps.next());
    if let Component::Normal(strfirst) = pfirst {
        let difftype = match strfirst.to_str() {
            Some("signature") => DiffType::Signature,
            Some("snapshot") => DiffType::Snapshot,
            Some("deleted") => DiffType::Deleted,
            _ => {
                return None;
            }
        };
        let realpath = pcomps.as_path();
        Some((difftype, realpath))
    } else {
        None
    }
}

#[cfg(test)]
mod test {
    use super::*;
    use backend::local::LocalBackend;
    use time_utils::to_pretty_local;

    #[test]
    fn open_from_local() {
        let backend = LocalBackend::new("tests/backups/single_vol").unwrap();
        let files = BackupFiles::new(&backend).unwrap();
        println!("Backup snapshots:");
        for snapshot in files.snapshots() {
            println!("Snapshot {}", to_pretty_local(snapshot.time()));
            for file in snapshot.files() {
                println!("    {:?}", file);
            }
        }
    }
}
<commit_msg>Added expected output for signatures test<commit_after>use std::io::{self, Read};
use std::iter::Iterator;
use std::path::{Component, Path, PathBuf};
use std::slice;

use flate2::read::GzDecoder;
use tar;
use time::Timespec;

use backend::Backend;
use collections::{CollectionsStatus, SignatureFile};

pub struct BackupFiles {
    chains: Vec<Chain>,
}

pub struct Snapshot<'a> {
    index: u8,
    chain: &'a Chain,
}

#[derive(Debug)]
pub struct File<'a> {
    pub path: &'a Path,
    pub last_modified: Timespec,
}

/// Iterator over a list of backup snapshots.
pub struct Snapshots<'a> {
    chain_iter: slice::Iter<'a, Chain>,
    chain: Option<&'a Chain>,
    snapshot_id: u8,
}

pub struct SnapshotFiles<'a> {
    index: u8,
    iter: slice::Iter<'a, PathSnapshots>,
}

enum DiffType {
    Signature,
    Snapshot,
    Deleted,
}

/// Store separately informations about the signatures and informations about the paths in the
/// signatures. This allows to reuse informations between snapshots and avoid duplicating them.
struct Chain {
    timestamps: Vec<Timespec>,
    files: Vec<PathSnapshots>,
}

struct PathSnapshots {
    // the directory or file path
    path: PathBuf,
    // all the snapshots for this path
    snapshots: Vec<PathSnapshot>,
}

struct PathSnapshot {
    // info are None if the snapshot has deleted this path
    info: Option<PathInfo>,
    // the index of the snapshot in the chain
    index: u8,
}

/// Informations about a path inside a snapshot.
struct PathInfo {
    mtime: Timespec,
}

impl BackupFiles {
    pub fn new<B: Backend>(backend: &B) -> io::Result<BackupFiles> {
        let collection = {
            let filenames = try!(backend.get_file_names());
            CollectionsStatus::from_filenames(&filenames)
        };
        let mut chains: Vec<Chain> = Vec::new();
        let coll_chains = collection.signature_chains();
        for coll_chain in coll_chains {
            // translate collections::SignatureChain into a Chain
            let mut chain = Chain {
                timestamps: Vec::new(),
                files: Vec::new(),
            };
            // add to the chain the full signature and all the incremental signatures
            // if an error occurs in the full signature exit
            let file = try!(backend.open_file(coll_chain.fullsig.file_name.as_ref()));
            try!(add_sigfile_to_chain(&mut chain, file, &coll_chain.fullsig));
            for inc in &coll_chain.inclist {
                // TODO: if an error occurs here, do not exit with an error, instead
                // break the iteration and store the error inside the chain
                let file = try!(backend.open_file(inc.file_name.as_ref()));
                try!(add_sigfile_to_chain(&mut chain, file, &inc));
            }
            chains.push(chain);
        }
        Ok(BackupFiles { chains: chains })
    }

    pub fn snapshots(&self) -> Snapshots {
        let mut iter = self.chains.iter();
        let first_chain = iter.next();
        Snapshots {
            chain_iter: iter,
            chain: first_chain,
            snapshot_id: 0,
        }
    }
}

impl<'a> Iterator for Snapshots<'a> {
    type Item = Snapshot<'a>;

    fn next(&mut self) -> Option<Snapshot<'a>> {
        loop {
            if let Some(chain) = self.chain {
                if let Some(_) = chain.timestamps.get(self.snapshot_id as usize) {
                    let result = Some(Snapshot {
                        index: self.snapshot_id,
                        chain: chain,
                    });
                    self.snapshot_id += 1;
                    return result;
                } else {
                    // this chain is completed
                    // go to next chain
                    self.chain = self.chain_iter.next();
                    self.snapshot_id = 0;
                }
            } else {
                // no other chains are present
                return None;
            }
        }
    }
}

impl<'a> Snapshot<'a> {
    pub fn time(&self) -> Timespec {
        self.chain.timestamps[self.index as usize]
    }

    pub fn files(&self) -> SnapshotFiles<'a> {
        SnapshotFiles {
            index: self.index,
            iter: self.chain.files.iter(),
        }
    }
}

impl<'a> Iterator for SnapshotFiles<'a> {
    type Item = File<'a>;

    fn next(&mut self) -> Option<File<'a>> {
        while let Some(path_snapshots) = self.iter.next() {
            if let Some(s) = path_snapshots.snapshots.iter().rev().find(|s| s.index <= self.index) {
                // now we have a path info present in this snapshot
                // if it is not deleted return it
                if let Some(ref info) = s.info {
                    return Some(File {
                        path: path_snapshots.path.as_ref(),
                        last_modified: info.mtime,
                    });
                }
            }
        }
        None
    }
}

fn add_sigfile_to_chain<R: Read>(chain: &mut Chain,
                                 file: R,
                                 sigfile: &SignatureFile)
                                 -> io::Result<()> {
    let result = {
        let snapshot_id = chain.files.len() as u8;
        if sigfile.compressed {
            let gz_decoder = try!(GzDecoder::new(file));
            add_sigtar_to_snapshots(&mut chain.files, tar::Archive::new(gz_decoder), snapshot_id)
        } else {
            add_sigtar_to_snapshots(&mut chain.files, tar::Archive::new(file), snapshot_id)
        }
    };
    if result.is_ok() {
        // add to the list of signatures only if everything is ok
        // we do not need to cleanup the chain if someting went wrong, because if the
        // list of signatures is not updated, the change is not observable
        chain.timestamps.push(sigfile.time);
    }
    result
}

fn add_sigtar_to_snapshots<R: Read>(snapshots: &mut Vec<PathSnapshots>,
                                    mut tar: tar::Archive<R>,
                                    snapshot_id: u8)
                                    -> io::Result<()> {
    let mut new_files: Vec<PathSnapshots> = Vec::new();
    {
        let mut old_snapshots = snapshots.iter_mut();
        for tarfile in try!(tar.files_mut()) {
            // we can ignore paths with errors
            // the only problem here is that we miss some change in the chain, but it is
            // better than abort the whole signature
            let tarfile = unwrap_or_continue!(tarfile);
            let header = tarfile.header();
            let path = unwrap_or_continue!(header.path());
            let (difftype, path) = unwrap_opt_or_continue!(parse_snapshot_path(&path));
            let info = match difftype {
                DiffType::Signature | DiffType::Snapshot => {
                    let time = Timespec::new(header.mtime().unwrap_or(0) as i64, 0);
                    Some(PathInfo { mtime: time })
                }
                _ => None,
            };
            let new_snapshot = PathSnapshot {
                info: info,
                index: snapshot_id,
            };
            // find the current path in the old snapshots
            // note: they are ordered
            let position = {
                let mut position: Option<&mut PathSnapshots> = None;
                while let Some(path_snapshots) = old_snapshots.next() {
                    if path_snapshots.path.as_path() < path {
                        continue;
                    }
                    if path_snapshots.path.as_path() == path {
                        // this path is already present in old snapshots: update them
                        position = Some(path_snapshots);
                    } else {
                        // we've already reached the first item next to the current path
                        // so, the path is not present in old snapshots
                        position = None;
                    }
                    break;
                }
                position
            };
            if let Some(path_snapshots) = position {
                path_snapshots.snapshots.push(new_snapshot);
            } else {
                // the path is not present in the old snapshots: add to new list
                new_files.push(PathSnapshots {
                    path: path.to_path_buf(),
                    snapshots: vec![new_snapshot],
                });
            }
        }
    }
    // merge the new files with old snapshots
    if !new_files.is_empty() {
        // TODO: Performance hurt here: we have two sorted arrays to merge,
        // better to use this algorithm: http://stackoverflow.com/a/4553321/1667955
        snapshots.extend(new_files.into_iter());
        snapshots.sort_by(|a, b| a.path.cmp(&b.path));
    }
    Ok(())
}

fn parse_snapshot_path(path: &Path) -> Option<(DiffType, &Path)> {
    // split the path in (first directory, the remaining path)
    // the first is the type, the remaining is the real path
    let mut pcomps = path.components();
    let pfirst = try_opt!(pcomps.next());
    if let Component::Normal(strfirst) = pfirst {
        let difftype = match strfirst.to_str() {
            Some("signature") => DiffType::Signature,
            Some("snapshot") => DiffType::Snapshot,
            Some("deleted") => DiffType::Deleted,
            _ => {
                return None;
            }
        };
        let realpath = pcomps.as_path();
        Some((difftype, realpath))
    } else {
        None
    }
}

#[cfg(test)]
mod test {
    use super::*;
    use backend::local::LocalBackend;
    use time_utils::to_pretty_local;

    #[test]
    fn open_from_local() {
        let backend = LocalBackend::new("tests/backups/single_vol").unwrap();
        let files = BackupFiles::new(&backend).unwrap();
        println!("Backup snapshots:");
        for snapshot in files.snapshots() {
            println!("Snapshot {}", to_pretty_local(snapshot.time()));
            for file in snapshot.files() {
                println!("    {:?}", file);
            }
        }

        // Expected output from duplicity (snapshot 1):
        //
        // Last full backup date: Wed Jun 17 18:25:45 2015
        // Sat Sep 28 18:30:59 2002 .
        // Tue Aug 28 18:23:30 2001 changeable_permission
        // Sat Jul 27 23:00:05 2002 deleted_file
        // Sat Jul 27 23:00:36 2002 directory_to_file
        // Sat Jul 27 23:00:36 2002 directory_to_file/file
        // Tue Aug 28 07:34:29 2001 executable
        // Tue Aug 28 18:19:27 2001 executable2
        // Tue Aug 28 07:32:46 2001 fifo
        // Sat Jul 27 23:23:54 2002 file_to_directory
        // Wed Jul 31 01:54:30 2002 largefile
        // Tue Aug 28 07:30:52 2001 regular_file
        // Thu Aug 30 00:40:37 2001 regular_file.sig
        // Fri Nov  1 04:44:47 2002 symbolic_link
        // Tue Aug 28 21:56:38 2001 test
        // Tue Aug 28 07:31:42 2001 two_hardlinked_files1
        // Tue Aug 28 07:31:42 2001 two_hardlinked_files2
        // Tue Aug 28 22:03:47 2001 <long-list-of-unicode-chars>
        //
        //
        // Expected output from duplicity (snapshot 2):
        //
        // Last full backup date: Wed Jun 17 18:25:45 2015
        // Wed Jul 31 01:55:32 2002 .
        // Tue Aug 28 18:23:30 2001 changeable_permission
        // Sat Jul 27 23:00:48 2002 directory_to_file
        // Tue Aug 28 07:34:29 2001 executable
        // Sat Jul 27 23:01:33 2002 executable2
        // Sat Jul 27 23:01:33 2002 executable2/another_file
        // Tue Aug 28 07:32:46 2001 fifo
        // Sat Jul 27 23:24:06 2002 file_to_directory
        // Wed Jul 31 01:55:24 2002 largefile
        // Sat Jul 27 23:00:18 2002 new_file
        // Sat Jul 27 22:59:32 2002 regular_file
        // Thu Aug 30 00:40:37 2001 regular_file.sig
        // Sat Jul 27 22:59:46 2002 symbolic_link
        // Tue Aug 28 21:56:38 2001 test
        // Tue Aug 28 07:31:42 2001 two_hardlinked_files1
        // Tue Aug 28 07:31:42 2001 two_hardlinked_files2
        // Tue Aug 28 22:03:47 2001 <long-list-of-unicode-chars>
        //
        //
        // Expected output from duplicity (snapshot 3):
        //
        // Last full backup date: Wed Jun 17 18:25:45 2015
        // Sat Sep 28 18:30:59 2002 .
        // Tue Aug 28 18:23:30 2001 changeable_permission
        // Tue Aug 28 07:34:29 2001 executable
        // Tue Aug 28 18:19:27 2001 executable2
        // Tue Aug 28 07:32:46 2001 fifo
        // Wed Jul 31 03:43:34 2002 largefile
        // Tue Aug 28 07:30:52 2001 regular_file
        // Thu Aug 30 00:40:37 2001 regular_file.sig
        // Fri Nov  1 04:44:48 2002 symbolic_link
        // Tue Aug 28 21:56:38 2001 test
        // Tue Aug 28 07:31:42 2001 two_hardlinked_files1
        // Tue Aug 28 07:31:42 2001 two_hardlinked_files2
        // Tue Aug 28 22:03:47 2001 <long-list-of-unicode-chars>
    }
}
<|endoftext|>


<commit_before>use super::BeaconStateError;
use crate::*;
use crate::{validator_registry::get_active_validator_indices, *};
use bls::create_proof_of_possession;
use rayon::prelude::*;
use ssz::TreeHash;

/// Builds a `BeaconState` for use in testing or benchmarking.
///
/// Building the `BeaconState` is a three step processes:
///
/// 1. Create a new `BeaconStateBuilder`.
/// 2. Call `Self::build()` or `Self::build_fast()` generate a  `BeaconState`.
/// 3. (Optional) Use builder functions to modify the `BeaconState`.
/// 4. Call `Self::cloned_state()` to obtain a `BeaconState` cloned from this struct.
///
/// Step (2) happens prior to step (3) because some functionality requires an existing
/// `BeaconState`.
///
/// Step (4) produces a clone of the BeaconState and doesn't consume the `BeaconStateBuilder` to
/// allow access to `self.keypairs` and `self.spec`.
pub struct BeaconStateBuilder {
    pub state: BeaconState,
}

impl BeaconStateBuilder {
    /// Create a new builder with the given number of validators.
    ///
    /// Spec v0.4.0
    pub fn new(genesis_time: u64, latest_eth1_data: Eth1Data, spec: &ChainSpec) -> Self {
        Self {
            state: BeaconState::genesis(genesis_time, latest_eth1_data, spec),
        }
    }

    /// Process deposit objects.
    ///
    /// Spec v0.4.0
    pub fn process_initial_deposits(
        &mut self,
        initial_validator_deposits: &[Deposit],
        spec: &ChainSpec,
    ) {
        let deposit_data = initial_validator_deposits
            .par_iter()
            .map(|deposit| &deposit.deposit_data)
            .collect();

        self.state.process_deposits(deposit_data, spec);

        self.activate_genesis_validators(spec);

        self.state.deposit_index = initial_validator_deposits.len() as u64;
    }

    fn activate_genesis_validators(&mut self, spec: &ChainSpec) {
        for validator_index in 0..self.state.validator_registry.len() {
            if self.state.get_effective_balance(validator_index, spec) >= spec.max_deposit_amount {
                self.state.activate_validator(validator_index, true, spec);
            }
        }
    }

    /// Instantiate the validator registry from a YAML file.
    ///
    /// This skips a lot of signing and verification, useful for fast test setups.
    ///
    /// Spec v0.4.0
    pub fn import_existing_validators(
        &mut self,
        validators: Vec<Validator>,
        initial_balances: Vec<u64>,
        deposit_index: u64,
        spec: &ChainSpec,
    ) {
        self.state.validator_registry = validators;

        assert_eq!(
            self.state.validator_registry.len(),
            initial_balances.len(),
            "Not enough balances for validators"
        );

        self.state.validator_balances = initial_balances;

        self.activate_genesis_validators(spec);

        self.state.deposit_index = deposit_index;
    }

    /// Updates the final state variables and returns a fully built genesis state.
    ///
    /// Spec v0.4.0
    pub fn build(mut self, spec: &ChainSpec) -> Result<BeaconState, BeaconStateError> {
        let genesis_active_index_root =
            get_active_validator_indices(&self.state.validator_registry, spec.genesis_epoch)
                .hash_tree_root();

        self.state.latest_active_index_roots = vec![
            Hash256::from_slice(&genesis_active_index_root);
            spec.latest_active_index_roots_length
        ];

        self.state.current_shuffling_seed = self.state.generate_seed(spec.genesis_epoch, spec)?;

        Ok(self.state)
    }

    /*
    /// Sets the `BeaconState` to be in the last slot of the given epoch.
    ///
    /// Sets all justification/finalization parameters to be be as "perfect" as possible (i.e.,
    /// highest justified and finalized slots, full justification bitfield, etc).
    pub fn teleport_to_end_of_epoch(&mut self, epoch: Epoch, spec: &ChainSpec) {
        let state = &mut self.state;

        let slot = epoch.end_slot(spec.slots_per_epoch);

        state.slot = slot;
        state.validator_registry_update_epoch = epoch - 1;

        state.previous_shuffling_epoch = epoch - 1;
        state.current_shuffling_epoch = epoch;

        state.previous_shuffling_seed = Hash256::from_low_u64_le(0);
        state.current_shuffling_seed = Hash256::from_low_u64_le(1);

        state.previous_justified_epoch = epoch - 2;
        state.justified_epoch = epoch - 1;
        state.justification_bitfield = u64::max_value();
        state.finalized_epoch = epoch - 1;
    }

    /// Creates a full set of attestations for the `BeaconState`. Each attestation has full
    /// participation from its committee and references the expected beacon_block hashes.
    ///
    /// These attestations should be fully conducive to justification and finalization.
    pub fn insert_attestations(&mut self) {
        let state = &mut self.state;

        state
            .build_epoch_cache(RelativeEpoch::Previous, &self.spec)
            .unwrap();
        state
            .build_epoch_cache(RelativeEpoch::Current, &self.spec)
            .unwrap();

        let current_epoch = state.current_epoch(&self.spec);
        let previous_epoch = state.previous_epoch(&self.spec);
        let current_epoch_depth =
            (state.slot - current_epoch.end_slot(self.spec.slots_per_epoch)).as_usize();

        let previous_epoch_slots = previous_epoch.slot_iter(self.spec.slots_per_epoch);
        let current_epoch_slots = current_epoch
            .slot_iter(self.spec.slots_per_epoch)
            .take(current_epoch_depth);

        for slot in previous_epoch_slots.chain(current_epoch_slots) {
            let committees = state
                .get_crosslink_committees_at_slot(slot, &self.spec)
                .unwrap()
                .clone();

            for (committee, shard) in committees {
                state
                    .latest_attestations
                    .push(committee_to_pending_attestation(
                        state, &committee, shard, slot, &self.spec,
                    ))
            }
        }
    }

    /// Returns a cloned `BeaconState`.
    pub fn cloned_state(&self) -> BeaconState {
        self.state.as_ref().expect("Genesis required").clone()
    }
    */
}

/*
/// Builds a valid PendingAttestation with full participation for some committee.
fn committee_to_pending_attestation(
    state: &BeaconState,
    committee: &[usize],
    shard: u64,
    slot: Slot,
    spec: &ChainSpec,
) -> PendingAttestation {
    let current_epoch = state.current_epoch(spec);
    let previous_epoch = state.previous_epoch(spec);

    let mut aggregation_bitfield = Bitfield::new();
    let mut custody_bitfield = Bitfield::new();

    for (i, _) in committee.iter().enumerate() {
        aggregation_bitfield.set(i, true);
        custody_bitfield.set(i, true);
    }

    let is_previous_epoch =
        state.slot.epoch(spec.slots_per_epoch) != slot.epoch(spec.slots_per_epoch);

    let justified_epoch = if is_previous_epoch {
        state.previous_justified_epoch
    } else {
        state.justified_epoch
    };

    let epoch_boundary_root = if is_previous_epoch {
        *state
            .get_block_root(previous_epoch.start_slot(spec.slots_per_epoch), spec)
            .unwrap()
    } else {
        *state
            .get_block_root(current_epoch.start_slot(spec.slots_per_epoch), spec)
            .unwrap()
    };

    let justified_block_root = *state
        .get_block_root(justified_epoch.start_slot(spec.slots_per_epoch), &spec)
        .unwrap();

    PendingAttestation {
        aggregation_bitfield,
        data: AttestationData {
            slot,
            shard,
            beacon_block_root: *state.get_block_root(slot, spec).unwrap(),
            epoch_boundary_root,
            crosslink_data_root: Hash256::zero(),
            latest_crosslink: Crosslink {
                epoch: slot.epoch(spec.slots_per_epoch),
                crosslink_data_root: Hash256::zero(),
            },
            justified_epoch,
            justified_block_root,
        },
        custody_bitfield,
        inclusion_slot: slot,
    }
}
*/
<commit_msg>Tidy BeaconStateBuilder struct<commit_after>use super::BeaconStateError;
use crate::validator_registry::get_active_validator_indices;
use crate::*;
use rayon::prelude::*;
use ssz::TreeHash;

/// Builds a `BeaconState` for use in production.
///
/// This struct should not be modified for use in testing scenarios. Use `TestingBeaconStateBuilder` for that purpose.
///
/// This struct should remain safe and sensible for production usage.
pub struct BeaconStateBuilder {
    pub state: BeaconState,
}

impl BeaconStateBuilder {
    /// Create a new builder with the given number of validators.
    ///
    /// Spec v0.4.0
    pub fn new(genesis_time: u64, latest_eth1_data: Eth1Data, spec: &ChainSpec) -> Self {
        Self {
            state: BeaconState::genesis(genesis_time, latest_eth1_data, spec),
        }
    }

    /// Process deposit objects.
    ///
    /// Spec v0.4.0
    pub fn process_initial_deposits(
        &mut self,
        initial_validator_deposits: &[Deposit],
        spec: &ChainSpec,
    ) {
        let deposit_data = initial_validator_deposits
            .par_iter()
            .map(|deposit| &deposit.deposit_data)
            .collect();

        self.state.process_deposits(deposit_data, spec);

        self.activate_genesis_validators(spec);

        self.state.deposit_index = initial_validator_deposits.len() as u64;
    }

    fn activate_genesis_validators(&mut self, spec: &ChainSpec) {
        for validator_index in 0..self.state.validator_registry.len() {
            if self.state.get_effective_balance(validator_index, spec) >= spec.max_deposit_amount {
                self.state.activate_validator(validator_index, true, spec);
            }
        }
    }

    /// Instantiate the validator registry from a YAML file.
    ///
    /// This skips a lot of signing and verification, useful for fast test setups.
    ///
    /// Spec v0.4.0
    pub fn import_existing_validators(
        &mut self,
        validators: Vec<Validator>,
        initial_balances: Vec<u64>,
        deposit_index: u64,
        spec: &ChainSpec,
    ) {
        self.state.validator_registry = validators;

        assert_eq!(
            self.state.validator_registry.len(),
            initial_balances.len(),
            "Not enough balances for validators"
        );

        self.state.validator_balances = initial_balances;

        self.activate_genesis_validators(spec);

        self.state.deposit_index = deposit_index;
    }

    /// Updates the final state variables and returns a fully built genesis state.
    ///
    /// Spec v0.4.0
    pub fn build(mut self, spec: &ChainSpec) -> Result<BeaconState, BeaconStateError> {
        let genesis_active_index_root =
            get_active_validator_indices(&self.state.validator_registry, spec.genesis_epoch)
                .hash_tree_root();

        self.state.latest_active_index_roots = vec![
            Hash256::from_slice(&genesis_active_index_root);
            spec.latest_active_index_roots_length
        ];

        self.state.current_shuffling_seed = self.state.generate_seed(spec.genesis_epoch, spec)?;

        Ok(self.state)
    }
}
<|endoftext|>


<commit_before>use super::*;

// This is the runtime cursor on iteration
// It hold a copy of the containing page next page lock guard
// These lock guards are preventing the node and their neighbourhoods been changed externally
// Ordering are specified that can also change lock pattern
pub struct RTCursor<KS, PS>
where
    KS: Slice<EntryKey> + Debug + 'static,
    PS: Slice<NodeCellRef> + 'static,
{
    pub index: usize,
    pub ordering: Ordering,
    pub page: Option<NodeCellRef>,
    pub deleted: DeletionSet,
    pub marker: PhantomData<(KS, PS)>,
    pub current: Option<EntryKey>,
}

impl<KS, PS> RTCursor<KS, PS>
where
    KS: Slice<EntryKey> + Debug + 'static,
    PS: Slice<NodeCellRef> + 'static,
{
    pub fn new(pos: usize, page: &NodeCellRef, ordering: Ordering, deleted: &DeletionSet) -> Self {
        let mut cursor = RTCursor {
            index: pos,
            ordering,
            page: Some(page.clone()),
            deleted: deleted.clone(),
            marker: PhantomData,
            current: None,
        };
        match ordering {
            Ordering::Forward
                if pos >= read_node(page, |node: &NodeReadHandler<KS, PS>| node.len()) =>
            {
                cursor.next();
            }
            _ => {
                cursor.current = Some(Self::read_current(page, pos));
            }
        }
        debug!(
            "Created cursor with pos {}, current {:?}, ordering: {:?}",
            cursor.index, cursor.current, cursor.ordering
        );
        cursor
    }
    fn boxed(self) -> Box<IndexCursor> {
        box self
    }

    fn read_current(node: &NodeCellRef, pos: usize) -> EntryKey {
        read_node(node, |node: &NodeReadHandler<KS, PS>| {
            node.extnode().keys.as_slice_immute()[pos].clone()
        })
    }

    fn next_candidate(&mut self) -> bool {
        if self.page.is_some() {
            let current_page = self.page.clone().unwrap();
            read_node(&current_page, |page: &NodeReadHandler<KS, PS>| {
                let ext_page = page.extnode();
                // debug!("Next id with index: {}, length: {}", self.index + 1, ext_page.len);
                match self.ordering {
                    Ordering::Forward => {
                        if page.is_empty() || self.index + 1 >= page.len() {
                            let next_node_ref = &ext_page.next;
                            return read_node(
                                next_node_ref,
                                |next_node: &NodeReadHandler<KS, PS>| {
                                    if next_node.is_none() {
                                        self.page = None;
                                        self.current = None;
                                        return false;
                                    } else if next_node.is_empty() {
                                        return self.next();
                                    } else if next_node.is_ext() {
                                        self.index = 0;
                                        self.page = Some(next_node_ref.clone());
                                        self.current =
                                            Some(Self::read_current(next_node_ref, self.index));
                                        return true;
                                    } else {
                                        unreachable!()
                                    }
                                },
                            );
                        } else {
                            self.index += 1;
                            // debug!("Advancing cursor to index {}", self.index);
                        }
                    }
                    Ordering::Backward => {
                        if page.is_empty() || self.index == 0 {
                            let prev_node_ref = &ext_page.prev;
                            return read_node(
                                prev_node_ref,
                                |prev_node: &NodeReadHandler<KS, PS>| {
                                    if prev_node.is_none() {
                                        self.page = None;
                                        self.current = None;
                                        return false;
                                    } else if prev_node.is_empty() {
                                        return self.next();
                                    } else if prev_node.is_ext() {
                                        self.index = prev_node.len() - 1;
                                        self.page = Some(prev_node_ref.clone());
                                        self.current =
                                            Some(Self::read_current(prev_node_ref, self.index));
                                        return true;
                                    } else {
                                        unreachable!()
                                    }
                                },
                            );
                        } else {
                            self.index -= 1;
                            // debug!("Advancing cursor to index {}", self.index);
                        }
                    }
                }
                self.current = Some(page.extnode().keys.as_slice_immute()[self.index].clone());
                true
            })
        } else {
            false
        }
    }
}

impl<KS, PS> IndexCursor for RTCursor<KS, PS>
where
    KS: Slice<EntryKey> + Debug + 'static,
    PS: Slice<NodeCellRef> + 'static,
{
    // TODO: Copy current after next
    fn next(&mut self) -> bool {
        loop {
            let has_candidate = self.next_candidate();
            if has_candidate {
                // search in deleted set and skip if exists
                if !self.deleted.read().contains(self.current().unwrap()) {
                    return true;
                }
            } else {
                return false;
            }
        }
    }

    // TODO: Use copied key reference
    fn current(&self) -> Option<&EntryKey> {
        self.current.as_ref()
    }
}
<commit_msg>cope with empty node for cursor<commit_after>use super::*;

// This is the runtime cursor on iteration
// It hold a copy of the containing page next page lock guard
// These lock guards are preventing the node and their neighbourhoods been changed externally
// Ordering are specified that can also change lock pattern
pub struct RTCursor<KS, PS>
where
    KS: Slice<EntryKey> + Debug + 'static,
    PS: Slice<NodeCellRef> + 'static,
{
    pub index: usize,
    pub ordering: Ordering,
    pub page: Option<NodeCellRef>,
    pub deleted: DeletionSet,
    pub marker: PhantomData<(KS, PS)>,
    pub current: Option<EntryKey>,
}

impl<KS, PS> RTCursor<KS, PS>
where
    KS: Slice<EntryKey> + Debug + 'static,
    PS: Slice<NodeCellRef> + 'static,
{
    pub fn new(pos: usize, page: &NodeCellRef, ordering: Ordering, deleted: &DeletionSet) -> Self {
        let mut cursor = RTCursor {
            index: pos,
            ordering,
            page: Some(page.clone()),
            deleted: deleted.clone(),
            marker: PhantomData,
            current: None,
        };
        match ordering {
            Ordering::Forward
                if pos >= read_node(page, |node: &NodeReadHandler<KS, PS>| node.len()) =>
            {
                cursor.next();
            }
            _ => {
                cursor.current = Some(Self::read_current(page, pos));
            }
        }
        debug!(
            "Created cursor with pos {}, current {:?}, ordering: {:?}",
            cursor.index, cursor.current, cursor.ordering
        );
        cursor
    }
    fn boxed(self) -> Box<IndexCursor> {
        box self
    }

    fn read_current(node: &NodeCellRef, pos: usize) -> EntryKey {
        read_node(node, |node: &NodeReadHandler<KS, PS>| {
            node.extnode().keys.as_slice_immute()[pos].clone()
        })
    }

    fn next_candidate(&mut self) -> bool {
        if self.page.is_some() {
            let current_page = self.page.clone().unwrap();
            read_node(&current_page, |page: &NodeReadHandler<KS, PS>| {
                // let ext_page = page.extnode();
                // debug!("Next id with index: {}, length: {}", self.index + 1, ext_page.len);
                match self.ordering {
                    Ordering::Forward => {
                        if page.is_empty() || self.index + 1 >= page.len() {
                            let next_node_ref = page.right_ref().unwrap();
                            return read_node(
                                next_node_ref,
                                |next_node: &NodeReadHandler<KS, PS>| {
                                    if next_node.is_none() {
                                        self.page = None;
                                        self.current = None;
                                        return false;
                                    } else if next_node.is_empty() {
                                        return self.next();
                                    } else if next_node.is_ext() {
                                        self.index = 0;
                                        self.page = Some(next_node_ref.clone());
                                        self.current =
                                            Some(Self::read_current(next_node_ref, self.index));
                                        return true;
                                    } else {
                                        unreachable!()
                                    }
                                },
                            );
                        } else {
                            self.index += 1;
                            // debug!("Advancing cursor to index {}", self.index);
                        }
                    }
                    Ordering::Backward => {
                        if page.is_empty() || self.index == 0 {
                            let prev_node_ref = page.left_ref().unwrap();
                            return read_node(
                                prev_node_ref,
                                |prev_node: &NodeReadHandler<KS, PS>| {
                                    if prev_node.is_none() {
                                        self.page = None;
                                        self.current = None;
                                        return false;
                                    } else if prev_node.is_empty() {
                                        return self.next();
                                    } else if prev_node.is_ext() {
                                        self.index = prev_node.len() - 1;
                                        self.page = Some(prev_node_ref.clone());
                                        self.current =
                                            Some(Self::read_current(prev_node_ref, self.index));
                                        return true;
                                    } else {
                                        unreachable!()
                                    }
                                },
                            );
                        } else {
                            self.index -= 1;
                            // debug!("Advancing cursor to index {}", self.index);
                        }
                    }
                }
                self.current = Some(Self::read_current(&current_page, self.index));
                true
            })
        } else {
            false
        }
    }
}

impl<KS, PS> IndexCursor for RTCursor<KS, PS>
where
    KS: Slice<EntryKey> + Debug + 'static,
    PS: Slice<NodeCellRef> + 'static,
{
    // TODO: Copy current after next
    fn next(&mut self) -> bool {
        loop {
            let has_candidate = self.next_candidate();
            if has_candidate {
                // search in deleted set and skip if exists
                if !self.deleted.read().contains(self.current().unwrap()) {
                    return true;
                }
            } else {
                return false;
            }
        }
    }

    // TODO: Use copied key reference
    fn current(&self) -> Option<&EntryKey> {
        self.current.as_ref()
    }
}
<|endoftext|>