/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This software may be used and distributed according to the terms of the
 * GNU General Public License version 2.
 */

use std::collections::HashSet;
use std::collections::VecDeque;
use std::num::NonZeroU64;

use anyhow::anyhow;
use anyhow::Error;
use anyhow::Result;
use async_trait::async_trait;
use blobstore::Loadable;
use bookmarks_types::BookmarkKey;
use bookmarks_types::BookmarkKind;
use changeset_fetcher::ChangesetFetcherRef;
use changesets::ChangesetsRef;
use chrono::DateTime;
use chrono::Utc;
use context::CoreContext;
use ephemeral_blobstore::BubbleId;
use futures::stream;
use futures::StreamExt;
use futures::TryStreamExt;
use logger_ext::Loggable;
use metaconfig_types::RepoConfigRef;
#[cfg(fbcode_build)]
use mononoke_new_commit_rust_logger::MononokeNewCommitLogger;
use mononoke_types::BonsaiChangeset;
use mononoke_types::ChangesetId;
use mononoke_types::Generation;
use once_cell::sync::Lazy;
use permission_checker::MononokeIdentitySet;
use phases::PhasesRef;
use regex::Regex;
use repo_blobstore::RepoBlobstoreRef;
use repo_identity::RepoIdentityRef;
use serde_derive::Serialize;
#[cfg(fbcode_build)]
use whence_logged::WhenceScribeLogged;

pub struct CommitInfo {
    changeset_id: ChangesetId,
    bubble_id: Option<NonZeroU64>,
    diff_id: Option<String>,
    changed_files_info: ChangedFilesInfo,
}

impl CommitInfo {
    pub fn new(bcs: &BonsaiChangeset, bubble_id: Option<BubbleId>) -> Self {
        CommitInfo {
            changeset_id: bcs.get_changeset_id(),
            bubble_id: bubble_id.map(Into::into),
            diff_id: extract_differential_revision(bcs.message()).map(ToString::to_string),
            changed_files_info: ChangedFilesInfo::new(bcs),
        }
    }

    pub fn update_changeset_id(
        &mut self,
        old_changeset_id: ChangesetId,
        new_changeset_id: ChangesetId,
    ) -> Result<()> {
        if self.changeset_id != old_changeset_id {
            return Err(anyhow!(
                concat!(
                    "programming error: attempting to update CommitInfo for incorrect changeset, ",
                    "expected {}, but modifying {}",
                ),
                old_changeset_id,
                self.changeset_id
            ));
        }
        self.changeset_id = new_changeset_id;
        Ok(())
    }
}

fn extract_differential_revision(message: &str) -> Option<&str> {
    static RE: Lazy<Regex> = Lazy::new(|| {
        Regex::new(r"(?m)^Differential Revision: [^\n]*/D([0-9]+)")
            .expect("Failed to compile differential revision regex")
    });

    Some(RE.captures(message)?.get(1)?.as_str())
}

pub struct ChangedFilesInfo {
    changed_files_count: u64,
    changed_files_size: u64,
}

impl ChangedFilesInfo {
    pub fn new(bcs: &BonsaiChangeset) -> Self {
        let changed_files_count = bcs.file_changes_map().len() as u64;
        let changed_files_size = bcs
            .file_changes_map()
            .values()
            .map(|fc| fc.size().unwrap_or(0))
            .sum::<u64>();

        Self {
            changed_files_count,
            changed_files_size,
        }
    }
}

#[derive(Serialize)]
struct PlainCommitInfo {
    // Repo ID is logged to legacy scuba for compatibility, but should be
    // considered deprecated and not logged to Logger.
    repo_id: i32,
    repo_name: String,
    is_public: bool,
    changeset_id: ChangesetId,
    #[serde(skip_serializing_if = "Option::is_none")]
    bubble_id: Option<NonZeroU64>,
    #[serde(skip_serializing_if = "Option::is_none")]
    diff_id: Option<String>,
    changed_files_count: u64,
    changed_files_size: u64,
    parents: Vec<ChangesetId>,
    generation: Generation,
    #[serde(skip_serializing_if = "Option::is_none")]
    bookmark: Option<String>,
    #[serde(skip_serializing_if = "Option::is_none")]
    user_unix_name: Option<String>,
    #[serde(skip_serializing_if = "MononokeIdentitySet::is_empty")]
    user_identities: MononokeIdentitySet,
    #[serde(skip_serializing_if = "Option::is_none")]
    source_hostname: Option<String>,
    #[serde(with = "::chrono::serde::ts_seconds")]
    received_timestamp: DateTime<Utc>,
}

impl PlainCommitInfo {
    async fn new(
        ctx: &CoreContext,
        repo: &(impl ChangesetsRef + RepoIdentityRef),
        received_timestamp: DateTime<Utc>,
        bookmark: Option<(&BookmarkKey, BookmarkKind)>,
        commit_info: CommitInfo,
    ) -> Result<PlainCommitInfo> {
        let CommitInfo {
            changeset_id,
            bubble_id,
            diff_id,
            changed_files_info:
                ChangedFilesInfo {
                    changed_files_count,
                    changed_files_size,
                },
        } = commit_info;
        let repo_id = repo.repo_identity().id().id();
        let repo_name = repo.repo_identity().name().to_string();
        let cs = repo
            .changesets()
            .get(ctx, changeset_id)
            .await?
            .ok_or_else(|| anyhow!("Changeset not found: {}", changeset_id))?;
        let parents = cs.parents;
        let generation = Generation::new(cs.gen);
        let user_unix_name = ctx.metadata().unix_name().map(|un| un.to_string());
        let user_identities = ctx.metadata().identities().clone();
        let source_hostname = ctx.metadata().client_hostname().map(|hn| hn.to_string());
        let (bookmark, is_public) = bookmark.map_or((None, false), |(name, kind)| {
            (Some(name.to_string()), kind.is_public())
        });

        Ok(PlainCommitInfo {
            repo_id,
            repo_name,
            is_public,
            changeset_id,
            bubble_id,
            diff_id,
            changed_files_count,
            changed_files_size,
            parents,
            generation,
            bookmark,
            user_unix_name,
            user_identities,
            source_hostname,
            received_timestamp,
        })
    }
}

#[async_trait]
impl Loggable for PlainCommitInfo {
    #[cfg(fbcode_build)]
    async fn log_to_logger(&self, ctx: &CoreContext) -> Result<()> {
        // Without override, WhenceScribeLogged is set to default which will cause
        // data being logged to "/sandbox" category if service is run from devserver.
        // But currently we use Logger only if we're in prod (as config implies), so
        // we should log to prod too, even from devserver.
        // For example, we can land a commit to prod from devserver, and logging for
        // this commit should go to prod, not to sandbox.
        MononokeNewCommitLogger::override_whence_scribe_logged(ctx.fb, WhenceScribeLogged::PROD);
        let mut logger = MononokeNewCommitLogger::new(ctx.fb);
        logger
            .set_repo_name(self.repo_name.clone())
            .set_is_public(self.is_public)
            .set_changeset_id(self.changeset_id.to_string())
            .set_parents(self.parents.iter().map(ToString::to_string).collect())
            .set_generation(self.generation.value() as i64)
            .set_changed_files_count(self.changed_files_count as i64)
            .set_changed_files_size(self.changed_files_size as i64)
            .set_pusher_identities(
                self.user_identities
                    .iter()
                    .map(ToString::to_string)
                    .collect(),
            )
            .set_received_timestamp(self.received_timestamp.timestamp());

        if let Some(bubble_id) = &self.bubble_id {
            logger.set_bubble_id(bubble_id.get() as i64);
        }
        if let Some(diff_id) = &self.diff_id {
            logger.set_diff_id(diff_id.clone());
        }
        if let Some(bookmark) = &self.bookmark {
            logger.set_bookmark_name(bookmark.to_string());
        }
        if let Some(source_hostname) = &self.source_hostname {
            logger.set_source_hostname(source_hostname.clone());
        }

        logger.attach_raw_scribe_write_cat()?;
        logger.log_async()?;

        Ok(())
    }
}

pub async fn log_new_commits(
    ctx: &CoreContext,
    repo: &(impl RepoIdentityRef + ChangesetsRef + RepoConfigRef),
    bookmark: Option<(&BookmarkKey, BookmarkKind)>,
    commit_infos: Vec<CommitInfo>,
) {
    let new_commit_logging_destination = repo
        .repo_config()
        .update_logging_config
        .new_commit_logging_destination
        .as_ref();

    // If nothing is going to be logged, we can exit early.
    if new_commit_logging_destination.is_none() {
        return;
    }

    let received_timestamp = Utc::now();

    let res = stream::iter(commit_infos)
        .map(Ok)
        .try_for_each_concurrent(100, |commit_info| async move {
            let plain_commit_info =
                PlainCommitInfo::new(ctx, repo, received_timestamp, bookmark, commit_info).await?;
            if let Some(new_commit_logging_destination) = new_commit_logging_destination {
                plain_commit_info
                    .log(ctx, new_commit_logging_destination)
                    .await;
            }
            anyhow::Ok(())
        })
        .await;

    if let Err(err) = res {
        ctx.scuba().clone().log_with_msg(
            "Failed to log new draft commit to scribe",
            Some(err.to_string()),
        );
    }
}

pub async fn log_new_bonsai_changesets(
    ctx: &CoreContext,
    repo: &(impl RepoIdentityRef + ChangesetsRef + RepoConfigRef),
    bookmark: &BookmarkKey,
    kind: BookmarkKind,
    commits_to_log: Vec<BonsaiChangeset>,
) {
    log_new_commits(
        ctx,
        repo,
        Some((bookmark, kind)),
        commits_to_log
            .iter()
            .map(|bcs| CommitInfo::new(bcs, None))
            .collect(),
    )
    .await;
}

/// Helper function for finding all the newly public commit that should be logged after
/// the bookmark moves to to_cs_id. For public commits we allow them to be logged twice:
/// once when they're actually created, second time when they become public.
pub async fn find_draft_ancestors(
    ctx: &CoreContext,
    repo: &(impl RepoIdentityRef + RepoConfigRef + PhasesRef + ChangesetFetcherRef + RepoBlobstoreRef),
    to_cs_id: ChangesetId,
) -> Result<Vec<BonsaiChangeset>, Error> {
    ctx.scuba()
        .clone()
        .log_with_msg("Started finding draft ancestors", None);

    let phases = repo.phases();
    let mut queue = VecDeque::new();
    let mut visited = HashSet::new();
    let mut drafts = vec![];
    queue.push_back(to_cs_id);
    visited.insert(to_cs_id);

    while let Some(cs_id) = queue.pop_front() {
        let public = phases
            .get_public(ctx, vec![cs_id], false /*ephemeral_derive*/)
            .await?;

        if public.contains(&cs_id) {
            continue;
        }
        drafts.push(cs_id);

        let parents = repo.changeset_fetcher().get_parents(ctx, cs_id).await?;
        for p in parents {
            if visited.insert(p) {
                queue.push_back(p);
            }
        }
    }

    let drafts = stream::iter(drafts)
        .map(Ok)
        .map_ok(|cs_id| async move { cs_id.load(ctx, repo.repo_blobstore()).await })
        .try_buffer_unordered(100)
        .try_collect::<Vec<_>>()
        .await?;

    ctx.scuba()
        .clone()
        .log_with_msg("Found draft ancestors", Some(format!("{}", drafts.len())));
    Ok(drafts)
}

#[cfg(test)]
mod test {
    use std::collections::HashSet;

    use blobrepo::AsBlobRepo;
    use fbinit::FacebookInit;
    use maplit::hashset;
    use mononoke_api_types::InnerRepo;
    use tests_utils::bookmark;
    use tests_utils::drawdag::create_from_dag;

    use super::*;

