/*
 * 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::BTreeSet;
use std::collections::HashSet;
use std::sync::Arc;

use anyhow::anyhow;
use anyhow::Result;
use bookmarks::BookmarkKey;
use futures::channel::mpsc;
use futures::channel::oneshot;
use futures::future::BoxFuture;
use futures::future::Shared;
use futures::stream::StreamExt;
use futures::FutureExt;
use futures_ext::future::FbFutureExt;
use futures_stats::FutureStats;
use futures_stats::TimedFutureExt;
use identity::Identity;
use land_service_if::types::*;
use mononoke_api::CoreContext;
use mononoke_api::Mononoke;
use mononoke_api::RepoContext;
use mononoke_types::BonsaiChangeset;
use permission_checker::MononokeIdentity;
use pushrebase_client::LocalPushrebaseClient;
use pushrebase_client::PushrebaseClient;
use repo_authorization::AuthorizationContext;
use stats::prelude::*;
use time_ext::DurationExt;

use crate::conversion_helpers;
use crate::errors;
use crate::errors::LandChangesetsError;
use crate::land_changeset_object::LandChangesetObject;
use crate::scuba_response::AddScubaResponse;

const LAND_CHANSET_BUFFER_SIZE: usize = 64;

define_stats! {
    prefix = "mononoke.land_service";
    total_request_start: timeseries(Rate, Sum),
    total_request_success: timeseries(Rate, Sum),
    total_request_internal_failure: timeseries(Rate, Sum),
    total_request_canceled: timeseries(Rate, Sum),

    total_batch_success: timeseries(Rate, Sum),
    total_batch_failures: timeseries(Rate, Sum),

    // Duration per changesets landed with and without batches
    method_completion_time_ms: dynamic_histogram("method.{}.completion_time_ms", (method: String); 10, 0, 1_000, Average, Sum, Count; P 5; P 50 ; P 90),
}

pub type EnqueueSender = mpsc::UnboundedSender<(
    oneshot::Sender<Result<LandChangesetsResponse, LandChangesetsError>>,
    LandChangesetObject,
)>;
pub fn setup_worker() -> (EnqueueSender, Shared<BoxFuture<'static, ()>>) {
    // The mpsc channel needed as a way to enqueue new requests while there is an
    // in-flight request.
    // - queue_sender will be used to add new requests to the queue (channel),
    // - receiver - to read a new batch of requests and land them.
    //
    // To notify the clients back that the request was successfully landed,
    // a oneshot channel is used. When the enqueued requests are processed, the clients
    // receive result of the operation:
    // error if something went wrong and nothing if it's ok.
    let (queue_sender, receiver) = mpsc::unbounded::<(
        oneshot::Sender<Result<LandChangesetsResponse, LandChangesetsError>>,
        LandChangesetObject,
    )>();
    let worker = async move {
        let enqueued_landed_changesets = receiver.ready_chunks(LAND_CHANSET_BUFFER_SIZE).for_each(
            move |batch /* (Sender, LandChangesetObject) */| async move {
                process_requests(batch).await;
            },
        );
        tokio::spawn(enqueued_landed_changesets);
    }
    .boxed()
    .shared();
    (queue_sender, worker)
}

async fn process_requests(
    requests: Vec<(
        oneshot::Sender<Result<LandChangesetsResponse, LandChangesetsError>>,
        LandChangesetObject,
    )>,
) {
    let mut changesets_batch: BTreeSet<Vec<u8>> = BTreeSet::new();
    let mut first_land_changeset_object_batched: Option<LandChangesetObject> = None;
    let mut backup_batch = Vec::new();

    for (sender, land_changeset_object) in requests.into_iter() {
        // If there is NO pushvars for a request, we batch it
        if land_changeset_object.request.pushvars.is_none() {
            if changesets_batch.is_empty() {
                first_land_changeset_object_batched = Some(land_changeset_object.clone());
            }
            changesets_batch.extend(land_changeset_object.request.changesets.clone());
            backup_batch.push((sender, land_changeset_object));
        //Otherwise, we just process it individually
        } else if let Err(err) =
            sender.send(impl_land_changesets(land_changeset_object.clone()).await)
        {
            let mut scuba = land_changeset_object.ctx.scuba().clone();
            scuba.log_with_msg(
                        "Failed to send individual response back without batching (i.e., request with pushvars)",
                        Some(format!("{:?}", err)),
                    );
        };
    }

    if let Some(mut land_changeset_object) = first_land_changeset_object_batched {
        land_changeset_object.request.changesets = changesets_batch;

        let (stats, result) = impl_land_changesets(land_changeset_object.clone())
            .timed()
            .await;
        log_batch_result(land_changeset_object.ctx, &stats, result, backup_batch).await;
        STATS::method_completion_time_ms.add_value(
            stats.completion_time.as_millis_unchecked() as i64,
            ("impl_land_changesets_with_batch".to_string(),),
        );
    }
}

//TODO: Once the batching is done, this method does not need to be public
pub async fn impl_land_changesets(
    land_changeset_object: LandChangesetObject,
) -> Result<LandChangesetsResponse, LandChangesetsError> {
    land_changeset_object
        .ctx
        .scuba()
        .clone()
        .log_with_msg("Request start", None);
    STATS::total_request_start.add_value(1);

    let (stats, res) = process_land_changesets_request(land_changeset_object.clone())
        .timed()
        .on_cancel_with_data(|stats| log_canceled(&land_changeset_object.ctx, &stats))
        .await;
    log_result(land_changeset_object.ctx, &stats, &res);
    STATS::method_completion_time_ms.add_value(
        stats.completion_time.as_millis_unchecked() as i64,
        ("impl_land_changesets".to_string(),),
    );
    res
}

/// Create a RepoContext
async fn get_repo_context(
    mononoke: Arc<Mononoke>,
    repo_name: String,
    ctx: CoreContext,
    authz: AuthorizationContext,
) -> Result<RepoContext, LandChangesetsError> {
    Ok(mononoke
        .repo(ctx, &repo_name)
        .await?
        .ok_or_else(|| {
            errors::internal_error(anyhow!("Not finding the repo: {}", repo_name).as_ref())
        })?
        .with_authorization_context(authz)
        .build()
        .await?)
}

// Check for the scm_service_identity
fn assert_internal_identity(
    identity: Identity,
    repo: &RepoContext,
) -> Result<(), LandChangesetsError> {
    let original_identities = repo.ctx().metadata().original_identities();
    if !original_identities.map_or(false, |ids| {
        ids.contains(&MononokeIdentity::from_identity(&identity))
    }) {
        return Err(errors::internal_error(
            anyhow!(
                "Insufficient permissions, internal options only. Identities: {}",
                original_identities
                    .map_or_else(|| "<none>".to_string(), permission_checker::pretty_print)
            )
            .as_ref(),
        )
        .into());
    }
    Ok(())
}

async fn process_land_changesets_request(
    land_changeset_object: LandChangesetObject,
) -> Result<LandChangesetsResponse, LandChangesetsError> {
    let LandChangesetObject {
        mononoke,
        identity,
        ctx,
        request,
    } = land_changeset_object;

    let authz = AuthorizationContext::new(&ctx);
    //TODO: Avoid using RepoContext, build a leaner Repo type if possible (T132600441)
    let repo: RepoContext =
        get_repo_context(mononoke, request.repo_name, ctx.clone(), authz.clone()).await?;

    assert_internal_identity(identity, &repo)?;

    let bookmark = BookmarkKey::new(request.bookmark)?;
    let changesets: HashSet<BonsaiChangeset> =
        conversion_helpers::convert_bonsai_changesets(request.changesets, &ctx, &repo).await?;

    let pushvars = conversion_helpers::convert_pushvars(request.pushvars.unwrap_or_default());

    let cross_repo_push_source =
        conversion_helpers::convert_cross_repo_push_source(request.cross_repo_push_source)?;

    let bookmark_restrictions =
        conversion_helpers::convert_bookmark_restrictions(request.bookmark_restrictions)?;

    let outcome = LocalPushrebaseClient {
        ctx: &ctx,
        authz: &authz,
        repo: &repo.inner_repo().clone(),
        hook_manager: repo.hook_manager().as_ref(),
    }
    .pushrebase(
        &bookmark,
        changesets,
        Some(&pushvars),
        cross_repo_push_source,
        bookmark_restrictions,
        false, // Currently mononoke server logs the new commits
    )
    .await?;

    Ok(LandChangesetsResponse {
        pushrebase_outcome: PushrebaseOutcome {
            head: outcome.head.as_ref().to_vec(),
            rebased_changesets: outcome
                .rebased_changesets
                .into_iter()
                .map(|rebased_changeset| {
                    conversion_helpers::convert_rebased_changesets_into_pairs(rebased_changeset)
                })
                .collect(),
            pushrebase_distance: conversion_helpers::convert_to_i64(outcome.pushrebase_distance.0)?,
            retry_num: conversion_helpers::convert_to_i64(outcome.retry_num.0)?,
            old_bookmark_value: outcome
                .old_bookmark_value
                .map(conversion_helpers::convert_changeset_id_to_vec_binary),
        },
    })
}

async fn log_batch_result(
    ctx: CoreContext,
    stats: &FutureStats,
    result: Result<LandChangesetsResponse, LandChangesetsError>,
    backup_batch: Vec<(
        oneshot::Sender<Result<LandChangesetsResponse, LandChangesetsError>>,
        LandChangesetObject,
    )>,
) {
    let mut scuba = ctx.scuba().clone();
    let batch_size = backup_batch.len();

    match result {
        Ok(ref response) => {
            // if batched request worked, send response back for each request
            for (sender, _) in backup_batch.into_iter() {
                if let Err(err) = sender.send(result.clone()) {
                    scuba.log_with_msg(
                        "Failed sending individual response back after batching completed",
                        Some(format!("{:?}", err)),
                    );
                };
            }
            response.add_scuba_response(&mut scuba);
            STATS::total_batch_success.add_value(1);
            scuba.log_with_msg(
                format!("Batching {} requests completed successfully", batch_size).as_str(),
                None,
            );
        }
        Err(err) => {
            STATS::total_batch_failures.add_value(1);
            scuba.log_with_msg(
                format!("Batching {} requests failed", batch_size).as_str(),
                None,
            );
            scuba.add("error batching", err.to_string());
            // if found error, process requests individually
            for (sender, land_changeset_object) in backup_batch.into_iter() {
                if let Err(err) = sender.send(impl_land_changesets(land_changeset_object).await) {
                    scuba.log_with_msg(
                        "Failed sending individual response back after batching failed",
                        Some(format!("{:?}", err)),
                    );
                };
            }
        }
    };

    ctx.perf_counters().insert_perf_counters(&mut scuba);
    scuba.add_future_stats(stats);
}

