/*
 * 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::Arguments;
use std::fmt::Write;
use std::sync::atomic::AtomicBool;
use std::sync::atomic::Ordering;
use std::sync::Arc;
use std::sync::Mutex;

use anyhow::Context;
use anyhow::Error;
use bytes::Bytes;
use cached_config::ConfigHandle;
use context::CoreContext;
use futures::future;
use futures::stream::Stream;
use futures::stream::StreamExt;
use futures::stream::TryStreamExt;
use gotham::state::FromState;
use gotham::state::State;
use gotham_derive::StateData;
use gotham_ext::body_ext::BodyExt;
use hostname::get_hostname;
use http::header::HeaderMap;
use http::uri::Authority;
use http::uri::Parts;
use http::uri::PathAndQuery;
use http::uri::Scheme;
use http::uri::Uri;
use hyper::client::HttpConnector;
use hyper::header;
use hyper::Body;
use hyper::Client;
use hyper::Request;
use hyper_openssl::HttpsConnector;
use lfs_protocol::RequestBatch;
use lfs_protocol::RequestObject;
use lfs_protocol::ResponseBatch;
use metaconfig_types::RepoConfigRef;
use mononoke_types::ContentId;
use qps::Qps;
use repo_authorization::AuthorizationContext;
use repo_permission_checker::RepoPermissionCheckerRef;
use slog::Logger;
use tokio::runtime::Handle;

use crate::config::ServerConfig;
use crate::errors::ErrorKind;
use crate::errors::LfsServerContextErrorKind;
use crate::middleware::LfsMethod;
use crate::middleware::RequestContext;
use crate::LfsRepos;
use crate::Repo;

pub type HttpsHyperClient = Client<HttpsConnector<HttpConnector>>;

// The user agent string presented to upstream
const CLIENT_USER_AGENT: &str = "mononoke-lfs-server/0.1.0 git/2.15.1";

struct LfsServerContextInner {
    repositories: LfsRepos,
    client: Arc<HttpsHyperClient>,
    server: Arc<ServerUris>,
    always_wait_for_upstream: bool,
    max_upload_size: Option<u64>,
    config_handle: ConfigHandle<ServerConfig>,
    logger: Logger,
    qps: Arc<Option<Qps>>,
    server_hostname: Arc<String>,
}

#[derive(Clone, StateData)]
pub struct LfsServerContext {
    inner: Arc<Mutex<LfsServerContextInner>>,
    will_exit: Arc<AtomicBool>,
}

impl LfsServerContext {
    pub fn new(
        repositories: LfsRepos,
        server: ServerUris,
        always_wait_for_upstream: bool,
        max_upload_size: Option<u64>,
        will_exit: Arc<AtomicBool>,
        config_handle: ConfigHandle<ServerConfig>,
        logger: Logger,
        qps: Option<Qps>,
    ) -> Result<Self, Error> {
        let connector = HttpsConnector::new()
            .map_err(Error::from)
            .context(ErrorKind::HttpClientInitializationFailed)?;
        let client = Client::builder().build(connector);
        let server_hostname =
            Arc::new(get_hostname().unwrap_or_else(|_| "UNKNOWN_HOSTNAME".to_string()));

        let inner = LfsServerContextInner {
            repositories,
            server: Arc::new(server),
            client: Arc::new(client),
            always_wait_for_upstream,
            max_upload_size,
            config_handle,
            logger,
            qps: Arc::new(qps),
            server_hostname,
        };

        Ok(LfsServerContext {
            inner: Arc::new(Mutex::new(inner)),
            will_exit,
        })
    }

    pub async fn request(
        &self,
        ctx: CoreContext,
        repository: String,
        host: String,
        method: LfsMethod,
    ) -> Result<RepositoryRequestContext, LfsServerContextErrorKind> {
        let (
            repo,
            client,
            server,
            always_wait_for_upstream,
            max_upload_size,
            config,
            server_hostname,
        ) = {
            let inner = self.inner.lock().expect("poisoned lock");

            match inner.repositories.get(&repository) {
                Some(repo) => (
                    repo,
                    inner.client.clone(),
                    inner.server.clone(),
                    inner.always_wait_for_upstream,
                    inner.max_upload_size,
                    inner.config_handle.get(),
                    inner.server_hostname.clone(),
                ),
                None => {
                    return Err(LfsServerContextErrorKind::RepositoryDoesNotExist(
                        repository,
                    ));
                }
            }
        };

        let enforce_acl_check =
            repo.repo_config().enforce_lfs_acl_check && config.enforce_acl_check();

        acl_check(&ctx, &repo, enforce_acl_check, method).await?;

        Ok(RepositoryRequestContext {
            ctx,
            repo,
            uri_builder: UriBuilder {
                repository,
                server,
                host,
                server_hostname,
            },
            client: HttpClient::Enabled(client),
            config,
            always_wait_for_upstream,
            max_upload_size,
        })
    }

    pub fn get_config_handle(&self) -> ConfigHandle<ServerConfig> {
        self.inner
            .lock()
            .expect("poisoned lock")
            .config_handle
            .clone()
    }
    pub fn qps(&self) -> Arc<Option<Qps>> {
        let inner = self.inner.lock().expect("poisoned lock");
        inner.qps.clone()
    }

    pub fn logger(&self) -> Logger {
        let inner = self.inner.lock().expect("poisoned lock");
        inner.logger.clone()
    }

    pub fn get_config(&self) -> Arc<ServerConfig> {
        let inner = self.inner.lock().expect("poisoned lock");
        inner.config_handle.get()
    }

    pub fn will_exit(&self) -> bool {
        self.will_exit.load(Ordering::Relaxed)
    }
}

async fn acl_check(
    ctx: &CoreContext,
    repo: &impl RepoPermissionCheckerRef,
    enforce_authorization: bool,
    method: LfsMethod,
) -> Result<(), LfsServerContextErrorKind> {
    let authz = AuthorizationContext::new(ctx);
    let acl_check = if method.is_read_only() {
        authz.check_full_repo_read(ctx, repo).await
    } else {
        authz.check_full_repo_draft(ctx, repo).await
    };

    if acl_check.is_denied() && enforce_authorization {
        Err(LfsServerContextErrorKind::Forbidden)
    } else {
        Ok(())
    }
}

#[derive(Clone)]
enum HttpClient {
    Enabled(Arc<HttpsHyperClient>),
    #[cfg(test)]
    Disabled,
}

#[derive(Clone)]
pub struct RepositoryRequestContext {
    pub ctx: CoreContext,
    pub repo: Arc<Repo>,
    pub uri_builder: UriBuilder,
    pub config: Arc<ServerConfig>,
    always_wait_for_upstream: bool,
    max_upload_size: Option<u64>,
    client: HttpClient,
}

pub struct HttpClientResponse<S: Stream<Item = Result<Bytes, Error>> + Send + 'static> {
    headers: HeaderMap,
    body: Option<S>,
    handle: Handle,
}

impl<S: Stream<Item = Result<Bytes, Error>> + Send + 'static> HttpClientResponse<S> {
    pub async fn concat(mut self) -> Result<Bytes, Error> {
        let body = self
            .body
            .take()
            .expect("Body cannot be missing")
            .try_concat_body(&self.headers)?
            .await?;
        Ok(body)
    }

    pub async fn discard(self) -> Result<(), Error> {
        self.concat().await?;
        Ok(())
    }

    pub fn into_inner(mut self) -> S {
        self.body.take().expect("Body cannot be missing")
    }
}

impl<S: Stream<Item = Result<Bytes, Error>> + Send + 'static> Drop for HttpClientResponse<S> {
    fn drop(&mut self) {
        if let Some(body) = self.body.take() {
            let discard = body.for_each(|_| future::ready(()));
            self.handle.spawn(discard);
        }
    }
}

fn host_maybe_port_to_host(host_maybe_port: &str) -> Result<String, LfsServerContextErrorKind> {
    Ok(host_maybe_port
        .parse::<Uri>()
        .map_err(|_e| LfsServerContextErrorKind::MissingHostHeader)?
        .authority()
        .ok_or(LfsServerContextErrorKind::MissingHostHeader)?
        .host()
        .to_string())
}

fn get_host_header(headers: &Option<&HeaderMap>) -> Result<String, LfsServerContextErrorKind> {
    let host_maybe_port = headers
        .ok_or(LfsServerContextErrorKind::MissingHostHeader)?
        .get(http::header::HOST)
        .ok_or(LfsServerContextErrorKind::MissingHostHeader)?
        .to_str()
        .map_err(|_e| LfsServerContextErrorKind::MissingHostHeader)?;

    host_maybe_port_to_host(host_maybe_port)
}

impl RepositoryRequestContext {
    pub async fn instantiate(
        state: &mut State,
        repository: String,
        method: LfsMethod,
    ) -> Result<Self, LfsServerContextErrorKind> {
        let req_ctx = state.borrow_mut::<RequestContext>();
        req_ctx.set_request(repository.clone(), method);

        let ctx = req_ctx.ctx.clone();

        let headers = HeaderMap::try_borrow_from(state);
        let host = get_host_header(&headers)?;

        let lfs_ctx = LfsServerContext::borrow_from(state);
        lfs_ctx.request(ctx, repository, host, method).await
    }

    pub fn logger(&self) -> &Logger {
        self.ctx.logger()
    }

    pub fn always_wait_for_upstream(&self) -> bool {
        self.always_wait_for_upstream
    }

    pub fn max_upload_size(&self) -> Option<u64> {
        self.max_upload_size
    }

    pub async fn dispatch(
        &self,
        mut request: Request<Body>,
    ) -> Result<HttpClientResponse<impl Stream<Item = Result<Bytes, Error>>>, Error> {
        #[allow(clippy::infallible_destructuring_match)]
        let client = match self.client {
            HttpClient::Enabled(ref client) => client,
            #[cfg(test)]
            HttpClient::Disabled => panic!("HttpClient is disabled in test"),
        };

        request.headers_mut().insert(
            header::USER_AGENT,
            header::HeaderValue::from_static(CLIENT_USER_AGENT),
        );
        let res = client.request(request);

        // NOTE: We spawn the request on an executor because we'd like to read the response even if
        // we drop the future returned here. The reason for that is that if we don't read a
        // response, Hyper will not reuse the conneciton for its pool (which makes sense for the
        // general case: if your server is sending you 5GB of data and you drop the future, you
        // don't want to read all that later just to reuse a connection).
        let fut = async move {
            let res = res.await.context(ErrorKind::UpstreamDidNotRespond)?;

            let (head, body) = res.into_parts();

            if !head.status.is_success() {
                let body = body.try_concat_body(&head.headers)?.await?;
                return Err(ErrorKind::UpstreamError(
                    head.status,
                    String::from_utf8_lossy(&body).to_string(),
                )
                .into());
            }

            // NOTE: This buffers the response here, since all our callsites need a concatenated
            // response. If we want to add callsites that need a streaming response, we should add
            // our own wrapper type that wraps the response and the headers.
            Ok(HttpClientResponse {
                headers: head.headers,
                body: Some(body.map_err(Error::from)),
                handle: Handle::current(),
            })
        };

        tokio::spawn(fut).await?
    }

    pub async fn upstream_batch(
        &self,
        batch: &RequestBatch,
    ) -> Result<Option<ResponseBatch>, ErrorKind> {
        let uri = match self.uri_builder.upstream_batch_uri()? {
            Some(uri) => uri,
            None => {
                return Ok(None);
            }
        };

        let body: Bytes = serde_json::to_vec(&batch)
            .map_err(|e| ErrorKind::SerializationFailed(e.into()))?
            .into();

        let req = Request::post(uri)
            .body(body.into())
            .map_err(|e| ErrorKind::Error(e.into()))?;

        let res = self
            .dispatch(req)
            .await
            .map_err(ErrorKind::UpstreamBatchNoResponse)?
            .concat()
            .await
            .map_err(ErrorKind::UpstreamBatchNoResponse)?;

        let batch = serde_json::from_slice::<ResponseBatch>(&res)
            .map_err(|e| ErrorKind::UpstreamBatchInvalid(e.into()))?;

        Ok(Some(batch))
    }
}

#[derive(Clone)]
pub struct UriBuilder {
    pub repository: String,
    pub server: Arc<ServerUris>,
    pub host: String,
    pub server_hostname: Arc<String>,
}

impl UriBuilder {
    fn pick_uri(&self) -> Result<&BaseUri, ErrorKind> {
        self.server
            .self_uris
            .iter()
            .find(|&x| x.authority.host() == self.host)
            .ok_or_else(|| ErrorKind::HostNotAllowlisted(self.host.clone()))
    }

    pub fn upload_uri(&self, object: &RequestObject) -> Result<Uri, ErrorKind> {
        self.pick_uri()?
            .build(format_args!(
                "{}/upload/{}/{}?server_hostname={}",
                &self.repository, object.oid, object.size, self.server_hostname,
            ))
            .map_err(|e| ErrorKind::UriBuilderFailed("upload_uri", e))
    }

    pub fn download_uri(&self, content_id: &ContentId) -> Result<Uri, ErrorKind> {
        self.pick_uri()?
            .build(format_args!(
                "{}/download/{}?server_hostname={}",
                &self.repository, content_id, self.server_hostname
            ))
            .map_err(|e| ErrorKind::UriBuilderFailed("download_uri", e))
    }

    pub fn consistent_download_uri(
        &self,
        content_id: &ContentId,
        routing_key: String,
    ) -> Result<Uri, ErrorKind> {
        self.pick_uri()?
            .build(format_args!(
                "{}/download/{}?routing={}&server_hostname={}",
                &self.repository, content_id, routing_key, self.server_hostname
            ))
            .map_err(|e| ErrorKind::UriBuilderFailed("consistent_download_uri", e))
    }

    pub fn upstream_batch_uri(&self) -> Result<Option<Uri>, ErrorKind> {
        self.server
            .upstream_uri
            .as_ref()
            .map(|uri| {
                uri.build(format_args!("objects/batch"))
                    .map_err(|e| ErrorKind::UriBuilderFailed("upstream_batch_uri", e))
            })
            .transpose()
    }
}

fn parse_and_check_uri(src: &str) -> Result<BaseUri, ErrorKind> {
    let uri = src
        .parse::<Uri>()
        .map_err(|_e| ErrorKind::InvalidUri(src.to_string(), "invalid uri"))?;

    let Parts {
        scheme,
        authority,
        path_and_query,
        ..
    } = uri.into_parts();

    Ok(BaseUri {
        scheme: scheme.ok_or_else(|| ErrorKind::InvalidUri(src.to_string(), "missing scheme"))?,
        authority: authority
            .ok_or_else(|| ErrorKind::InvalidUri(src.to_string(), "missing authority"))?,
        path_and_query,
    })
}

#[derive(Debug)]
pub struct ServerUris {
    /// The list of allowlisted root URLs to use when composing URLs for this LFS server
    self_uris: Vec<BaseUri>,
    /// The URL for an upstream LFS server
    upstream_uri: Option<BaseUri>,
}

impl ServerUris {
    pub fn new(self_uris: Vec<String>, upstream_uri: Option<String>) -> Result<Self, Error> {
        Ok(Self {
            self_uris: self_uris
                .into_iter()
                .map(|v| parse_and_check_uri(&v))
                .collect::<Result<Vec<_>, _>>()?,
            upstream_uri: upstream_uri.map(|v| parse_and_check_uri(&v)).transpose()?,
        })
    }
}

#[derive(Debug)]
struct BaseUri {
    scheme: Scheme,
    authority: Authority,
    path_and_query: Option<PathAndQuery>,
}

impl BaseUri {
    pub fn build(&self, args: Arguments) -> Result<Uri, Error> {
        let mut p = String::new();
        if let Some(ref path_and_query) = self.path_and_query {
            write!(&mut p, "{}", path_and_query)?;
            if !path_and_query.path().ends_with('/') {
                write!(&mut p, "/")?;
            }
        }
        p.write_fmt(args)?;

        Uri::builder()
            .scheme(self.scheme.clone())
            .authority(self.authority.clone())
            .path_and_query(&p[..])
            .build()
            .map_err(Error::from)
    }
}

#[cfg(test)]
mod test {
    use std::str::FromStr;

    use fbinit::FacebookInit;
    use lfs_protocol::Sha256 as LfsSha256;
    use mononoke_types::hash::Sha256;
    use mononoke_types::ContentId;
    use repo_permission_checker::AlwaysAllowRepoPermissionChecker;
    use repo_permission_checker::MockRepoPermissionChecker;
    use test_repo_factory::TestRepoFactory;

    use super::*;

    const ONES_HASH: &str = "1111111111111111111111111111111111111111111111111111111111111111";
    const TWOS_HASH: &str = "2222222222222222222222222222222222222222222222222222222222222222";
    const SIZE: u64 = 123;
    const SERVER_HOSTNAME: &str = "mzr.re";

    pub fn uri_builder(
        self_uris: Vec<&str>,
        upstream_uri: Option<&str>,
        host: String,
    ) -> Result<UriBuilder, Error> {
        let server = ServerUris::new(
            self_uris.into_iter().map(|v| v.to_string()).collect(),
            upstream_uri.map(|v| v.to_string()),
        )?;
        Ok(UriBuilder {
            repository: "repo123".to_string(),
            server: Arc::new(server),
            host,
            server_hostname: Arc::new(SERVER_HOSTNAME.to_string()),
        })
    }

    pub struct TestContextBuilder<'a> {
        fb: FacebookInit,
        repo: Repo,
        self_uris: Vec<&'a str>,
        upstream_uri: Option<String>,
        config: ServerConfig,
        host: String,
    }

    impl TestContextBuilder<'_> {
        pub fn repo(mut self, repo: Repo) -> Self {
            self.repo = repo;
            self
        }

        pub fn upstream_uri(mut self, upstream_uri: Option<String>) -> Self {
            self.upstream_uri = upstream_uri;
            self
        }

        pub fn config(mut self, config: ServerConfig) -> Self {
            self.config = config;
            self
        }

        pub fn build(self) -> Result<RepositoryRequestContext, Error> {
            let Self {
                fb,
                repo,
                self_uris,
                upstream_uri,
                config,
                host,
            } = self;

            let uri_builder = uri_builder(self_uris, upstream_uri.as_deref(), host)?;

            Ok(RepositoryRequestContext {
                ctx: CoreContext::test_mock(fb),
                repo: Arc::new(repo),
                config: Arc::new(config),
                uri_builder,
                always_wait_for_upstream: false,
                max_upload_size: None,
                client: HttpClient::Disabled,
            })
        }
    }

    impl RepositoryRequestContext {
        pub async fn test_builder(fb: FacebookInit) -> Result<TestContextBuilder<'static>, Error> {
            let repo = TestRepoFactory::new(fb)?.build().await?;

            Self::test_builder_with_repo(fb, repo)
        }

        pub fn test_builder_with_repo(
            fb: FacebookInit,
            repo: Repo,
        ) -> Result<TestContextBuilder<'static>, Error> {
            Ok(TestContextBuilder {
                fb,
                repo,
                self_uris: vec!["http://foo.com/"],
                upstream_uri: Some("http://bar.com".to_string()),
                config: ServerConfig::default(),
                host: "foo.com".to_string(),
            })
        }
    }

    fn obj() -> Result<RequestObject, Error> {
        Ok(RequestObject {
            oid: LfsSha256::from_str(ONES_HASH)?,
            size: SIZE,
        })
    }

    fn content_id() -> Result<ContentId, Error> {
        ContentId::from_str(ONES_HASH)
    }

    fn oid() -> Result<Sha256, Error> {
        Sha256::from_str(TWOS_HASH)
    }

    fn build_url(base: &str, hash: &str, size: Option<u64>) -> String {
        if let Some(size) = size {
            format!("{base}/{hash}/{size}?server_hostname={SERVER_HOSTNAME}")
        } else {
            format!("{base}/{hash}?server_hostname={SERVER_HOSTNAME}")
        }
    }

    #[test]
    fn test_basic_upload_uri() -> Result<(), Error> {
        let b = uri_builder(
            vec!["http://foo.com"],
            Some("http://bar.com"),
            "foo.com".to_string(),
        )?;
        assert_eq!(
            b.upload_uri(&obj()?)?.to_string(),
            build_url("http://foo.com/repo123/upload", ONES_HASH, Some(SIZE)),
        );
        Ok(())
    }

    #[test]
    fn test_basic_upload_uri_slash() -> Result<(), Error> {
        let b = uri_builder(
            vec!["http://foo.com/"],
            Some("http://bar.com"),
            "foo.com".to_string(),
        )?;
        assert_eq!(
            b.upload_uri(&obj()?)?.to_string(),
            build_url("http://foo.com/repo123/upload", ONES_HASH, Some(SIZE)),
        );
        Ok(())
    }

    #[test]
    fn test_prefix_upload_uri() -> Result<(), Error> {
        let b = uri_builder(
            vec!["http://foo.com/bar"],
            Some("http://bar.com"),
            "foo.com".to_string(),
        )?;
        assert_eq!(
            b.upload_uri(&obj()?)?.to_string(),
            build_url("http://foo.com/bar/repo123/upload", ONES_HASH, Some(SIZE)),
        );
        Ok(())
    }

    #[test]
    fn test_prefix_slash_upload_uri() -> Result<(), Error> {
        let b = uri_builder(
            vec!["http://foo.com/bar/"],
            Some("http://bar.com"),
            "foo.com".to_string(),
        )?;
        assert_eq!(
            b.upload_uri(&obj()?)?.to_string(),
            //format!("http://foo.com/bar/repo123/upload/{}/{}", ONES_HASH, SIZE),
            build_url("http://foo.com/bar/repo123/upload", ONES_HASH, Some(SIZE)),
        );
        Ok(())
    }

    #[test]
    fn test_basic_download_uri() -> Result<(), Error> {
        let b = uri_builder(
            vec!["http://foo.com"],
            Some("http://bar.com"),
            "foo.com".to_string(),
        )?;
        assert_eq!(
            b.download_uri(&content_id()?)?.to_string(),
            build_url("http://foo.com/repo123/download", ONES_HASH, None),
        );
        Ok(())
    }

    #[test]
    fn test_basic_download_uri_slash() -> Result<(), Error> {
        let b = uri_builder(
            vec!["http://foo.com/"],
            Some("http://bar.com"),
            "foo.com".to_string(),
        )?;
        assert_eq!(
            b.download_uri(&content_id()?)?.to_string(),
            build_url("http://foo.com/repo123/download", ONES_HASH, None),
        );
        Ok(())
    }

    #[test]
    fn test_prefix_download_uri() -> Result<(), Error> {
        let b = uri_builder(
            vec!["http://foo.com/bar"],
            Some("http://bar.com"),
            "foo.com".to_string(),
        )?;
        assert_eq!(
            b.download_uri(&content_id()?)?.to_string(),
            build_url("http://foo.com/bar/repo123/download", ONES_HASH, None),
        );
        Ok(())
    }

    #[test]
    fn test_prefix_slash_download_uri() -> Result<(), Error> {
        let b = uri_builder(
            vec!["http://foo.com/bar/"],
            Some("http://bar.com"),
            "foo.com".to_string(),
        )?;
        assert_eq!(
            b.download_uri(&content_id()?)?.to_string(),
            build_url("http://foo.com/bar/repo123/download", ONES_HASH, None),
        );
        Ok(())
    }

    #[test]
    fn test_basic_consistent_download_uri() -> Result<(), Error> {
        let b = uri_builder(
            vec!["http://foo.com"],
            Some("http://bar.com"),
            "foo.com".to_string(),
        )?;
        assert_eq!(
            b.consistent_download_uri(&content_id()?, format!("{}", oid()?))?
                .to_string(),
            format!(
                "http://foo.com/repo123/download/{}?routing={}&server_hostname={}",
                ONES_HASH, TWOS_HASH, SERVER_HOSTNAME
            ),
        );
        Ok(())
    }

    #[test]
    fn test_basic_upstream_batch_uri() -> Result<(), Error> {
        let b = uri_builder(
            vec!["http://foo.com"],
            Some("http://bar.com"),
            "foo.com".to_string(),
        )?;
        assert_eq!(
            b.upstream_batch_uri()?.map(|uri| uri.to_string()),
            Some("http://bar.com/objects/batch".to_string()),
        );
        Ok(())
    }

    #[test]
    fn test_basic_upstream_batch_uri_slash() -> Result<(), Error> {
        let b = uri_builder(
            vec!["http://foo.com/"],
            Some("http://bar.com"),
            "foo.com".to_string(),
        )?;
        assert_eq!(
            b.upstream_batch_uri()?.map(|uri| uri.to_string()),
            Some("http://bar.com/objects/batch".to_string()),
        );
        Ok(())
    }

    #[test]
    fn test_prefix_upstream_batch_uri() -> Result<(), Error> {
        let b = uri_builder(
            vec!["http://foo.com"],
            Some("http://bar.com/foo"),
            "foo.com".to_string(),
        )?;
        assert_eq!(
            b.upstream_batch_uri()?.map(|uri| uri.to_string()),
            Some("http://bar.com/foo/objects/batch".to_string()),
        );
        Ok(())
    }

    #[test]
    fn test_prefix_slash_upstream_batch_uri() -> Result<(), Error> {
        let b = uri_builder(
            vec!["http://foo.com"],
            Some("http://bar.com/foo/"),
            "foo.com".to_string(),
        )?;
        assert_eq!(
            b.upstream_batch_uri()?.map(|uri| uri.to_string()),
            Some("http://bar.com/foo/objects/batch".to_string()),
        );
        Ok(())
    }

    #[fbinit::test]
    async fn test_acl_check_no_certificates(fb: FacebookInit) -> Result<(), Error> {
        let aclchecker = AlwaysAllowRepoPermissionChecker::new();
        let repo: Repo = test_repo_factory::TestRepoFactory::new(fb)?
            .with_permission_checker(Arc::new(aclchecker))
            .build()
            .await?;

        let ctx = CoreContext::test_mock(fb);
        // No certificates + ACL enforcement off = free pass
        acl_check(&ctx, &repo, false, LfsMethod::Download).await?;

        Ok(())
    }

    #[fbinit::test]
    async fn test_acl_check_action(fb: FacebookInit) -> Result<(), Error> {
        let mut aclchecker = MockRepoPermissionChecker::new();
        aclchecker
            .expect_check_if_read_access_allowed()
            .return_const(true)
            .times(1);
        aclchecker
            .expect_check_if_draft_access_allowed_with_tunable_enforcement()
            .return_const(true)
            .times(1);

        let repo: Repo = test_repo_factory::TestRepoFactory::new(fb)?
            .with_permission_checker(Arc::new(aclchecker))
            .build()
            .await?;

        let ctx = CoreContext::test_mock(fb);
        acl_check(&ctx, &repo, false, LfsMethod::Download).await?;
        acl_check(&ctx, &repo, false, LfsMethod::Upload).await?;

        Ok(())
    }

    #[test]
    fn test_host_maybe_port_to_host() -> Result<(), Error> {
        assert_eq!(host_maybe_port_to_host("example.com")?, "example.com");
        assert_eq!(host_maybe_port_to_host("example.com:90")?, "example.com");
        assert_eq!(host_maybe_port_to_host("[::1]")?, "[::1]");
        assert_eq!(host_maybe_port_to_host("[::1]:80")?, "[::1]");
        assert_eq!(host_maybe_port_to_host("127.0.0.1:80")?, "127.0.0.1");
        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.
 */