    #[fbinit::test]
    async fn test_find_draft_ancestors_simple(fb: FacebookInit) -> Result<(), Error> {
        let ctx = CoreContext::test_mock(fb);
        let repo: InnerRepo = test_repo_factory::build_empty(fb).await?;
        let mapping = create_from_dag(
            &ctx,
            repo.as_blob_repo(),
            r##"
            A-B-C-D
            "##,
        )
        .await?;

        let cs_id = mapping.get("A").unwrap();
        let to_cs_id = mapping.get("D").unwrap();
        bookmark(&ctx, repo.as_blob_repo(), "book")
            .set_to(*cs_id)
            .await?;
        let drafts = find_draft_ancestors(&ctx, &repo, *to_cs_id).await?;

        let drafts = drafts
            .into_iter()
            .map(|bcs| bcs.get_changeset_id())
            .collect::<HashSet<_>>();

        assert_eq!(
            drafts,
            hashset! {
                *mapping.get("B").unwrap(),
                *mapping.get("C").unwrap(),
                *mapping.get("D").unwrap(),
            }
        );

        bookmark(&ctx, repo.as_blob_repo(), "book")
            .set_to(*mapping.get("B").unwrap())
            .await?;
        let drafts = find_draft_ancestors(&ctx, &repo, *to_cs_id).await?;

        let drafts = drafts
            .into_iter()
            .map(|bcs| bcs.get_changeset_id())
            .collect::<HashSet<_>>();

        assert_eq!(
            drafts,
            hashset! {
                *mapping.get("C").unwrap(),
                *mapping.get("D").unwrap(),
            }
        );
        Ok(())
    }

    #[fbinit::test]
    async fn test_find_draft_ancestors_merge(fb: FacebookInit) -> Result<(), Error> {
        let ctx = CoreContext::test_mock(fb);
        let repo: InnerRepo = test_repo_factory::build_empty(fb).await?;
        let mapping = create_from_dag(
            &ctx,
            repo.as_blob_repo(),
            r##"
              B
             /  \
            A    D
             \  /
               C
            "##,
        )
        .await?;

        let cs_id = mapping.get("B").unwrap();
        let to_cs_id = mapping.get("D").unwrap();
        bookmark(&ctx, repo.as_blob_repo(), "book")
            .set_to(*cs_id)
            .await?;
        let drafts = find_draft_ancestors(&ctx, &repo, *to_cs_id).await?;

        let drafts = drafts
            .into_iter()
            .map(|bcs| bcs.get_changeset_id())
            .collect::<HashSet<_>>();

        assert_eq!(
            drafts,
            hashset! {
                *mapping.get("C").unwrap(),
                *mapping.get("D").unwrap(),
            }
        );

        Ok(())
    }
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This software may be used and distributed according to the terms of the
 * GNU General Public License version 2.
 */

//! Log changes to the repository (new commits and bookmark updates) to
//! external telemetry.

mod bookmark_logger;
mod commit_logger;

pub use crate::bookmark_logger::log_bookmark_operation;
pub use crate::bookmark_logger::BookmarkInfo;
pub use crate::bookmark_logger::BookmarkOperation;
pub use crate::commit_logger::find_draft_ancestors;
pub use crate::commit_logger::log_new_bonsai_changesets;
pub use crate::commit_logger::log_new_commits;
pub use crate::commit_logger::CommitInfo;


// @generated by autocargo
use std::env;
use std::fs;
use std::path::Path;

use thrift_compiler::Config;
use thrift_compiler::GenContext;

#[rustfmt::skip]
fn main() {
    // Rerun if this gets rewritten.
    println!("cargo:rerun-if-changed=thrift_build.rs");

    let out_dir = env::var_os("OUT_DIR").expect("OUT_DIR env not provided");
    let out_dir: &Path = out_dir.as_ref();
    fs::write(
        out_dir.join("cratemap"),
        "filenodes crate
mercurial_thrift mercurial_thrift
mononoke_types_thrift mononoke_types_thrift",
    ).expect("Failed to write cratemap");

    let conf = {
        let mut conf = Config::from_env(GenContext::Lib).expect("Failed to instantiate thrift_compiler::Config");

        let path_from_manifest_to_base: &Path = "../../../..".as_ref();
        let cargo_manifest_dir =
            env::var_os("CARGO_MANIFEST_DIR").expect("CARGO_MANIFEST_DIR not provided");
        let cargo_manifest_dir: &Path = cargo_manifest_dir.as_ref();
        let base_path = cargo_manifest_dir
            .join(path_from_manifest_to_base)
            .canonicalize()
            .expect("Failed to canonicalize base_path");
        // TODO: replace canonicalize() with std::path::absolute() when
        // https://github.com/rust-lang/rust/pull/91673 is available (~Rust 1.60)
        // and remove this block.
        #[cfg(windows)]
        let base_path = Path::new(
            base_path
                .as_path()
                .to_string_lossy()
                .trim_start_matches(r"\\?\"),
            )
            .to_path_buf();

        conf.base_path(base_path);

        let options = "types_crate=filenodes-if__types";
        if !options.is_empty() {
            conf.options(options);
        }

        let include_srcs = vec![

        ];
        conf.include_srcs(include_srcs);

        conf
    };

    conf
        .run(&[
            "filenodes.thrift"
        ])
        .expect("Failed while running thrift compilation");
}


// @generated by autocargo
::codegen_includer_proc_macro::include!();


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This software may be used and distributed according to the terms of the
 * GNU General Public License version 2.
 */

use abomonation_derive::Abomonation;
use anyhow::anyhow;
use anyhow::Result;
use async_trait::async_trait;
use context::CoreContext;
use mercurial_types::HgChangesetId;
use mercurial_types::HgFileNodeId;
use mercurial_types::HgNodeHash;
use mercurial_types::RepoPath;
use mononoke_types::hash;
use quickcheck_arbitrary_derive::Arbitrary;

pub mod thrift {
    pub use filenodes_if::*;
}

pub fn blake2_path_hash(data: &Vec<u8>) -> hash::Blake2 {
    let mut hash_content = hash::Context::new("path".as_bytes());
    hash_content.update(data);
    hash_content.finish()
}

#[derive(Clone, Debug, Eq, PartialEq)]
pub struct PreparedFilenode {
    pub path: RepoPath,
    pub info: FilenodeInfo,
}

#[derive(Abomonation, Arbitrary, Clone, Debug, Eq, PartialEq)]
pub struct FilenodeInfo {
    pub filenode: HgFileNodeId,
    pub p1: Option<HgFileNodeId>,
    pub p2: Option<HgFileNodeId>,
    pub copyfrom: Option<(RepoPath, HgFileNodeId)>,
    pub linknode: HgChangesetId,
}

#[derive(Abomonation, Clone, Debug, Eq, PartialEq)]
pub enum FilenodeRange {
    /// A range of filenodes.
    Filenodes(Vec<FilenodeInfo>),

    /// This range is too large to be fetched.
    TooBig,
}

// The main purpose of FilenodeResult is to force callers to deal with situation
// when filenodes are disabled. This shouldn't happen normally, but it
// might happen in exceptional situation like e.g. filenodes db being
// unavailable.
//
// The guideline here is the following - if the code might affect the critical
// read path i.e. serving "hg pull"/"hg update" then it should not rely on
// filenodes being available and it needs to have a workaround.
#[derive(Debug)]
#[must_use]
pub enum FilenodeResult<T> {
    Present(T),
    Disabled,
}

impl<T> FilenodeResult<T> {
    pub fn map<U>(self, func: impl Fn(T) -> U) -> FilenodeResult<U> {
        match self {
            FilenodeResult::Present(t) => FilenodeResult::Present(func(t)),
            FilenodeResult::Disabled => FilenodeResult::Disabled,
        }
    }

    pub fn do_not_handle_disabled_filenodes(self) -> Result<T> {
        match self {
            FilenodeResult::Present(t) => Ok(t),
            FilenodeResult::Disabled => Err(anyhow!("filenodes are disabled")),
        }
    }
}

impl FilenodeInfo {
    pub fn from_thrift(info: thrift::FilenodeInfo) -> Result<Self> {
        let catch_block = || {
            let copyfrom = match info.copyfrom {
                None => None,
                Some(copyfrom) => Some((
                    RepoPath::from_thrift(copyfrom.path)?,
                    HgFileNodeId::new(HgNodeHash::from_thrift(copyfrom.filenode)?),
                )),
            };

            Ok(Self {
                filenode: HgFileNodeId::new(HgNodeHash::from_thrift(info.filenode)?),
                p1: HgNodeHash::from_thrift_opt(info.p1)?.map(HgFileNodeId::new),
                p2: HgNodeHash::from_thrift_opt(info.p2)?.map(HgFileNodeId::new),
                copyfrom,
                linknode: HgChangesetId::new(HgNodeHash::from_thrift(info.linknode)?),
            })
        };

        catch_block()
    }

    pub fn into_thrift(self) -> thrift::FilenodeInfo {
        thrift::FilenodeInfo {
            filenode: self.filenode.into_nodehash().into_thrift(),
            p1: self.p1.map(|p| p.into_nodehash().into_thrift()),
            p2: self.p2.map(|p| p.into_nodehash().into_thrift()),
            copyfrom: self.copyfrom.map(|copyfrom| thrift::FilenodeCopyFrom {
                path: copyfrom.0.into_thrift(),
                filenode: copyfrom.1.into_nodehash().into_thrift(),
            }),
            linknode: self.linknode.into_nodehash().into_thrift(),
        }
    }
}

#[facet::facet]
#[async_trait]
pub trait Filenodes: Send + Sync {
    async fn add_filenodes(
        &self,
        ctx: &CoreContext,
        info: Vec<PreparedFilenode>,
    ) -> Result<FilenodeResult<()>>;

    async fn add_or_replace_filenodes(
        &self,
        ctx: &CoreContext,
        info: Vec<PreparedFilenode>,
    ) -> Result<FilenodeResult<()>>;

    async fn get_filenode(
        &self,
        ctx: &CoreContext,
        path: &RepoPath,
        filenode: HgFileNodeId,
    ) -> Result<FilenodeResult<Option<FilenodeInfo>>>;

    async fn get_all_filenodes_maybe_stale(
        &self,
        ctx: &CoreContext,
        path: &RepoPath,
        limit: Option<u64>,
    ) -> Result<FilenodeResult<FilenodeRange>>;

    fn prime_cache(&self, ctx: &CoreContext, filenodes: &[PreparedFilenode]);
}

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

    use super::*;

    quickcheck! {
        fn filenodes_info_thrift_roundtrip(obj: FilenodeInfo) -> bool {
            let thrift_struct = obj.clone().into_thrift();
            let obj2 = FilenodeInfo::from_thrift(thrift_struct)
                .expect("converting a valid Thrift structure should always work");
            obj == obj2
        }
    }
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This software may be used and distributed according to the terms of the
 * GNU General Public License version 2.
 */

use anyhow::Error;
use anyhow::Result;
use futures::future;
use futures::future::Future;
use futures::future::TryFutureExt;
use mononoke_types::hash;

use crate::expected_size::ExpectedSize;
use crate::incremental_hash::Blake3IncrementalHasher;
use crate::incremental_hash::GitSha1IncrementalHasher;
use crate::incremental_hash::Sha1IncrementalHasher;
use crate::incremental_hash::Sha256IncrementalHasher;
use crate::multiplexer::Multiplexer;
use crate::streamhash::hash_stream;

type Aliases = (hash::Sha1, hash::Sha256, hash::RichGitSha1, hash::Blake3);

/// Hashes are dependent on the ExpectedSize we received, so we want to make sure callers verify
/// that the size they gave us is actually the one they observed before giving them access to the
/// resulting hashes. To do so, we require callers to provide the effective size to retrieve
/// aliases.
#[derive(Debug)]
pub struct RedeemableAliases {
    expected_size: ExpectedSize,
    aliases: Aliases,
}

impl RedeemableAliases {
    pub fn new(expected_size: ExpectedSize, aliases: Aliases) -> Self {
        Self {
            expected_size,
            aliases,
        }
    }

    pub fn redeem(self, size: u64) -> Result<Aliases> {
        let Self {
            expected_size,
            aliases,
        } = self;
        expected_size.check_equals(size).map(move |_| aliases)
    }
}

// Given a multiplexer, attach new aliases computations to it.
pub fn add_aliases_to_multiplexer<T: AsRef<[u8]> + Send + Sync + Clone + 'static>(
    multiplexer: &mut Multiplexer<T>,
    expected_size: ExpectedSize,
) -> impl Future<Output = Result<RedeemableAliases, Error>> + std::marker::Unpin {
    let sha1 = multiplexer.add(|stream| hash_stream(Sha1IncrementalHasher::new(), stream));
    let sha256 = multiplexer.add(|stream| hash_stream(Sha256IncrementalHasher::new(), stream));
    let git_sha1 = multiplexer
        .add(move |stream| hash_stream(GitSha1IncrementalHasher::new(expected_size), stream));
    let seeded_blake3 =
        multiplexer.add(|stream| hash_stream(Blake3IncrementalHasher::new_seeded(), stream));

    future::try_join4(sha1, sha256, git_sha1, seeded_blake3)
        .map_ok(move |aliases| RedeemableAliases::new(expected_size, aliases))
        .map_err(Error::from)
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This software may be used and distributed according to the terms of the
 * GNU General Public License version 2.
 */

use std::fmt;
use std::fmt::Debug;
use std::pin::Pin;

use anyhow::Error;
use anyhow::Result;
use bytes::Bytes;
use bytes::BytesMut;
use futures::future::BoxFuture;
use futures::future::FutureExt;
use futures::future::TryFutureExt;
use futures::stream::BoxStream;
use futures::stream::Stream;
use futures::stream::StreamExt;
use futures::stream::TryStreamExt;
use futures::task::Context;
use futures::task::Poll;

use crate::expected_size::ExpectedSize;

#[must_use = "streams do nothing unless polled"]
#[pin_project::pin_project]
#[derive(Debug)]
pub struct ChunkStream<S> {
    #[pin]
    stream: S,
    state: ChunkStreamState,
}

#[derive(Debug)]
struct ChunkStreamState {
    chunk_size: usize,
    buff: BytesMut,
    emitted: bool,
    had_data: bool,
    done: bool,
}

impl<S> ChunkStream<S> {
    pub fn new(stream: S, chunk_size: usize) -> ChunkStream<S> {
        assert!(chunk_size > 0);

        ChunkStream {
            stream,
            state: ChunkStreamState {
                chunk_size,
                buff: BytesMut::with_capacity(chunk_size),
                emitted: false,
                had_data: false,
                done: false,
            },
        }
    }
}

impl<S, E> Stream for ChunkStream<S>
where
    S: Stream<Item = Result<Bytes, E>>,
{
    type Item = S::Item;

    fn poll_next(self: Pin<&mut Self>, ctx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
        let mut proj = self.project();

        if proj.state.done {
            return Poll::Ready(None);
        }

        loop {
            if proj.state.buff.len() >= proj.state.chunk_size {
                // We've buffered more data than we need. Emit some.
                proj.state.emitted = true;
                let chunk = proj.state.buff.split_to(proj.state.chunk_size).freeze();
                return Poll::Ready(Some(Ok(chunk)));
            }

            // We need more data. Poll for some! Note the as_mut() here is used to reborrow the
            // stream and avoid moving it into the loop iteration.

            match futures::ready!(proj.stream.as_mut().poll_next(ctx)) {
                Some(Ok(bytes)) => {
                    // We got more data. Extend our buffer, then see if that is enough to return. Note
                    // that extend_from slice implicitly extends our BytesMut.
                    proj.state.had_data = true;
                    proj.state.buff.extend_from_slice(&bytes);
                    continue;
                }
                Some(Err(e)) => {
                    return Poll::Ready(Some(Err(e)));
                }
                None => {
                    // Fallthrough
                }
            };

            // No more data is coming.

            proj.state.done = true;

            // Return whatever we have left. However, we need to be a little careful to handle
            // empty data here.
            //
            // If our buffer happens to just be empty, but we emitted data, that just means our
            // data was disivible by our chunk size, and we are done.
            //
            // However, if our buffer is empty, but we never emitted, then we have two possible
            // cases to handle:
            //
            // - Our underlying stream was empty Bytes. In this case, we should return empty Bytes
            // too (we're returning a representation of the underlying content, chunked).
            //
            // - Our underlying stream was empty. In this case, we shouldn't return anything.

            let out = if !proj.state.buff.is_empty() || (proj.state.had_data && !proj.state.emitted)
            {
                // We did have some buffered data. Emit that.
                proj.state.emitted = true;
                let chunk = std::mem::replace(&mut proj.state.buff, BytesMut::new()).freeze();
                Poll::Ready(Some(Ok(chunk)))
            } else {
                // We have no more buffered data. We're done.
                Poll::Ready(None)
            };

            return out;
        }
    }
}

pub enum Chunks<'a> {
    Inline(BoxFuture<'a, Result<Bytes, Error>>),
    Chunked(ExpectedSize, BoxStream<'a, Result<Bytes, Error>>),
}

impl Debug for Chunks<'_> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match self {
            Chunks::Inline(_) => write!(f, "Chunks::Inline(..)"),
            Chunks::Chunked(size, _) => write!(f, "Chunks::Chunked({:?}, ...)", size),
        }
    }
}