fn log_result<T: AddScubaResponse>(
    ctx: CoreContext,
    stats: &FutureStats,
    result: &Result<T, LandChangesetsError>,
) {
    let mut scuba = ctx.scuba().clone();

    match result {
        Ok(response) => {
            response.add_scuba_response(&mut scuba);
            STATS::total_request_success.add_value(1);
            scuba.add("status", "SUCCESS");
        }
        Err(err) => {
            STATS::total_request_internal_failure.add_value(1);
            scuba.add("status", "INTERNAL_ERROR");
            scuba.add("error", err.to_string());
        }
    };

    ctx.perf_counters().insert_perf_counters(&mut scuba);
    scuba.add_future_stats(stats);
    scuba.log_with_msg("Request complete", None);
}

fn log_canceled(ctx: &CoreContext, stats: &FutureStats) {
    STATS::total_request_canceled.add_value(1);
    let mut scuba = ctx.scuba().clone();
    ctx.perf_counters().insert_perf_counters(&mut scuba);
    scuba.add_future_stats(stats);
    scuba.add("status", "CANCELED");
    scuba.log_with_msg("Request canceled", None);
}


/*
 * 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::process::Stdio;

use anyhow::Context as _;
use anyhow::Error;
use anyhow::Result;
use bytes::Bytes;
use bytes::BytesMut;
use context::CoreContext;
use filestore::Alias;
use filestore::FetchKey;
use filestore::FilestoreConfigRef;
use filestore::StoreRequest;
use futures::future;
use futures::future::TryFutureExt;
use futures::stream::Stream;
use futures::stream::TryStreamExt;
use mercurial_types::blobs::LFSContent;
use mononoke_types::ContentMetadataV2;
use repo_blobstore::RepoBlobstoreRef;
use slog::info;
use tokio::io::BufReader;
use tokio::process::Child;
use tokio::process::Command;
use tokio_util::codec::BytesCodec;
use tokio_util::codec::FramedRead;

fn lfs_stream(
    lfs_helper: &str,
    lfs: &LFSContent,
) -> Result<(Child, impl Stream<Item = Result<Bytes, Error>>)> {
    let mut cmd = Command::new(lfs_helper)
        .arg(format!("{}", lfs.oid().to_hex()))
        .arg(format!("{}", lfs.size()))
        .stdin(Stdio::null())
        .stdout(Stdio::piped())
        .spawn()
        .with_context(|| format!("Error starting lfs_helper: {:?}", lfs_helper))?;

    let stdout = cmd
        .stdout
        .take()
        .expect("stdout was piped earlier and is missing here");
    let stdout = BufReader::new(stdout);
    let stream = FramedRead::new(stdout, BytesCodec::new())
        .map_ok(BytesMut::freeze)
        .map_err(Error::from);

    Ok((cmd, stream))
}

async fn do_lfs_upload(
    ctx: &CoreContext,
    repo: &(impl RepoBlobstoreRef + FilestoreConfigRef),
    lfs_helper: &str,
    lfs: &LFSContent,
) -> Result<ContentMetadataV2, Error> {
    let metadata = filestore::get_metadata(
        repo.repo_blobstore(),
        ctx,
        &FetchKey::Aliased(Alias::Sha256(lfs.oid())),
    )
    .await?;

    if let Some(metadata) = metadata {
        info!(
            ctx.logger(),
            "lfs_upload: reusing blob {:?}", metadata.sha256
        );
        return Ok(metadata);
    }

    info!(ctx.logger(), "lfs_upload: importing blob {:?}", lfs.oid());
    let req = StoreRequest::with_sha256(lfs.size(), lfs.oid());

    let (mut child, stream) = lfs_stream(lfs_helper, lfs)?;

    let upload = filestore::store(
        repo.repo_blobstore(),
        *repo.filestore_config(),
        ctx,
        &req,
        stream,
    );

    // NOTE: We ignore the child exit code here. Since the Filestore validates the object
    // we're uploading by SHA256, that's indeed fine (it doesn't matter if the Child failed
    // if it gave us exactly the content we wanted).
    let (_, meta) = future::try_join(child.wait().map_err(Error::from), upload).await?;

    info!(ctx.logger(), "lfs_upload: imported blob {:?}", meta.sha256);

    Ok(meta)
}

pub async fn lfs_upload(
    ctx: &CoreContext,
    repo: &(impl RepoBlobstoreRef + FilestoreConfigRef),
    lfs_helper: &str,
    lfs: &LFSContent,
) -> Result<ContentMetadataV2, Error> {
    let max_attempts = 5;
    let mut attempt = 0;

    loop {
        let res = do_lfs_upload(ctx, repo, lfs_helper, lfs).await;

        if res.is_ok() || attempt > max_attempts {
            break res;
        }

        attempt += 1;
    }
}


/*
 * 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.
 */

mod protocol;
mod str_serialized;

pub use protocol::git_lfs_mime;
pub use protocol::ObjectAction;
pub use protocol::ObjectError;
pub use protocol::ObjectStatus;
pub use protocol::Operation;
pub use protocol::RequestBatch;
pub use protocol::RequestObject;
pub use protocol::ResponseBatch;
pub use protocol::ResponseError;
pub use protocol::ResponseObject;
pub use protocol::Sha256;
pub use protocol::Transfer;