#![recursion_limit = "256"]
#![feature(never_type)]
#![feature(let_chains)]

use std::fs::File;
use std::io::Write;
use std::net::ToSocketAddrs;
use std::sync::atomic::AtomicBool;
use std::sync::atomic::Ordering;
use std::sync::Arc;

use anyhow::anyhow;
use anyhow::bail;
use anyhow::Context;
use anyhow::Error;
use anyhow::Result;
use cached_config::ConfigHandle;
use clap::Parser;
use cloned::cloned;
use cmdlib_caching::CachelibSettings;
use connection_security_checker::ConnectionSecurityChecker;
use fbinit::FacebookInit;
use filestore::FilestoreConfig;
use futures::channel::oneshot;
use futures::future::try_select;
use futures::pin_mut;
use futures::TryFutureExt;
use gotham_ext::handler::MononokeHttpHandler;
use gotham_ext::middleware::LoadMiddleware;
use gotham_ext::middleware::LogMiddleware;
use gotham_ext::middleware::MetadataMiddleware;
use gotham_ext::middleware::PostResponseMiddleware;
use gotham_ext::middleware::ScubaMiddleware;
use gotham_ext::middleware::ServerIdentityMiddleware;
use gotham_ext::middleware::TimerMiddleware;
use gotham_ext::middleware::TlsSessionDataMiddleware;
use gotham_ext::serve;
use hyper::header::HeaderValue;
use metaconfig_types::RepoConfig;
use metaconfig_types::ShardedService;
use mononoke_app::args::parse_config_spec_to_path;
use mononoke_app::args::ReadonlyArgs;
use mononoke_app::args::RepoFilterAppExtension;
use mononoke_app::args::ShutdownTimeoutArgs;
use mononoke_app::args::TLSArgs;
use mononoke_app::fb303::AliveService;
use mononoke_app::fb303::Fb303AppExtension;
use mononoke_app::MononokeApp;
use mononoke_app::MononokeAppBuilder;
use mononoke_repos::MononokeRepos;
use qps::Qps;
use repo_blobstore::RepoBlobstore;
use repo_identity::RepoIdentity;
use repo_permission_checker::RepoPermissionChecker;
use slog::info;
use tokio::net::TcpListener;