/// Chunk a stream of incoming data for storage. We use the incoming size hint to decide whether
/// to chunk.
pub fn make_chunks<'a, S>(
    data: S,
    expected_size: ExpectedSize,
    chunk_size: Option<u64>,
) -> Chunks<'a>
where
    S: Stream<Item = Result<Bytes, Error>> + Send + 'a,
{
    // NOTE: We stop reading if the stream we are provided exceeds the expected_size we were given.
    // While we do check later that the stream matches *exactly* the size we were given, doing this
    // check lets us bail early (and e.g. ensures that if we are told something is 1 byte but it
    // actually is 1TB, we don't try to buffer the whole 1TB).
    let limit = {
        let mut observed_size: u64 = 0; // This moves into the closure below and serves as its state.
        move |chunk: Result<Bytes, Error>| {
            // NOTE: unwrap() will fail if we have a Bytes whose length is too large to fit in a u64.
            // We presumably don't have such Bytes in memory!
            let chunk = chunk?;
            observed_size += u64::try_from(chunk.len()).unwrap();
            expected_size.check_less(observed_size)?;
            Result::<_, Error>::Ok(chunk)
        }
    };

    let data = data.map(limit);

    match chunk_size {
        Some(chunk_size) if expected_size.should_chunk(chunk_size) => {
            let stream = ChunkStream::new(data, chunk_size as usize);
            Chunks::Chunked(expected_size, stream.boxed())
        }
        _ => {
            let fut = data
                .try_fold(
                    expected_size.new_buffer(),
                    |mut bytes, incoming| async move {
                        bytes.extend_from_slice(incoming.as_ref());
                        Result::<_, Error>::Ok(bytes)
                    },
                )
                .map_ok(BytesMut::freeze)
                .boxed();
            Chunks::Inline(fut)
        }
    }
}

#[cfg(test)]
mod test {
    use assert_matches::assert_matches;
    use futures::stream;
    use quickcheck::quickcheck;
    use tokio::runtime::Runtime;

    use super::*;

    #[test]
    fn test_make_chunks_no_chunk_size() {
        let in_stream = stream::empty();

        match make_chunks(in_stream, ExpectedSize::new(10), None) {
            Chunks::Inline(_) => {}
            c => panic!("Did not expect {:?}", c),
        };
    }

    #[test]
    fn test_make_chunks_no_chunking() {
        let in_stream = stream::empty();

        match make_chunks(in_stream, ExpectedSize::new(10), Some(100)) {
            Chunks::Inline(_) => {}
            c => panic!("Did not expect {:?}", c),
        };
    }

    #[test]
    fn test_make_chunks_no_chunking_limit() {
        let in_stream = stream::empty();

        match make_chunks(in_stream, ExpectedSize::new(100), Some(100)) {
            Chunks::Inline(_) => {}
            c => panic!("Did not expect {:?}", c),
        };
    }

    #[test]
    fn test_make_chunks_chunking() {
        let in_stream = stream::empty();

        match make_chunks(in_stream, ExpectedSize::new(1000), Some(100)) {
            Chunks::Chunked(h, _) if h.check_equals(1000).is_ok() => {}
            c => panic!("Did not expect {:?}", c),
        };
    }

    #[tokio::test]
    async fn test_make_chunks_overflow_inline() {
        // Make chunks buffers if we expect content that is small enough to fit the chunk size.
        // However, if we get more content than that, we should stop.

        let chunks = vec![
            Bytes::from(vec![1; 5]),
            Bytes::from(vec![1; 5]),
            Bytes::from(vec![1; 5]),
            Bytes::from(vec![1; 5]),
        ];
        let in_stream = stream::iter(chunks).map(Ok);

        let fut = match make_chunks(in_stream, ExpectedSize::new(10), Some(100)) {
            c @ Chunks::Chunked(..) => panic!("Did not expect {:?}", c),
            Chunks::Inline(fut) => fut,
        };

        fut.await
            .expect_err("make_chunks should abort if the content does not end as advertised");
    }

    #[tokio::test]
    async fn test_make_chunks_overflow_chunked() {
        // If we get more content than advertises, abort.

        let chunks = vec![
            Bytes::from(vec![1; 5]),
            Bytes::from(vec![1; 5]),
            Bytes::from(vec![1; 5]),
            Bytes::from(vec![1; 5]),
        ];
        let in_stream = stream::iter(chunks).map(Ok);

        let fut = match make_chunks(in_stream, ExpectedSize::new(10), Some(1)) {
            Chunks::Chunked(_, stream) => stream.try_collect::<Vec<_>>(),
            c @ Chunks::Inline(..) => panic!("Did not expect {:?}", c),
        };

        fut.await
            .expect_err("make_chunks should abort if the content does not end as advertised");
    }

    #[tokio::test]
    async fn test_stream_of_empty_bytes() {
        // If we give ChunkStream a stream that contains empty bytes, then we should return one
        // chunk of empty bytes.
        let chunks = vec![Bytes::new()];
        let in_stream = stream::iter(chunks).map(Result::<_, ()>::Ok);
        let mut stream = ChunkStream::new(in_stream, 1);

        assert_eq!(stream.try_next().await, Ok(Some(Bytes::new())));
        assert_eq!(stream.try_next().await, Ok(None));
    }

    #[tokio::test]
    async fn test_stream_of_repeated_empty_bytes() {
        // If we give ChunkStream a stream that contains however many empty bytes, then we should
        // return a single chunk of empty bytes.

        let chunks = vec![Bytes::new(), Bytes::new()];
        let in_stream = stream::iter(chunks).map(Result::<_, ()>::Ok);
        let mut stream = ChunkStream::new(in_stream, 1);

        assert_eq!(stream.try_next().await, Ok(Some(Bytes::new())));
        assert_eq!(stream.try_next().await, Ok(None));
    }

    #[tokio::test]
    async fn test_empty_stream() {
        // If we give ChunkStream an empty stream, it should retun an empty stream.

        let in_stream = stream::iter(vec![]).map(Result::<_, ()>::Ok);
        let mut stream = ChunkStream::new(in_stream, 1);

        assert_eq!(stream.next().await, None);
    }

    #[tokio::test]
    async fn test_bigger_incoming_chunks() {
        // Explicitly test that ChunkStream handles splitting chunks.
        let chunks = vec![vec![1; 10], vec![1; 10]];
        assert!(do_check_chunk_stream(chunks, 5).await)
    }