/*
 * 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::HashMap;
use std::fmt;
use std::fmt::Debug;
use std::fmt::Display;
use std::mem;
use std::str::FromStr;

use anyhow::bail;
use anyhow::Error;
use anyhow::Result;
use faster_hex::hex_decode;
use faster_hex::hex_string;
use http::Uri;
use mime::Mime;
use once_cell::sync::Lazy;
use quickcheck::Arbitrary;
use quickcheck::Gen;
use serde::Deserialize;
use serde::Serialize;

use crate::str_serialized;

// This module provides types conforming to the Git-LFS protocol specification:
// https://github.com/git-lfs/git-lfs/blob/master/docs/api/batch.md

static GIT_LFS_MIME: Lazy<Mime> = Lazy::new(|| "application/vnd.git-lfs+json".parse().unwrap());

pub fn git_lfs_mime() -> Mime {
    GIT_LFS_MIME.clone()
}

#[derive(Copy, Clone, Serialize, Debug, Deserialize, Eq, PartialEq, Hash)]
pub enum Operation {
    #[serde(rename = "download")]
    Download,
    #[serde(rename = "upload")]
    Upload,
}

impl Display for Operation {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Self::Download => write!(f, "download"),
            Self::Upload => write!(f, "upload"),
        }
    }
}

impl Arbitrary for Operation {
    fn arbitrary(g: &mut Gen) -> Self {
        if bool::arbitrary(g) {
            Operation::Download
        } else {
            Operation::Upload
        }
    }
}

#[derive(Clone, Serialize, Debug, Deserialize, Eq, PartialEq, Hash)]
pub enum Transfer {
    #[serde(rename = "basic")]
    Basic,
    #[serde(other)]
    Unknown,
}

impl Arbitrary for Transfer {
    fn arbitrary(_g: &mut Gen) -> Self {
        // We don't generate invalid Transfer instances for testing.
        Transfer::Basic
    }
}

impl Default for Transfer {
    fn default() -> Self {
        Self::Basic
    }
}

#[derive(Clone, Serialize, Debug, Deserialize, Hash, Eq, PartialEq)]
pub struct Ref {
    pub name: String,
}

impl Arbitrary for Ref {
    fn arbitrary(g: &mut Gen) -> Self {
        Self {
            name: String::arbitrary(g),
        }
    }
}

#[derive(Clone, Serialize, Deserialize, Hash, Eq, PartialEq, Copy)]
pub struct Sha256(pub [u8; 32]);

impl Sha256 {
    fn to_hex(&self) -> String {
        hex_string(&self.0)
    }
}

impl FromStr for Sha256 {
    type Err = Error;

    fn from_str(s: &str) -> Result<Self> {
        if s.len() != mem::size_of::<Sha256>() * 2 {
            bail!("invalid sha256 length: {}", s);
        }

        let mut ret = [0; mem::size_of::<Sha256>()];
        hex_decode(s.as_bytes(), &mut ret)?;
        Ok(Sha256(ret))
    }
}

impl Debug for Sha256 {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "Sha256({})", self.to_hex())
    }
}

impl Display for Sha256 {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        Display::fmt(&self.to_hex(), fmt)
    }
}

#[derive(Clone, Serialize, Debug, Deserialize, Hash, Eq, PartialEq, Copy)]
pub struct RequestObject {
    #[serde(with = "str_serialized")]
    pub oid: Sha256,
    pub size: u64,
}

impl Arbitrary for RequestObject {
    fn arbitrary(g: &mut Gen) -> Self {
        let mut oid = [0u8; mem::size_of::<Sha256>()];
        for b in oid.iter_mut() {
            *b = u8::arbitrary(g);
        }
        Self {
            oid: Sha256(oid),
            size: u64::arbitrary(g),
        }
    }
}

fn default_client_transfers() -> Vec<Transfer> {
    vec![Transfer::default()]
}

#[derive(Clone, Serialize, Debug, Deserialize, PartialEq)]
pub struct RequestBatch {
    pub operation: Operation,
    #[serde(default = "default_client_transfers")]
    pub transfers: Vec<Transfer>,
    pub r#ref: Option<Ref>,
    pub objects: Vec<RequestObject>,
}

impl Arbitrary for RequestBatch {
    fn arbitrary(g: &mut Gen) -> Self {
        Self {
            operation: Operation::arbitrary(g),
            transfers: Vec::arbitrary(g),
            r#ref: Option::arbitrary(g),
            objects: Vec::arbitrary(g),
        }
    }
}

#[derive(Clone, Serialize, Debug, Deserialize, PartialEq)]
pub struct ObjectAction {
    #[serde(with = "str_serialized")]
    pub href: Uri,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub header: Option<HashMap<String, String>>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub expires_in: Option<i32>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub expires_at: Option<String>,
}

impl Arbitrary for ObjectAction {
    fn arbitrary(g: &mut Gen) -> Self {
        // We generate a basic URL here. Nothing very fancy.
        let proto = if bool::arbitrary(g) { "http" } else { "https" };

        let domain = if bool::arbitrary(g) {
            "foo.com"
        } else {
            "bar.com"
        };

        let path = if bool::arbitrary(g) { "" } else { "/123" };

        let uri: Uri = format!("{}://{}{}", proto, domain, path).parse().unwrap();

        Self {
            href: uri,
            header: Option::arbitrary(g),
            expires_in: Option::arbitrary(g),
            expires_at: Option::arbitrary(g),
        }
    }
}

impl ObjectAction {
    pub fn new(href: Uri) -> Self {
        Self {
            href,
            header: None,
            expires_in: None,
            expires_at: None,
        }
    }
}

#[derive(Clone, Serialize, Debug, Deserialize, Hash, PartialEq, Eq)]
pub struct ObjectError {
    pub code: u16,
    pub message: String,
}

impl Arbitrary for ObjectError {
    fn arbitrary(g: &mut Gen) -> Self {
        Self {
            code: u16::arbitrary(g),
            message: String::arbitrary(g),
        }
    }
}

#[derive(Clone, Serialize, Debug, Deserialize, PartialEq)]
#[serde(untagged)]
pub enum ObjectStatus {
    Ok {
        #[serde(default)]
        authenticated: bool,
        actions: HashMap<Operation, ObjectAction>,
    },
    Err {
        error: ObjectError,
    },
}

impl Arbitrary for ObjectStatus {
    fn arbitrary(g: &mut Gen) -> Self {
        if bool::arbitrary(g) {
            let mut actions = HashMap::new();

            if bool::arbitrary(g) {
                actions.insert(Operation::Download, ObjectAction::arbitrary(g));
            }

            if bool::arbitrary(g) {
                actions.insert(Operation::Upload, ObjectAction::arbitrary(g));
            }

            Self::Ok {
                authenticated: bool::arbitrary(g),
                actions,
            }
        } else {
            Self::Err {
                error: ObjectError::arbitrary(g),
            }
        }
    }
}

#[derive(Clone, Serialize, Debug, Deserialize, PartialEq)]
pub struct ResponseObject {
    #[serde(flatten)]
    pub object: RequestObject,
    #[serde(flatten)]
    pub status: ObjectStatus,
}

impl Arbitrary for ResponseObject {
    fn arbitrary(g: &mut Gen) -> Self {
        Self {
            object: RequestObject::arbitrary(g),
            status: ObjectStatus::arbitrary(g),
        }
    }
}

#[derive(Clone, Serialize, Debug, Deserialize, PartialEq)]
pub struct ResponseBatch {
    #[serde(default)]
    pub transfer: Transfer,
    pub objects: Vec<ResponseObject>,
}

impl Arbitrary for ResponseBatch {
    fn arbitrary(g: &mut Gen) -> Self {
        Self {
            transfer: Transfer::arbitrary(g),
            objects: Vec::arbitrary(g),
        }
    }
}

#[derive(Clone, Serialize, Debug, Deserialize)]
pub struct ResponseError {
    pub message: String,
    #[serde(skip_serializing, skip_deserializing)]
    pub documentation_url: Option<Uri>,
    pub request_id: Option<String>,
}

impl Arbitrary for ResponseError {
    fn arbitrary(g: &mut Gen) -> Self {
        Self {
            message: String::arbitrary(g),
            // TODO: It'd be nice to generate those too.
            documentation_url: None,
            request_id: Option::arbitrary(g),
        }
    }
}

#[cfg(test)]
mod test {
    use assert_matches::assert_matches;
    use maplit::hashmap;
    use quickcheck::quickcheck;
    use serde_json::json;

    use super::*;

    const ONES_SHA256: &str = "1111111111111111111111111111111111111111111111111111111111111111";

    #[test]
    pub fn test_deserialize_ok_object() {
        let j = json!({
            "oid": ONES_SHA256,
            "size": 123,
            "actions": {
                "download": {
                    "href": "https://some-download.com",
                    "header": {
                        "Key": "value"
                    },
                    "expires_at": "2016-11-10T15:29:07Z",
                }
            }
        });

        assert_matches!(
            serde_json::from_str::<ResponseObject>(&j.to_string()),
            Ok(ResponseObject {
                object: RequestObject { oid: _, size: 123 },
                status: ObjectStatus::Ok {
                    authenticated: false,
                    actions: _,
                },
            })
        )
    }

    #[test]
    pub fn test_deserialize_err_object() {
        let j = json!({
            "oid": ONES_SHA256,
            "size": 123,
            "error": {
                "code": 404,
                "message": "Object does not exist"
            }
        });

        assert_matches!(
            serde_json::from_str::<ResponseObject>(&j.to_string()),
            Ok(ResponseObject {
                object: RequestObject { oid: _, size: 123 },
                status: ObjectStatus::Err {
                    error: ObjectError {
                        code: 404,
                        message: _,
                    },
                },
            })
        )
    }

    #[test]
    pub fn test_deserialize_action() {
        let j = json!({
            "href": "https://some-download.com",
            "header": {
                "Key": "value"
            },
            "expires_at": "2016-11-10T15:29:07Z",
        });

        let res = serde_json::from_str::<ObjectAction>(&j.to_string()).unwrap();
        assert_eq!(
            res.href,
            "https://some-download.com".parse::<Uri>().unwrap()
        );
        assert_eq!(
            res.header,
            Some(hashmap! { "Key".to_string() => "value".to_string() })
        );
        assert_eq!(res.expires_at, Some("2016-11-10T15:29:07Z".to_string()));
    }

    quickcheck! {
        fn request_batch_roundtrip(batch: RequestBatch) -> bool {
            let json = serde_json::to_string(&batch).unwrap();
            let rt = serde_json::from_str::<RequestBatch>(&json).unwrap();
            rt == batch
        }

        fn response_batch_roundtrip(batch: ResponseBatch) -> bool {
            let json = serde_json::to_string(&batch).unwrap();
            let rt = serde_json::from_str::<ResponseBatch>(&json).unwrap();
            rt == batch
        }
    }
}


/*
 * 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::marker::PhantomData;
use std::str::FromStr;

use serde::de;
use serde::ser;

struct FromStrVisitor<T> {
    _phantom: PhantomData<T>,
}

impl<'de, T> de::Visitor<'de> for FromStrVisitor<T>
where
    T: FromStr,
    <T as std::str::FromStr>::Err: fmt::Display,
{
    type Value = T;

    fn expecting(&self, formatter: &mut fmt::Formatter) -> fmt::Result {
        formatter.write_str("a string")
    }

    fn visit_str<E>(self, v: &str) -> Result<Self::Value, E>
    where
        E: de::Error,
    {
        v.parse::<T>().map_err(E::custom)
    }
}

pub fn deserialize<'de, T, D>(deserializer: D) -> Result<T, D::Error>
where
    D: de::Deserializer<'de>,
    T: FromStr,
    <T as std::str::FromStr>::Err: fmt::Display,
{
    deserializer.deserialize_any(FromStrVisitor::<T> {
        _phantom: PhantomData,
    })
}

pub fn serialize<T, S>(t: &T, serializer: S) -> Result<S::Ok, S::Error>
where
    S: ser::Serializer,
    T: ToString,
{
    serializer.serialize_str(&t.to_string())
}


/*
 * 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::HashMap;
use std::num::NonZeroU16;
use std::time::Instant;

use anyhow::Context;
use anyhow::Error;
use blobstore::Blobstore;
use blobstore::Loadable;
use blobstore::LoadableError;
use filestore::Alias;
use futures::future;
use futures::future::FutureExt;
use futures::pin_mut;
use futures::select;
use gotham::state::FromState;
use gotham::state::State;
use gotham_derive::StateData;
use gotham_derive::StaticResponseExtender;
use gotham_ext::body_ext::BodyExt;
use gotham_ext::error::HttpError;
use gotham_ext::middleware::RequestStartTime;
use gotham_ext::middleware::ScubaMiddlewareState;
use gotham_ext::response::BytesBody;
use gotham_ext::response::TryIntoResponse;
use http::header::HeaderMap;
use hyper::Body;
use hyper::StatusCode;
use lfs_protocol::git_lfs_mime;
use lfs_protocol::ObjectAction;
use lfs_protocol::ObjectError;
use lfs_protocol::ObjectStatus;
use lfs_protocol::Operation;
use lfs_protocol::RequestBatch;
use lfs_protocol::RequestObject;
use lfs_protocol::ResponseBatch;
use lfs_protocol::ResponseObject;
use lfs_protocol::Transfer;
use maplit::hashmap;
use mononoke_types::hash::Sha256;
use mononoke_types::typed_hash::ContentId;
use mononoke_types::BlobstoreKey;
use rand::Rng;
use redactedblobstore::has_redaction_root_cause;
use repo_blobstore::RepoBlobstoreRef;
use serde::Deserialize;
use slog::debug;
use stats::prelude::*;
use time_ext::DurationExt;
use time_window_counter::GlobalTimeWindowCounterBuilder;

use crate::errors::ErrorKind;
use crate::lfs_server_context::RepositoryRequestContext;
use crate::lfs_server_context::UriBuilder;
use crate::middleware::LfsMethod;
use crate::popularity::consistent_routing;
use crate::scuba::LfsScubaKey;

define_stats! {
    prefix ="mononoke.lfs.batch";
    download_redirect_internal: timeseries(Rate, Sum),
    download_redirect_upstream: timeseries(Rate, Sum),
    download_unknown: timeseries(Rate, Sum),
    upload_redirect: timeseries(Rate, Sum),
    upload_no_redirect: timeseries(Rate, Sum),
    upload_rejected: timeseries(Rate, Sum),
}

enum Source {
    Internal,
    Upstream,
}

#[derive(Deserialize, StateData, StaticResponseExtender)]
pub struct BatchParams {
    repository: String,
}

/// A collection of objects available in a specific server (internal, or upstream).
struct ServerObjects {
    objects: HashMap<lfs_protocol::Sha256, (u64, ObjectAction)>,
}

impl ServerObjects {
    fn empty() -> Self {
        Self {
            objects: HashMap::new(),
        }
    }

    fn get(&self, oid: &lfs_protocol::Sha256) -> Option<(RequestObject, ObjectAction)> {
        let (size, action) = self.objects.get(oid)?;
        Some((
            RequestObject {
                oid: *oid,
                size: *size,
            },
            action.clone(),
        ))
    }

    fn contains(&self, oid: &lfs_protocol::Sha256) -> bool {
        self.objects.contains_key(oid)
    }
}

impl FromIterator<(RequestObject, ObjectAction)> for ServerObjects {
    fn from_iter<I: IntoIterator<Item = (RequestObject, ObjectAction)>>(iter: I) -> Self {
        let mut objects = HashMap::new();

        for (obj, action) in iter {
            objects.insert(obj.oid, (obj.size, action));
        }

        Self { objects }
    }
}

impl FromIterator<(InternalObject, ObjectAction)> for ServerObjects {
    fn from_iter<I: IntoIterator<Item = (InternalObject, ObjectAction)>>(iter: I) -> Self {
        let mut objects = HashMap::new();

        for (obj, action) in iter {
            objects.insert(obj.oid.into(), (obj.download_size(), action));
        }

        Self { objects }
    }
}

enum UpstreamObjects {
    UpstreamPresence(ServerObjects),
    NoUpstream,
}

impl UpstreamObjects {
    fn should_upload(&self, oid: &lfs_protocol::Sha256) -> bool {
        match self {
            // Upload to upstream if the object is missing there.
            UpstreamObjects::UpstreamPresence(objs) => !objs.contains(oid),
            // Without an upstream, we never need to upload there.
            UpstreamObjects::NoUpstream => false,
        }
    }

    fn download_action(&self, oid: &lfs_protocol::Sha256) -> Option<(RequestObject, ObjectAction)> {
        match self {
            // Passthrough download actions from upstream.
            UpstreamObjects::UpstreamPresence(map) => map.get(oid),
            // In the absence of an upstream, we cannot download from there.
            UpstreamObjects::NoUpstream => None,
        }
    }
}

// TODO: Unit tests for this. We could use a client that lets us do stub things.
async fn upstream_objects(
    ctx: &RepositoryRequestContext,
    objects: &[RequestObject],
) -> Result<UpstreamObjects, ErrorKind> {
    let objects = objects.to_vec();

    let batch = RequestBatch {
        operation: Operation::Download,
        r#ref: None,
        transfers: vec![Transfer::Basic],
        objects,
    };

    let res = ctx.upstream_batch(&batch).await?;

    let ResponseBatch { transfer, objects } = match res {
        Some(res) => res,
        None => {
            return Ok(UpstreamObjects::NoUpstream);
        }
    };

    let objects = match transfer {
        Transfer::Basic => {
            // Extract valid download actions from upstream. Those are the objects upstream has.
            objects
                .into_iter()
                .filter_map(|object| {
                    let ResponseObject { object, status } = object;

                    let mut actions = match status {
                        ObjectStatus::Ok {
                            authenticated: false,
                            actions,
                        } => actions,
                        _ => HashMap::new(),
                    };

                    actions
                        .remove(&Operation::Download)
                        .map(|action| (object, action))
                })
                .collect()
        }
        Transfer::Unknown => ServerObjects::empty(),
    };

    Ok(UpstreamObjects::UpstreamPresence(objects))
}

/// An object available internally (i.e. on this server).
#[derive(Copy, Clone, PartialEq, Eq, Debug)]
pub struct InternalObject {
    id: ContentId,
    oid: Sha256,
    size: Option<u64>,
}

impl InternalObject {
    pub fn new(id: ContentId, oid: Sha256, size: Option<u64>) -> Self {
        Self { id, oid, size }
    }

    pub fn id(&self) -> ContentId {
        self.id
    }

    pub fn download_size(&self) -> u64 {
        self.size.unwrap_or(0)
    }
}

async fn resolve_internal_object(
    ctx: &RepositoryRequestContext,
    oid: Sha256,
) -> Result<Option<InternalObject>, Error> {
    let blobstore = ctx.repo.repo_blobstore();

    let content_id = Alias::Sha256(oid).load(&ctx.ctx, blobstore).await;

    let content_id = match content_id {
        Ok(content_id) => content_id,
        Err(LoadableError::Missing(_)) => return Ok(None),
        Err(e) => return Err(Error::from(e).context(ErrorKind::LocalAliasLoadError)),
    };

    // The filestore may allow aliases to be created before the contents are created (the creation
    // of the content is what makes it logically exists), so we should check for the content's
    // existence before we proceed here. Querying metadata does this implicitly, since metadata is
    // written after content. This wouldn't matter if we didn't have an upstream, but it does
    // matter for now to handle the (very much edge-y) case of the content existing in the
    // upstream, its alias existing locally, but not its content (T57777060).

    let meta = filestore::get_metadata(&blobstore, &ctx.ctx, &(content_id.into()))
        .await
        .with_context(|| format!("Failed fetching content metadata for {:?}", content_id));

    match meta {
        Ok(Some(meta)) => {
            return Ok(Some(InternalObject::new(
                meta.content_id,
                meta.sha256,
                Some(meta.total_size),
            )));
        }
        Ok(None) => {
            return Ok(None);
        }
        Err(e) if has_redaction_root_cause(&e) => {
            // Fallthrough to an existence check.
        }
        Err(e) => {
            return Err(e);
        }
    };

    // If a load error was caused by redaction, then check for existence instead, which isn't
    // subject to redaction (only the content is).

    let exists = blobstore
        .is_present(&ctx.ctx, &content_id.blobstore_key())
        .await
        .with_context(|| format!("Failed to check for existence of: {:?}", content_id))?
        .assume_not_found_if_unsure();

    if exists {
        return Ok(Some(InternalObject::new(content_id, oid, None)));
    }

    Ok(None)
}

fn generate_routing_key(tasks_per_content: NonZeroU16, oid: Sha256) -> String {
    // Randomly generate task number to send to.
    let task_n = rand::thread_rng().gen_range(0..tasks_per_content.get());
    // For the base task, no extension is added to routing key.
    let mut routing_key = format!("{}", oid);
    if task_n > 0 {
        // All other tasks have tailing number in routing key.
        routing_key = format!("{}-{}", routing_key, task_n);
    }
    routing_key
}

async fn internal_objects(
    ctx: &RepositoryRequestContext,
    objects: &[RequestObject],
) -> Result<ServerObjects, ErrorKind> {
    let futs = objects.iter().map(|req| async move {
        let obj = resolve_internal_object(ctx, req.oid.into())
            .await
            .map_err(ErrorKind::Error)?;

        let consistent_routing = match obj {
            Some(obj) => consistent_routing(ctx, obj, GlobalTimeWindowCounterBuilder).await,
            None => Some(NonZeroU16::new(1).unwrap()),
        };

        Result::<_, ErrorKind>::Ok((obj, consistent_routing))
    });

    let objs = future::try_join_all(futs).await?;

    objs.into_iter()
        .filter_map(|(maybe_obj, consistent_routing)| match maybe_obj {
            // Map the objects we have locally into an action routing to a Mononoke LFS server.
            Some(obj) => {
                let uri = if let Some(consistent_routing) = consistent_routing && ctx.config.enable_consistent_routing() {
                    let routing_key = generate_routing_key(consistent_routing, obj.oid);
                    ctx.uri_builder
                        .consistent_download_uri(&obj.id, routing_key)
                } else {
                    ctx.uri_builder.download_uri(&obj.id)
                };

                Some(uri.map(|uri| (obj, ObjectAction::new(uri))))
            }
            None => None,
        })
        .collect::<Result<ServerObjects, ErrorKind>>()
}

fn batch_upload_response_objects(
    uri_builder: &UriBuilder,
    max_upload_size: Option<u64>,
    objects: &[RequestObject],
    upstream: &UpstreamObjects,
    internal: &ServerObjects,
) -> Result<Vec<ResponseObject>, ErrorKind> {
    let objects: Result<Vec<ResponseObject>, ErrorKind> = objects
        .iter()
        .map(|object| {
            let status = match (
                upstream.should_upload(&object.oid),
                internal.contains(&object.oid),
                max_upload_size,
            ) {
                (false, true, _) => {
                    // Object doesn't need to be uploaded anywhere: move on.
                    STATS::upload_no_redirect.add_value(1);

                    ObjectStatus::Ok {
                        authenticated: false,
                        actions: hashmap! {},
                    }
                }
                (_, _, Some(max_upload_size)) if object.size > max_upload_size => {
                    // Object is too large and upload is required: reject it (note: this doesn't
                    // enforce that uploads cannot be done: the upload endpoint has its own
                    // validation too).
                    STATS::upload_rejected.add_value(1);

                    ObjectStatus::Err {
                        error: ObjectError {
                            code: StatusCode::BAD_REQUEST.as_u16(),
                            message: ErrorKind::UploadTooLarge(object.size, max_upload_size)
                                .to_string(),
                        },
                    }
                }
                _ => {
                    // Object is missing in at least one location. Require uploading it.
                    STATS::upload_redirect.add_value(1);
                    let uri = uri_builder.upload_uri(object)?;
                    let action = ObjectAction::new(uri);

                    ObjectStatus::Ok {
                        authenticated: false,
                        actions: hashmap! { Operation::Upload => action },
                    }
                }
            };

            Ok(ResponseObject {
                object: *object,
                status,
            })
        })
        .collect();
    objects
}

async fn batch_upload(
    ctx: &RepositoryRequestContext,
    batch: RequestBatch,
) -> Result<ResponseBatch, ErrorKind> {
    let (upstream, internal) = future::try_join(
        upstream_objects(ctx, &batch.objects),
        internal_objects(ctx, &batch.objects),
    )
    .await?;

    let objects = batch_upload_response_objects(
        &ctx.uri_builder,
        ctx.max_upload_size(),
        &batch.objects,
        &upstream,
        &internal,
    )?;

    Ok(ResponseBatch {
        transfer: Transfer::Basic,
        objects,
    })
}

/// This method peforms the routing logic for a given object being requested, given what's
/// available in upstream and downstream.
fn route_download_for_object<'a>(
    upstream: &'a Result<UpstreamObjects, Error>,
    internal: &'a ServerObjects,
    oid: &'_ lfs_protocol::Sha256,
) -> Result<Option<(Source, RequestObject, ObjectAction)>, Error> {
    // If we have the object internally, then get it from there.
    if let Some((obj, action)) = internal.get(oid) {
        return Ok(Some((Source::Internal, obj, action)));
    }

    match upstream {
        // If our upstream succeded, then we try to get the object from there.
        Ok(ref upstream) => {
            if let Some((obj, action)) = upstream.download_action(oid) {
                return Ok(Some((Source::Upstream, obj, action)));
            }
        }
        // If our upstream failed, then we don't know whether the object is available anywhere, so
        // that's an error.
        Err(ref cause) => {
            let err =
                Error::new(ErrorKind::ObjectNotInternallyAvailableAndUpstreamUnavailable(*oid))
                    .context(cause.to_string());
            return Err(err);
        }
    }

    // If upstream succeeded and we found the result neither locally nor remotely, then that's
    // just a missing object.
    Ok(None)
}

fn batch_download_response_objects(
    objects: &[RequestObject],
    upstream: &Result<UpstreamObjects, Error>,
    internal: &ServerObjects,
    scuba: &mut Option<&mut ScubaMiddlewareState>,
) -> Result<Vec<ResponseObject>, Error> {
    let mut upstream_blobs = vec![];
    let responses = objects
        .iter()
        .map(|request_object| {
            let location = route_download_for_object(upstream, internal, &request_object.oid)?;

            let res = match location {
                Some((source, server_object, action)) => {
                    match source {
                        Source::Internal => STATS::download_redirect_internal.add_value(1),
                        Source::Upstream => {
                            upstream_blobs.push(server_object.oid.to_string());
                            STATS::download_redirect_upstream.add_value(1);
                        }
                    };

                    let status = ObjectStatus::Ok {
                        authenticated: false,
                        actions: hashmap! { Operation::Download => action },
                    };

                    ResponseObject {
                        object: server_object,
                        status,
                    }
                }
                None => {
                    STATS::download_unknown.add_value(1);

                    let status = ObjectStatus::Err {
                        error: ObjectError {
                            code: StatusCode::NOT_FOUND.as_u16(),
                            message: "Object does not exist".to_string(),
                        },
                    };

                    ResponseObject {
                        object: *request_object,
                        status,
                    }
                }
            };

            Result::<_, Error>::Ok(res)
        })
        .collect::<Result<Vec<_>, _>>()?;

    ScubaMiddlewareState::maybe_add(
        scuba,
        LfsScubaKey::BatchInternalMissingBlobs,
        upstream_blobs,
    );

    Ok(responses)
}

/// Try to prepare a batch response with only internal objects. Returns None if any are missing.
fn batch_download_internal_only_response_objects(
    objects: &[RequestObject],
    internal: &ServerObjects,
) -> Option<Vec<ResponseObject>> {
    let res = objects
        .iter()
        .map(|request_object| {
            let (server_object, action) = internal.get(&request_object.oid)?;

            let status = ObjectStatus::Ok {
                authenticated: false,
                actions: hashmap! { Operation::Download => action },
            };

            Some(ResponseObject {
                object: server_object,
                status,
            })
        })
        .collect::<Option<Vec<_>>>()?;

    // Record stats only if all have succeeded.
    STATS::download_redirect_internal.add_value(res.len() as i64);

    Some(res)
}

async fn batch_download(
    ctx: &RepositoryRequestContext,
    batch: RequestBatch,
    scuba: &mut Option<&mut ScubaMiddlewareState>,
) -> Result<ResponseBatch, ErrorKind> {
    let upstream = upstream_objects(ctx, &batch.objects).fuse();
    let internal = internal_objects(ctx, &batch.objects).fuse();
    pin_mut!(upstream, internal);

    let mut update_batch_order = |status| {
        ScubaMiddlewareState::maybe_add(scuba, LfsScubaKey::BatchOrder, status);
    };

    update_batch_order("error");

    let objects = select! {
        upstream_objects = upstream => {
            update_batch_order("upstream");
            debug!(ctx.logger(), "batch: upstream ready");
            let internal_objects = internal.await?;
            debug!(ctx.logger(), "batch: internal ready");
            batch_download_response_objects(&batch.objects, &upstream_objects.map_err(Error::from), &internal_objects, scuba)
        }
        internal_objects = internal => {
            debug!(ctx.logger(), "batch: internal ready");
            let internal_objects = internal_objects?;

            let objects = if ctx.always_wait_for_upstream() {
                None
            } else {
                batch_download_internal_only_response_objects(&batch.objects, &internal_objects)
            };

            if let Some(objects) = objects {
                // We were able to return with just internal, don't wait for upstream.
                update_batch_order("internal");
                debug!(ctx.logger(), "batch: skip upstream");
                Ok(objects)
            } else {
                // We don't have all the objects: wait for upstream.
                update_batch_order("both");
                let upstream_objects = upstream.await;
                debug!(ctx.logger(), "batch: upstream ready");
                batch_download_response_objects(&batch.objects, &upstream_objects.map_err(Error::from), &internal_objects, scuba)
            }
        }
    }.map_err(ErrorKind::Error)?;

    Ok(ResponseBatch {
        transfer: Transfer::Basic,
        objects,
    })
}

// TODO: Do we want to validate the client's Accept & Content-Type headers here?
pub async fn batch(state: &mut State) -> Result<impl TryIntoResponse, HttpError> {
    let BatchParams { repository } = state.take();
    let start_time = state
        .try_borrow::<RequestStartTime>()
        .map_or_else(Instant::now, |t| t.0);

    let ctx =
        RepositoryRequestContext::instantiate(state, repository.clone(), LfsMethod::Batch).await?;

    ScubaMiddlewareState::maybe_add(
        &mut state.try_borrow_mut::<ScubaMiddlewareState>(),
        LfsScubaKey::BatchRequestContextReadyUs,
        start_time.elapsed().as_micros_unchecked(),
    );

    let body = Body::take_from(state);
    let headers = HeaderMap::try_borrow_from(state);

    let body = body
        .try_concat_body_opt(headers)
        .map_err(HttpError::e400)?
        .await
        .context(ErrorKind::ClientCancelled)
        .map_err(HttpError::e400)?;

    let mut scuba = state.try_borrow_mut::<ScubaMiddlewareState>();

    ScubaMiddlewareState::maybe_add(
        &mut scuba,
        LfsScubaKey::BatchRequestReceivedUs,
        start_time.elapsed().as_micros_unchecked(),
    );

    let request_batch = serde_json::from_slice::<RequestBatch>(&body)
        .context(ErrorKind::InvalidBatch)
        .map_err(HttpError::e400)?;

    ScubaMiddlewareState::maybe_add(
        &mut scuba,
        LfsScubaKey::BatchObjectCount,
        request_batch.objects.len(),
    );

    ScubaMiddlewareState::maybe_add(
        &mut scuba,
        LfsScubaKey::BatchRequestParsedUs,
        start_time.elapsed().as_micros_unchecked(),
    );

    let res = match request_batch.operation {
        Operation::Upload => batch_upload(&ctx, request_batch).await,
        Operation::Download => batch_download(&ctx, request_batch, &mut scuba).await,
    };

    ScubaMiddlewareState::maybe_add(
        &mut scuba,
        LfsScubaKey::BatchResponseReadyUs,
        start_time.elapsed().as_micros_unchecked(),
    );

    let res = res.map_err(|e| match e {
        ErrorKind::HostNotAllowlisted(_) => HttpError::e400(e),
        _ => HttpError::e500(e),
    })?;
    let body = serde_json::to_string(&res).map_err(HttpError::e500)?;

    Ok(BytesBody::new(body, git_lfs_mime()))
}

#[cfg(test)]
mod test {
    use std::collections::HashSet;
    use std::sync::atomic::AtomicU64;
    use std::sync::atomic::Ordering;
    use std::sync::Arc;

    use async_trait::async_trait;
    use blobstore::BlobstoreBytes;
    use blobstore::BlobstoreGetData;
    use bytes::Bytes;
    use context::CoreContext;
    use fbinit::FacebookInit;
    use filestore::FilestoreConfigRef;
    use filestore::StoreRequest;
    use futures::stream;
    use hyper::Uri;
    use memblob::Memblob;
    use mononoke_types::ContentMetadataV2Id;
    use mononoke_types_mocks::hash::FOURS_SHA256;
    use mononoke_types_mocks::hash::ONES_SHA256;
    use mononoke_types_mocks::hash::THREES_SHA256;
    use mononoke_types_mocks::hash::TWOS_SHA256;
    use pretty_assertions::assert_eq;
    use redactedblobstore::RedactedBlobs;
    use redactedblobstore::RedactedMetadata;
    use test_repo_factory::TestRepoFactory;

    use super::*;
    use crate::lfs_server_context::ServerUris;
    use crate::Repo;

    const SERVER_HOSTNAME: &str = "mzr.re";

    fn obj(oid: Sha256, size: u64) -> RequestObject {
        RequestObject {
            oid: oid.into(),
            size,
        }
    }

    #[test]
    fn test_download() -> Result<(), Error> {
        let o1 = obj(ONES_SHA256, 111);
        let o2 = obj(TWOS_SHA256, 222);
        let o3 = obj(THREES_SHA256, 333);
        let o4 = obj(FOURS_SHA256, 444);

        let upstream = hashmap! {
            o1 => ObjectAction::new("http://foo.com/1".parse()?),
            o2 => ObjectAction::new("http://foo.com/2".parse()?),
        }
        .into_iter()
        .collect();

        let upstream = Ok(UpstreamObjects::UpstreamPresence(upstream));

        let internal = hashmap! {
            o2 => ObjectAction::new("http://bar.com/2".parse()?),
            o3 => ObjectAction::new("http://bar.com/3".parse()?),
        }
        .into_iter()
        .collect();

        let res =
            batch_download_response_objects(&[o1, o2, o3, o4], &upstream, &internal, &mut None)?;

        assert_eq!(
            vec![
                ResponseObject {
                    object: o1,
                    status: ObjectStatus::Ok {
                        authenticated: false,
                        // This is in upstream only
                        actions: hashmap! { Operation::Download =>  ObjectAction::new("http://foo.com/1".parse()?) }
                    }
                },
                ResponseObject {
                    object: o2,
                    status: ObjectStatus::Ok {
                        authenticated: false,
                        // This is in both, so it'll go to internal
                        actions: hashmap! { Operation::Download =>  ObjectAction::new("http://bar.com/2".parse()?) }
                    }
                },
                ResponseObject {
                    object: o3,
                    status: ObjectStatus::Ok {
                        authenticated: false,
                        // This is in internal only
                        actions: hashmap! { Operation::Download =>  ObjectAction::new("http://bar.com/3".parse()?) }
                    }
                },
                ResponseObject {
                    object: o4,
                    status: ObjectStatus::Err {
                        error: ObjectError {
                            code: 404,
                            message: "Object does not exist".to_string(),
                        }
                    }
                }
            ],
            res
        );

        // Test sending the wrong size. Ensure we still find the object here.

        let res2 = batch_download_response_objects(
            &[obj(ONES_SHA256, 0), obj(THREES_SHA256, 0)],
            &upstream,
            &internal,
            &mut None,
        )?;

        assert_eq!(
            vec![
                ResponseObject {
                    object: o1,
                    status: ObjectStatus::Ok {
                        authenticated: false,
                        actions: hashmap! { Operation::Download =>  ObjectAction::new("http://foo.com/1".parse()?) }
                    }
                },
                ResponseObject {
                    object: o3,
                    status: ObjectStatus::Ok {
                        authenticated: false,
                        actions: hashmap! { Operation::Download =>  ObjectAction::new("http://bar.com/3".parse()?) }
                    }
                },
            ],
            res2
        );

        Ok(())
    }

    #[test]
    fn test_routing_keys() -> Result<(), Error> {
        // allowed keys
        let allowed_routing_key_base: String = format!("{}", ONES_SHA256);
        let allowed_routing_key_one: String = format!("{}-1", allowed_routing_key_base);
        // base case
        let routing_key_base = generate_routing_key(NonZeroU16::new(1).unwrap(), ONES_SHA256);
        assert_eq!(&routing_key_base, &allowed_routing_key_base);

        // random key case
        let allowed_routing_keys = vec![allowed_routing_key_base, allowed_routing_key_one];
        for _ in 0..5 {
            let routing_key = generate_routing_key(NonZeroU16::new(2).unwrap(), ONES_SHA256);
            assert!(allowed_routing_keys.contains(&routing_key))
        }

        Ok(())
    }

    #[test]
    fn test_download_upstream_failed_and_its_ok() -> Result<(), Error> {
        let o1 = obj(ONES_SHA256, 111);

        let req = vec![o1];

        let upstream = Err(Error::msg("Oops"));

        let internal = hashmap! {
            o1 => ObjectAction::new("http://foo.com/1".parse()?),
        }
        .into_iter()
        .collect();

        let res = batch_download_response_objects(&req, &upstream, &internal, &mut None)?;

        assert_eq!(
            vec![ResponseObject {
                object: o1,
                status: ObjectStatus::Ok {
                    authenticated: false,
                    actions: hashmap! { Operation::Download =>  ObjectAction::new("http://foo.com/1".parse()?) }
                }
            },],
            res
        );

        Ok(())
    }

    #[test]
    fn test_download_upstream_failed_and_its_not_ok() -> Result<(), Error> {
        let o1 = obj(ONES_SHA256, 111);

        let req = vec![o1];

        let upstream = Err(Error::msg("Oops"));

        let internal = ServerObjects::empty();

        let res = batch_download_response_objects(&req, &upstream, &internal, &mut None);
        assert!(res.is_err());

        Ok(())
    }

    fn upload_uri(object: &RequestObject) -> Result<Uri, Error> {
        let r = format!(
            "http://foo.com/repo123/upload/{}/{}?server_hostname={}",
            object.oid, object.size, SERVER_HOSTNAME
        )
        .parse()?;
        Ok(r)
    }

    #[test]
    fn test_upload() -> Result<(), Error> {
        let o1 = obj(ONES_SHA256, 123);
        let o2 = obj(TWOS_SHA256, 456);
        let o3 = obj(THREES_SHA256, 789);
        let o4 = obj(FOURS_SHA256, 1111);

        let req = vec![o1, o2, o3, o4];

        let upstream = hashmap! {
            o1 => ObjectAction::new("http://foo.com/1".parse()?),
            o2 => ObjectAction::new("http://foo.com/2".parse()?),
        }
        .into_iter()
        .collect();

        let internal = hashmap! {
            o2 => ObjectAction::new("http://bar.com/2".parse()?),
            o3 => ObjectAction::new("http://bar.com/3".parse()?),
        }
        .into_iter()
        .collect();

        let server = ServerUris::new(
            vec!["http://foo.com".to_string()],
            Some("http://bar.com".to_string()),
        )?;
        let uri_builder = UriBuilder {
            repository: "repo123".to_string(),
            server: Arc::new(server),
            host: "foo.com".to_string(),
            server_hostname: Arc::new(SERVER_HOSTNAME.to_string()),
        };

        let res = batch_upload_response_objects(
            &uri_builder,
            Some(1000),
            &req,
            &UpstreamObjects::UpstreamPresence(upstream),
            &internal,
        )?;

        assert_eq!(
            vec![
                ResponseObject {
                    object: o1,
                    status: ObjectStatus::Ok {
                        authenticated: false,
                        // This is in upstream only, so it needs uploading
                        actions: hashmap! { Operation::Upload =>  ObjectAction::new(upload_uri(&o1)?) }
                    }
                },
                ResponseObject {
                    object: o2,
                    status: ObjectStatus::Ok {
                        authenticated: false,
                        // This is in both, so no actions are required.
                        actions: hashmap! {}
                    }
                },
                ResponseObject {
                    object: o3,
                    status: ObjectStatus::Ok {
                        authenticated: false,
                        // This is in internal only, so it needs uploading
                        actions: hashmap! { Operation::Upload =>  ObjectAction::new(upload_uri(&o3)?) }
                    }
                },
                ResponseObject {
                    object: o4,
                    status: ObjectStatus::Err {
                        error: ObjectError {
                            code: 400,
                            message: "Object size (1111) exceeds max allowed size (1000)"
                                .to_string(),
                        }
                    }
                },
            ],
            res
        );

        Ok(())
    }

    #[fbinit::test]
    async fn test_resolve_missing(fb: FacebookInit) -> Result<(), Error> {
        let ctx = RepositoryRequestContext::test_builder(fb).await?.build()?;
        assert_eq!(resolve_internal_object(&ctx, ONES_SHA256).await?, None);
        Ok(())
    }

    #[fbinit::test]
    async fn test_resolve_present(fb: FacebookInit) -> Result<(), Error> {
        let ctx = RepositoryRequestContext::test_builder(fb).await?.build()?;

        let meta = filestore::store(
            ctx.repo.repo_blobstore(),
            *ctx.repo.filestore_config(),
            &ctx.ctx,
            &StoreRequest::new(6),
            stream::once(async move { Ok(Bytes::from("foobar")) }),
        )
        .await?;

        assert_eq!(
            resolve_internal_object(&ctx, meta.sha256)
                .await?
                .map(|o| o.id),
            Some(meta.content_id)
        );
        Ok(())
    }

    #[derive(Debug)]
    struct HideBlob<B> {
        inner: B,
        hide: HashSet<String>,
        hits: Arc<AtomicU64>,
    }

    impl<B: std::fmt::Display> std::fmt::Display for HideBlob<B> {
        fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
            write!(f, "HideBlob<{}>", &self.inner)
        }
    }

    #[async_trait]
    impl<B: Blobstore> Blobstore for HideBlob<B> {
        async fn get<'a>(
            &'a self,
            ctx: &'a CoreContext,
            key: &'a str,
        ) -> Result<Option<BlobstoreGetData>, Error> {
            if self.hide.contains(key) {
                self.hits.fetch_add(1, Ordering::SeqCst);
                return Ok(None);
            }
            self.inner.get(ctx, key).await
        }

        async fn put<'a>(
            &'a self,
            ctx: &'a CoreContext,
            key: String,
            value: BlobstoreBytes,
        ) -> Result<(), Error> {
            self.inner.put(ctx, key, value).await
        }
    }

    #[fbinit::test]
    async fn test_resolve_redacted(fb: FacebookInit) -> Result<(), Error> {
        // Create a single factory with a blobstore we've created.
        let mut factory = TestRepoFactory::new(fb)?;
        let stub_blobstore = Arc::new(Memblob::default());
        factory.with_blobstore(stub_blobstore.clone());

        // First, have the filestore tell us what the hash for this blob would be, so we can create
        // a new repo and redact it.
        let stub: Repo = factory.build().await?;

        let meta = filestore::store(
            stub.repo_blobstore(),
            *stub.filestore_config(),
            &CoreContext::test_mock(fb),
            &StoreRequest::new(6),
            stream::once(async move { Ok(Bytes::from("foobar")) }),
        )
        .await?;

        // Now, create a new blob repo with the same blob store, but with some data redacted.
        let repo = factory
            .redacted(Some(RedactedBlobs::FromSql(Arc::new(hashmap! {
                meta.content_id.blobstore_key() => RedactedMetadata {
                    task: "test".to_string(),
                    log_only: false,
                }
            }))))
            .build()
            .await?;

        let ctx = RepositoryRequestContext::test_builder(fb)
            .await?
            .repo(repo)
            .build()?;

        assert_eq!(
            resolve_internal_object(&ctx, meta.sha256)
                .await?
                .map(|o| o.id),
            Some(meta.content_id)
        );

        // Now, create another one with the same redaction, but this time pretend the metadata does
        // not exist. Building the key is a bit hacky here, but we have an assertion on the number
        // of hits later to make sure it works.

        let key = format!(
            "repo0000.{}",
            ContentMetadataV2Id::from(meta.content_id).blobstore_key()
        );

        let hits = Arc::new(AtomicU64::new(0));

        let hide_blobstore = HideBlob {
            inner: stub_blobstore,
            hide: HashSet::from_iter(Some(key)),
            hits: hits.clone(),
        };

        let repo = factory
            .with_blobstore(Arc::new(hide_blobstore))
            .build()
            .await?;

        let ctx = RepositoryRequestContext::test_builder(fb)
            .await?
            .repo(repo)
            .build()?;

        assert_eq!(
            resolve_internal_object(&ctx, meta.sha256)
                .await?
                .map(|o| o.id),
            Some(meta.content_id)
        );

        assert_eq!(hits.load(Ordering::SeqCst), 1);

        Ok(())
    }

    #[fbinit::test]
    async fn test_resolve_size(fb: FacebookInit) -> Result<(), Error> {
        let repo: Repo = test_repo_factory::build_empty(fb).await?;

        let meta = filestore::store(
            repo.repo_blobstore(),
            *repo.filestore_config(),
            &CoreContext::test_mock(fb),
            &StoreRequest::new(6),
            stream::once(async move { Ok(Bytes::from("foobar")) }),
        )
        .await?;

        let ctx = RepositoryRequestContext::test_builder(fb)
            .await?
            .repo(repo)
            .build()?;

        let obj = RequestObject {
            oid: meta.sha256.into(),
            size: 6,
        };

        let (obj1, action1) = internal_objects(&ctx, &[obj])
            .await?
            .get(&meta.sha256.into())
            .context("Missing v1")?;

        // Note: give the server the wrong size here.
        let (obj2, action2) = internal_objects(&ctx, &[RequestObject { size: 0, ..obj }])
            .await?
            .get(&meta.sha256.into())
            .context("Missing v2")?;

        assert_eq!(obj1, obj);
        assert_eq!(obj2, obj);
        assert_eq!(action1, action2);

        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::collections::BTreeSet;
use std::num::NonZeroU16;
use std::str::FromStr;

use anyhow::bail;
use anyhow::Context;
use anyhow::Error;
use gotham_ext::middleware::PostResponseConfig;
use permission_checker::MononokeIdentitySet;
use rate_limiting::LoadShedLimit;
use serde::de::Deserializer;
use serde::de::Error as _;
use serde::ser::Serializer;
use serde::Deserialize;
use serde::Serialize;

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ObjectPopularity {
    /// SCS counter category to use for blob popularity.
    pub category: String,
    /// How long (in seconds) to lookback
    pub window: u32,
    /// Objects whose sum of downloads exceeds the threshold during the window will not be
    /// consistently-routed. This ensures the full pool of servers can be used to serve very
    /// popular blobs.
    pub thresholds: Vec<ConsistentRoutingRing>,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ConsistentRoutingRing {
    pub threshold: u64,
    pub mode: ConsistentRoutingRingMode,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum ConsistentRoutingRingMode {
    Num { tasks_per_content: NonZeroU16 },
    All,
}

impl TryFrom<lfs_server_config::ConsistentRoutingRingMode> for ConsistentRoutingRingMode {
    type Error = Error;

    fn try_from(value: lfs_server_config::ConsistentRoutingRingMode) -> Result<Self, Self::Error> {
        let mode = match value {
            lfs_server_config::ConsistentRoutingRingMode::num(t) => {
                let tasks_per_content = t
                    .tasks_per_content
                    .try_into()
                    .with_context(|| "tasks_per_content is < 0")?;
                ConsistentRoutingRingMode::Num {
                    tasks_per_content: NonZeroU16::new(tasks_per_content)
                        .with_context(|| "tasks_per_content is 0")?,
                }
            }
            lfs_server_config::ConsistentRoutingRingMode::all(_) => ConsistentRoutingRingMode::All,
            lfs_server_config::ConsistentRoutingRingMode::UnknownField(_) => {
                bail!("Unknown field in ConsistentRoutingRingMode")
            }
        };
        Ok(mode)
    }
}

impl TryFrom<lfs_server_config::ConsistentRoutingRing> for ConsistentRoutingRing {
    type Error = Error;

    fn try_from(value: lfs_server_config::ConsistentRoutingRing) -> Result<Self, Self::Error> {
        let threshold = value
            .threshold
            .try_into()
            .with_context(|| format!("Invalid threshold: {:?}", value.threshold))?;

        let mode = value.mode.try_into().context("Invalid mode")?;

        Ok(Self { threshold, mode })
    }
}

impl TryFrom<lfs_server_config::ObjectPopularity> for ObjectPopularity {
    type Error = Error;

    fn try_from(value: lfs_server_config::ObjectPopularity) -> Result<Self, Self::Error> {
        let window = value
            .window
            .try_into()
            .with_context(|| format!("Invalid window: {:?}", value.window))?;

        let thresholds = value
            .thresholds
            .clone()
            .into_iter()
            .map(|l| l.try_into())
            .collect::<Result<Vec<_>, _>>()
            .context("Invalid routing thresholds")?;

        Ok(Self {
            category: value.category,
            window,
            thresholds,
        })
    }
}

#[derive(Debug, Clone)]
pub struct ServerConfig {
    pub raw_server_config: lfs_server_config::LfsServerConfig,
    loadshedding_limits: Vec<LoadShedLimit>,
    object_popularity: Option<ObjectPopularity>,
    disable_compression_identities: Vec<MononokeIdentitySet>,
}

impl TryFrom<lfs_server_config::LfsServerConfig> for ServerConfig {
    type Error = Error;

    fn try_from(value: lfs_server_config::LfsServerConfig) -> Result<Self, Error> {
        let loadshedding_limits = value
            .loadshedding_limits
            .clone()
            .into_iter()
            .map(|l| l.try_into())
            .collect::<Result<Vec<_>, _>>()
            .context("Invalid loadshedding config")?;

        let object_popularity = value
            .object_popularity
            .as_ref()
            .map(|o| o.clone().try_into())
            .transpose()
            .with_context(|| "Invalid object popularity")?;

        let mut disable_compression_identities: Vec<MononokeIdentitySet> = Vec::new();
        for list in value.disable_compression_identities.iter() {
            let idents = list
                .iter()
                .map(|i| FromStr::from_str(i))
                .collect::<Result<BTreeSet<_>, _>>()?;
            disable_compression_identities.push(idents);
        }

        Ok(Self {
            raw_server_config: value,
            loadshedding_limits,
            object_popularity,
            disable_compression_identities,
        })
    }
}

impl<'de> Deserialize<'de> for ServerConfig {
    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
    where
        D: Deserializer<'de>,
    {
        let raw = lfs_server_config::LfsServerConfig::deserialize(deserializer)?;
        let config = Self::try_from(raw).map_err(|e| D::Error::custom(format!("{:?}", e)))?;
        Ok(config)
    }
}

impl Serialize for ServerConfig {
    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
    where
        S: Serializer,
    {
        lfs_server_config::LfsServerConfig::serialize(&self.raw_server_config, serializer)
    }
}

impl Default for ServerConfig {
    fn default() -> Self {
        let raw_server_config = lfs_server_config::LfsServerConfig {
            track_bytes_sent: false,
            enable_consistent_routing: false,
            disable_hostname_logging: false,
            enforce_acl_check: false,
            loadshedding_limits: vec![],
            object_popularity: None,
            disable_compression: false,
            disable_compression_identities: vec![],
            enforce_authentication: false,
        };

        Self {
            raw_server_config,
            loadshedding_limits: vec![],
            object_popularity: None,
            disable_compression_identities: vec![],
        }
    }
}

impl ServerConfig {
    pub fn track_bytes_sent(&self) -> bool {
        self.raw_server_config.track_bytes_sent
    }
    pub fn enable_consistent_routing(&self) -> bool {
        self.raw_server_config.enable_consistent_routing
    }
    pub fn disable_hostname_logging(&self) -> bool {
        self.raw_server_config.disable_hostname_logging
    }
    pub fn loadshedding_limits(&self) -> Vec<LoadShedLimit> {
        self.loadshedding_limits.clone()
    }
    pub fn enforce_acl_check(&self) -> bool {
        self.raw_server_config.enforce_acl_check
    }
    pub fn enforce_authentication(&self) -> bool {
        self.raw_server_config.enforce_authentication
    }
    pub fn object_popularity(&self) -> Option<&ObjectPopularity> {
        self.object_popularity.as_ref()
    }
    #[cfg(test)]
    pub fn object_popularity_mut(&mut self) -> &mut Option<ObjectPopularity> {
        &mut self.object_popularity
    }
    pub fn disable_compression(&self) -> bool {
        self.raw_server_config.disable_compression
    }
    pub fn disable_compression_identities(&self) -> &Vec<MononokeIdentitySet> {
        &self.disable_compression_identities
    }
    #[cfg(test)]
    pub fn disable_compression_identities_mut(&mut self) -> &mut Vec<MononokeIdentitySet> {
        &mut self.disable_compression_identities
    }
}

impl PostResponseConfig for ServerConfig {
    fn resolve_hostname(&self) -> bool {
        !self.disable_hostname_logging()
    }
}


/*
 * 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::str::FromStr;

use anyhow::Context;
use anyhow::Error;
use filestore::Alias;
use filestore::FetchKey;
use filestore::Range;
use futures::stream::StreamExt;
use futures::stream::TryStreamExt;
use gotham::state::FromState;
use gotham::state::State;
use gotham_derive::StateData;
use gotham_derive::StaticResponseExtender;
use gotham_ext::content_encoding::ContentEncoding;
use gotham_ext::error::HttpError;
use gotham_ext::middleware::ScubaMiddlewareState;
use gotham_ext::response::CompressedResponseStream;
use gotham_ext::response::ResponseStream;
use gotham_ext::response::ResponseTryStreamExt;
use gotham_ext::response::StreamBody;
use gotham_ext::response::TryIntoResponse;
use http::header::HeaderMap;
use http::header::RANGE;
use mononoke_types::hash::Sha256;
use mononoke_types::ContentId;
use permission_checker::MononokeIdentitySet;
use redactedblobstore::has_redaction_root_cause;
use repo_blobstore::RepoBlobstoreRef;
use serde::Deserialize;
use stats::prelude::*;

use crate::config::ServerConfig;
use crate::errors::ErrorKind;
use crate::lfs_server_context::RepositoryRequestContext;
use crate::middleware::LfsMethod;
use crate::scuba::LfsScubaKey;
use crate::util::is_identity_subset;

define_stats! {
    prefix = "mononoke.lfs.download";
    size_bytes_sent: timeseries(
        "size_bytes_sent";
        Sum;
        Duration::from_secs(5), Duration::from_secs(15), Duration::from_secs(60)
    ),
}
#[derive(Deserialize, StateData, StaticResponseExtender)]
pub struct DownloadParamsContentId {
    repository: String,
    content_id: String,
}

#[derive(Deserialize, StateData, StaticResponseExtender)]
pub struct DownloadParamsSha256 {
    repository: String,
    oid: String,
}

fn parse_range(header: &str) -> Result<Range, Error> {
    static RE: once_cell::sync::Lazy<regex::Regex> =
        once_cell::sync::Lazy::new(|| regex::Regex::new(r"bytes=(\d+)-(\d+)").unwrap());

    let caps = RE
        .captures(header)
        .with_context(|| format!("Unsupported range: {}", header))?;

    let range = Range::range_inclusive(
        caps[1]
            .parse()
            .with_context(|| format!("Invalid range start: {}", &caps[1]))?,
        caps[2]
            .parse()
            .with_context(|| format!("Invalid range end: {}", &caps[2]))?,
    )?
    .strict();

    Ok(range)
}

fn extract_range(state: &State) -> Result<Option<Range>, Error> {
    let header = match HeaderMap::try_borrow_from(state).and_then(|h| h.get(RANGE)) {
        Some(h) => h,
        None => return Ok(None),
    };

    let header = std::str::from_utf8(header.as_bytes()).context("Invalid range")?;

    Ok(Some(parse_range(header)?))
}

fn should_disable_compression(
    config: &ServerConfig,
    client_idents: Option<&MononokeIdentitySet>,
) -> bool {
    if config.disable_compression() {
        return true;
    }

    is_identity_subset(config.disable_compression_identities(), client_idents)
}

async fn fetch_by_key(
    ctx: RepositoryRequestContext,
    key: FetchKey,
    content_encoding: ContentEncoding,
    range: Option<Range>,
    scuba: &mut Option<&mut ScubaMiddlewareState>,
) -> Result<impl TryIntoResponse, HttpError> {
    // Query a stream out of the Filestore
    let fetched = filestore::fetch_range_with_size(
        ctx.repo.repo_blobstore().clone(),
        ctx.ctx.clone(),
        &key,
        range.unwrap_or_else(Range::all),
    )
    .await
    .map_err(|e| {
        if has_redaction_root_cause(&e) {
            HttpError::e410(e)
        } else {
            HttpError::e500(e.context(ErrorKind::FilestoreReadFailure))
        }
    })?;

    // Return a 404 if the stream doesn't exist.
    let (stream, size) = fetched
        .ok_or(ErrorKind::ObjectDoesNotExist(key))
        .map_err(HttpError::e404)?;

    ScubaMiddlewareState::maybe_add(scuba, LfsScubaKey::DownloadContentSize, size);

    let stream = match content_encoding {
        ContentEncoding::Identity => ResponseStream::new(stream)
            .set_content_length(size)
            .left_stream(),
        ContentEncoding::Compressed(c) => CompressedResponseStream::new(stream, c).right_stream(),
    };

    let stream = if ctx.config.track_bytes_sent() {
        stream
            .inspect_ok(|bytes| STATS::size_bytes_sent.add_value(bytes.len() as i64))
            .left_stream()
    } else {
        stream.right_stream()
    };

    let stream = stream.end_on_err();

    let mut body = StreamBody::new(stream, mime::APPLICATION_OCTET_STREAM);
    if range.is_some() {
        body.partial = true;
    }
    Ok(body)
}

async fn download_inner(
    state: &mut State,
    repository: String,
    key: FetchKey,
    method: LfsMethod,
) -> Result<impl TryIntoResponse, HttpError> {
    let range = extract_range(state).map_err(HttpError::e400)?;

    let ctx = RepositoryRequestContext::instantiate(state, repository.clone(), method).await?;

    let disable_compression =
        should_disable_compression(&ctx.config, Some(ctx.ctx.metadata().identities()));

    let content_encoding = if disable_compression {
        ContentEncoding::Identity
    } else {
        ContentEncoding::from_state(state)
    };

    let mut scuba = state.try_borrow_mut::<ScubaMiddlewareState>();

    fetch_by_key(ctx, key, content_encoding, range, &mut scuba).await
}

pub async fn download(state: &mut State) -> Result<impl TryIntoResponse, HttpError> {
    let DownloadParamsContentId {
        repository,
        content_id,
    } = state.take();

    let content_id = ContentId::from_str(&content_id)
        .context(ErrorKind::InvalidContentId)
        .map_err(HttpError::e400)?;

    let key = FetchKey::Canonical(content_id);

    download_inner(state, repository, key, LfsMethod::Download).await
}

pub async fn download_sha256(state: &mut State) -> Result<impl TryIntoResponse, HttpError> {
    let DownloadParamsSha256 { repository, oid } = state.take();

    let oid = Sha256::from_str(&oid)
        .context(ErrorKind::InvalidOid)
        .map_err(HttpError::e400)?;

    let key = FetchKey::Aliased(Alias::Sha256(oid));

    download_inner(state, repository, key, LfsMethod::DownloadSha256).await
}

#[cfg(test)]
mod test {
    use std::sync::Arc;

    use anyhow::Error;
    use fbinit::FacebookInit;
    use http::StatusCode;
    use maplit::hashmap;
    use mononoke_types::typed_hash::BlobstoreKey;
    use mononoke_types_mocks::contentid::ONES_CTID;
    use permission_checker::MononokeIdentity;
    use redactedblobstore::RedactedBlobs;
    use redactedblobstore::RedactedMetadata;
    use test_repo_factory::TestRepoFactory;

    use super::*;

    #[fbinit::test]
    async fn test_redacted_fetch(fb: FacebookInit) -> Result<(), Error> {
        let content_id = ONES_CTID;
        let reason = "test reason";

        let repo = TestRepoFactory::new(fb)?
            .redacted(Some(RedactedBlobs::FromSql(Arc::new(
                hashmap! { content_id.blobstore_key() => RedactedMetadata {
                   task: reason.to_string(),
                   log_only: false,
                }},
            ))))
            .build()
            .await?;

        let ctx = RepositoryRequestContext::test_builder(fb)
            .await?
            .repo(repo)
            .build()?;

        let key = FetchKey::Canonical(content_id);

        let err = fetch_by_key(ctx, key, ContentEncoding::Identity, None, &mut None)
            .await
            .map(|_| ())
            .unwrap_err();
        assert_eq!(err.status_code, StatusCode::GONE);
        assert!(err.error.to_string().contains(reason));
        Ok(())
    }

    #[test]
    fn test_parse_range() -> Result<(), Error> {
        // NOTE: This range is inclusive, so here we want bytes 1, 2, 3, 5 (a 5-byte range starting
        // at byte 1).
        assert_eq!(parse_range("bytes=1-5")?, Range::sized(1, 5).strict());
        assert!(parse_range("1-5").is_err());
        assert!(parse_range("foo=1-5").is_err());
        Ok(())
    }

    #[test]
    fn test_should_disable_compression() -> Result<(), Error> {
        let mut config = ServerConfig::default();
        let test_ident = MononokeIdentity::new("USER", "test");
        let test_ident_2 = MononokeIdentity::new("USER", "test2");
        let mut client_idents = MononokeIdentitySet::new();
        client_idents.insert(test_ident.clone());

        assert!(!should_disable_compression(&config, None));
        assert!(!should_disable_compression(
            &config,
            Some(&MononokeIdentitySet::new())
        ));
        assert!(!should_disable_compression(&config, Some(&client_idents)));

        config
            .disable_compression_identities_mut()
            .push(client_idents.clone());
        assert!(!should_disable_compression(&config, None));
        assert!(should_disable_compression(&config, Some(&client_idents)));

        // A client must match all idents in a MononokeIdentitySet for compression to be disabled.
        config.disable_compression_identities_mut().clear();

        let mut and_set = MononokeIdentitySet::new();
        and_set.insert(test_ident);
        and_set.insert(test_ident_2);

        config.disable_compression_identities_mut().push(and_set);
        assert!(!should_disable_compression(&config, Some(&client_idents)));

        config.raw_server_config.disable_compression = true;
        assert!(should_disable_compression(&config, None));
        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 filestore::FetchKey;
use gotham_ext::error::HttpError;
use hyper::StatusCode;
use lfs_protocol::RequestObject;
use lfs_protocol::ResponseObject;
use thiserror::Error;

#[derive(Debug, Error)]
pub enum ErrorKind {
    #[error("Client cancelled the request")]
    ClientCancelled,
    #[error("An error occurred forwarding the request to upstream")]
    UpstreamDidNotRespond,
    #[error("An error ocurred receiving a response from upstream ({0}): {1}")]
    UpstreamError(StatusCode, String),
    #[error("Could not serialize")]
    SerializationFailed(#[source] anyhow::Error),
    #[error("Could not initialize HTTP client")]
    HttpClientInitializationFailed,
    #[error("Could not build {0}: {1}")]
    UriBuilderFailed(&'static str, anyhow::Error),
    #[error("Invalid Uri {0}: {1}")]
    InvalidUri(String, &'static str),
    #[error("Host {0} is not allowlisted")]
    HostNotAllowlisted(String),
    #[error("Object does not exist: {0:?}")]
    ObjectDoesNotExist(FetchKey),
    #[error("Could not dispatch batch request to upstream")]
    UpstreamBatchNoResponse(#[source] anyhow::Error),
    #[error("Upstream batch response is invalid")]
    UpstreamBatchInvalid(#[source] anyhow::Error),
    #[error("Could not fetch upstream batch")]
    UpstreamBatchError,
    #[error("Could not perform upstream upload")]
    UpstreamUploadError,
    #[error("Upstream batch response included an invalid transfer")]
    UpstreamInvalidTransfer,
    #[error("Upstream batch response did not include requested object: {0:?}")]
    UpstreamMissingObject(RequestObject),
    #[error("Upstream batch response included an invalid object: {0:?}")]
    UpstreamInvalidObject(ResponseObject),
    #[error("Could not load local alias")]
    LocalAliasLoadError,
    #[error("Could not parse Request Batch")]
    InvalidBatch,
    #[error("Could not parse Content ID")]
    InvalidContentId,
    #[error("Could not parse SHA256")]
    InvalidOid,
    #[error("Could not access Filestore for reads")]
    FilestoreReadFailure,
    #[error("Could not access Filestore for writes")]
    FilestoreWriteFailure,
    #[error("Object size ({0}) exceeds max allowed size ({1})")]
    UploadTooLarge(u64, u64),
    #[error("Object is not internally available, and upstream is not available: {0}")]
    ObjectNotInternallyAvailableAndUpstreamUnavailable(lfs_protocol::Sha256),
    #[error("Object could not be synced from upstream: {0:?}")]
    ObjectCannotBeSynced(RequestObject),

    /// A generic error occurred, and we'd like to propagate it.
    #[error(transparent)]
    Error(anyhow::Error),
}

#[derive(Debug, Error)]
pub enum LfsServerContextErrorKind {
    #[error("Operated not permitted")]
    Forbidden,
    #[error("Client not authenticated")]
    NotAuthenticated,
    #[error("Repository does not exist: {0}")]
    RepositoryDoesNotExist(String),
    #[error("Missing host header")]
    MissingHostHeader,
}

impl From<LfsServerContextErrorKind> for HttpError {
    fn from(e: LfsServerContextErrorKind) -> HttpError {
        use LfsServerContextErrorKind::*;
        match e {
            Forbidden => HttpError::e403(e),
            RepositoryDoesNotExist(_) => HttpError::e400(e),
            MissingHostHeader => HttpError::e400(e),
            NotAuthenticated => HttpError::e403(e),
        }
    }
}