use crate::lfs_server_context::LfsServerContext;
use crate::lfs_server_context::ServerUris;
use crate::middleware::OdsMiddleware;
use crate::middleware::RequestContextMiddleware;
use crate::scuba::LfsScubaHandler;
use crate::service::build_router;

mod batch;
mod config;
mod download;
mod errors;
mod git_upload;
mod lfs_server_context;
mod middleware;
mod popularity;
mod scuba;
mod service;
mod upload;
mod util;

const SERVICE_NAME: &str = "mononoke_lfs_server";

// Used to determine how many entries are in cachelib's HashTable. A smaller
// object size results in more entries and possibly higher idle memory usage.
// More info: https://fburl.com/wiki/i78i3uzk
const CACHE_OBJECT_SIZE: usize = 256 * 1024;

#[facet::container]
#[derive(Clone)]
pub struct Repo {
    #[facet]
    repo_identity: RepoIdentity,

    #[init(repo_identity.name().to_string())]
    name: String,

    #[facet]
    repo_config: RepoConfig,

    #[facet]
    filestore_config: FilestoreConfig,

    #[facet]
    repo_blobstore: RepoBlobstore,

    #[facet]
    repo_permission_checker: dyn RepoPermissionChecker,
}

/// Mononoke LFS Server
#[derive(Parser)]
struct LfsServerArgs {
    /// Shutdown timeout args for this service
    #[clap(flatten)]
    shutdown_timeout_args: ShutdownTimeoutArgs,
    /// TLS parameters for this service
    #[clap(flatten)]
    tls_params: TLSArgs,
    /// The host to listen on locally
    #[clap(long, default_value = "127.0.0.1")]
    listen_host: String,
    /// The port to listen on locally
    #[clap(long, default_value = "8001")]
    listen_port: String,
    /// Path for file in which to write the bound tcp address in rust std::net::SocketAddr format
    #[clap(long)]
    bound_address_file: Option<String>,
    /// The base URLs for this server
    #[clap(value_delimiter = ',')]
    self_urls: Vec<String>,
    /// The base URL for an upstream server
    #[clap(long)]
    upstream_url: Option<String>,
    /// Whether to always wait for an upstream response (primarily useful in testing)
    #[clap(long)]
    always_wait_for_upstream: bool,
    /// Path to config in configerator
    #[clap(long)]
    live_config: Option<String>,
    /// Whether or not to use test-friendly logging
    #[clap(long)]
    test_friendly_logging: bool,
    // A file to which to log TLS session data, including master secrets.
    // Use this for debugging with tcpdump.
    // Note that this compromises the secrecy of TLS sessions.
    #[clap(long)]
    tls_session_data_log_file: Option<String>,
    /// Whether to enable Mononoke-specific small git blob uploads
    #[clap(long)]
    git_blob_upload_allowed: bool,
    /// A limit (in bytes) to enforce for uploads.
    #[clap(long)]
    max_upload_size: Option<u64>,
    #[clap(flatten)]
    readonly: ReadonlyArgs,
    /// Path to config
    #[clap(long)]
    cslb_config: Option<String>,
}

#[derive(Clone)]
pub struct LfsRepos {
    pub(crate) repos: Arc<MononokeRepos<Repo>>,
}

impl LfsRepos {
    pub(crate) async fn new(app: &MononokeApp) -> Result<Self> {
        let repos_mgr = app
            .open_managed_repos(Some(ShardedService::LargeFilesService))
            .await?;
        let repos = repos_mgr.repos().clone();
        Ok(Self { repos })
    }

    pub(crate) fn get(&self, repo_name: &str) -> Option<Arc<Repo>> {
        self.repos.get_by_name(repo_name)
    }
}