    #[tokio::test]
    async fn test_smaller_incoming_chunks() {
        // Explicitly test that ChunkStream handles putting chunks together.
        let chunks = vec![vec![1; 10], vec![1; 10]];
        assert!(do_check_chunk_stream(chunks, 15).await)
    }

    #[tokio::test]
    async fn test_stream_exhaustion() {
        #[pin_project::pin_project]
        struct StrictStream {
            chunks: Vec<Bytes>,
            done: bool,
        }

        impl Stream for StrictStream {
            type Item = Result<Bytes, ()>;

            fn poll_next(
                mut self: Pin<&mut Self>,
                _: &mut Context<'_>,
            ) -> Poll<Option<Self::Item>> {
                if self.done {
                    panic!("StrictStream was done");
                }

                match self.chunks.pop() {
                    Some(b) => Poll::Ready(Some(Ok(b))),
                    None => {
                        self.done = true;
                        Poll::Ready(None)
                    }
                }
            }
        }

        let chunks = vec![
            Bytes::from(vec![1; 5]),
            Bytes::from(vec![1; 5]),
            Bytes::from(vec![1; 5]),
            Bytes::from(vec![1; 5]),
        ];

        let mut stream = ChunkStream::new(
            StrictStream {
                chunks,
                done: false,
            },
            10,
        );

        assert_matches!(stream.try_next().await, Ok(Some(_)));
        assert_matches!(stream.try_next().await, Ok(Some(_)));
        assert_matches!(stream.try_next().await, Ok(None));
        assert_matches!(stream.try_next().await, Ok(None));
    }

    async fn do_check_chunk_stream(in_chunks: Vec<Vec<u8>>, size: usize) -> bool {
        let in_chunks: Vec<Bytes> = in_chunks.into_iter().map(Bytes::from).collect();
        let chunk_stream = ChunkStream::new(
            stream::iter(in_chunks.clone()).map(Result::<_, ()>::Ok),
            size,
        );
        let out_chunks = chunk_stream.try_collect::<Vec<_>>().await.unwrap();

        let expected_bytes = in_chunks
            .iter()
            .fold(BytesMut::new(), |mut bytes, chunk| {
                bytes.extend_from_slice(chunk);
                bytes
            })
            .freeze();

        let got_bytes = out_chunks
            .iter()
            .fold(BytesMut::new(), |mut bytes, chunk| {
                bytes.extend_from_slice(chunk);
                bytes
            })
            .freeze();

        // The contents should be the same
        if expected_bytes != got_bytes {
            return false;
        }

        // If there were no contents, then just return that.
        if expected_bytes.is_empty() {
            return true;
        }

        // All chunks except for the last one must equal chunk size
        for chunk in out_chunks[0..out_chunks.len() - 1].iter() {
            if chunk.len() != size {
                return false;
            }
        }

        // The last chunk must smaller than the chunk size
        if out_chunks[out_chunks.len() - 1].len() > size {
            return false;
        }

        true
    }

    quickcheck! {
        fn check_chunk_stream(in_chunks: Vec<Vec<u8>>, size: u8) -> bool {
            let size = (size as usize) + 1; // Don't allow 0 as the size.
            let rt = Runtime::new().unwrap();
            rt.block_on(do_check_chunk_stream(in_chunks, size))
        }

        fn check_make_chunks_fut_joins(in_chunks: Vec<Vec<u8>>) -> bool {
            let rt = Runtime::new().unwrap();

            let in_chunks: Vec<Bytes> = in_chunks.into_iter().map(Bytes::from).collect();
            let in_stream = stream::iter(in_chunks.clone()).map(Ok);

            let expected_bytes = in_chunks.iter().fold(BytesMut::new(), |mut bytes, chunk| {
                bytes.extend_from_slice(chunk);
                bytes
            }).freeze();

            let len = expected_bytes.len() as u64;

            let fut = match make_chunks(in_stream, ExpectedSize::new(len), Some(len)) {
                Chunks::Inline(fut) => fut,
                c => panic!("Did not expect {:?}", c),
            };

            let out_bytes = rt.block_on(fut).unwrap();
            out_bytes == expected_bytes
        }
    }
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This software may be used and distributed according to the terms of the
 * GNU General Public License version 2.
 */

use anyhow::Context;
use anyhow::Result;
use blobstore::BlobCopier;
use blobstore::Blobstore;
use blobstore::Loadable;
use blobstore::Storable;
use context::CoreContext;
use futures::future;
use futures::stream;
use futures::stream::StreamExt;
use futures::stream::TryStreamExt;
use mononoke_types::BlobstoreKey;
use mononoke_types::BlobstoreValue;
use mononoke_types::ContentAlias;
use mononoke_types::ContentMetadataV2;
use slog::info;
use strum::IntoEnumIterator;

use crate::Alias;
use crate::AliasBlob;
use crate::FileContents;
use crate::FilestoreConfig;

pub async fn copy(
    original_blobstore: &impl Blobstore,
    copier: &(impl BlobCopier + Sync),
    config: FilestoreConfig,
    ctx: &CoreContext,
    data: &ContentMetadataV2,
) -> Result<()> {
    // See reasoning about order of writes in ./finalize.rs::finalize (https://fburl.com/code/3w8dncr3)

    // Ensure that all aliases are covered, and missing out an alias gives a compile time error.
    future::try_join_all(Alias::iter().map(|alias| {
        match alias {
            Alias::Sha1(_) => copier.copy(ctx, Alias::Sha1(data.sha1).blobstore_key()),
            Alias::GitSha1(_) => {
                copier.copy(ctx, Alias::GitSha1(data.git_sha1.sha1()).blobstore_key())
            }
            Alias::Sha256(_) => copier.copy(ctx, Alias::Sha256(data.sha256).blobstore_key()),
            Alias::SeededBlake3(_) => Box::pin(async move {
                let blake3 = Alias::SeededBlake3(data.seeded_blake3);
                match copier.copy(ctx, blake3.blobstore_key()).await
                {
                    resp @ Ok(_) => resp,
                    // The backfilling for ContentMetadataV2 has happened in different stages so the alias#
                    // might be missing. Regenerate it.
                    Err(_) => {
                        info!(
                            ctx.logger(),
                            "Failure in copying seeded blake3 ({:?}) alias for content ID {:?}. Generating Alias.",
                            data.seeded_blake3,
                            data.content_id
                        );
                        let content_alias = ContentAlias::from_content_id(data.content_id);
                        AliasBlob(blake3, content_alias).store(ctx, original_blobstore).await?;
                        // Now that the alias has been regenerated, try copying again.
                        copier.copy(ctx, blake3.blobstore_key()).await
                    }
                }
            })
        }
    }))
    .await
    .with_context(|| {
        format!(
            "Failure in copying alias for content id {:?}",
            data.content_id
        )
    })?;

    // Files are stored inline or in chunks, depending on their size. If they're chunked,
    // we need to copy all chunks. Unfortunately, the only way to know how they're stored is
    // by loading FileContents, which might be large-ish if the file is actually inlined.
    let file_contents = data.content_id.load(ctx, original_blobstore).await?;
    match file_contents {
        FileContents::Chunked(chunked) => {
            stream::iter(
                chunked
                    .into_chunks()
                    .into_iter()
                    .map(|c| copier.copy(ctx, c.chunk_id().blobstore_key())),
            )
            .buffer_unordered(config.concurrency)
            .try_collect()
            .await?
        }
        FileContents::Bytes(_) => {}
    }

    copier.copy(ctx, data.content_id.blobstore_key()).await?;

    copier
        .copy(ctx, data.clone().into_blob().id().blobstore_key())
        .await?;
    Ok(())
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This software may be used and distributed according to the terms of the
 * GNU General Public License version 2.
 */

use std::fmt::Debug;

use mononoke_types::hash::Blake3;
use mononoke_types::hash::RichGitSha1;
use mononoke_types::hash::Sha1;
use mononoke_types::hash::Sha256;
use mononoke_types::ContentId;
use thiserror::Error;

use crate::expected_size::ExpectedSize;
use crate::FetchKey;

#[derive(Debug)]
pub struct InvalidHash<T: Debug> {
    pub expected: T,
    pub effective: T,
}

#[derive(Debug, Error)]
pub enum ErrorKind {
    #[error("Invalid size: {0:?} was expected, {1:?} was observed")]
    InvalidSize(ExpectedSize, u64),

    #[error("Invalid ContentId: {0:?}")]
    InvalidContentId(InvalidHash<ContentId>),

    #[error("Invalid Sha1: {0:?}")]
    InvalidSha1(InvalidHash<Sha1>),

    #[error("Invalid Sha256: {0:?}")]
    InvalidSha256(InvalidHash<Sha256>),

    #[error("Invalid RichGitSha1: {0:?}")]
    InvalidGitSha1(InvalidHash<RichGitSha1>),

    #[error("Invalid Blake3: {0:?}")]
    InvalidBlake3(InvalidHash<Blake3>),

    #[error("Missing content: {0:?}")]
    MissingContent(FetchKey),
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This software may be used and distributed according to the terms of the
 * GNU General Public License version 2.
 */

use anyhow::Result;
use bytes::BytesMut;

use crate::errors::ErrorKind;
use crate::incremental_hash::AdvisorySize;

/// ExpectedSize is an opaque struct that lets us encapsulate the incoming size for a Filestore
/// operation and make meaningful comparisons with it, but doesn't let us easily read it back. This
/// ensures that we cannot accidentally use the size hint (we shouldn't be trusted) when we need an
/// actual observed size.
#[derive(Debug, Copy, Clone)]
pub struct ExpectedSize(u64);

impl ExpectedSize {
    pub fn new(size: u64) -> Self {
        Self(size)
    }

    pub fn should_chunk(&self, chunk_size: u64) -> bool {
        self.0 > chunk_size
    }

    pub fn check_equals(&self, size: u64) -> Result<()> {
        if size == self.0 {
            return Ok(());
        }
        Err(ErrorKind::InvalidSize(self.clone(), size).into())
    }

    pub fn check_less(&self, size: u64) -> Result<()> {
        if size <= self.0 {
            return Ok(());
        }
        Err(ErrorKind::InvalidSize(self.clone(), size).into())
    }

    pub fn new_buffer(&self) -> BytesMut {
        // NOTE: This will panic if we can't fit an u64 into usize. That's expected.
        BytesMut::with_capacity(self.0.try_into().unwrap())
    }
}

/// The incremental_hash crate does need access to the internal u64 to create its hashes, so we
/// expose it through a Trait defined here. Don't use that trait if you're not incremental_hash.
impl AdvisorySize for ExpectedSize {
    fn advise(&self) -> u64 {
        self.0
    }
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This software may be used and distributed according to the terms of the
 * GNU General Public License version 2.
 */

use std::borrow::Borrow;
use std::cmp::max;
use std::cmp::min;

use anyhow::Error;
use blobstore::Blobstore;
use blobstore::Loadable;
use blobstore::LoadableError;
use bytes::Bytes;
use cloned::cloned;
use context::CoreContext;
use futures::future;
use futures::stream;
use futures::stream::Stream;
use futures::stream::StreamExt;
use itertools::Either;
use mononoke_types::ContentChunk;
use mononoke_types::ContentChunkId;
use mononoke_types::ContentId;
use mononoke_types::FileContents;
use thiserror::Error;

// TODO: Make this configurable? Perhaps as a global, since it's something that only makes sense at
// the program level (as opposed to e;g. chunk size, which makes sense at the repo level).
const BUFFER_MEMORY_BUDGET: u64 = 16 * 1024 * 1024; // 16MB.

#[derive(Debug, Error)]
pub enum ErrorKind {
    #[error("Chunk not found: {0:?}")]
    ChunkNotFound(ContentChunkId),
}

#[derive(Debug, Copy, Clone, PartialEq, Eq)]
pub struct Range {
    inner: RangeInner,
    strict: bool,
}

impl Range {
    pub fn all() -> Self {
        Self {
            inner: RangeInner::All,
            strict: false,
        }
    }

    pub fn sized(start: u64, size: u64) -> Self {
        Self {
            inner: RangeInner::Span {
                start,
                end: start.saturating_add(size),
            },
            strict: false,
        }
    }

    pub fn strict(self) -> Self {
        Self {
            strict: true,
            ..self
        }
    }

    pub fn range_inclusive(start: u64, end: u64) -> Result<Self, Error> {
        if start > end {
            return Err(anyhow::anyhow!("Invalid range bounds: {}-{}", start, end));
        }

        let end = end + 1;

        Ok(Self {
            inner: RangeInner::Span { start, end },
            strict: false,
        })
    }
}

#[derive(Debug, Copy, Clone, PartialEq, Eq)]
enum RangeInner {
    All,
    Span { start: u64, end: u64 },
}

impl RangeInner {
    /// Returns how many byts stream_file_bytes will return when this range is applied to a file of
    /// size file_size.
    fn real_size(&self, file_size: u64) -> u64 {
        match self {
            RangeInner::All => file_size,
            RangeInner::Span { start, end } => {
                let end = std::cmp::min(*end, file_size);
                end.saturating_sub(*start)
            }
        }
    }

    fn exceeds_file_size(&self, file_size: u64) -> bool {
        match self {
            RangeInner::All => true,
            RangeInner::Span { start: _, end } => *end > file_size,
        }
    }
}

pub fn stream_file_bytes<'a, B: Blobstore + Clone + 'a>(
    blobstore: B,
    ctx: impl Borrow<CoreContext> + Clone + Send + Sync + 'a,
    file_contents: FileContents,
    range: Range,
) -> Result<impl Stream<Item = Result<Bytes, Error>> + 'a, Error> {
    let Range {
        inner: range,
        strict,
    } = range;

    if strict && range.exceeds_file_size(file_contents.size()) {
        return Err(Error::msg("Range exceeds file size"));
    }

    let stream = match file_contents {
        FileContents::Bytes(bytes) => {
            // File is just a single chunk of bytes. Return the correct
            // slice, based on the requested range.
            let bytes = match range {
                RangeInner::Span {
                    start: range_start,
                    end: range_end,
                } => {
                    let len = bytes.len() as u64;
                    let slice_start = min(range_start, len);
                    let slice_end = min(range_end, len);
                    bytes.slice((slice_start as usize)..(slice_end as usize))
                }
                RangeInner::All => bytes,
            };
            stream::once(future::ready(Ok(bytes))).left_stream()
        }
        FileContents::Chunked(chunked) => {
            // File is split into multiple chunks. Dispatch fetches for the chunks that overlap the
            // range, and buffer them. We know all chunks are the same size (same possibly for the
            // last one) so we use that to get our buffer size.
            let chunks = chunked.into_chunks();

            let max_chunk_size = chunks.iter().map(|c| c.size()).max();
            let buffer_size = match max_chunk_size {
                Some(size) if size > 0 => BUFFER_MEMORY_BUDGET / size,
                _ => 1,
            };
            let buffer_size = std::cmp::max(buffer_size, 1);

            // NOTE: buffer_size cannot be greater than our memory budget given how it's computed,
            // so that's safe.
            let buffer_size = buffer_size.try_into().unwrap();

            let chunk_iter = match range {
                RangeInner::Span {
                    start: range_start,
                    end: range_end,
                } => {
                    let iter = chunks
                        .into_iter()
                        .map({
                            // Compute chunk start and end within the file.
                            let mut stream_offset = 0;
                            move |chunk| {
                                let chunk_start = stream_offset.clone();
                                let chunk_size = chunk.size();
                                stream_offset += chunk_size;
                                (chunk_start, chunk_start + chunk_size, chunk)
                            }
                        })
                        .skip_while(move |(_chunk_start, chunk_end, _chunk)| {
                            // Skip chunks from before the range.
                            *chunk_end <= range_start
                        })
                        .take_while(move |(chunk_start, _chunk_end, _chunk)| {
                            // Take chunks that overlap the range.
                            *chunk_start < range_end
                        })
                        .map(move |(chunk_start, chunk_end, chunk)| {
                            // Compute the slice within this chunk that we need.
                            let slice_start = max(chunk_start, range_start);
                            let slice_end = min(chunk_end, range_end);
                            if (slice_start, slice_end) == (chunk_start, chunk_end) {
                                (RangeInner::All, chunk)
                            } else {
                                (
                                    RangeInner::Span {
                                        start: slice_start - chunk_start,
                                        end: slice_end - chunk_start,
                                    },
                                    chunk,
                                )
                            }
                        });
                    Either::Left(iter)
                }
                RangeInner::All => {
                    let iter = chunks.into_iter().map(|chunk| (RangeInner::All, chunk));
                    Either::Right(iter)
                }
            };

            stream::iter(chunk_iter.map(move |(chunk_range, chunk)| {
                // Send some (maybe all) of this chunk.
                let chunk_id = chunk.chunk_id();
                cloned!(ctx, blobstore);
                async move {
                    let bytes = chunk_id
                        .load(ctx.borrow(), &blobstore)
                        .await
                        .map_err(move |err| match err {
                            LoadableError::Error(err) => err,
                            LoadableError::Missing(_) => ErrorKind::ChunkNotFound(chunk_id).into(),
                        })
                        .map(ContentChunk::into_bytes)?;

                    let bytes = match chunk_range {
                        RangeInner::Span { start, end } => {
                            bytes.slice((start as usize)..(end as usize))
                        }
                        RangeInner::All => bytes,
                    };

                    Ok(bytes)
                }
            }))
            .buffered(buffer_size)
            .right_stream()
        }
    };

    Ok(stream)
}

pub async fn fetch_with_size<'a, B: Blobstore + Clone + 'a>(
    blobstore: B,
    ctx: impl Borrow<CoreContext> + Clone + Send + Sync + 'a,
    content_id: ContentId,
    range: Range,
) -> Result<Option<(impl Stream<Item = Result<Bytes, Error>> + 'a, u64)>, Error> {
    let maybe_file_contents = {
        cloned!(ctx, blobstore);
        async move { content_id.load(ctx.borrow(), &blobstore).await }.await
    }
    .map(Some)
    .or_else(|err| match err {
        LoadableError::Error(err) => Err(err),
        LoadableError::Missing(_) => Ok(None),
    })?;

    let file_contents = match maybe_file_contents {
        Some(file_contents) => file_contents,
        None => return Ok(None),
    };

    let file_size = file_contents.size();
    let stream = stream_file_bytes(blobstore, ctx, file_contents, range)?;
    Ok(Some((stream, range.inner.real_size(file_size))))
}

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

    #[test]
    fn test_real_size() {
        assert_eq!(RangeInner::All.real_size(5), 5);

        // Bytes 0 to 9 in a 5 bytes file
        assert_eq!(RangeInner::Span { start: 0, end: 10 }.real_size(5), 5);

        // Bytes 0 to 4 in a 5 bytes file
        assert_eq!(RangeInner::Span { start: 0, end: 5 }.real_size(5), 5);

        // Bytes 0 to 3 in a 5 bytes file
        assert_eq!(RangeInner::Span { start: 0, end: 4 }.real_size(5), 4);

        // Bytes 1 to 3 in a 5 bytes file
        assert_eq!(RangeInner::Span { start: 1, end: 4 }.real_size(5), 3);

        // Bytes 3 to 3 in a 5 bytes file
        assert_eq!(RangeInner::Span { start: 3, end: 4 }.real_size(5), 1);

        // Bytes 10 to 10 in a 5 bytes file
        assert_eq!(RangeInner::Span { start: 10, end: 11 }.real_size(5), 0);

        // Bytes 10 to 14 in a 5 bytes file
        assert_eq!(RangeInner::Span { start: 10, end: 15 }.real_size(5), 0);

        // Nothing
        assert_eq!(RangeInner::Span { start: 10, end: 10 }.real_size(5), 0);

        // Nothing
        assert_eq!(RangeInner::Span { start: 0, end: 0 }.real_size(5), 0);
    }
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This software may be used and distributed according to the terms of the
 * GNU General Public License version 2.
 */

use std::fmt::Display;

use anyhow::Context;
use anyhow::Result;
use async_trait::async_trait;
use blobstore::Blobstore;
use blobstore::Loadable;
use blobstore::LoadableError;
use blobstore::Storable;
use context::CoreContext;
use edenapi_types::AnyFileContentId;
use mononoke_types::errors::ErrorKind;
use mononoke_types::hash;
use mononoke_types::BlobstoreKey;
use mononoke_types::ContentAlias;
use mononoke_types::ContentId;
use strum::EnumIter;

/// Key for fetching - we can access with any of the supported key types
#[derive(Debug, Copy, Clone)]
pub enum FetchKey {
    Canonical(ContentId),
    Aliased(Alias),
}

impl From<ContentId> for FetchKey {
    fn from(content_id: ContentId) -> Self {
        FetchKey::Canonical(content_id)
    }
}

impl From<Alias> for FetchKey {
    fn from(alias: Alias) -> Self {
        FetchKey::Aliased(alias)
    }
}

impl From<hash::Sha256> for FetchKey {
    fn from(hash: hash::Sha256) -> Self {
        FetchKey::Aliased(Alias::Sha256(hash))
    }
}

impl From<hash::Sha1> for FetchKey {
    fn from(hash: hash::Sha1) -> Self {
        FetchKey::Aliased(Alias::Sha1(hash))
    }
}

impl From<hash::Blake3> for FetchKey {
    fn from(hash: hash::Blake3) -> Self {
        FetchKey::Aliased(Alias::SeededBlake3(hash))
    }
}

impl From<AnyFileContentId> for FetchKey {
    fn from(id: AnyFileContentId) -> Self {
        match id {
            AnyFileContentId::ContentId(id) => Self::from(ContentId::from(id)),
            AnyFileContentId::Sha1(id) => Self::from(hash::Sha1::from(id)),
            AnyFileContentId::Sha256(id) => Self::from(hash::Sha256::from(id)),
            AnyFileContentId::SeededBlake3(id) => Self::from(hash::Blake3::from(id)),
        }
    }
}

impl FetchKey {
    pub fn blobstore_key(&self) -> String {
        match self {
            Self::Canonical(cid) => cid.blobstore_key(),
            Self::Aliased(alias) => alias.blobstore_key(),
        }
    }
}

#[derive(Debug, Copy, Clone, PartialEq, Eq, Hash, EnumIter)]
pub enum Alias {
    Sha1(hash::Sha1),
    Sha256(hash::Sha256),
    GitSha1(hash::GitSha1),
    SeededBlake3(hash::Blake3),
}

impl Display for Alias {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Self::Sha1(sha1) => f.write_fmt(format_args!("{:?}", sha1)),
            Self::Sha256(sha256) => f.write_fmt(format_args!("{:?}", sha256)),
            Self::GitSha1(gitsha1) => f.write_fmt(format_args!("{:?}", gitsha1)),
            Self::SeededBlake3(seeded_blake3) => f.write_fmt(format_args!("{:?}", seeded_blake3)),
        }
    }
}

#[async_trait]
impl Loadable for FetchKey {
    type Value = ContentId;

    /// Return the canonical ID for a key. It doesn't check if the corresponding content actually
    /// exists:
    /// - When called with content_id, it doesn't check the content id is stored in the blobstore
    /// - It is possible for an alias to exist before the ID if there was an interrupted store
    /// operation.
    async fn load<'a, B: Blobstore>(
        &'a self,
        ctx: &'a CoreContext,
        blobstore: &'a B,
    ) -> Result<Self::Value, LoadableError> {
        match self {
            FetchKey::Canonical(content_id) => Ok(*content_id),
            FetchKey::Aliased(alias) => alias.load(ctx, blobstore).await,
        }
    }
}