#[fbinit::main]
fn main(fb: FacebookInit) -> Result<(), Error> {
    let cachelib_settings = CachelibSettings {
        expected_item_size_bytes: Some(CACHE_OBJECT_SIZE),
        ..Default::default()
    };

    let app = MononokeAppBuilder::new(fb)
        .with_app_extension(Fb303AppExtension {})
        .with_app_extension(RepoFilterAppExtension {})
        .with_default_cachelib_settings(cachelib_settings)
        .build::<LfsServerArgs>()?;

    let args: LfsServerArgs = app.args()?;

    let logger = app.logger().clone();
    let config_store = app.config_store();
    let acl_provider = app.environment().acl_provider.clone();

    let listen_host = args.listen_host.clone();
    let listen_port = args.listen_port.clone();
    let bound_addr_path = args.bound_address_file.clone();

    let git_blob_upload_allowed = args.git_blob_upload_allowed;

    let addr = format!("{}:{}", listen_host, listen_port);

    let tls_certificate = args.tls_params.tls_certificate.clone();
    let tls_private_key = args.tls_params.tls_private_key.clone();
    let tls_ca = args.tls_params.tls_ca.clone();
    let tls_ticket_seeds = args.tls_params.tls_ticket_seeds.clone();

    let tls_acceptor = match (tls_certificate, tls_private_key, tls_ca, tls_ticket_seeds) {
        (Some(tls_certificate), Some(tls_private_key), Some(tls_ca), tls_ticket_seeds) => {
            let acceptor = secure_utils::SslConfig::new(
                tls_ca,
                tls_certificate,
                tls_private_key,
                tls_ticket_seeds,
            )
            .build_tls_acceptor(logger.clone())?;
            Some(acceptor)
        }
        (None, None, None, None) => None,
        _ => bail!("TLS flags must be passed together"),
    };

    let tls_session_data_log = args.tls_session_data_log_file.clone();

    let scuba_logger = app.environment().scuba_sample_builder.clone();

    let will_exit = Arc::new(AtomicBool::new(false));

    let config_handle = match &args.live_config {
        Some(spec) => config_store.get_config_handle_DEPRECATED(parse_config_spec_to_path(spec)?),
        None => Ok(ConfigHandle::default()),
    };

    let config_handle = config_handle.context(Error::msg("Failed to load configuration"))?;

    let cslb_config = args.cslb_config;

    let qps = match cslb_config {
        Some(config) => {
            Some(Qps::new(fb, config, config_store).with_context(|| "Failed to initialize QPS")?)
        }
        None => None,
    };
    let max_upload_size: Option<u64> = args.max_upload_size;

    let self_urls = args.self_urls;
    let upstream_url = args.upstream_url;
    let always_wait_for_upstream = args.always_wait_for_upstream;
    let log_middleware = if args.test_friendly_logging {
        LogMiddleware::test_friendly()
    } else {
        LogMiddleware::slog(logger.clone())
    };

    app.start_monitoring(SERVICE_NAME, AliveService)?;
    app.start_stats_aggregation()?;

    let common = &app.repo_configs().common;
    let internal_identity = common.internal_identity.clone();
    let (shutdown_tx, shutdown_rx) = oneshot::channel::<()>();
    let server = {
        cloned!(acl_provider, common, logger, will_exit);
        move |app| async move {
            let repos = LfsRepos::new(&app)
                .await
                .context(Error::msg("Error opening repos"))?;

            let addr = addr
                .to_socket_addrs()
                .context(Error::msg("Invalid Listener Address"))?
                .next()
                .ok_or_else(|| Error::msg("Invalid Socket Address"))?;

            let listener = TcpListener::bind(&addr)
                .await
                .context(Error::msg("Could not start TCP listener"))?;

            // We use the listen_host rather than the ip of listener.local_addr()
            // because the certs user passed will be referencing listen_host
            let bound_addr = format!("{}:{}", listen_host, listener.local_addr()?.port());

            // For tests we use one empty string self_url, map it to None
            let self_urls =
                if self_urls.is_empty() || (self_urls.len() == 1 && self_urls[0].is_empty()) {
                    None
                } else {
                    Some(self_urls)
                };

            let self_urls = if let Some(self_urls) = self_urls {
                self_urls
            } else {
                let protocol = if tls_acceptor.is_some() {
                    "https://"
                } else {
                    "http://"
                };
                vec![protocol.to_owned() + &bound_addr]
            };

            let server_uris = ServerUris::new(self_urls, upstream_url)?;

            let ctx = LfsServerContext::new(
                repos,
                server_uris,
                always_wait_for_upstream,
                max_upload_size,
                will_exit,
                config_handle.clone(),
                logger.clone(),
                qps,
            )?;
            let enforce_authentication = ctx.get_config().enforce_authentication();

            let router = build_router(fb, ctx, git_blob_upload_allowed);

            let capture_session_data = tls_session_data_log.is_some();

            let handler = MononokeHttpHandler::builder()
                .add(TlsSessionDataMiddleware::new(tls_session_data_log)?)
                .add(MetadataMiddleware::new(
                    fb,
                    logger.clone(),
                    internal_identity,
                ))
                .add(PostResponseMiddleware::with_config(config_handle))
                .add(RequestContextMiddleware::new(
                    fb,
                    logger.clone(),
                    enforce_authentication,
                    args.readonly.readonly,
                ))
                .add(LoadMiddleware::new())
                .add(log_middleware)
                .add(ServerIdentityMiddleware::new(HeaderValue::from_static(
                    "mononoke-lfs",
                )))
                .add(<ScubaMiddleware<LfsScubaHandler>>::new(scuba_logger))
                .add(OdsMiddleware::new())
                .add(TimerMiddleware::new())
                .build(router);

            info!(&logger, "Listening on {}", bound_addr);

            // Write out the bound address if requested, this is helpful in tests when using automatic binding with :0
            if let Some(bound_addr_path) = bound_addr_path {
                let mut writer = File::create(bound_addr_path)?;
                writer.write_all(bound_addr.as_bytes())?;
                writer.write_all(b"\n")?;
            }

            let serve = async move {
                if let Some(tls_acceptor) = tls_acceptor {
                    let connection_security_checker =
                        ConnectionSecurityChecker::new(acl_provider.as_ref(), &common).await?;

                    serve::https(
                        logger,
                        listener,
                        tls_acceptor,
                        capture_session_data,
                        connection_security_checker,
                        handler,
                    )
                    .await
                } else {
                    serve::http(logger, listener, handler).await
                }
            };
            pin_mut!(serve);
            try_select(
                serve,
                shutdown_rx.map_err(|err| anyhow!("Cancelled channel: {}", err)),
            )
            .await
            .map_err(|e| futures::future::Either::factor_first(e).0)?;
            Ok(())
        }
    };

    app.run_until_terminated(
        server,
        move || will_exit.store(true, Ordering::Relaxed),
        args.shutdown_timeout_args.shutdown_grace_period,
        async move {
            let _ = shutdown_tx.send(());
            // Currently we kill off in-flight requests as soon as we've closed the listener.
            // If this is a problem in prod, this would be the point at which to wait
            // for all connections to shut down.
            // To do this properly, we'd need to track the `Connection` futures that Gotham
            // gets from Hyper, tell them to gracefully shutdown, then wait for them to complete
        },
        args.shutdown_timeout_args.shutdown_timeout,
    )?;

    info!(&logger, "Exiting...");
    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::num::NonZeroU16;
use std::time::Duration;

use anyhow::Error;
use fbinit::FacebookInit;
use stats::prelude::*;
use time_window_counter::BoxGlobalTimeWindowCounter;
use time_window_counter::GlobalTimeWindowCounterBuilder;
use tokio::time;

use crate::batch::InternalObject;
use crate::config::ConsistentRoutingRingMode;
use crate::config::ObjectPopularity;
use crate::lfs_server_context::RepositoryRequestContext;

define_stats! {
    prefix = "mononoke.lfs.popularity";
    success: timeseries(Rate, Sum),
    error: timeseries(Rate, Sum),
    timeout: timeseries(Rate, Sum),
}

const OBJECT_POPULARITY_TIMEOUT: Duration = Duration::from_millis(10);

pub trait PopularityBuilder {
    fn build(
        &self,
        fb: FacebookInit,
        category: impl AsRef<str>,
        key: impl AsRef<str>,
        min_time_window: u32,
        max_time_window: u32,
    ) -> BoxGlobalTimeWindowCounter;
}

impl PopularityBuilder for GlobalTimeWindowCounterBuilder {
    fn build(
        &self,
        fb: FacebookInit,
        category: impl AsRef<str>,
        key: impl AsRef<str>,
        min_time_window: u32,
        max_time_window: u32,
    ) -> BoxGlobalTimeWindowCounter {
        GlobalTimeWindowCounterBuilder::build(fb, category, key, min_time_window, max_time_window)
    }
}

async fn increment_and_fetch_object_popularity<B: PopularityBuilder>(
    ctx: &RepositoryRequestContext,
    obj: InternalObject,
    config: &ObjectPopularity,
    builder: B,
) -> Result<u64, Error> {
    let key = format!("{}/{}", &ctx.repo.name, obj.id());
    let ctr = builder.build(
        ctx.ctx.fb,
        &config.category,
        &key,
        config.window,
        config.window,
    );

    ctr.bump(obj.download_size() as f64);
    let v = ctr.get(config.window).await? as i64;
    let v = v.try_into()?;
    Ok(v)
}

/// None -> no consistent routing - routing randomly to all tasks
/// Some(n) -> consistently route to n tasks
///
/// By default, route all blobs consistently to one task - Some(1)
pub async fn consistent_routing<B: PopularityBuilder>(
    ctx: &RepositoryRequestContext,
    obj: InternalObject,
    builder: B,
) -> Option<NonZeroU16> {
    let config = match ctx.config.object_popularity() {
        Some(r) => r,
        None => return Some(NonZeroU16::new(1).unwrap()),
    };

    let popularity = time::timeout(
        OBJECT_POPULARITY_TIMEOUT,
        increment_and_fetch_object_popularity(ctx, obj, config, builder),
    )
    .await;

    match popularity {
        Ok(Ok(popularity)) => {
            STATS::success.add_value(1);
            // The bigger the popularity, the bigger tasks_per_content.
            // Thresholds are sorted in increasing order.
            // [pop_thresh = 0,    tasks_per_content = 1]
            // [pop_thresh = 100,  tasks_per_content = 10]
            // [pop_thresh = 1000, tasks_per_content = All]
            for ring in config.thresholds.iter().rev() {
                if popularity >= ring.threshold {
                    match ring.mode {
                        ConsistentRoutingRingMode::Num { tasks_per_content } => {
                            return Some(tasks_per_content);
                        }
                        ConsistentRoutingRingMode::All => return None,
                    }
                }
            }
            // Shouldn't be reached as we require to have a ring with popularity
            // threshold 0. Let's default to 1 anyway.
        }
        Ok(Err(_)) => {
            // Errored
            STATS::error.add_value(1);
        }
        Err(..) => {
            // Timed out
            STATS::timeout.add_value(1);
        }
    };

    // Default to one task per blob
    Some(NonZeroU16::new(1).unwrap())
}

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

    use async_trait::async_trait;
    use futures::future;
    use mononoke_types_mocks::contentid::ONES_CTID;
    use mononoke_types_mocks::hash::ONES_SHA256;
    use time_window_counter::GlobalTimeWindowCounter;

    use super::*;
    use crate::config::ConsistentRoutingRing;
    use crate::config::ConsistentRoutingRingMode;
    use crate::config::ObjectPopularity;
    use crate::config::ServerConfig;

    fn dummy(size: impl Into<Option<u64>>) -> InternalObject {
        InternalObject::new(ONES_CTID, ONES_SHA256, size.into())
    }

    #[derive(Clone, Default)]
    struct DummyCounter(Arc<AtomicU64>);

    impl PopularityBuilder for DummyCounter {
        fn build(
            &self,
            _: FacebookInit,
            _: impl AsRef<str>,
            _: impl AsRef<str>,
            _: u32,
            _: u32,
        ) -> BoxGlobalTimeWindowCounter {
            Box::new(self.clone())
        }
    }

    #[async_trait]
    impl GlobalTimeWindowCounter for DummyCounter {
        async fn get(&self, _time_window: u32) -> Result<f64, Error> {
            Ok(self.0.load(Ordering::Relaxed) as f64)
        }

        fn bump(&self, value: f64) {
            self.0.fetch_add(value as u64, Ordering::Relaxed);
        }
    }

    #[derive(Copy, Clone)]
    struct TimeoutCounter;

    impl PopularityBuilder for TimeoutCounter {
        fn build(
            &self,
            _: FacebookInit,
            _: impl AsRef<str>,
            _: impl AsRef<str>,
            _: u32,
            _: u32,
        ) -> BoxGlobalTimeWindowCounter {
            Box::new(*self)
        }
    }

    #[async_trait]
    impl GlobalTimeWindowCounter for TimeoutCounter {
        async fn get(&self, _time_window: u32) -> Result<f64, Error> {
            future::pending().await
        }

        fn bump(&self, _: f64) {
            // noop
        }
    }

    #[derive(Copy, Clone)]
    struct ErrorCounter;

    impl PopularityBuilder for ErrorCounter {
        fn build(
            &self,
            _: FacebookInit,
            _: impl AsRef<str>,
            _: impl AsRef<str>,
            _: u32,
            _: u32,
        ) -> BoxGlobalTimeWindowCounter {
            Box::new(*self)
        }
    }

    #[async_trait]
    impl GlobalTimeWindowCounter for ErrorCounter {
        async fn get(&self, _time_window: u32) -> Result<f64, Error> {
            Err(Error::msg("error counter"))
        }

        fn bump(&self, _: f64) {
            // noop
        }
    }

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

        assert!(consistent_routing(&ctx, dummy(None), ctr).await.is_some());

        Ok(())
    }

    #[fbinit::test]
    async fn test_popularity(fb: FacebookInit) -> Result<(), Error> {
        let mut config = ServerConfig::default();
        *config.object_popularity_mut() = Some(ObjectPopularity {
            category: "foo".into(),
            window: 100,
            thresholds: vec![
                ConsistentRoutingRing {
                    threshold: 0,
                    mode: ConsistentRoutingRingMode::Num {
                        tasks_per_content: std::num::NonZeroU16::new(1).unwrap(),
                    },
                },
                ConsistentRoutingRing {
                    threshold: 10,
                    mode: ConsistentRoutingRingMode::All,
                },
            ],
        });

        let ctx = RepositoryRequestContext::test_builder(fb)
            .await?
            .config(config)
            .build()?;
        let ctr = DummyCounter::default();

        assert!(
            consistent_routing(&ctx, dummy(4), ctr.clone())
                .await
                .is_some()
        );

        assert!(
            consistent_routing(&ctx, dummy(4), ctr.clone())
                .await
                .is_some()
        );

        assert!(
            consistent_routing(&ctx, dummy(4), ctr.clone())
                .await
                .is_none()
        );

        Ok(())
    }

    #[fbinit::test]
    async fn test_popularity_rings(fb: FacebookInit) -> Result<(), Error> {
        let mut config = ServerConfig::default();
        *config.object_popularity_mut() = Some(ObjectPopularity {
            category: "foo".into(),
            window: 100,
            thresholds: vec![
                ConsistentRoutingRing {
                    threshold: 0,
                    mode: ConsistentRoutingRingMode::Num {
                        tasks_per_content: std::num::NonZeroU16::new(1).unwrap(),
                    },
                },
                ConsistentRoutingRing {
                    threshold: 10,
                    mode: ConsistentRoutingRingMode::Num {
                        tasks_per_content: std::num::NonZeroU16::new(10).unwrap(),
                    },
                },
                ConsistentRoutingRing {
                    threshold: 15,
                    mode: ConsistentRoutingRingMode::All,
                },
            ],
        });

        let ctx = RepositoryRequestContext::test_builder(fb)
            .await?
            .config(config)
            .build()?;
        let ctr = DummyCounter::default();

        assert_eq!(
            consistent_routing(&ctx, dummy(4), ctr.clone()).await,
            Some(std::num::NonZeroU16::new(1).unwrap())
        );

        assert_eq!(
            consistent_routing(&ctx, dummy(4), ctr.clone()).await,
            Some(std::num::NonZeroU16::new(1).unwrap())
        );

        assert_eq!(
            consistent_routing(&ctx, dummy(4), ctr.clone()).await,
            Some(std::num::NonZeroU16::new(10).unwrap())
        );

        assert!(
            consistent_routing(&ctx, dummy(4), ctr.clone())
                .await
                .is_none()
        );

        Ok(())
    }

    #[fbinit::test]
    async fn test_timeout(fb: FacebookInit) -> Result<(), Error> {
        let mut config = ServerConfig::default();
        *config.object_popularity_mut() = Some(ObjectPopularity {
            category: "foo".into(),
            window: 100,
            thresholds: vec![
                ConsistentRoutingRing {
                    threshold: 0,
                    mode: ConsistentRoutingRingMode::Num {
                        tasks_per_content: std::num::NonZeroU16::new(1).unwrap(),
                    },
                },
                ConsistentRoutingRing {
                    threshold: 10,
                    mode: ConsistentRoutingRingMode::All,
                },
            ],
        });

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

        assert!(
            consistent_routing(&ctx, dummy(None), TimeoutCounter)
                .await
                .is_some()
        );

        Ok(())
    }

    #[fbinit::test]
    async fn test_error(fb: FacebookInit) -> Result<(), Error> {
        let mut config = ServerConfig::default();
        *config.object_popularity_mut() = Some(ObjectPopularity {
            category: "foo".into(),
            window: 100,
            thresholds: vec![
                ConsistentRoutingRing {
                    threshold: 0,
                    mode: ConsistentRoutingRingMode::Num {
                        tasks_per_content: std::num::NonZeroU16::new(1).unwrap(),
                    },
                },
                ConsistentRoutingRing {
                    threshold: 10,
                    mode: ConsistentRoutingRingMode::All,
                },
            ],
        });

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

        assert!(
            consistent_routing(&ctx, dummy(None), ErrorCounter)
                .await
                .is_some()
        );

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

use clientinfo::ClientInfo;
use clientinfo::CLIENT_INFO_HEADER;
use gotham::state::State;
use gotham_ext::middleware::PostResponseInfo;
use gotham_ext::middleware::ScubaHandler;
use scuba_ext::MononokeScubaSampleBuilder;

use crate::middleware::RequestContext;
use crate::util::read_header_value_ignore_err;

struct ClientInfoHeader(ClientInfo);

impl FromStr for ClientInfoHeader {
    type Err = serde_json::Error;

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        let client_info = serde_json::from_str(s)?;
        Ok(Self(client_info))
    }
}