impl Alias {
    pub fn blobstore_key(&self) -> String {
        match self {
            Alias::GitSha1(git_sha1) => format!("alias.gitsha1.{}", git_sha1.to_hex()),
            Alias::Sha1(sha1) => format!("alias.sha1.{}", sha1.to_hex()),
            Alias::Sha256(sha256) => format!("alias.sha256.{}", sha256.to_hex()),
            Alias::SeededBlake3(blake3) => format!("alias.seeded_blake3.{}", blake3.to_hex()),
        }
    }

    #[inline]
    pub fn sampling_fingerprint(&self) -> u64 {
        match self {
            Alias::GitSha1(git_sha1) => git_sha1.sampling_fingerprint(),
            Alias::Sha1(sha1) => sha1.sampling_fingerprint(),
            Alias::Sha256(sha256) => sha256.sampling_fingerprint(),
            Alias::SeededBlake3(blake3) => blake3.sampling_fingerprint(),
        }
    }
}

#[async_trait]
impl Loadable for Alias {
    type Value = ContentId;

    async fn load<'a, B: Blobstore>(
        &'a self,
        ctx: &'a CoreContext,
        blobstore: &'a B,
    ) -> Result<Self::Value, LoadableError> {
        let key = self.blobstore_key();
        let get = blobstore.get(ctx, &key);
        let maybe_alias = get.await?;
        let blob = maybe_alias.ok_or_else(|| LoadableError::Missing(key.clone()))?;

        ContentAlias::from_bytes(blob.into_raw_bytes())
            .map(|alias| alias.content_id())
            .with_context(|| ErrorKind::BlobKeyError(key.clone()))
            .map_err(LoadableError::Error)
    }
}

pub struct AliasBlob(pub Alias, pub ContentAlias);

#[async_trait]
impl Storable for AliasBlob {
    type Key = ();

    async fn store<'a, B: Blobstore>(
        self,
        ctx: &'a CoreContext,
        blobstore: &'a B,
    ) -> Result<Self::Key> {
        blobstore
            .put(ctx, self.0.blobstore_key(), self.1.into_blob())
            .await
    }
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This software may be used and distributed according to the terms of the
 * GNU General Public License version 2.
 */

use anyhow::Error;
use blobstore::Blobstore;
use blobstore::Storable;
use context::CoreContext;
use futures::future;
use mononoke_types::BlobstoreValue;
use mononoke_types::ContentAlias;
use mononoke_types::ContentMetadataV2;
use strum::IntoEnumIterator;

use crate::errors::ErrorKind;
use crate::errors::InvalidHash;
use crate::fetch_key::Alias;
use crate::fetch_key::AliasBlob;
use crate::prepare::Prepared;
use crate::StoreRequest;

fn check_hash<T: std::fmt::Debug + PartialEq + Copy>(
    expected: Option<T>,
    effective: T,
) -> Result<(), InvalidHash<T>> {
    if let Some(expected) = expected {
        if expected != effective {
            return Err(InvalidHash {
                expected,
                effective,
            });
        }
    }

    Ok(())
}