#[derive(Copy, Clone, Debug)]
pub enum LfsScubaKey {
    /// The repository this request was for.
    Repository,
    /// The method this request matched for in our handlers.
    Method,
    /// If an error was encountered during processing, the error message.
    ErrorMessage,
    /// Total count fo errors that occurred during processing.
    ErrorCount,
    /// The order in which the response to a batch request was produced.
    BatchOrder,
    /// The number of objects in a batch request
    BatchObjectCount,
    /// The objects that could not be serviced by this LFS server in a batch request
    BatchInternalMissingBlobs,
    /// Timing checkpoints in batch requests
    BatchRequestContextReadyUs,
    BatchRequestReceivedUs,
    BatchRequestParsedUs,
    BatchResponseReadyUs,
    /// Whether the upload was a sync
    UploadSync,
    /// The actual size of the content being sent
    DownloadContentSize,
    /// The attempt information reported by the client
    ClientAttempt,
    ClientAttemptsLeft,
    ClientThrottleAttemptsLeft,
    /// Fields of ClientInfo
    SandcastleNonce,
    SandcastleAlias,
    ClientTwJob,
    ClientTwTask,
}

impl AsRef<str> for LfsScubaKey {
    fn as_ref(&self) -> &'static str {
        use LfsScubaKey::*;

        match self {
            Repository => "repository",
            Method => "method",
            ErrorMessage => "error_msg",
            ErrorCount => "error_count",
            BatchOrder => "batch_order",
            BatchObjectCount => "batch_object_count",
            BatchInternalMissingBlobs => "batch_internal_missing_blobs",
            BatchRequestContextReadyUs => "batch_context_ready_us",
            BatchRequestReceivedUs => "batch_request_received_us",
            BatchRequestParsedUs => "batch_request_parsed_us",
            BatchResponseReadyUs => "batch_response_ready_us",
            UploadSync => "upload_sync",
            DownloadContentSize => "download_content_size",
            ClientAttempt => "client_attempt",
            ClientAttemptsLeft => "client_attempts_left",
            ClientThrottleAttemptsLeft => "client_throttle_attempts_left",
            SandcastleNonce => "sandcastle_nonce",
            SandcastleAlias => "sandcastle_alias",
            ClientTwJob => "client_tw_job",
            ClientTwTask => "client_tw_task",
        }
    }
}

impl From<LfsScubaKey> for String {
    fn from(k: LfsScubaKey) -> String {
        k.as_ref().to_string()
    }
}

#[derive(Clone)]
pub struct LfsScubaHandler {
    ctx: Option<RequestContext>,
    client_attempt: Option<u64>,
    client_attempts_left: Option<u64>,
    client_throttle_attempts_left: Option<u64>,
    client_info: Option<ClientInfo>,
}

impl ScubaHandler for LfsScubaHandler {
    fn from_state(state: &State) -> Self {
        let client_attempt = read_header_value_ignore_err(state, "X-Attempt");

        let client_attempts_left = read_header_value_ignore_err(state, "X-Attempts-Left");

        let client_throttle_attempts_left =
            read_header_value_ignore_err(state, "X-Throttle-Attempts-Left");

        let client_info: Option<ClientInfo> =
            read_header_value_ignore_err(state, CLIENT_INFO_HEADER)
                .map(|ci: ClientInfoHeader| ci.0);

        Self {
            ctx: state.try_borrow::<RequestContext>().cloned(),
            client_attempt,
            client_attempts_left,
            client_throttle_attempts_left,
            client_info,
        }
    }

    fn populate_scuba(self, info: &PostResponseInfo, scuba: &mut MononokeScubaSampleBuilder) {
        scuba.add_opt(LfsScubaKey::ClientAttempt, self.client_attempt);
        scuba.add_opt(LfsScubaKey::ClientAttemptsLeft, self.client_attempts_left);
        scuba.add_opt(
            LfsScubaKey::ClientThrottleAttemptsLeft,
            self.client_throttle_attempts_left,
        );

        if let Some(ctx) = self.ctx {
            if let Some(repository) = ctx.repository {
                scuba.add(LfsScubaKey::Repository, repository.as_ref());
            }

            if let Some(method) = ctx.method {
                scuba.add(LfsScubaKey::Method, method.to_string());
            }

            if let Some(err) = info.first_error() {
                scuba.add(LfsScubaKey::ErrorMessage, format!("{:?}", err));
            }

            scuba.add(LfsScubaKey::ErrorCount, info.error_count());

            ctx.ctx.perf_counters().insert_perf_counters(scuba);
        }

        if let Some(client_info) = self.client_info {
            scuba.add_opt(
                LfsScubaKey::SandcastleNonce,
                client_info.fb.sandcastle_nonce(),
            );

            scuba.add_opt(
                LfsScubaKey::SandcastleAlias,
                client_info.fb.sandcastle_alias(),
            );

            scuba.add_opt(LfsScubaKey::ClientTwJob, client_info.fb.tw_job());

            scuba.add_opt(LfsScubaKey::ClientTwTask, client_info.fb.tw_task());
        }
    }
}


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

use anyhow::Context;
use anyhow::Error;
use bytes::Bytes;
use filestore::Alias;
use filestore::FetchKey;
use filestore::FilestoreConfigRef;
use filestore::StoreRequest;
use futures::channel::mpsc::channel;
use futures::SinkExt;
use futures::Stream;
use futures::StreamExt;
use futures::TryStreamExt;
use futures_util::try_join;
use gotham::state::FromState;
use gotham::state::State;
use gotham_derive::StateData;
use gotham_derive::StaticResponseExtender;
use gotham_ext::error::HttpError;
use gotham_ext::middleware::HttpScubaKey;
use gotham_ext::middleware::ScubaMiddlewareState;
use gotham_ext::response::EmptyBody;
use gotham_ext::response::TryIntoResponse;
use http::header::CONTENT_LENGTH;
use hyper::Body;
use hyper::Request;
use lfs_protocol::ObjectAction;
use lfs_protocol::ObjectStatus;
use lfs_protocol::Operation;
use lfs_protocol::RequestBatch;
use lfs_protocol::RequestObject;
use lfs_protocol::ResponseBatch;
use lfs_protocol::Sha256 as LfsSha256;
use lfs_protocol::Transfer;
use mononoke_types::hash::Sha256;
use repo_blobstore::RepoBlobstoreRef;
use serde::Deserialize;
use stats::prelude::*;

use crate::errors::ErrorKind;
use crate::lfs_server_context::RepositoryRequestContext;
use crate::middleware::LfsMethod;
use crate::scuba::LfsScubaKey;
use crate::util::read_header_value;

define_stats! {
    prefix ="mononoke.lfs.upload";
    upstream_uploads: timeseries(Rate, Sum),
    upstream_success: timeseries(Rate, Sum),
    internal_uploads: timeseries(Rate, Sum),
    internal_success: timeseries(Rate, Sum),
    size_bytes: histogram(1_500_000, 0, 150_000_000, Average, Sum, Count; P 5; P 25; P 50; P 75; P 95; P 97; P 99),
}

// Small buffers for Filestore & Dewey
const BUFFER_SIZE: usize = 5;

mod closeable_sender {
    use std::pin::Pin;

    use futures::channel::mpsc::SendError;
    use futures::channel::mpsc::Sender;
    use futures::sink::Sink;
    use futures::task::Context;
    use futures::task::Poll;
    use pin_project::pin_project;

    #[pin_project]
    pub struct CloseableSender<T> {
        #[pin]
        inner: Sender<T>,
    }

    impl<T> CloseableSender<T> {
        pub fn new(inner: Sender<T>) -> Self {
            Self { inner }
        }
    }

    impl<T> Sink<T> for CloseableSender<T> {
        type Error = SendError;

        fn poll_ready(
            self: Pin<&mut Self>,
            ctx: &mut Context<'_>,
        ) -> Poll<Result<(), Self::Error>> {
            match self.project().inner.poll_ready(ctx) {
                Poll::Ready(Err(e)) if e.is_disconnected() => Poll::Ready(Ok(())),
                x => x,
            }
        }

        fn start_send(self: Pin<&mut Self>, msg: T) -> Result<(), Self::Error> {
            match self.project().inner.start_send(msg) {
                Err(e) if e.is_disconnected() => Ok(()),
                x => x,
            }
        }

        fn poll_flush(
            self: Pin<&mut Self>,
            ctx: &mut Context<'_>,
        ) -> Poll<Result<(), Self::Error>> {
            self.project().inner.poll_flush(ctx)
        }

        fn poll_close(
            self: Pin<&mut Self>,
            ctx: &mut Context<'_>,
        ) -> Poll<Result<(), Self::Error>> {
            self.project().inner.poll_close(ctx)
        }
    }
}

use closeable_sender::CloseableSender;

// NOTE: We don't deserialize things beyond a String form, in order to report errors in our
// controller, not in routing.
#[derive(Deserialize, StateData, StaticResponseExtender)]
pub struct UploadParams {
    repository: String,
    oid: String,
    size: String,
}

fn find_actions(
    batch: ResponseBatch,
    object: &RequestObject,
) -> Result<HashMap<Operation, ObjectAction>, Error> {
    let ResponseBatch { transfer, objects } = batch;

    match transfer {
        Transfer::Basic => objects
            .into_iter()
            .find(|o| o.object == *object)
            .ok_or_else(|| ErrorKind::UpstreamMissingObject(*object).into())
            .and_then(|o| match o.status {
                ObjectStatus::Ok {
                    authenticated: false,
                    actions,
                } => Ok(actions),
                _ => Err(ErrorKind::UpstreamInvalidObject(o).into()),
            }),
        Transfer::Unknown => Err(ErrorKind::UpstreamInvalidTransfer.into()),
    }
}

async fn internal_upload<S>(
    ctx: &RepositoryRequestContext,
    oid: Sha256,
    size: u64,
    data: S,
) -> Result<(), Error>
where
    S: Stream<Item = Result<Bytes, Error>> + Unpin + Send + 'static,
{
    STATS::internal_uploads.add_value(1);

    filestore::store(
        ctx.repo.repo_blobstore(),
        *ctx.repo.filestore_config(),
        &ctx.ctx,
        &StoreRequest::with_sha256(size, oid),
        data,
    )
    .await
    .context(ErrorKind::FilestoreWriteFailure)?;

    STATS::internal_success.add_value(1);

    Ok(())
}

async fn upstream_upload<S>(
    ctx: &RepositoryRequestContext,
    oid: Sha256,
    size: u64,
    data: S,
) -> Result<(), Error>
where
    S: Stream<Item = Result<Bytes, Error>> + Unpin + Send + 'static,
{
    let object = RequestObject {
        oid: LfsSha256(oid.into_inner()),
        size,
    };

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

    let res = ctx
        .upstream_batch(&batch)
        .await
        .context(ErrorKind::UpstreamBatchError)?;

    let batch = match res {
        Some(res) => res,
        None => {
            // We have no upstream: discard this copy.
            return Ok(());
        }
    };

    let action = find_actions(batch, &object)?.remove(&Operation::Upload);

    if let Some(action) = action {
        // TODO: We are discarding expiry and headers here. We probably shouldn't.
        // TODO: We are discarding verify actions.
        STATS::upstream_uploads.add_value(1);
        let ObjectAction { href, .. } = action;

        let body = Body::wrap_stream(data);
        let req = Request::put(href)
            .header("Content-Length", &size.to_string())
            .body(body)?;

        // NOTE: We read the response body here, otherwise Hyper will not allow this connection to
        // be reused.
        ctx.dispatch(req)
            .await
            .context(ErrorKind::UpstreamUploadError)?
            .discard()
            .await?;

        STATS::upstream_success.add_value(1);
        return Ok(());
    }

    Ok(())
}

async fn upload_from_client<S>(
    ctx: &RepositoryRequestContext,
    oid: Sha256,
    size: u64,
    body: S,
    scuba: &mut Option<&mut ScubaMiddlewareState>,
) -> Result<(), Error>
where
    S: Stream<Item = Result<Bytes, ()>> + Unpin + Send + 'static,
{
    let (internal_send, internal_recv) = channel::<Result<Bytes, ()>>(BUFFER_SIZE);
    let (upstream_send, upstream_recv) = channel::<Result<Bytes, ()>>(BUFFER_SIZE);

    // CloseableSender lets us allow a stream to close without breaking the sender. This is useful
    // if e.g. upstream already has the data we want to send, so we decide not to send it. This
    // gives us better error messages, since if one of our uploads fail, we're guaranteed that the
    // consume_stream future isn't the one that'll return an error.
    let internal_send = CloseableSender::new(internal_send);
    let upstream_send = CloseableSender::new(upstream_send);

    let mut sink = internal_send.fanout(upstream_send);

    let internal_recv = internal_recv
        .map_err(|()| ErrorKind::ClientCancelled)
        .err_into();

    let upstream_recv = upstream_recv
        .map_err(|()| ErrorKind::ClientCancelled)
        .err_into();

    let internal_upload = internal_upload(ctx, oid, size, internal_recv);
    let upstream_upload = upstream_upload(ctx, oid, size, upstream_recv);

    let mut received: usize = 0;

    let mut data = body
        .map_ok(|chunk| {
            received += chunk.len();
            chunk
        })
        .map(Ok);

    // Note: this closure simply creates a single future that sends all data then closes the sink.
    // It needs to be a single future because all 3 futures below need to make progress
    // concurrently for the upload to succeed (if the destinations aren't making progress, we'll
    // deadlock, if the source isn't making progress, we'll deadlock too, and if the sink doesn't
    // close, we'll never finish the uploads).
    let consume_stream = async {
        sink.send_all(&mut data)
            .await
            .map_err(|_| ErrorKind::ClientCancelled)
            .map_err(Error::from)?;

        sink.close().await?;

        Ok(())
    };

    let res = try_join!(internal_upload, upstream_upload, consume_stream);

    ScubaMiddlewareState::maybe_add(scuba, HttpScubaKey::RequestBytesReceived, received);

    res.map(|_| ())
}

async fn sync_internal_and_upstream(
    ctx: &RepositoryRequestContext,
    oid: Sha256,
    size: u64,
    scuba: &mut Option<&mut ScubaMiddlewareState>,
) -> Result<(), Error> {
    let key = FetchKey::Aliased(Alias::Sha256(oid));

    let res = filestore::fetch(ctx.repo.repo_blobstore().clone(), ctx.ctx.clone(), &key).await?;

    match res {
        Some(stream) => {
            // We have the data, so presumably upstream does not have it.
            ScubaMiddlewareState::maybe_add(scuba, LfsScubaKey::UploadSync, "internal_to_upstream");
            upstream_upload(ctx, oid, size, stream).await?;
        }
        None => {
            // We do not have the data. Get it from upstream.
            ScubaMiddlewareState::maybe_add(scuba, LfsScubaKey::UploadSync, "upstream_to_internal");
            let object = RequestObject {
                oid: LfsSha256(oid.into_inner()),
                size,
            };

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

            let batch = ctx
                .upstream_batch(&batch)
                .await
                .context(ErrorKind::UpstreamBatchError)?
                .ok_or(ErrorKind::ObjectCannotBeSynced(object))?;

            let action = find_actions(batch, &object)?
                .remove(&Operation::Download)
                .ok_or(ErrorKind::ObjectCannotBeSynced(object))?;

            let req = Request::get(action.href).body(Body::empty())?;

            let stream = ctx
                .dispatch(req)
                .await
                .context(ErrorKind::ObjectCannotBeSynced(object))?
                .into_inner();

            internal_upload(ctx, oid, size, stream).await?;
        }
    }

    Ok(())
}

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

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

    let oid = Sha256::from_str(&oid).map_err(HttpError::e400)?;
    let size = size.parse().map_err(Error::from).map_err(HttpError::e400)?;
    let content_length: Option<u64> = read_header_value(state, CONTENT_LENGTH)
        .transpose()
        .map_err(HttpError::e400)?;

    if let Some(content_length) = content_length {
        ScubaMiddlewareState::try_borrow_add(
            state,
            HttpScubaKey::RequestContentLength,
            content_length,
        );
    }

    STATS::size_bytes.add_value(size as i64);

    if let Some(max_upload_size) = ctx.max_upload_size() {
        if size > max_upload_size {
            Err(HttpError::e400(ErrorKind::UploadTooLarge(
                size,
                max_upload_size,
            )))?;
        }
    }

    // The key invariant of our proxy design is that if you upload to this LFS server, then the
    // content will be present in both this LFS server and its upstream. To do so, we've
    // historically asked the client to upload whatever data either server is missing. However,
    // this only works if the client has the data, but it doesn't always. Indeed, if the client
    // only holds a LFS pointer (but knows the data is available on the server), then we can find
    // ourselves in a situation where the client needs to upload to us, but cannot (because it does
    // not have the data). However, the client still needs a mechanism to assert that the data is
    // in both this LFS server and its upstream. So, to support this mechanism, we let the client
    // send an empty request when uploading. When this happens, we assume we must have the content
    // somewhere, and try to sync it as necessary (to upstream if we have it internally, and to
    // internal if we don't).

    match content_length {
        Some(0) if size > 0 => {
            let mut scuba = state.try_borrow_mut::<ScubaMiddlewareState>();
            sync_internal_and_upstream(&ctx, oid, size, &mut scuba)
                .await
                .map_err(HttpError::e500)?;
        }
        _ => {
            // TODO: More appropriate status codes here
            let body = Body::take_from(state).map_err(|_| ());
            let mut scuba = state.try_borrow_mut::<ScubaMiddlewareState>();
            upload_from_client(&ctx, oid, size, body, &mut scuba)
                .await
                .map_err(HttpError::e500)?;
        }
    }

    Ok(EmptyBody::new())
}

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

    use chaosblob::ChaosBlobstore;
    use chaosblob::ChaosOptions;
    use fbinit::FacebookInit;
    use futures::future;
    use futures::stream;
    use memblob::Memblob;
    use test_repo_factory::TestRepoFactory;

    use super::*;

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

        let body = stream::once(future::ready(Ok(Bytes::from("foobar"))));
        let oid =
            Sha256::from_str("c3ab8ff13720e8ad9047dd39466b3c8974e592c2fa383d4a3960714caef0c4f2")?;
        let size = 6;

        upload_from_client(&ctx, oid, size, body, &mut None).await?;

        Ok(())
    }

    #[fbinit::test]
    async fn test_upload_from_client_failing_internal(fb: FacebookInit) -> Result<(), Error> {
        // Create a test repo with a blobstore that fails all reads and writes.
        let repo = TestRepoFactory::new(fb)?
            .with_blobstore(Arc::new(ChaosBlobstore::new(
                Memblob::default(),
                ChaosOptions::new(NonZeroU32::new(1), NonZeroU32::new(1)),
            )))
            .build()
            .await?;

        let ctx = RepositoryRequestContext::test_builder_with_repo(fb, repo)?
            .upstream_uri(None)
            .build()?;

        let body = stream::once(future::ready(Ok(Bytes::from("foobar"))));
        let oid =
            Sha256::from_str("c3ab8ff13720e8ad9047dd39466b3c8974e592c2fa383d4a3960714caef0c4f2")?;
        let size = 6;

        let r = upload_from_client(&ctx, oid, size, body, &mut None).await;

        assert!(r.is_err());

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

use anyhow::Error;
use gotham::state::FromState;
use gotham::state::State;
use http::header::AsHeaderName;
use http::header::HeaderMap;
use permission_checker::MononokeIdentitySet;

pub fn read_header_value<K, T>(state: &State, header: K) -> Option<Result<T, Error>>
where
    K: AsHeaderName,
    T: FromStr,
    <T as FromStr>::Err: std::error::Error + Send + Sync + 'static,
{
    let headers = HeaderMap::try_borrow_from(state)?;
    let val = headers.get(header)?;
    let val = std::str::from_utf8(val.as_bytes())
        .map_err(Error::from)
        .and_then(|val| T::from_str(val).map_err(Error::from));
    Some(val)
}

pub fn read_header_value_ignore_err<K, T>(state: &State, header: K) -> Option<T>
where
    K: AsHeaderName,
    T: FromStr,
    <T as FromStr>::Err: std::error::Error + Send + Sync + 'static,
{
    let headers = HeaderMap::try_borrow_from(state)?;
    let val = headers.get(header)?;
    let val = std::str::from_utf8(val.as_bytes()).ok()?;

    T::from_str(val).ok()
}

pub fn is_identity_subset<'a>(
    subset_idents: impl IntoIterator<Item = &'a MononokeIdentitySet>,
    client_idents: Option<&MononokeIdentitySet>,
) -> bool {
    let client_idents = match client_idents {
        Some(idents) => idents,
        None => return false,
    };

    subset_idents
        .into_iter()
        .any(|subset_ids| subset_ids.is_subset(client_idents))
}