pub async fn finalize<B: Blobstore>(
    blobstore: &B,
    ctx: &CoreContext,
    req: Option<&StoreRequest>,
    outcome: Prepared,
) -> Result<ContentMetadataV2, Error> {
    let Prepared {
        sha1,
        sha256,
        git_sha1,
        seeded_blake3,
        contents,
        is_ascii,
        is_binary,
        is_utf8,
        ends_in_newline,
        newline_count,
        first_line,
        is_generated,
        is_partially_generated,
    } = outcome;

    let total_size = contents.size();

    let blob = contents.into_blob();
    let content_id = *blob.id();

    if let Some(req) = req {
        let StoreRequest {
            expected_size,
            canonical: req_content_id,
            sha1: req_sha1,
            sha256: req_sha256,
            git_sha1: req_git_sha1,
            seeded_blake3: req_seeded_blake3,
        } = req;

        expected_size.check_equals(total_size)?;

        {
            use ErrorKind::*;
            check_hash(*req_content_id, content_id).map_err(InvalidContentId)?;
            check_hash(*req_sha1, sha1).map_err(InvalidSha1)?;
            check_hash(*req_sha256, sha256).map_err(InvalidSha256)?;
            check_hash(*req_git_sha1, git_sha1).map_err(InvalidGitSha1)?;
            check_hash(*req_seeded_blake3, seeded_blake3).map_err(InvalidBlake3)?
        }
    }

    let alias = ContentAlias::from_content_id(content_id);
    // Ensure that all aliases are covered, and missing out an alias gives a compile time error.
    future::try_join_all(Alias::iter().map(|alias_type| match alias_type {
        Alias::Sha1(_) => AliasBlob(Alias::Sha1(sha1), alias.clone()).store(ctx, blobstore),
        Alias::GitSha1(_) => {
            AliasBlob(Alias::GitSha1(git_sha1.sha1()), alias.clone()).store(ctx, blobstore)
        }
        Alias::Sha256(_) => AliasBlob(Alias::Sha256(sha256), alias.clone()).store(ctx, blobstore),
        Alias::SeededBlake3(_) => {
            AliasBlob(Alias::SeededBlake3(seeded_blake3), alias.clone()).store(ctx, blobstore)
        }
    }))
    .await?;

    // Since we don't have atomicity for multiple puts, we need to make sure they're ordered
    // correctly:
    //
    // - write the forward-mapping aliases
    // - write the data blob
    // - write the metadata blob
    //
    // Rationale for this order: since we can't guarantee the aliases are written atomically,
    // on failure we could end up writing some but not others. If the underlying blob exists
    // at that point, we've got an inconsistency. However writing the data blob is atomic,
    // and the aliases are only meaningful as references to that blob (in other words, an
    // alias referring to an absent blob is itself considered to be absent, so logically all
    // all the aliases come into existence atomically when the data blob is written).
    // Once the data blob is written we can write the metadata object. This is just a
    // cache, as everything in it can be computed from the content id. Therefore, in principle,
    // if it doesn't get written we can fix it up later.

    blob.store(ctx, blobstore).await?;

    let metadata = ContentMetadataV2 {
        total_size,
        content_id,
        sha1,
        git_sha1,
        sha256,
        seeded_blake3,
        is_ascii,
        is_binary,
        is_utf8,
        ends_in_newline,
        newline_count,
        first_line,
        is_generated,
        is_partially_generated,
    };
    metadata.clone().into_blob().store(ctx, blobstore).await?;

    Ok(metadata)
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This software may be used and distributed according to the terms of the
 * GNU General Public License version 2.
 */

use blake3::Hasher as Blake3;
use bytes::Bytes;
use digest::Digest;
use mononoke_types::hash;
use mononoke_types::typed_hash;
use mononoke_types::ContentId;
use sha1::Sha1;
use sha2::Sha256;

pub fn hash_bytes<H>(mut hasher: impl Hasher<H>, bytes: &Bytes) -> H {
    hasher.update(bytes);
    hasher.finish()
}

pub trait AdvisorySize {
    fn advise(&self) -> u64;
}

impl AdvisorySize for &Bytes {
    fn advise(&self) -> u64 {
        // NOTE: This will panic if the size of the Bytes buffer we have doesn't fit in a u64.
        self.len().try_into().unwrap()
    }
}

pub trait Hasher<H> {
    /// Update the Hasher with new bytes
    fn update<T: AsRef<[u8]>>(&mut self, bytes: T);

    /// Turn the Hasher into the actual Hash.
    fn finish(self) -> H;
}

pub struct ContentIdIncrementalHasher(typed_hash::ContentIdContext);

impl ContentIdIncrementalHasher {
    pub fn new() -> Self {
        Self(typed_hash::ContentIdContext::new())
    }
}

impl Hasher<ContentId> for ContentIdIncrementalHasher {
    fn update<T: AsRef<[u8]>>(&mut self, bytes: T) {
        self.0.update(bytes)
    }

    fn finish(self) -> ContentId {
        self.0.finish()
    }
}

pub struct Sha1IncrementalHasher(Sha1);

impl Sha1IncrementalHasher {
    pub fn new() -> Self {
        Self(Sha1::new())
    }
}

impl Hasher<hash::Sha1> for Sha1IncrementalHasher {
    fn update<T: AsRef<[u8]>>(&mut self, bytes: T) {
        self.0.update(bytes.as_ref())
    }

    fn finish(self) -> hash::Sha1 {
        let hash = self.0.finalize().into();
        hash::Sha1::from_byte_array(hash)
    }
}

pub struct Sha256IncrementalHasher(Sha256);

impl Sha256IncrementalHasher {
    pub fn new() -> Self {
        Self(Sha256::new())
    }
}

impl Hasher<hash::Sha256> for Sha256IncrementalHasher {
    fn update<T: AsRef<[u8]>>(&mut self, bytes: T) {
        self.0.update(bytes.as_ref())
    }

    fn finish(self) -> hash::Sha256 {
        let hash = self.0.finalize().into();
        hash::Sha256::from_byte_array(hash)
    }
}

pub struct GitSha1IncrementalHasher(Sha1, u64);

impl GitSha1IncrementalHasher {
    pub fn new<A: AdvisorySize>(size: A) -> Self {
        let size = size.advise();
        let mut sha1 = Sha1::new();
        let prototype = hash::RichGitSha1::from_byte_array([0; 20], "blob", size);
        sha1.update(&prototype.prefix());
        Self(sha1, size)
    }
}

impl Hasher<hash::RichGitSha1> for GitSha1IncrementalHasher {
    fn update<T: AsRef<[u8]>>(&mut self, bytes: T) {
        self.0.update(bytes.as_ref())
    }

    fn finish(self) -> hash::RichGitSha1 {
        let hash = self.0.finalize().into();
        hash::RichGitSha1::from_byte_array(hash, "blob", self.1)
    }
}

/// Incremental hasher for non-seeded and seeded Blake3 hash
pub struct Blake3IncrementalHasher(Blake3);

impl Blake3IncrementalHasher {
    /// Creates a non-seeded Blake3 incremental hasher
    pub fn _new() -> Self {
        Self(Blake3::new())
    }

    /// Creates a seeded Blake3 incremental hasher
    pub fn new_seeded() -> Self {
        #[cfg(fbcode_build)]
        let bytes = blake3_constant::BLAKE3_HASH_KEY.as_bytes();
        #[cfg(not(fbcode_build))]
        let bytes = "20220728-2357111317192329313741#".as_bytes();
        if bytes.len() != blake3::KEY_LEN {
            panic!(
                "Seed for Blake3 hash needs exactly {} bytes",
                blake3::KEY_LEN
            )
        } else {
            let mut ret = [0; blake3::KEY_LEN];
            ret.copy_from_slice(bytes);
            Self(Blake3::new_keyed(&ret))
        }
    }
}

impl Hasher<hash::Blake3> for Blake3IncrementalHasher {
    fn update<T: AsRef<[u8]>>(&mut self, bytes: T) {
        self.0.update(bytes.as_ref());
    }

    fn finish(self) -> hash::Blake3 {
        let hash = self.0.finalize().into();
        hash::Blake3::from_byte_array(hash)
    }
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This software may be used and distributed according to the terms of the
 * GNU General Public License version 2.
 */

#![feature(never_type)]
#![type_length_limit = "2000000"]

use std::borrow::Borrow;

use anyhow::Error;
use blobstore::Blobstore;
use blobstore::Loadable;
use blobstore::LoadableError;
use bytes::Bytes;
use bytes::BytesMut;
use cloned::cloned;
use context::CoreContext;
use futures::future::Future;
use futures::future::TryFutureExt;
use futures::stream;
use futures::stream::Stream;
use futures::stream::TryStreamExt;
use mononoke_types::hash;
use mononoke_types::BlobstoreKey;
use mononoke_types::ContentId;
use mononoke_types::ContentMetadataV2;
use mononoke_types::FileContents;

mod alias;
mod chunk;
mod copy;
mod errors;
mod expected_size;
mod fetch;
mod fetch_key;
mod finalize;
mod incremental_hash;
mod metadata;
mod multiplexer;
mod prepare;
mod rechunk;
mod streamhash;

pub use alias::add_aliases_to_multiplexer;
pub use copy::copy;
pub use fetch::Range;
pub use fetch_key::Alias;
pub use fetch_key::AliasBlob;
pub use fetch_key::FetchKey;
pub use incremental_hash::hash_bytes;
pub use incremental_hash::Blake3IncrementalHasher;
pub use incremental_hash::GitSha1IncrementalHasher;
pub use incremental_hash::Sha1IncrementalHasher;
pub use incremental_hash::Sha256IncrementalHasher;
pub use rechunk::force_rechunk;
pub use rechunk::rechunk;

#[cfg(test)]
mod test;

/// File storage.
///
/// This is a specialized wrapper around a blobstore specifically for user data files (
/// rather than metadata, trees, etc). Its primary (initial) goals are:
/// - providing a streaming interface for file access
/// - maintain multiple aliases for each file using different key schemes
/// - maintain reverse mapping from primary key to aliases
///
/// Secondary:
/// - Implement chunking at this level
/// - Compression
/// - Range access (initially fetch, and later store)
///
/// Implementation notes:
/// This code takes over the management of file content in a backwards compatible way - it uses
/// the same blobstore key structure and the same encoding schemes for existing files.
/// Extensions (compression, chunking) will change this, but it will still allow backwards
/// compatibility.
#[facet::facet]
#[derive(Debug, Copy, Clone)]
pub struct FilestoreConfig {
    pub chunk_size: Option<u64>,
    pub concurrency: usize,
}

impl FilestoreConfig {
    pub fn no_chunking_filestore() -> Self {
        Self {
            chunk_size: None,
            concurrency: 1,
        }
    }
}

/// Key for storing. We'll compute any missing keys, but we must have the total size.
#[derive(Debug, Clone)]
pub struct StoreRequest {
    pub expected_size: expected_size::ExpectedSize,
    pub canonical: Option<ContentId>,
    pub sha1: Option<hash::Sha1>,
    pub sha256: Option<hash::Sha256>,
    pub git_sha1: Option<hash::RichGitSha1>,
    pub seeded_blake3: Option<hash::Blake3>,
}

impl StoreRequest {
    pub fn new(size: u64) -> Self {
        use expected_size::ExpectedSize;

        Self {
            expected_size: ExpectedSize::new(size),
            canonical: None,
            sha1: None,
            sha256: None,
            git_sha1: None,
            seeded_blake3: None,
        }
    }

    pub fn with_canonical(size: u64, canonical: ContentId) -> Self {
        use expected_size::ExpectedSize;

        Self {
            expected_size: ExpectedSize::new(size),
            canonical: Some(canonical),
            sha1: None,
            sha256: None,
            git_sha1: None,
            seeded_blake3: None,
        }
    }

    pub fn with_sha1(size: u64, sha1: hash::Sha1) -> Self {
        use expected_size::ExpectedSize;

        Self {
            expected_size: ExpectedSize::new(size),
            canonical: None,
            sha1: Some(sha1),
            sha256: None,
            git_sha1: None,
            seeded_blake3: None,
        }
    }

    pub fn with_sha256(size: u64, sha256: hash::Sha256) -> Self {
        use expected_size::ExpectedSize;

        Self {
            expected_size: ExpectedSize::new(size),
            canonical: None,
            sha1: None,
            sha256: Some(sha256),
            git_sha1: None,
            seeded_blake3: None,
        }
    }

    pub fn with_git_sha1(size: u64, git_sha1: hash::RichGitSha1) -> Self {
        use expected_size::ExpectedSize;

        Self {
            expected_size: ExpectedSize::new(size),
            canonical: None,
            sha1: None,
            sha256: None,
            git_sha1: Some(git_sha1),
            seeded_blake3: None,
        }
    }

    pub fn with_seeded_blake3(size: u64, seeded_blake3: hash::Blake3) -> Self {
        Self {
            expected_size: expected_size::ExpectedSize::new(size),
            canonical: None,
            sha1: None,
            sha256: None,
            git_sha1: None,
            seeded_blake3: Some(seeded_blake3),
        }
    }

    pub fn with_fetch_key(size: u64, key: FetchKey) -> Self {
        use FetchKey::*;
        match key {
            Canonical(id) => Self::with_canonical(size, id),
            Aliased(Alias::Sha1(id)) => Self::with_sha1(size, id),
            Aliased(Alias::GitSha1(id)) => {
                Self::with_git_sha1(size, hash::RichGitSha1::from_sha1(id, "blob", size))
            }
            Aliased(Alias::Sha256(id)) => Self::with_sha256(size, id),
            Aliased(Alias::SeededBlake3(id)) => Self::with_seeded_blake3(size, id),
        }
    }
}

/// Fetch the metadata for the underlying content. This will return None if the content does
/// not exist. It might recompute metadata on the fly if the content exists but the metadata does
/// not.
pub async fn get_metadata<B: Blobstore>(
    blobstore: &B,
    ctx: &CoreContext,
    key: &FetchKey,
) -> Result<Option<ContentMetadataV2>, Error> {
    let maybe_id = key
        .load(ctx, blobstore)
        .await
        .map(Some)
        .or_else(|err| match err {
            LoadableError::Error(err) => Err(err),
            LoadableError::Missing(_) => Ok(None),
        })?;

    match maybe_id {
        Some(id) => metadata::get_metadata(blobstore, ctx, id).await,
        None => Ok(None),
    }
}

/// Fetch the metadata for the underlying content. This will return None if the content does
/// not exist, Some(None) if the metadata does not exist, and Some(Some(ContentMetadata))
/// when metadata found. It will not recompute metadata on the fly
pub async fn get_metadata_readonly<B: Blobstore>(
    blobstore: &B,
    ctx: &CoreContext,
    key: &FetchKey,
) -> Result<Option<Option<ContentMetadataV2>>, Error> {
    let maybe_id = key
        .load(ctx, blobstore)
        .await
        .map(Some)
        .or_else(|err| match err {
            LoadableError::Error(err) => Err(err),
            LoadableError::Missing(_) => Ok(None),
        })?;

    match maybe_id {
        Some(id) => metadata::get_metadata_readonly(blobstore, ctx, id)
            .await
            .map(Some),
        None => Ok(Some(None)),
    }
}

/// Return true if the given key exists. A successful return means the key definitely
/// either exists or doesn't; an error means the existence could not be determined.
pub async fn exists<B: Blobstore>(
    blobstore: &B,
    ctx: &CoreContext,
    key: &FetchKey,
) -> Result<bool, Error> {
    let maybe_id = key
        .load(ctx, blobstore)
        .await
        .map(Some)
        .or_else(|err| match err {
            LoadableError::Error(err) => Err(err),
            LoadableError::Missing(_) => Ok(None),
        })?;

    match maybe_id {
        Some(id) => blobstore
            .is_present(ctx, &id.blobstore_key())
            .await?
            .fail_if_unsure(),
        None => Ok(false),
    }
}

/// Fetch a file as a stream. This returns either success with a stream of data and file size if
/// the file exists, success with None if it does not exist, or an Error if either existence can't
/// be determined or if opening the file failed. File contents are returned in chunks configured by
/// FilestoreConfig::read_chunk_size - this defines the max chunk size, but they may be shorter
/// (not just the final chunks - any of them). Chunks are guaranteed to have non-zero size.
pub async fn fetch_with_size<'a, B: Blobstore + Clone + 'a>(
    blobstore: B,
    ctx: impl Borrow<CoreContext> + Clone + Send + Sync + 'a,
    key: &FetchKey,
) -> Result<Option<(impl Stream<Item = Result<Bytes, Error>> + 'a, u64)>, Error> {
    let content_id =
        key.load(ctx.borrow(), &blobstore)
            .await
            .map(Some)
            .or_else(|err| match err {
                LoadableError::Error(err) => Err(err),
                LoadableError::Missing(_) => Ok(None),
            })?;

    match content_id {
        Some(content_id) => {
            fetch::fetch_with_size(blobstore, ctx, content_id, fetch::Range::all()).await
        }
        None => Ok(None),
    }
}

/// This function has the same functionality as fetch_with_size, but only
/// returns data for a range within the file.
///
/// Requests for data beyond the end of the file will return only the part of
/// the file that overlaps with the requested range, if any.
pub async fn fetch_range_with_size<'a, B: Blobstore + Clone + 'a>(
    blobstore: B,
    ctx: impl Borrow<CoreContext> + Clone + Send + Sync + 'a,
    key: &FetchKey,
    range: Range,
) -> Result<Option<(impl Stream<Item = Result<Bytes, Error>> + 'a, u64)>, Error> {
    let content_id =
        key.load(ctx.borrow(), &blobstore)
            .await
            .map(Some)
            .or_else(|err| match err {
                LoadableError::Error(err) => Err(err),
                LoadableError::Missing(_) => Ok(None),
            })?;

    match content_id {
        Some(content_id) => fetch::fetch_with_size(blobstore, ctx, content_id, range).await,
        None => Ok(None),
    }
}

/// This function has the same functionality as fetch_with_size, but doesn't return the file size.
pub async fn fetch<'a, B: Blobstore + Clone + 'a>(
    blobstore: B,
    ctx: impl Borrow<CoreContext> + Clone + Send + Sync + 'a,
    key: &FetchKey,
) -> Result<Option<impl Stream<Item = Result<Bytes, Error>> + 'a>, Error> {
    let res = fetch_with_size(blobstore, ctx, key).await?;
    Ok(res.map(|(stream, _len)| stream))
}

/// Fetch the contents of a blob concatenated together. This bad for buffering, and you shouldn't
/// add new callsites. This is only for compatibility with existin callsites.
pub async fn fetch_concat_opt<B: Blobstore>(
    blobstore: &B,
    ctx: &CoreContext,
    key: &FetchKey,
) -> Result<Option<Bytes>, Error> {
    let res = fetch_with_size(blobstore, ctx, key).await?;

    match res {
        Some((stream, len)) => {
            let len = len
                .try_into()
                .map_err(|_| anyhow::format_err!("Cannot fetch file with length {}", len))?;

            let buf = BytesMut::with_capacity(len);

            let bytes = stream
                .try_fold(buf, |mut buffer, chunk| async move {
                    buffer.extend_from_slice(&chunk);
                    Result::<_, Error>::Ok(buffer)
                })
                .await?;

            Ok(Some(bytes.freeze()))
        }
        None => Ok(None),
    }
}

/// Similar to `fetch_concat_opt`, but requires the blob to be present, or errors out.
pub async fn fetch_concat<B: Blobstore>(
    blobstore: &B,
    ctx: &CoreContext,
    key: impl Into<FetchKey>,
) -> Result<Bytes, Error> {
    let key: FetchKey = key.into();
    let bytes = fetch_concat_opt(blobstore, ctx, &key).await?;
    bytes.ok_or_else(|| errors::ErrorKind::MissingContent(key).into())
}

/// Fetch content associated with the key as a stream
///
/// Moslty behaves as the `fetch`, except it is pushing missing content error into stream if
/// data associated with the key was not found.
pub fn fetch_stream<'a, B: Blobstore + Clone + 'a>(
    blobstore: B,
    ctx: impl Borrow<CoreContext> + Clone + Send + Sync + 'a,
    key: impl Into<FetchKey>,
) -> impl Stream<Item = Result<Bytes, Error>> + 'a {
    let key: FetchKey = key.into();

    async move {
        let stream = fetch(blobstore, ctx, &key)
            .await?
            .ok_or(errors::ErrorKind::MissingContent(key))?;
        Result::<_, Error>::Ok(stream)
    }
    .try_flatten_stream()
}