/*
 * 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 bytes::Bytes;
use filestore::FilestoreConfigRef;
use filestore::StoreRequest;
use futures::Stream;
use futures::TryStreamExt;
use gotham::state::FromState;
use gotham::state::State;
use gotham_derive::StateData;
use gotham_derive::StaticResponseExtender;
use gotham_ext::error::HttpError;
use gotham_ext::middleware::HttpScubaKey;
use gotham_ext::middleware::ScubaMiddlewareState;
use gotham_ext::response::EmptyBody;
use gotham_ext::response::TryIntoResponse;
use http::header::CONTENT_LENGTH;
use hyper::Body;
use mononoke_types::hash::GitSha1;
use mononoke_types::hash::RichGitSha1;
use repo_blobstore::RepoBlobstoreRef;
use serde::Deserialize;
use stats::prelude::*;

use crate::errors::ErrorKind;
use crate::lfs_server_context::RepositoryRequestContext;
use crate::middleware::LfsMethod;
use crate::util::read_header_value;

define_stats! {
    prefix ="mononoke.lfs.git_upload_blob";
    total_uploads: timeseries(Rate, Sum),
    upload_success: timeseries(Rate, Sum),
    size_bytes: histogram(100_000, 0, 10_000_000, Average, Sum, Count; P 5; P 25; P 50; P 75; P 95; P 97; P 99),
}

// NOTE: We don't deserialize things beyond a String form, in order to report errors in our
// controller, not in routing.
#[derive(Deserialize, StateData, StaticResponseExtender)]
pub struct GitBlobParams {
    repository: String,
    oid: String,
    size: String,
}

async fn upload_blob<S>(
    ctx: &RepositoryRequestContext,
    oid: RichGitSha1,
    size: u64,
    body: S,
) -> Result<(), Error>
where
    S: Stream<Item = Result<Bytes, ()>> + Unpin + Send + 'static,
{
    STATS::total_uploads.add_value(1);

    filestore::store(
        ctx.repo.repo_blobstore(),
        *ctx.repo.filestore_config(),
        &ctx.ctx,
        &StoreRequest::with_git_sha1(size, oid),
        body.map_err(|()| ErrorKind::ClientCancelled).err_into(),
    )
    .await
    .context(ErrorKind::FilestoreWriteFailure)?;

    STATS::upload_success.add_value(1);

    Ok(())
}

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

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

    let oid = GitSha1::from_str(&oid).map_err(HttpError::e400)?;
    let size = size.parse().map_err(Error::from).map_err(HttpError::e400)?;
    let content_length: Option<u64> = read_header_value(state, CONTENT_LENGTH)
        .transpose()
        .map_err(HttpError::e400)?;

    if let Some(content_length) = content_length {
        ScubaMiddlewareState::try_borrow_add(
            state,
            HttpScubaKey::RequestContentLength,
            content_length,
        );
    }

    STATS::size_bytes.add_value(size as i64);

    if let Some(max_upload_size) = ctx.max_upload_size() {
        if size > max_upload_size {
            return Err(HttpError::e400(ErrorKind::UploadTooLarge(
                size,
                max_upload_size,
            )));
        }
    }

    let oid = RichGitSha1::from_sha1(oid, "blob", size);
    let body = Body::take_from(state).map_err(|_| ());
    upload_blob(&ctx, oid, size, body)
        .await
        .map_err(HttpError::e500)?;

    Ok(EmptyBody::new())
}


/*
 * 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 blob;

pub use blob::git_upload_blob;
pub use blob::GitBlobParams;


/*
 * 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 ods;
mod request_context;

pub use self::ods::OdsMiddleware;
pub use self::request_context::LfsMethod;
pub use self::request_context::RequestContext;
pub use self::request_context::RequestContextMiddleware;


/*
 * 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 gotham::state::State;
use gotham_ext::middleware::Middleware;
use gotham_ext::middleware::PostResponseCallbacks;
use hyper::Body;
use hyper::Response;
use hyper::StatusCode;
use stats::prelude::*;
use time_ext::DurationExt;

use super::LfsMethod;
use super::RequestContext;

define_stats! {
    prefix = "mononoke.lfs.request";
    success: timeseries(Rate, Sum),
    requests: timeseries(Rate, Sum),
    failure_4xx: timeseries(Rate, Sum),
    failure_429: timeseries(Rate, Sum),
    failure_404: timeseries(Rate, Sum),
    failure_5xx: timeseries(Rate, Sum),

    repo_requests: dynamic_timeseries("{}.requests", (repo_and_method: String); Rate, Sum),
    repo_success: dynamic_timeseries("{}.success", (repo_and_method: String); Rate, Sum),
    repo_failure_4xx: dynamic_timeseries("{}.failure_4xx", (repo_and_method: String); Rate, Sum),
    repo_failure_5xx: dynamic_timeseries("{}.failure_5xx", (repo_and_method: String); Rate, Sum),
    git_upload_blob_duration: dynamic_histogram("{}.git_upload_blob_ms", (repo: String); 100, 0, 5000, Average, Sum, Count; P 5; P 25; P 50; P 75; P 95; P 97; P 99),
    upload_duration: dynamic_histogram("{}.upload_ms", (repo: String); 100, 0, 5000, Average, Sum, Count; P 5; P 25; P 50; P 75; P 95; P 97; P 99),
    download_duration: dynamic_histogram("{}.download_ms", (repo: String); 100, 0, 5000, Average, Sum, Count; P 5; P 25; P 50; P 75; P 95; P 97; P 99),
    download_sha256_duration: dynamic_histogram("{}.download_sha256_ms", (repo: String); 100, 0, 5000, Average, Sum, Count; P 5; P 25; P 50; P 75; P 95; P 97; P 99),
    batch_duration: dynamic_histogram("{}.batch_ms", (repo: String); 10, 0, 500, Average, Sum, Count; P 5; P 25; P 50; P 75; P 95; P 97; P 99),
    response_bytes_sent: dynamic_histogram("{}.response_bytes_sent", (repo_and_method: String); 1_500_000, 0, 150_000_000, Average, Sum, Count; P 5; P 25; P 50; P 75; P 95; P 97; P 99),
}

fn log_stats(state: &mut State, status: StatusCode) -> Option<()> {
    // Not all requests have a valid method and repo, so calculate the top level HTTP stats first.
    STATS::requests.add_value(1);
    if status.is_success() {
        STATS::success.add_value(1);
    } else if status.is_client_error() {
        if status == StatusCode::TOO_MANY_REQUESTS {
            STATS::failure_429.add_value(1);
        } else if status == StatusCode::NOT_FOUND {
            STATS::failure_404.add_value(1);
        }

        STATS::failure_4xx.add_value(1);
    } else if status.is_server_error() {
        STATS::failure_5xx.add_value(1);
    }

    let ctx = state.try_borrow::<RequestContext>()?;
    let method = ctx.method?;
    let repo = ctx.repository.clone()?;
    let repo_and_method = format!("{}.{}", &repo, method);

    let callbacks = state.try_borrow_mut::<PostResponseCallbacks>()?;

    callbacks.add(move |info| {
        if let Some(duration) = info.duration {
            match method {
                LfsMethod::Upload => {
                    STATS::upload_duration.add_value(duration.as_millis_unchecked() as i64, (repo,))
                }
                LfsMethod::Download => STATS::download_duration
                    .add_value(duration.as_millis_unchecked() as i64, (repo,)),
                LfsMethod::DownloadSha256 => STATS::download_sha256_duration
                    .add_value(duration.as_millis_unchecked() as i64, (repo,)),
                LfsMethod::Batch => {
                    STATS::batch_duration.add_value(duration.as_millis_unchecked() as i64, (repo,))
                }
                LfsMethod::GitBlob => STATS::git_upload_blob_duration
                    .add_value(duration.as_millis_unchecked() as i64, (repo,)),
            }
        }

        STATS::repo_requests.add_value(1, (repo_and_method.clone(),));

        if status.is_success() {
            STATS::repo_success.add_value(1, (repo_and_method.clone(),));
        } else if status.is_client_error() {
            STATS::repo_failure_4xx.add_value(1, (repo_and_method.clone(),));
        } else if status.is_server_error() {
            STATS::repo_failure_5xx.add_value(1, (repo_and_method.clone(),));
        }

        if let Some(response_bytes_sent) = info.meta.as_ref().map(|m| m.body().bytes_sent) {
            STATS::response_bytes_sent.add_value(response_bytes_sent as i64, (repo_and_method,))
        }
    });

    Some(())
}

pub struct OdsMiddleware {}

impl OdsMiddleware {
    pub fn new() -> Self {
        Self {}
    }
}

#[async_trait::async_trait]
impl Middleware for OdsMiddleware {
    async fn outbound(&self, state: &mut State, response: &mut Response<Body>) {
        if let Some(ctx) = state.try_borrow::<RequestContext>() {
            if ctx.should_log {
                log_stats(state, response.status());
            }
        }
    }
}


/*
 * 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::default::Default;
use std::fmt;
use std::sync::Arc;

use context::CoreContext;
use context::SessionContainer;
use fbinit::FacebookInit;
use gotham::state::FromState;
use gotham::state::State;
use gotham_derive::StateData;
use gotham_ext::middleware::MetadataState;
use gotham_ext::middleware::Middleware;
use gotham_ext::state_ext::StateExt;
use hyper::body::Body;
use hyper::Response;
use hyper::StatusCode;
use metadata::Metadata;
use permission_checker::MononokeIdentitySetExt;
use scuba_ext::MononokeScubaSampleBuilder;
use slog::error;
use slog::o;
use slog::Logger;

#[derive(Copy, Clone)]
pub enum LfsMethod {
    Upload,
    Download,
    DownloadSha256,
    Batch,
    // Methods below this are for pushing git objects, not for LFS
    // They do not correspond to any LFS protocol
    GitBlob,
}

impl fmt::Display for LfsMethod {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        let name = match self {
            Self::Upload => "upload",
            Self::Download => "download",
            Self::DownloadSha256 => "download_sha256",
            Self::Batch => "batch",
            Self::GitBlob => "git_blob_upload",
        };
        write!(f, "{}", name)
    }
}

impl LfsMethod {
    pub fn is_read_only(&self) -> bool {
        match self {
            Self::Download | Self::DownloadSha256 | Self::Batch => true,
            Self::Upload | Self::GitBlob => false,
        }
    }
}

#[derive(StateData, Clone)]
pub struct RequestContext {
    pub ctx: CoreContext,
    pub repository: Option<String>,
    pub method: Option<LfsMethod>,
    pub error_msg: Option<String>,
    pub should_log: bool,
}

impl RequestContext {
    fn new(ctx: CoreContext, should_log: bool) -> Self {
        Self {
            ctx,
            repository: None,
            method: None,
            error_msg: None,
            should_log,
        }
    }

    pub fn set_request(&mut self, repository: String, method: LfsMethod) {
        self.repository = Some(repository);
        self.method = Some(method);
    }
}

#[derive(Clone)]
pub struct RequestContextMiddleware {
    fb: FacebookInit,
    logger: Logger,
    enforce_authentication: bool,
    readonly: bool,
}

impl RequestContextMiddleware {
    pub fn new(
        fb: FacebookInit,
        logger: Logger,
        enforce_authentication: bool,
        readonly: bool,
    ) -> Self {
        Self {
            fb,
            logger,
            enforce_authentication,
            readonly,
        }
    }
}

#[async_trait::async_trait]
impl Middleware for RequestContextMiddleware {
    async fn inbound(&self, state: &mut State) -> Option<Response<Body>> {
        let request_id = state.short_request_id();

        let logger = self.logger.new(o!("request_id" => request_id.to_string()));
        let metadata = if let Some(metadata_state) = MetadataState::try_borrow_from(state) {
            metadata_state.metadata().clone()
        } else {
            Metadata::default()
        };

        let should_log = !metadata.identities().is_proxygen_test_identity();

        if metadata.identities().is_empty() && self.enforce_authentication {
            let msg = "Client not authenticated".to_string();
            error!(self.logger, "{}", &msg);
            let response = Response::builder()
                .status(StatusCode::FORBIDDEN)
                .body(
                    format!(
                        "{{\"message:\"{}\", \"request_id\":\"{}\"}}",
                        msg,
                        state.short_request_id()
                    )
                    .into(),
                )
                .expect("Couldn't build http response");

            return Some(response);
        }

        let session = SessionContainer::builder(self.fb)
            .metadata(Arc::new(metadata))
            .readonly(self.readonly)
            .build();
        let ctx = session.new_context(logger, MononokeScubaSampleBuilder::with_discard());

        state.put(RequestContext::new(ctx, should_log));

        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 anyhow::Context;
use anyhow::Error;
use gotham::state::State;
use gotham_ext::error::ErrorFormatter;
use gotham_ext::state_ext::StateExt;
use lfs_protocol::git_lfs_mime;
use lfs_protocol::ResponseError;
use mime::Mime;

pub struct LfsErrorFormatter;

impl ErrorFormatter for LfsErrorFormatter {
    type Body = Vec<u8>;

    fn format(&self, error: &Error, state: &State) -> Result<(Self::Body, Mime), Error> {
        let message = format!("{:#}", error);

        let res = ResponseError {
            message,
            documentation_url: None,
            request_id: Some(state.short_request_id().to_string()),
        };

        let body = serde_json::to_vec(&res).context("Failed to serialize error")?;

        Ok((body, git_lfs_mime()))
    }
}


/*
 * 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::pin::Pin;

use anyhow::anyhow;
use anyhow::Result;
use cached_config::ConfigHandle;
use fbinit::FacebookInit;
use futures::future::FutureExt;
use gotham::handler::HandlerFuture;
use gotham::middleware::Middleware;
use gotham::state::FromState;
use gotham::state::State;
use gotham_derive::NewMiddleware;
use gotham_ext::error::HttpError;
use gotham_ext::middleware::MetadataState;
use gotham_ext::response::build_error_response;
use http::HeaderMap;
use hyper::Uri;
use qps::Qps;
use slog::trace;

use super::error_formatter::LfsErrorFormatter;
use crate::config::ServerConfig;
use crate::LfsServerContext;

const HEADER_REVPROXY_REGION: &str = "x-fb-revproxy-region";

// NOTE: Our Throttling middleware is implemented as Gotham middleware for 3 reasons:
// - It needs to replace responses.
// - It needs to do asynchronously.
// - It only needs to run if we're going to serve a request.

#[derive(Clone, NewMiddleware)]
pub struct ThrottleMiddleware {
    fb: FacebookInit,
    handle: ConfigHandle<ServerConfig>,
}

impl ThrottleMiddleware {
    pub fn new(fb: FacebookInit, handle: ConfigHandle<ServerConfig>) -> Self {
        Self { fb, handle }
    }
}

impl Middleware for ThrottleMiddleware {
    fn call<Chain>(self, state: State, chain: Chain) -> Pin<Box<HandlerFuture>>
    where
        Chain: FnOnce(State) -> Pin<Box<HandlerFuture>>,
    {
        if let Some(uri) = Uri::try_borrow_from(&state) {
            if uri.path() == "/health_check" {
                return chain(state);
            }
        }
        let identities = state
            .try_borrow::<MetadataState>()
            .map(|metadata_state| metadata_state.metadata().identities());

        for limit in self.handle.get().loadshedding_limits().iter() {
            if let Err(err) = limit.should_load_shed(self.fb, identities) {
                let err = HttpError::e429(err);

                let res =
                    async move { build_error_response(err, state, &LfsErrorFormatter) }.boxed();

                return res;
            }
        }

        chain(state)
    }
}

#[derive(Clone, NewMiddleware)]
pub struct QpsMiddleware {
    lfs_ctx: LfsServerContext,
}

impl QpsMiddleware {
    pub fn new(lfs_ctx: LfsServerContext) -> Self {
        Self { lfs_ctx }
    }
}

impl Middleware for QpsMiddleware {
    fn call<Chain>(self, state: State, chain: Chain) -> Pin<Box<HandlerFuture>>
    where
        Chain: FnOnce(State) -> Pin<Box<HandlerFuture>>,
    {
        if let Some(uri) = Uri::try_borrow_from(&state) {
            if uri.path() == "/health_check" {
                return chain(state);
            }
        }

        let headers = HeaderMap::try_borrow_from(&state).expect("No headers in the request");

        if let Err(e) = { bump_qps(headers, &self.lfs_ctx.qps()) } {
            trace!(self.lfs_ctx.logger(), "Failed to bump QPS: {:?}", e);
        }

        chain(state)
    }
}

fn bump_qps(headers: &HeaderMap, qps: &Option<Qps>) -> Result<()> {
    let qps = match qps {
        Some(qps) => qps,
        None => return Ok(()),
    };
    match headers.get(HEADER_REVPROXY_REGION) {
        Some(proxy_region) => {
            qps.bump(proxy_region.to_str()?)?;
            Ok(())
        }
        None => Err(anyhow!("No {:?} header.", HEADER_REVPROXY_REGION)),
    }
}


/*
 * 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 error_formatter;
mod middleware;
mod router;

pub use router::build_router;


/*
 * 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::pin::Pin;

use fbinit::FacebookInit;
use futures::FutureExt;
use gotham::handler::HandlerFuture;
use gotham::helpers::http::response::create_empty_response;
use gotham::helpers::http::response::create_response;
use gotham::middleware::state::StateMiddleware;
use gotham::pipeline::new_pipeline;
use gotham::pipeline::single_pipeline;
use gotham::router::builder::build_router as gotham_build_router;
use gotham::router::builder::DefineSingleRoute;
use gotham::router::builder::DrawRoutes;
use gotham::router::Router;
use gotham::state::FromState;
use gotham::state::State;
use gotham_ext::response::build_response;
use hyper::Body;
use hyper::Response;
use hyper::StatusCode;

use super::error_formatter::LfsErrorFormatter;
use super::middleware::QpsMiddleware;
use super::middleware::ThrottleMiddleware;
use crate::batch;
use crate::download;
use crate::git_upload;
use crate::lfs_server_context::LfsServerContext;
use crate::upload;

// These 3 methods are wrappers to go from async fn's to the implementations Gotham expects,
// as well as creating HTTP responses using build_response().
fn batch_handler(mut state: State) -> Pin<Box<HandlerFuture>> {
    async move {
        let res = batch::batch(&mut state).await;
        build_response(res, state, &LfsErrorFormatter)
    }
    .boxed()
}

fn download_handler(mut state: State) -> Pin<Box<HandlerFuture>> {
    async move {
        let res = download::download(&mut state).await;
        build_response(res, state, &LfsErrorFormatter)
    }
    .boxed()
}

fn download_sha256_handler(mut state: State) -> Pin<Box<HandlerFuture>> {
    async move {
        let res = download::download_sha256(&mut state).await;
        build_response(res, state, &LfsErrorFormatter)
    }
    .boxed()
}

fn upload_handler(mut state: State) -> Pin<Box<HandlerFuture>> {
    async move {
        let res = upload::upload(&mut state).await;
        build_response(res, state, &LfsErrorFormatter)
    }
    .boxed()
}

fn git_upload_blob_handler(mut state: State) -> Pin<Box<HandlerFuture>> {
    async move {
        let res = git_upload::git_upload_blob(&mut state).await;
        build_response(res, state, &LfsErrorFormatter)
    }
    .boxed()
}

fn health_handler(state: State) -> (State, &'static str) {
    let lfs_ctx = LfsServerContext::borrow_from(&state);
    let res = if lfs_ctx.will_exit() {
        "EXITING"
    } else {
        "I_AM_ALIVE"
    };
    (state, res)
}

fn config_handler(state: State) -> (State, Response<Body>) {
    let lfs_ctx = LfsServerContext::borrow_from(&state);

    let res = match serde_json::to_string(&*lfs_ctx.get_config()) {
        Ok(json) => create_response(&state, StatusCode::OK, mime::APPLICATION_JSON, json),
        Err(_) => create_empty_response(&state, StatusCode::INTERNAL_SERVER_ERROR),
    };

    (state, res)
}

pub fn build_router(
    fb: FacebookInit,
    lfs_ctx: LfsServerContext,
    allow_git_blob_upload: bool,
) -> Router {
    let pipeline = new_pipeline()
        .add(ThrottleMiddleware::new(fb, lfs_ctx.get_config_handle()))
        .add(QpsMiddleware::new(lfs_ctx.clone()))
        .add(StateMiddleware::new(lfs_ctx))
        .build();

    let (chain, pipelines) = single_pipeline(pipeline);

    gotham_build_router(chain, pipelines, |route| {
        route
            .post("/:repository/objects/batch")
            .with_path_extractor::<batch::BatchParams>()
            .to(batch_handler);

        route
            .get("/:repository/download/:content_id")
            .with_path_extractor::<download::DownloadParamsContentId>()
            .to(download_handler);

        route
            .get("/:repository/download_sha256/:oid")
            .with_path_extractor::<download::DownloadParamsSha256>()
            .to(download_sha256_handler);

        route
            .put("/:repository/upload/:oid/:size")
            .with_path_extractor::<upload::UploadParams>()
            .to(upload_handler);

        if allow_git_blob_upload {
            route
                .put("/git_blob_upload/:repository/:oid/:size")
                .with_path_extractor::<git_upload::GitBlobParams>()
                .to(git_upload_blob_handler);
        }

        route.get("/health_check").to(health_handler);
        route.get("/config").to(config_handler);
    })
}