/// This function has the same functionality as fetch_range_with_size, but doesn't return the file size.
pub async fn fetch_range<'a, B: Blobstore>(
    blobstore: &'a B,
    ctx: &'a CoreContext,
    key: &FetchKey,
    range: Range,
) -> Result<Option<impl Stream<Item = Result<Bytes, Error>> + 'a>, Error> {
    let res = fetch_range_with_size(blobstore, ctx, key, range).await?;
    Ok(res.map(|(stream, _len)| stream))
}

/// Fetch the start of a file. Returns a Future that resolves with Some(Bytes) if the file was
/// found, and None otherwise. If Bytes are found, this function is guaranteed to return as many
/// Bytes as requested unless the file is shorter than that.
pub async fn peek<B: Blobstore>(
    blobstore: &B,
    ctx: &CoreContext,
    key: &FetchKey,
    size: usize,
) -> Result<Option<Bytes>, Error> {
    let maybe_stream = fetch(blobstore, ctx, key).await?;

    match maybe_stream {
        None => Ok(None),
        Some(stream) => {
            let mut stream = chunk::ChunkStream::new(stream, size);
            stream.try_next().await
        }
    }
}

/// Store a file from a stream. This is guaranteed atomic - either the store will succeed
/// for the entire file, or it will fail and the file will logically not exist (however
/// there's no guarantee that any partially written parts will be cleaned up).
pub async fn store<B: Blobstore + Clone + 'static>(
    blobstore: &B,
    config: FilestoreConfig,
    ctx: &CoreContext,
    req: &StoreRequest,
    data: impl Stream<Item = Result<Bytes, Error>> + Send,
) -> Result<ContentMetadataV2, Error> {
    use chunk::Chunks;

    let prepared = match chunk::make_chunks(data, req.expected_size, config.chunk_size) {
        Chunks::Inline(fut) => prepare::prepare_bytes(fut.await?).await,
        Chunks::Chunked(expected_size, chunks) => {
            prepare::prepare_chunked(
                ctx.clone(),
                blobstore.clone(),
                expected_size,
                chunks,
                config.concurrency,
            )
            .await?
        }
    };

    finalize::finalize(blobstore, ctx, Some(req), prepared).await
}

/// Store a set of bytes, and immediately return their Contentid and size. This function is
/// inefficient for large files, since it will hash the file twice if it's larger than the chunk
/// size. This function is intended as a transition function while we convert writers to streams
/// and refactor them to not expect to be able to obtain the ContentId for the content they
/// uploaded immediately. Do NOT add new callsites.
pub fn store_bytes<B: Blobstore + Clone + 'static>(
    blobstore: &B,
    config: FilestoreConfig,
    ctx: &CoreContext,
    bytes: Bytes,
) -> (
    (ContentId, u64),
    impl Future<Output = Result<(), Error>> + Send,
) {
    // NOTE: Like in other places in the Filestore, we assume that the size of buffers being passed
    // in can be represented in 64 bits (which is OK for the world we live in).

    let content_id = FileContents::content_id_for_bytes(&bytes);
    let size: u64 = bytes.len().try_into().unwrap();

    cloned!(ctx, blobstore);
    let upload = async move {
        store(
            &blobstore,
            config,
            &ctx,
            &StoreRequest::with_canonical(size, content_id),
            stream::once(async move { Ok(bytes) }),
        )
        .await?;

        Ok(())
    };

    ((content_id, size), upload)
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This software may be used and distributed according to the terms of the
 * GNU General Public License version 2.
 */

use anyhow::Error;
use blobstore::Blobstore;
use blobstore::Loadable;
use blobstore::LoadableError;
use blobstore::Storable;
use bytes::Bytes;
use context::CoreContext;
use futures::future;
use futures::task::Poll;
use futures::TryFutureExt;
use mononoke_types::private::ErrorKind;
use mononoke_types::BlobstoreKey;
use mononoke_types::BlobstoreValue;
use mononoke_types::ContentAlias;
use mononoke_types::ContentId;
use mononoke_types::ContentMetadataV2;
use mononoke_types::ContentMetadataV2Id;
use slog::warn;
use strum::IntoEnumIterator;
use thiserror::Error;

use crate::alias::add_aliases_to_multiplexer;
use crate::expected_size::ExpectedSize;
use crate::fetch;
use crate::multiplexer::Multiplexer;
use crate::multiplexer::MultiplexerError;
use crate::prepare::add_partial_metadata_to_multiplexer;
use crate::Alias;
use crate::AliasBlob;

#[derive(Debug, Error)]
pub enum RebuildBackmappingError {
    #[error("Not found: {0:?}")]
    NotFound(ContentId),

    #[error("Error computing metadata for {0:?}: {1:?}")]
    InternalError(ContentId, #[source] Error),
}

/// Finds the metadata for a ContentId. Returns None if the content does not exist, and returns
/// the metadata otherwise. This might recompute the metadata on the fly if it is found to
/// be missing but the content exists.
pub async fn get_metadata<B: Blobstore>(
    blobstore: &B,
    ctx: &CoreContext,
    content_id: ContentId,
) -> Result<Option<ContentMetadataV2>, Error> {
    let metadata = get_metadata_readonly(blobstore, ctx, content_id).await;
    // We found the metadata, return it.
    if let Ok(Some(_)) = metadata {
        return metadata;
    } else if let Err(e) = metadata {
        match e.downcast_ref::<ErrorKind>() {
            // The backfilling for ContentMetadataV2 has happened in different stages.
            // If any of the later fields are missing, we get invalid thrift error. In
            // that case we need to rebuild the metadata, so do not return.
            Some(ErrorKind::InvalidThrift(..)) => {
                let key = ContentMetadataV2Id::from(content_id).blobstore_key();
                let msg = format!(
                    "Invalid ContentMetadataV2 format exists in blobstore for key {}. Error: {}",
                    key, e
                );
                warn!(ctx.logger(), "{}", &msg);
                let mut scuba = ctx.scuba().clone();
                scuba.add("blobstore_key", key);
                scuba.log_with_msg("ContentMetadataV2 backfill repair", msg);
            }
            _ => return Err(e),
        }
    }

    // We didn't find the metadata. Try to recompute it. This might fail if the
    // content doesn't exist, or due to an internal error.
    rebuild_metadata(blobstore, ctx, content_id)
        .await
        .map(Some)
        .or_else({
            use RebuildBackmappingError::*;
            |e| match e {
                // If we didn't find the ContentId we're rebuilding the metadata for,
                // then there is nothing else to do but indicate this metadata does not
                // exist.
                NotFound(_) => Ok(None),
                // If we ran into some error rebuilding the metadata that isn't not
                // having found the content, then we pass it up.
                e @ InternalError(..) => Err(e.into()),
            }
        })
}

/// Finds the metadata for a ContentId. Returns None if the content metadata does not exist
/// and returns Some(metadata) if it already exists. Does not recompute it on the fly.
pub async fn get_metadata_readonly<B: Blobstore>(
    blobstore: &B,
    ctx: &CoreContext,
    content_id: ContentId,
) -> Result<Option<ContentMetadataV2>, Error> {
    ContentMetadataV2Id::from(content_id)
        .load(ctx, blobstore)
        .await
        .map(Some)
        .or_else(|err| match err {
            LoadableError::Error(err) => Err(err),
            LoadableError::Missing(_) => Ok(None),
        })
}

/// If the metadata is missing, we can rebuild it on the fly, since all that's needed to do so
/// is the file contents. This can happen if we successfully stored a file, but failed to store
/// its metadata. To rebuild the metadata, we peek at the content in the blobstore to get
/// its size, then produce a stream of its contents and compute aliases over it. Finally, store
/// the metadata, and return it.
async fn rebuild_metadata<B: Blobstore>(
    blobstore: &B,
    ctx: &CoreContext,
    content_id: ContentId,
) -> Result<ContentMetadataV2, RebuildBackmappingError> {
    use RebuildBackmappingError::*;

    let file_contents = content_id
        .load(ctx, blobstore)
        .await
        .map_err(|err| match err {
            LoadableError::Error(err) => InternalError(content_id, err),
            LoadableError::Missing(_) => NotFound(content_id),
        })?;

    // NOTE: We implicitly trust data from the Filestore here. We do not validate
    // the size, nor the ContentId.
    let total_size = file_contents.size();
    let content_stream =
        fetch::stream_file_bytes(blobstore, ctx, file_contents, fetch::Range::all())
            .map_err(|e| InternalError(content_id, e))?;

    let mut multiplexer = Multiplexer::<Bytes>::new();
    let aliases = add_aliases_to_multiplexer(&mut multiplexer, ExpectedSize::new(total_size))
        .map_err(|e| InternalError(content_id, e));
    let metadata = add_partial_metadata_to_multiplexer(&mut multiplexer)
        .map_err(|e| InternalError(content_id, e));

    let futs = future::try_join(aliases, metadata);
    let res = multiplexer.drain(content_stream).await;
    match res {
        // All is well - get the results when our futures complete.
        Ok(_) => {
            let (redeemable, metadata) = futs.await?;
            let (sha1, sha256, git_sha1, seeded_blake3) = redeemable
                .redeem(total_size)
                .map_err(|e| InternalError(content_id, e))?;
            // To maintain consistency, rebuild the aliases to the content along with rebuilding
            // content metadata.
            let alias = ContentAlias::from_content_id(content_id);
            // Ensure that all aliases are covered, and missing out an alias gives a compile time error.
            future::try_join_all(
                Alias::iter()
                    .map(|alias_type| match alias_type {
                        Alias::Sha1(_) => AliasBlob(Alias::Sha1(sha1), alias.clone()),
                        Alias::GitSha1(_) => {
                            AliasBlob(Alias::GitSha1(git_sha1.sha1()), alias.clone())
                        }

                        Alias::Sha256(_) => AliasBlob(Alias::Sha256(sha256), alias.clone()),
                        Alias::SeededBlake3(_) => {
                            AliasBlob(Alias::SeededBlake3(seeded_blake3), alias.clone())
                        }
                    })
                    .map(|alias_blob| alias_blob.store(ctx, blobstore)),
            )
            .await
            .map_err(|e| InternalError(content_id, e))?;

            let metadata = ContentMetadataV2 {
                total_size,
                content_id,
                sha1,
                sha256,
                git_sha1,
                seeded_blake3,
                is_binary: metadata.is_binary,
                is_ascii: metadata.is_ascii,
                is_utf8: metadata.is_utf8,
                ends_in_newline: metadata.ends_in_newline,
                newline_count: metadata.newline_count,
                first_line: metadata.first_line,
                is_generated: metadata.is_generated,
                is_partially_generated: metadata.is_partially_generated,
            };

            let blob = metadata.clone().into_blob();

            blob.store(ctx, blobstore)
                .await
                .map_err(|e| InternalError(content_id, e))?;

            Ok(metadata)
        }
        // If the Multiplexer hit an error, then it's worth handling the Cancelled case
        // separately: Cancelled means our Multiplexer noted that one of its readers
        // stopped reading. Usually, this will be because one of the readers failed. So,
        // let's just poll the readers once to see if they have an error value ready, and
        // if so, let's return that Error (because it'll be a more usable one). If not,
        // we'll passthrough the cancellation (but, we do have a unit test to make sure we
        // hit the happy path that prettifies the error).
        Err(m @ MultiplexerError::Cancelled) => match futures::poll!(futs) {
            Poll::Ready(Err(e)) => Err(e),
            _ => Err(InternalError(content_id, m.into())),
        },

        Err(m @ MultiplexerError::InputError(..)) => Err(InternalError(content_id, m.into())),
    }
}