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

#![feature(try_blocks)]

mod git_reader;
mod gitimport_objects;
mod gitlfs;

use std::collections::BTreeMap;
use std::collections::HashSet;
use std::path::Path;
use std::process::Stdio;
use std::sync::RwLock;

use anyhow::bail;
use anyhow::format_err;
use anyhow::Context;
use anyhow::Error;
use bytes::Bytes;
use cloned::cloned;
use context::CoreContext;
use futures::Stream;
use futures::StreamExt;
use futures::TryFutureExt;
use futures::TryStreamExt;
use git_hash::ObjectId;
use git_object::Object;
use linked_hash_map::LinkedHashMap;
use manifest::bonsai_diff;
use manifest::BonsaiDiffFileChange;
use mononoke_types::ChangesetId;
use mononoke_types::MPath;
use slog::debug;
use slog::info;
use sorted_vector_map::SortedVectorMap;
use tokio::io::AsyncBufReadExt;
use tokio::io::BufReader;
use tokio::process::Command;
use tokio::task;

pub use crate::git_reader::GitRepoReader;
pub use crate::gitimport_objects::convert_time_to_datetime;
pub use crate::gitimport_objects::oid_to_sha1;
use crate::gitimport_objects::read_raw_object;
pub use crate::gitimport_objects::CommitMetadata;
pub use crate::gitimport_objects::ExtractedCommit;
pub use crate::gitimport_objects::GitLeaf;
pub use crate::gitimport_objects::GitManifest;
pub use crate::gitimport_objects::GitTree;
pub use crate::gitimport_objects::GitUploader;
pub use crate::gitimport_objects::GitimportPreferences;
pub use crate::gitimport_objects::GitimportTarget;
pub use crate::gitimport_objects::TagMetadata;
pub use crate::gitlfs::GitImportLfs;
pub use crate::gitlfs::LfsMetaData;

pub const HGGIT_MARKER_EXTRA: &str = "hg-git-rename-source";
pub const HGGIT_MARKER_VALUE: &[u8] = b"git";
pub const HGGIT_COMMIT_ID_EXTRA: &str = "convert_revision";

// TODO: Try to produce copy-info?
async fn find_file_changes<S, U>(
    ctx: &CoreContext,
    lfs: &GitImportLfs,
    reader: &GitRepoReader,
    uploader: U,
    changes: S,
) -> Result<SortedVectorMap<MPath, U::Change>, Error>
where
    S: Stream<Item = Result<BonsaiDiffFileChange<GitLeaf>, Error>>,
    U: GitUploader,
{
    changes
        .map_ok(|change| async {
            task::spawn({
                cloned!(ctx, reader, uploader, lfs);
                async move {
                    match change {
                        BonsaiDiffFileChange::Changed(path, ty, GitLeaf(oid))
                        | BonsaiDiffFileChange::ChangedReusedId(path, ty, GitLeaf(oid)) => {
                            let git_bytes = {
                                let object = reader.get_object(&oid).await?;
                                let blob = object
                                    .parsed
                                    .try_into_blob()
                                    .map_err(|_| format_err!("{} is not a blob", oid))?;
                                Bytes::from(blob.data)
                            };

                            uploader
                                .upload_file(&ctx, &lfs, &path, ty, oid, git_bytes)
                                .await
                                .map(|change| (path, change))
                        }
                        BonsaiDiffFileChange::Deleted(path) => Ok((path, U::deleted())),
                    }
                }
            })
            .await?
        })
        .try_buffer_unordered(100)
        .try_collect()
        .await
}

pub struct GitimportAccumulator {
    inner: LinkedHashMap<ObjectId, ChangesetId>,
}

impl GitimportAccumulator {
    pub fn new() -> Self {
        Self {
            inner: LinkedHashMap::new(),
        }
    }

    pub fn len(&self) -> usize {
        self.inner.len()
    }

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

    pub fn insert(&mut self, oid: ObjectId, cs_id: ChangesetId) {
        self.inner.insert(oid, cs_id);
    }

    pub fn get(&self, oid: &git_hash::oid) -> Option<ChangesetId> {
        self.inner.get(oid).copied()
    }
}

pub async fn is_annotated_tag(
    path: &Path,
    prefs: &GitimportPreferences,
    object_id: &ObjectId,
) -> Result<bool, Error> {
    let reader = GitRepoReader::new(&prefs.git_command_path, path).await?;
    Ok(reader
        .get_object(object_id)
        .await
        .with_context(|| {
            format_err!(
                "Failed to fetch git object {} for checking if its a tag",
                object_id,
            )
        })?
        .parsed
        .as_tag()
        .is_some())
}

pub async fn create_changeset_for_annotated_tag<Uploader: GitUploader>(
    ctx: &CoreContext,
    uploader: &Uploader,
    path: &Path,
    prefs: &GitimportPreferences,
    tag_id: &ObjectId,
    original_changeset_id: &ChangesetId,
) -> Result<ChangesetId, Error> {
    let reader = GitRepoReader::new(&prefs.git_command_path, path).await?;
    // Get the parsed Git Tag
    let tag_metadata = TagMetadata::new(ctx, *tag_id, &reader)
        .await
        .with_context(|| format_err!("Failed to create TagMetadata from git tag {}", tag_id))?;
    // Create the corresponding changeset for the Git Tag at Mononoke end
    let changeset_id = uploader
        .generate_changeset_for_annotated_tag(ctx, *original_changeset_id, tag_metadata)
        .await
        .with_context(|| format_err!("Failed to generate changeset for git tag {}", tag_id))?;
    Ok(changeset_id)
}

pub async fn upload_git_tag<Uploader: GitUploader>(
    ctx: &CoreContext,
    uploader: &Uploader,
    path: &Path,
    prefs: &GitimportPreferences,
    tag_id: &ObjectId,
) -> Result<(), Error> {
    let reader = GitRepoReader::new(&prefs.git_command_path, path).await?;
    let tag_bytes = read_raw_object(&reader, tag_id)
        .await
        .with_context(|| format_err!("Failed to fetch git tag {}", tag_id))?;
    // Upload Git Tag
    uploader
        .upload_object(ctx, *tag_id, tag_bytes)
        .await
        .with_context(|| format_err!("Failed to upload raw git tag {}", tag_id,))
}

pub async fn gitimport_acc<Uploader: GitUploader>(
    ctx: &CoreContext,
    path: &Path,
    uploader: &Uploader,
    target: &GitimportTarget,
    prefs: &GitimportPreferences,
) -> Result<GitimportAccumulator, Error> {
    let repo_name = if let Some(name) = &prefs.gitrepo_name {
        String::from(name)
    } else {
        let name_path = if path.ends_with(".git") {
            path.parent().unwrap_or(path)
        } else {
            path
        };
        String::from(name_path.to_string_lossy())
    };
    let dry_run = prefs.dry_run;

    let reader = GitRepoReader::new(&prefs.git_command_path, path).await?;
    let roots = target.get_roots();
    // TODO - remove this
    let nb_commits_to_import = target.get_nb_commits(&prefs.git_command_path, path).await?;
    if 0 == nb_commits_to_import {
        info!(ctx.logger(), "Nothing to import for repo {}.", repo_name);
        return Ok(GitimportAccumulator::new());
    }

    let acc = RwLock::new(GitimportAccumulator::new());

    // Kick off a stream that consumes the walk and prepared commits. Then, produce the Bonsais.
    target
        .list_commits(&prefs.git_command_path, path)
        .await?
        .try_filter_map({
            let acc = &acc;
            let uploader = &uploader;
            let repo_name = &repo_name;
            move |oid| async move {
                if let Some(bcs_id) = uploader.check_commit_uploaded(ctx, &oid).await? {
                    acc.write().expect("lock poisoned").insert(oid, bcs_id);
                    let git_sha1 = oid_to_sha1(&oid)?;
                    info!(
                        ctx.logger(),
                        "GitRepo:{} commit {} of {} - Oid:{} => Bid:{} (already exists)",
                        repo_name,
                        acc.read().expect("lock poisoned").len(),
                        nb_commits_to_import,
                        git_sha1.to_brief(),
                        bcs_id.to_brief()
                    );
                    Ok(None)
                } else {
                    Ok(Some(oid))
                }
            }
        })
        .map_ok(|oid| {
            cloned!(ctx, reader, uploader, prefs.lfs);
            async move {
                task::spawn({
                    async move {
                        let ExtractedCommit {
                            metadata,
                            tree,
                            parent_trees,
                            original_commit,
                        } = ExtractedCommit::new(&ctx, oid, &reader)
                            .await
                            .with_context(|| format!("While extracting {}", oid))?;

                        let file_changes = find_file_changes(
                            &ctx,
                            &lfs,
                            &reader,
                            uploader,
                            bonsai_diff(ctx.clone(), reader.clone(), tree, parent_trees),
                        )
                        .await?;

                        Result::<_, Error>::Ok((metadata, file_changes, original_commit, tree))
                    }
                })
                .await?
            }
        })
        .try_buffered(prefs.concurrency)
        .and_then(|(metadata, file_changes, original_commit, tree)| {
            let acc = &acc;
            let uploader = &uploader;
            let repo_name = &repo_name;
            let reader = &reader;
            async move {
                let oid = metadata.oid;
                let bonsai_parents = metadata
                    .parents
                    .iter()
                    .map(|p| {
                        roots
                            .get(p)
                            .copied()
                            .or_else(|| acc.read().expect("lock poisoned").get(p))
                            .ok_or_else(|| {
                                format_err!(
                                    "Couldn't find parent: {} in local list of imported commits",
                                    p
                                )
                            })
                    })
                    .collect::<Result<Vec<_>, _>>()
                    .with_context(|| format_err!("While looking for parents of {}", oid))?;

                // Before generating the corresponding changeset at Mononoke end, upload the raw git commit
                // and the git tree pointed to by the git commit.
                let tree_for_commit =
                    read_raw_object(reader, &tree.0).await.with_context(|| {
                        format_err!("Failed to fetch git tree {} for commit {}", tree.0, oid)
                    })?;
                // Upload Git Tree
                uploader
                    .upload_object(ctx, tree.0, tree_for_commit)
                    .await
                    .with_context(|| {
                        format_err!(
                            "Failed to upload raw git tree {} for commit {}",
                            tree.0,
                            oid
                        )
                    })?;
                // Upload Git commit
                uploader
                    .upload_object(ctx, oid, original_commit)
                    .await
                    .with_context(|| format_err!("Failed to upload raw git commit {}", oid))?;
                let (int_cs, bcs_id) = uploader
                    .generate_changeset_for_commit(
                        ctx,
                        bonsai_parents,
                        metadata,
                        file_changes,
                        dry_run,
                    )
                    .await?;
                acc.write().expect("lock poisoned").insert(oid, bcs_id);

                let git_sha1 = oid_to_sha1(&oid)?;
                info!(
                    ctx.logger(),
                    "GitRepo:{} commit {} of {} - Oid:{} => Bid:{}",
                    &repo_name,
                    acc.read().expect("lock poisoned").len(),
                    nb_commits_to_import,
                    git_sha1.to_brief(),
                    bcs_id.to_brief()
                );
                Ok((int_cs, git_sha1))
            }
        })
        // Chunk together into Vec<std::result::Result<(bcs, oid), Error> >
        .chunks(prefs.concurrency)
        // Go from Vec<Result<X,Y>> -> Result<Vec<X>,Y>
        .map(|v| v.into_iter().collect::<Result<Vec<_>, Error>>())
        .try_for_each(|v| async {
            cloned!(ctx, uploader);
            task::spawn(async move { uploader.finalize_batch(&ctx, dry_run, v).await }).await?
        })
        .await?;

    debug!(ctx.logger(), "Completed git import for repo {}.", repo_name);
    Ok(acc.into_inner().expect("lock poisoned"))
}

pub async fn gitimport(
    ctx: &CoreContext,
    path: &Path,
    uploader: &impl GitUploader,
    target: &GitimportTarget,
    prefs: &GitimportPreferences,
) -> Result<LinkedHashMap<ObjectId, ChangesetId>, Error> {
    Ok(gitimport_acc(ctx, path, uploader, target, prefs)
        .await?
        .inner)
}

/// Object representing Git refs. maybe_tag_id will only
/// have a value if the ref is a tag pointing to a commit.
#[derive(Debug, Hash, PartialEq, Eq, PartialOrd, Ord)]
pub struct GitRef {
    pub name: Vec<u8>,
    pub maybe_tag_id: Option<ObjectId>,
}

impl GitRef {
    fn new(name: Vec<u8>) -> Self {
        Self {
            name,
            maybe_tag_id: None,
        }
    }
}

pub async fn read_git_refs(
    path: &Path,
    prefs: &GitimportPreferences,
) -> Result<BTreeMap<GitRef, ObjectId>, Error> {
    let reader = GitRepoReader::new(&prefs.git_command_path, path).await?;

    let mut command = Command::new(&prefs.git_command_path)
        .current_dir(path)
        .env_clear()
        .kill_on_drop(false)
        .stdout(Stdio::piped())
        .arg("for-each-ref")
        .arg("--format=%(objectname) %(refname)")
        .spawn()
        .with_context(|| format!("failed to run git with {:?}", prefs.git_command_path))?;
    let stdout = BufReader::new(command.stdout.take().context("stdout not set up")?);
    let mut lines = stdout.lines();

    let mut refs = BTreeMap::new();

    while let Some(line) = lines
        .next_line()
        .await
        .context("git command didn't output anything")?
    {
        if let Some((oid_str, ref_name)) = line.split_once(' ') {
            let mut oid: ObjectId = oid_str.parse().context("reading refs")?;
            let mut git_ref = GitRef::new(ref_name.into());
            loop {
                let object = reader.get_object(&oid).await.with_context(|| {
                    format!("unable to read git object: {oid} for ref: {ref_name}")
                })?;
                match object.parsed {
                    Object::Tree(_) => {
                        // This happens in the Linux kernel repo, because Linus was being clever - a commit and a tree
                        // are both treeish for the purposes of things like checkout and diff.
                        break;
                    }
                    Object::Blob(_) => {
                        bail!("ref {} points to a blob", ref_name);
                    }
                    Object::Commit(_) => {
                        refs.insert(git_ref, oid);
                        break;
                    }
                    // If the ref is a tag, then we capture the object id of the tag.
                    // The loop is designed to peel the tag but we want the outermost
                    // tag object so only get the ID if we haven't already done it before.
                    Object::Tag(tag) => {
                        if git_ref.maybe_tag_id.is_none() {
                            git_ref.maybe_tag_id = Some(oid);
                        }
                        oid = tag.target;
                    }
                }
            }
        }
    }
    Ok(refs)
}

pub async fn import_tree_as_single_bonsai_changeset(
    ctx: &CoreContext,
    path: &Path,
    uploader: impl GitUploader,
    git_cs_id: ObjectId,
    prefs: &GitimportPreferences,
) -> Result<ChangesetId, Error> {
    let reader = GitRepoReader::new(&prefs.git_command_path, path).await?;

    let sha1 = oid_to_sha1(&git_cs_id)?;

    let ExtractedCommit {
        tree,
        metadata,
        original_commit,
        ..
    } = ExtractedCommit::new(ctx, git_cs_id, &reader)
        .await
        .with_context(|| format!("While extracting {}", git_cs_id))?;

    let file_changes = find_file_changes(
        ctx,
        &prefs.lfs,
        &reader.clone(),
        uploader.clone(),
        bonsai_diff(ctx.clone(), reader, tree, HashSet::new()),
    )
    .await?;

    // Before generating the corresponding changeset at Mononoke end, upload the raw git commit.
    uploader
        .upload_object(ctx, git_cs_id, original_commit)
        .await
        .with_context(|| format_err!("Failed to upload raw git commit {}", git_cs_id))?;

    uploader
        .generate_changeset_for_commit(ctx, vec![], metadata, file_changes, prefs.dry_run)
        .and_then(|(cs, id)| {
            uploader
                .finalize_batch(ctx, prefs.dry_run, vec![(cs, sha1)])
                .map_ok(move |_| id)
        })
        .await
}


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

use std::str;

use anyhow::Error;
use bytes::Bytes;
use bytes::BytesMut;
use futures::future;
use futures::future::MapOk;
use futures::future::Ready;
use futures::future::TryFutureExt;
use futures::stream::Stream;
use futures::stream::TryFold;
use futures::stream::TryStreamExt;
use http::header::HeaderMap;
use http::header::CONTENT_LENGTH;

type BodyFuture<S, E> = MapOk<
    TryFold<
        S,
        Ready<Result<BytesMut, E>>,
        BytesMut,
        fn(BytesMut, Bytes) -> Ready<Result<BytesMut, E>>,
    >,
    fn(BytesMut) -> Bytes,
>;

fn extend_buff<E>(mut buff: BytesMut, chunk: Bytes) -> Ready<Result<BytesMut, E>> {
    buff.extend_from_slice(chunk.as_ref());
    future::ready(Ok(buff))
}

pub trait BodyExt<E>: Stream<Item = Result<Bytes, E>> + Sized {
    fn try_concat_body_opt(
        self,
        headers: Option<&HeaderMap>,
    ) -> Result<BodyFuture<Self, E>, Error> {
        match headers {
            Some(headers) => self.try_concat_body(headers),
            None => {
                let headers = HeaderMap::default();
                self.try_concat_body(&headers)
            }
        }
    }

    fn try_concat_body(self, headers: &HeaderMap) -> Result<BodyFuture<Self, E>, Error> {
        let buff = if let Some(val) = headers.get(CONTENT_LENGTH) {
            let size = str::from_utf8(val.as_bytes())?.parse()?;
            BytesMut::with_capacity(size)
        } else {
            BytesMut::new()
        };

        Ok(self
            .try_fold(
                buff,
                extend_buff as fn(BytesMut, Bytes) -> Ready<Result<BytesMut, E>>,
            )
            .map_ok(BytesMut::freeze as fn(BytesMut) -> Bytes))
    }
}

impl<S, E> BodyExt<E> for S where S: Stream<Item = Result<Bytes, E>> {}

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

    use super::*;

    fn make_stream() -> impl Stream<Item = Result<Bytes, Error>> {
        stream::iter(vec![
            Result::<_, Error>::Ok(Bytes::from("1")),
            Result::<_, Error>::Ok(Bytes::from("2")),
        ])
    }

    #[tokio::test]
    async fn test_no_content_length() -> Result<(), Error> {
        let h = HeaderMap::new();
        let v = make_stream().try_concat_body(&h)?.await?;
        assert_eq!(v, Bytes::from("12"));
        Ok(())
    }

    #[tokio::test]
    async fn test_correct_content_length() -> Result<(), Error> {
        let mut h = HeaderMap::new();
        h.insert(CONTENT_LENGTH, "2".parse()?);

        let v = make_stream().try_concat_body(&h)?.await?;
        assert_eq!(v, Bytes::from("12"));

        Ok(())
    }

    #[tokio::test]
    async fn test_incorrect_content_length() -> Result<(), Error> {
        let mut h = HeaderMap::new();
        h.insert(CONTENT_LENGTH, "123".parse()?);

        let v = make_stream().try_concat_body(&h)?.await?;
        assert_eq!(v, Bytes::from("12"));

        Ok(())
    }

    #[tokio::test]
    async fn test_invalid_content_length() -> Result<(), Error> {
        let mut h = HeaderMap::new();
        h.insert(CONTENT_LENGTH, "foobar".parse()?);

        let v = make_stream().try_concat_body(&h);
        assert!(v.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 anyhow::Error;
use gotham::state::FromState;
use gotham::state::State;
use http::header::HeaderMap;
use http::header::HeaderValue;
use http::header::ACCEPT_ENCODING;

const GZIP: &str = "gzip";
const ZSTD: &str = "zstd";

#[derive(Copy, Clone, Debug, Eq, PartialEq)]
pub enum ContentCompression {
    Gzip,
    Zstd,
}

impl ContentCompression {
    pub fn as_str(&self) -> &'static str {
        match self {
            Self::Gzip => GZIP,
            Self::Zstd => ZSTD,
        }
    }
}

#[derive(Copy, Clone, Debug, Eq, PartialEq)]
pub enum ContentEncoding {
    Compressed(ContentCompression),
    Identity,
}

impl ContentEncoding {
    pub fn as_str(&self) -> &'static str {
        match self {
            Self::Compressed(c) => c.as_str(),
            Self::Identity => "identity",
        }
    }
}

impl From<ContentEncoding> for HeaderValue {
    fn from(e: ContentEncoding) -> HeaderValue {
        HeaderValue::from_static(e.as_str())
    }
}

impl ContentEncoding {
    pub fn from_state(state: &State) -> Self {
        let header = HeaderMap::try_borrow_from(state).and_then(|h| h.get(ACCEPT_ENCODING));

        match header {
            Some(h) => Self::from_header(h.as_bytes()).unwrap_or(Self::Identity),
            None => Self::Identity,
        }
    }

    /// Parse an [Accept-Encoding header] and provide a ContentEncoding
    /// representing it. This ignores client preferences entirely, since we
    /// control client & server, and instead uses our own preference ordering,
    /// restricted to what the client allows.
    ///
    /// [Accept-Encoding header]: https://developer.mozilla.org/en-US/docs/Web/HTTP/Headers/Accept-Encoding
    fn from_header(header: &[u8]) -> Result<Self, Error> {
        let mut gzip = false;
        let mut zstd = false;

        let encodings = std::str::from_utf8(header)?.split(' ');

        for encoding in encodings {
            let encoding = match encoding.split(';').next() {
                Some(encoding) => encoding.trim(),
                None => continue,
            };

            if encoding == ZSTD {
                zstd = true;
            }

            if encoding == GZIP {
                gzip = true;
            }
        }

        if zstd {
            return Ok(Self::Compressed(ContentCompression::Zstd));
        }

        if gzip {
            return Ok(Self::Compressed(ContentCompression::Gzip));
        }

        Ok(Self::Identity)
    }
}

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

    #[test]
    fn test_decode_content_encoding() -> Result<(), Error> {
        assert_eq!(
            ContentEncoding::from_header(b"identity, foobar")?,
            ContentEncoding::Identity
        );

        assert_eq!(
            ContentEncoding::from_header(b"foobar")?,
            ContentEncoding::Identity
        );

        assert_eq!(
            ContentEncoding::from_header(b"foobar, identity, zstd")?,
            ContentEncoding::Compressed(ContentCompression::Zstd),
        );

        assert_eq!(
            ContentEncoding::from_header(b"gzip, zstd")?,
            ContentEncoding::Compressed(ContentCompression::Zstd),
        );

        assert_eq!(
            ContentEncoding::from_header(b"deflate, gzip;q=1.0, *;q=0.5")?,
            ContentEncoding::Compressed(ContentCompression::Gzip),
        );

        assert_eq!(
            ContentEncoding::from_header(b"")?,
            ContentEncoding::Identity
        );

        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 anyhow::Error;
use gotham::state::State;
use hyper::Body;
use hyper::StatusCode;
use mime::Mime;
use rate_limiting::RateLimitReason;

pub trait ErrorFormatter {
    type Body: Into<Body>;

    fn format(&self, error: &Error, state: &State) -> Result<(Self::Body, Mime), Error>;
}

/// Wrapper around an anyhow::Error to indicate which
/// HTTP status code should be returned to the client.
#[derive(Debug)]
pub struct HttpError {
    pub error: Error,
    pub status_code: StatusCode,
}

impl HttpError {
    pub fn e400<E: Into<Error>>(err: E) -> Self {
        Self {
            error: err.into(),
            status_code: StatusCode::BAD_REQUEST,
        }
    }

    pub fn e401<E: Into<Error>>(err: E) -> Self {
        Self {
            error: err.into(),
            status_code: StatusCode::UNAUTHORIZED,
        }
    }

    pub fn e403<E: Into<Error>>(err: E) -> Self {
        Self {
            error: err.into(),
            status_code: StatusCode::FORBIDDEN,
        }
    }

    pub fn e404<E: Into<Error>>(err: E) -> Self {
        Self {
            error: err.into(),
            status_code: StatusCode::NOT_FOUND,
        }
    }

    pub fn e410<E: Into<Error>>(err: E) -> Self {
        Self {
            error: err.into(),
            status_code: StatusCode::GONE,
        }
    }

    pub fn e429<E: Into<Error>>(err: E) -> Self {
        Self {
            error: err.into(),
            status_code: StatusCode::TOO_MANY_REQUESTS,
        }
    }

    pub fn e500<E: Into<Error>>(err: E) -> Self {
        Self {
            error: err.into(),
            status_code: StatusCode::INTERNAL_SERVER_ERROR,
        }
    }

    pub fn e503<E: Into<Error>>(err: E) -> Self {
        Self {
            error: err.into(),
            status_code: StatusCode::SERVICE_UNAVAILABLE,
        }
    }
}

impl From<RateLimitReason> for HttpError {
    fn from(r: RateLimitReason) -> Self {
        Self::e429(r)
    }
}


/*
 * 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::net::SocketAddr;
use std::panic::RefUnwindSafe;
use std::pin::Pin;
use std::sync::Arc;

use futures::future::BoxFuture;
use futures::future::FutureExt;
use futures::task;
use gotham::handler::Handler;
use gotham::handler::HandlerFuture;
use gotham::handler::IntoResponse;
use gotham::handler::NewHandler;
use gotham::state::State;
use hyper::service::Service;
use hyper::Body;
use hyper::Request;
use hyper::Response;

use crate::middleware::Middleware;
use crate::socket_data::TlsSocketData;

#[derive(Clone)]
pub struct MononokeHttpHandler<H> {
    inner: H,
    middleware: Arc<Vec<Box<dyn Middleware>>>,
}

impl<H> MononokeHttpHandler<H> {
    pub fn into_service(
        self,
        addr: SocketAddr,
        tls_socket_data: Option<TlsSocketData>,
    ) -> MononokeHttpHandlerAsService<H> {
        MononokeHttpHandlerAsService {
            handler: self,
            addr,
            tls_socket_data,
        }
    }
}

// This trait is boilerplate to let us bind this in Gotham as the service. This is essentially
// patterned match off of how the Router Handler works in Gotham.
impl<H: Handler + Clone + Send + Sync + 'static + RefUnwindSafe> NewHandler
    for MononokeHttpHandler<H>
{
    type Instance = MononokeHttpHandler<H>;

    fn new_handler(&self) -> Result<Self::Instance, anyhow::Error> {
        Ok(self.clone())
    }
}

/// Executes a stack middleware. If any Middleware instance preempts the request, it returns the
/// index of the first middleware that didn't run and the response.
async fn run_middleware(
    middleware: &[Box<dyn Middleware>],
    state: &mut State,
) -> Option<(usize, Response<Body>)> {
    for (i, m) in middleware.iter().enumerate() {
        if let Some(response) = m.inbound(state).await {
            return Some((i + 1, response));
        }
    }

    None
}

// Gotham's router only runs middleware on route matches, which means we don't get any visibility
// into e.g. 404's. So, we use this handler to wrap Gotham's router and run our own middleware
// stack. Another reason not to use Gotham's middleware stack is that it places all middleware
// instances in the poll chain, so for e.g. an upload, we'd chain poll() calls through every single
// instance of middleware. In contrast, our middleware stack is completely out of the way once it
// has finished executing.
impl<H: Handler + Send + Sync + 'static> Handler for MononokeHttpHandler<H> {
    fn handle(self, mut state: State) -> Pin<Box<HandlerFuture>> {
        let fut = async move {
            // On request, middleware is called in order, then called the other way around on
            // response (this is what regular Router middleware in Gotham would do).
            let (idx, mut state, mut response) =
                match run_middleware(self.middleware.as_ref(), &mut state).await {
                    Some((idx, r)) => (idx, state, r),
                    None => {
                        // NOTE: It's a bit unfortunate that we have to return a HandlerFuture here
                        // when really we'd rather be working with just (State, HttpResponse<Body>)
                        // everywhere, but that's how it is in Gotham.
                        let (state, res) = match self.inner.handle(state).await {
                            Ok((state, res)) => (state, res),
                            Err((state, err)) => {
                                let response = err.into_response(&state);
                                (state, response)
                            }
                        };

                        let idx = self.middleware.len();
                        (idx, state, res)
                    }
                };

            // On outbound, only run middleware that did run in inbound.
            for m in self.middleware[0..idx].iter().rev() {
                m.outbound(&mut state, &mut response).await;
            }

            Ok((state, response))
        };

        fut.boxed()
    }
}

// NOTE: This is just syntactic sugar for MononokeHttpHandler::builder(), which is why the
// "handler" type is () here.
impl MononokeHttpHandler<()> {
    pub fn builder() -> MononokeHttpHandlerBuilder {
        MononokeHttpHandlerBuilder { middleware: vec![] }
    }
}

pub struct MononokeHttpHandlerBuilder {
    middleware: Vec<Box<dyn Middleware>>,
}

impl MononokeHttpHandlerBuilder {
    #[allow(clippy::should_implement_trait)]
    pub fn add<T: Middleware>(mut self, middleware: T) -> Self {
        self.middleware.push(Box::new(middleware));
        self
    }

    pub fn build<H: Handler + Send + Sync + 'static>(self, inner: H) -> MononokeHttpHandler<H> {
        MononokeHttpHandler {
            inner,
            middleware: Arc::new(self.middleware),
        }
    }
}

/// This is an instance of MononokeHttpHandlerAsService that is connected to a client. We can use
/// it to call into Gotham explicitly, or use it as a Hyper service.
#[derive(Clone)]
pub struct MononokeHttpHandlerAsService<H> {
    handler: MononokeHttpHandler<H>,
    addr: SocketAddr,
    tls_socket_data: Option<TlsSocketData>,
}

impl<H: Handler + Clone + Send + Sync + 'static + RefUnwindSafe> MononokeHttpHandlerAsService<H> {
    pub async fn call_gotham(self, req: Request<Body>) -> Response<Body> {
        let mut state = State::from_request(req, self.addr);
        if let Some(tls_socket_data) = self.tls_socket_data {
            tls_socket_data.populate_state(&mut state);
        }

        match self.handler.handle(state).await {
            Ok((_state, res)) => res,
            Err((state, err)) => err.into_response(&state),
        }
    }
}

impl<H: Handler + Clone + Send + Sync + 'static + RefUnwindSafe> Service<Request<Body>>
    for MononokeHttpHandlerAsService<H>
{
    type Response = Response<Body>;
    type Error = anyhow::Error;
    type Future = BoxFuture<'static, Result<Self::Response, Self::Error>>;

    fn poll_ready(&mut self, _cx: &mut task::Context<'_>) -> task::Poll<Result<(), Self::Error>> {
        task::Poll::Ready(Ok(()))
    }

    fn call(&mut self, req: Request<Body>) -> Self::Future {
        let this = self.clone();

        async move {
            let res = this.call_gotham(req).await;
            Ok(res)
        }
        .boxed()
    }
}

#[cfg(test)]
mod test {
    use futures::future;
    use gotham::test::TestServer;
    use gotham_derive::StateData;
    use hyper::http::StatusCode;
    use hyper::Body;

    use super::*;

    // Basic response handler for tests

    #[derive(Clone)]
    struct TestHandler;

    impl Handler for TestHandler {
        fn handle(self, state: State) -> Pin<Box<HandlerFuture>> {
            let response = Response::builder()
                .status(StatusCode::OK)
                .body(Body::empty())
                .unwrap();

            future::ready(Ok((state, response))).boxed()
        }
    }

    // Handler that expects to not be called

    #[derive(Clone)]
    struct PanicHandler;

    impl Handler for PanicHandler {
        fn handle(self, _state: State) -> Pin<Box<HandlerFuture>> {
            panic!("PanicHandler::handle was called")
        }
    }

    // Basic middleware for tests

    struct NoopMiddleware;

    #[async_trait::async_trait]
    impl Middleware for NoopMiddleware {}

    // MiddlewareValueMiddleware is used below for asserting that middleware is called in the right
    // order.

    #[derive(StateData)]
    pub struct MiddlewareValue(u64);

    struct MiddlewareValueMiddleware(Option<u64>, u64);

    #[async_trait::async_trait]
    impl Middleware for MiddlewareValueMiddleware {
        async fn inbound(&self, state: &mut State) -> Option<Response<Body>> {
            assert_eq!(state.try_borrow::<MiddlewareValue>().map(|v| v.0), self.0);
            state.put(MiddlewareValue(self.1));
            None
        }

        async fn outbound(&self, state: &mut State, _response: &mut Response<Body>) {
            assert_eq!(state.take::<MiddlewareValue>().0, self.1);

            if let Some(v) = self.0 {
                state.put(MiddlewareValue(v))
            }
        }
    }

    // InterceptMiddleware is used to check that middleware can intercept requests.

    struct InterceptMiddleware;

    #[async_trait::async_trait]
    impl Middleware for InterceptMiddleware {
        async fn inbound(&self, _state: &mut State) -> Option<Response<Body>> {
            let response = Response::builder()
                .status(StatusCode::NOT_FOUND)
                .body(Body::empty())
                .unwrap();
            Some(response)
        }

        async fn outbound(&self, _state: &mut State, response: &mut Response<Body>) {
            assert_eq!(response.status(), StatusCode::NOT_FOUND);
        }
    }

    // PanicMiddleware is used to check that intercepting a request interrupts the call chain.

    struct PanicMiddleware;

    #[async_trait::async_trait]
    impl Middleware for PanicMiddleware {
        async fn inbound(&self, _state: &mut State) -> Option<Response<Body>> {
            panic!("PanicMiddleware::inbound was called");
        }

        async fn outbound(&self, _state: &mut State, _response: &mut Response<Body>) {
            panic!("PanicMiddleware::outbound was called");
        }
    }

    #[test]
    fn test_empty() -> Result<(), anyhow::Error> {
        let handler = MononokeHttpHandler::builder().build(TestHandler);
        let server = TestServer::new(handler)?;
        let res = server.client().get("http://host/").perform()?;
        assert_eq!(res.status(), StatusCode::OK);
        Ok(())
    }

    #[test]
    fn test_noop() -> Result<(), anyhow::Error> {
        let handler = MononokeHttpHandler::builder()
            .add(NoopMiddleware)
            .add(NoopMiddleware)
            .build(TestHandler);
        let server = TestServer::new(handler)?;
        let res = server.client().get("http://host/").perform()?;
        assert_eq!(res.status(), StatusCode::OK);
        Ok(())
    }

    #[test]
    fn test_chain() -> Result<(), anyhow::Error> {
        let handler = MononokeHttpHandler::builder()
            .add(MiddlewareValueMiddleware(None, 1))
            .add(MiddlewareValueMiddleware(Some(1), 2))
            .build(TestHandler);
        let server = TestServer::new(handler)?;
        let res = server.client().get("http://host/").perform()?;
        assert_eq!(res.status(), StatusCode::OK);
        Ok(())
    }

    #[test]
    fn test_intercept_alone() -> Result<(), anyhow::Error> {
        let handler = MononokeHttpHandler::builder()
            .add(InterceptMiddleware)
            .build(TestHandler);
        let server = TestServer::new(handler)?;
        let res = server.client().get("http://host/").perform()?;
        assert_eq!(res.status(), StatusCode::NOT_FOUND);
        Ok(())
    }

    #[test]
    fn test_intercept_chain() -> Result<(), anyhow::Error> {
        let handler = MononokeHttpHandler::builder()
            .add(MiddlewareValueMiddleware(None, 1))
            .add(MiddlewareValueMiddleware(Some(1), 2))
            .add(InterceptMiddleware)
            .add(PanicMiddleware)
            .build(PanicHandler);
        let server = TestServer::new(handler)?;
        let res = server.client().get("http://host/").perform()?;
        assert_eq!(res.status(), StatusCode::NOT_FOUND);
        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.
 */

pub mod body_ext;
pub mod content_encoding;
pub mod error;
pub mod handler;
pub mod middleware;
pub mod response;
pub mod serve;
pub mod socket_data;
pub mod state_ext;


/*
 * 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::panic::RefUnwindSafe;
use std::sync::Arc;

use anyhow::Context as _;
use anyhow::Error;
use cloned::cloned;
use connection_security_checker::ConnectionSecurityChecker;
use futures::future::TryFutureExt;
use gotham::handler::Handler;
use hyper::server::conn::Http;
use openssl::ssl::Ssl;
use openssl::ssl::SslAcceptor;
use quiet_stream::QuietShutdownStream;
use slog::warn;
use slog::Logger;
use tokio::net::TcpListener;
use tokio_openssl::SslStream;

use crate::handler::MononokeHttpHandler;
use crate::socket_data::TlsSocketData;

pub async fn https<H>(
    logger: Logger,
    listener: TcpListener,
    acceptor: SslAcceptor,
    capture_session_data: bool,
    connection_security_checker: ConnectionSecurityChecker,
    handler: MononokeHttpHandler<H>,
) -> Result<(), Error>
where
    H: Handler + Clone + Send + Sync + 'static + RefUnwindSafe,
{
    let connection_security_checker = Arc::new(connection_security_checker);
    let acceptor = Arc::new(acceptor);

    loop {
        let (socket, peer_addr) = listener
            .accept()
            .await
            .context("Error accepting connections")?;

        cloned!(acceptor, logger, handler, connection_security_checker);

        let task = async move {
            let ssl = Ssl::new(acceptor.context()).context("Error creating Ssl")?;
            let ssl_socket = SslStream::new(ssl, socket).context("Error creating SslStream")?;
            let mut ssl_socket = Box::pin(ssl_socket);

            ssl_socket
                .as_mut()
                .accept()
                .await
                .context("Error performing TLS handshake")?;

            let tls_socket_data = TlsSocketData::from_ssl(
                ssl_socket.ssl(),
                connection_security_checker.as_ref(),
                capture_session_data,
            )
            .await;

            let service = handler
                .clone()
                .into_service(peer_addr, Some(tls_socket_data));

            let ssl_socket = QuietShutdownStream::new(ssl_socket);

            Http::new()
                .serve_connection(ssl_socket, service)
                .await
                .context("Error serving connection")?;

            Result::<_, Error>::Ok(())
        };

        tokio::spawn(task.map_err(move |e| {
            warn!(&logger, "HTTPS Server error: {:?}", e);
        }));
    }
}

pub async fn http<H>(
    logger: Logger,
    listener: TcpListener,
    handler: MononokeHttpHandler<H>,
) -> Result<(), Error>
where
    H: Handler + Clone + Send + Sync + 'static + RefUnwindSafe,
{
    loop {
        let (socket, peer_addr) = listener
            .accept()
            .await
            .context("Error accepting connections")?;

        cloned!(logger, handler);

        let task = async move {
            let service = handler.clone().into_service(peer_addr, None);

            let socket = QuietShutdownStream::new(socket);

            Http::new()
                .serve_connection(socket, service)
                .await
                .context("Error serving connection")?;

            Result::<_, Error>::Ok(())
        };

        tokio::spawn(task.map_err(move |e| {
            warn!(&logger, "HTTP Server error: {:?}", e);
        }));
    }
}


/*
 * 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 bytes::Bytes;
use connection_security_checker::ConnectionSecurityChecker;
use gotham::state::State;
use gotham_derive::StateData;
use openssl::ssl::SslRef;
use permission_checker::MononokeIdentity;
use permission_checker::MononokeIdentitySet;

#[derive(Clone)]
pub struct TlsSocketData {
    identities: Option<TlsCertificateIdentities>,
    session_data: Option<TlsSessionData>,
}

impl TlsSocketData {
    pub async fn from_ssl(
        ssl: &SslRef,
        connection_security_checker: &ConnectionSecurityChecker,
        capture_session_data: bool,
    ) -> Self {
        let identities = TlsCertificateIdentities::from_ssl(ssl, connection_security_checker).await;

        let session_data = if capture_session_data {
            TlsSessionData::from_ssl(ssl)
        } else {
            None
        };

        Self {
            identities,
            session_data,
        }
    }

    pub fn trusted_proxy(proxy_identites: MononokeIdentitySet) -> Self {
        Self {
            identities: Some(TlsCertificateIdentities::TrustedProxy(proxy_identites)),
            session_data: None,
        }
    }

    pub fn authenticated_identities(identities: MononokeIdentitySet) -> Self {
        Self {
            identities: Some(TlsCertificateIdentities::Authenticated(identities)),
            session_data: None,
        }
    }

    pub fn populate_state(self, state: &mut State) {
        if let Some(identities) = self.identities {
            state.put(identities);
        }

        if let Some(session_data) = self.session_data {
            state.put(session_data);
        }
    }
}

#[derive(Clone, StateData)]
pub struct TlsSessionData {
    pub client_random: Bytes,
    pub master_key: Bytes,
}

impl TlsSessionData {
    pub fn from_ssl(ssl: &SslRef) -> Option<Self> {
        let session = ssl.session()?;

        // NOTE: The OpenSSL API for getting session data is that you pass a zero-sized destination
        // to get the size. This is why we do this here.
        let mut empty: [u8; 0] = [];

        // NOTE: We use assert_eq! below, because it would be a programming error to receive less
        // than the proper size back here.

        let client_random_len = ssl.client_random(&mut empty);
        let mut client_random = vec![0; client_random_len];
        assert_eq!(client_random_len, ssl.client_random(&mut client_random[..]));

        let master_key_len = session.master_key(&mut empty);
        let mut master_key = vec![0; master_key_len];
        assert_eq!(master_key_len, session.master_key(&mut master_key[..]));

        Some(Self {
            client_random: client_random.into(),
            master_key: master_key.into(),
        })
    }
}

#[derive(Clone, StateData)]
pub enum TlsCertificateIdentities {
    TrustedProxy(MononokeIdentitySet),
    Authenticated(MononokeIdentitySet),
}

impl TlsCertificateIdentities {
    pub async fn from_ssl(
        ssl: &SslRef,
        connection_security_checker: &ConnectionSecurityChecker,
    ) -> Option<Self> {
        let peer_certificate = ssl.peer_certificate()?;
        let identities = MononokeIdentity::try_from_x509(&peer_certificate).ok()?;

        if connection_security_checker
            .check_if_trusted(&identities)
            .await
        {
            Some(TlsCertificateIdentities::TrustedProxy(identities))
        } else {
            Some(TlsCertificateIdentities::Authenticated(identities))
        }
    }
}


/*
 * 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::request_id;
use gotham::state::State;

pub trait StateExt {
    fn short_request_id(&self) -> &str;
}

const SHORT_ID_LEN: usize = 8;

impl StateExt for State {
    fn short_request_id(&self) -> &str {
        let req = request_id(self);
        &req[0..SHORT_ID_LEN.min(req.len())]
    }
}


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

use std::fmt;
use std::fmt::Debug;
use std::fmt::Display;
use std::sync::atomic::AtomicI64;
use std::sync::atomic::Ordering;
use std::sync::Arc;

use gotham::state::State;
use gotham_derive::StateData;
use hyper::Body;
use hyper::Response;

use super::Middleware;
use super::PostResponseCallbacks;

#[derive(StateData, Debug, Copy, Clone)]
pub struct RequestLoad(pub i64);

impl Display for RequestLoad {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        write!(fmt, "l")?;
        fmt::Display::fmt(&self.0, fmt)
    }
}

#[derive(Clone)]
pub struct LoadMiddleware {
    // NOTE: This should always be >0 but considering this is for monitoring right now and
    // therefore isn't super sensitive, let's be a bit conservative and try to do something
    // reasonable on underflow by using an AtomicI64.
    requests: Arc<AtomicI64>,
}

impl LoadMiddleware {
    pub fn new() -> Self {
        Self {
            requests: Arc::new(AtomicI64::new(0)),
        }
    }
}

#[async_trait::async_trait]
impl Middleware for LoadMiddleware {
    async fn inbound(&self, state: &mut State) -> Option<Response<Body>> {
        let old_request_count = self.requests.fetch_add(1, Ordering::Relaxed);
        state.put(RequestLoad(old_request_count + 1));
        None
    }

    async fn outbound(&self, state: &mut State, response: &mut Response<Body>) {
        if let Some(request_load) = state.try_take::<RequestLoad>() {
            let headers = response.headers_mut();
            headers.insert("X-Load", request_load.0.into());
        }

        if let Some(callbacks) = state.try_borrow_mut::<PostResponseCallbacks>() {
            callbacks.add({
                let requests = self.requests.clone();
                move |_| {
                    requests.fetch_sub(1, Ordering::Relaxed);
                }
            });
        }
    }
}


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

use std::fmt;
use std::fmt::Debug;
use std::fmt::Display;
use std::time::Duration;

use gotham::state::FromState;
use gotham::state::State;
use hyper::Body;
use hyper::Method;
use hyper::Response;
use hyper::StatusCode;
use hyper::Uri;
use hyper::Version;
use slog::info;
use slog::o;
use slog::Logger;
use time_ext::DurationExt;

use super::MetadataState;
use super::Middleware;
use super::PostResponseCallbacks;
use super::RequestLoad;
use crate::state_ext::StateExt;

const DIRECTION_REQUEST_IN: &str = "IN  >";
const DIRECTION_RESPONSE_OUT: &str = "OUT <";

// We have to turn out formats into macros to avoid duplicating them:

macro_rules! SLOG_FORMAT {
    () => {
        "{} {} {} \"{} {} {:?}\" {} {} {} {}"
    };
}

macro_rules! TEST_FRIENDLY_FORMAT {
    () => {
        "{} {} {} {}"
    };
}

/// We use DurationForDisplay to append ms on non-empty durations.
#[derive(Debug)]
struct DurationForDisplay(Option<u64>);

impl From<Option<Duration>> for DurationForDisplay {
    fn from(duration: Option<Duration>) -> Self {
        Self(duration.map(|d| d.as_millis_unchecked()))
    }
}

impl Display for DurationForDisplay {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        match self.0 {
            Some(duration) => {
                fmt::Display::fmt(&duration, fmt)?;
                write!(fmt, "ms")
            }
            None => write!(fmt, "-"),
        }
    }
}

#[derive(Clone)]
pub enum LogMiddleware {
    TestFriendly,
    Slog(Logger),
}

enum LogEntry {
    RequestIn,
    ResponseOut(StatusCode),
}

impl LogMiddleware {
    pub fn test_friendly() -> Self {
        Self::TestFriendly
    }

    pub fn slog(logger: Logger) -> Self {
        Self::Slog(logger)
    }
}

fn log_request_slog(logger: &Logger, state: &mut State, entry: LogEntry) -> Option<()> {
    let uri = Uri::try_borrow_from(state)?;
    if uri.path() == "/health_check" {
        return None;
    }
    let uri = uri.to_string();

    let load = *RequestLoad::borrow_from(state);
    let method = Method::borrow_from(state).clone();
    let version = *Version::borrow_from(state);
    let request_id = state.short_request_id().to_string();
    let address = MetadataState::try_borrow_from(state)
        .and_then(|metadata| metadata.metadata().client_ip())
        .map(|addr| addr.to_string());

    let callbacks = state.try_borrow_mut::<PostResponseCallbacks>()?;
    let logger = logger.new(o!("request_id" => request_id));

    match entry {
        LogEntry::RequestIn => {
            info!(
                &logger,
                SLOG_FORMAT!(),
                DIRECTION_REQUEST_IN,
                address.as_ref().map_or("-", String::as_ref),
                "-",
                method,
                uri,
                version,
                "-",
                "-",
                "-",
                load,
            );
        }
        LogEntry::ResponseOut(status) => {
            callbacks.add(move |info| {
                info!(
                    &logger,
                    SLOG_FORMAT!(),
                    DIRECTION_RESPONSE_OUT,
                    address.as_ref().map_or("-", String::as_ref),
                    info.client_hostname.as_ref().map_or("-", String::as_ref),
                    method,
                    uri,
                    version,
                    status.as_u16(),
                    info.meta.as_ref().map_or(0, |m| m.body().bytes_sent),
                    DurationForDisplay::from(info.duration),
                    load,
                );
            });
        }
    }

    None
}

fn log_request_test_friendly(state: &mut State, entry: LogEntry) -> Option<()> {
    let method = Method::try_borrow_from(state)?;
    let uri = Uri::try_borrow_from(state)?;

    match entry {
        LogEntry::RequestIn => {
            eprintln!(
                TEST_FRIENDLY_FORMAT!(),
                DIRECTION_REQUEST_IN, method, uri, "-"
            );
        }
        LogEntry::ResponseOut(status) => {
            eprintln!(
                TEST_FRIENDLY_FORMAT!(),
                DIRECTION_RESPONSE_OUT, method, uri, status
            );
        }
    };

    None
}

impl LogMiddleware {
    fn log(&self, state: &mut State, entry: LogEntry) {
        match self {
            Self::TestFriendly => {
                log_request_test_friendly(state, entry);
            }
            Self::Slog(ref logger) => {
                log_request_slog(logger, state, entry);
            }
        }
    }
}

#[async_trait::async_trait]
impl Middleware for LogMiddleware {
    async fn inbound(&self, state: &mut State) -> Option<Response<Body>> {
        let entry = LogEntry::RequestIn;
        self.log(state, entry);
        None
    }

    async fn outbound(&self, state: &mut State, response: &mut Response<Body>) {
        let entry = LogEntry::ResponseOut(response.status());
        self.log(state, entry);
    }
}


/*
 * 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::net::IpAddr;
use std::net::SocketAddr;

use cats::try_get_cats_idents;
use fbinit::FacebookInit;
use gotham::state::client_addr;
use gotham::state::FromState;
use gotham::state::State;
use gotham_derive::StateData;
use hyper::header::HeaderMap;
use hyper::Body;
use hyper::Response;
use hyper::StatusCode;
use metaconfig_types::Identity;
use metadata::Metadata;
use percent_encoding::percent_decode;
use permission_checker::MononokeIdentity;
use permission_checker::MononokeIdentitySet;
use slog::error;
use slog::Logger;

use super::Middleware;
use crate::socket_data::TlsCertificateIdentities;
use crate::state_ext::StateExt;

const ENCODED_CLIENT_IDENTITY: &str = "x-fb-validated-client-encoded-identity";
const CLIENT_IP: &str = "tfb-orig-client-ip";

#[derive(StateData, Default)]
pub struct MetadataState(Metadata);

impl MetadataState {
    pub fn metadata(&self) -> &Metadata {
        &self.0
    }
}

pub struct MetadataMiddleware {
    fb: FacebookInit,
    logger: Logger,
    internal_identity: Identity,
}

impl MetadataMiddleware {
    pub fn new(fb: FacebookInit, logger: Logger, internal_identity: Identity) -> Self {
        Self {
            fb,
            logger,
            internal_identity,
        }
    }

    fn extract_client_identities(
        &self,
        tls_certificate_identities: TlsCertificateIdentities,
        headers: &HeaderMap,
    ) -> Option<MononokeIdentitySet> {
        match tls_certificate_identities {
            TlsCertificateIdentities::TrustedProxy(idents) => {
                match request_identities_from_headers(headers) {
                    Some(identies_from_headers) => Some(identies_from_headers),
                    None => Some(idents),
                }
            }
            TlsCertificateIdentities::Authenticated(idents) => Some(idents),
        }
    }
}

fn request_ip_from_headers(headers: &HeaderMap) -> Option<IpAddr> {
    let header = headers.get(CLIENT_IP)?;
    let header = header.to_str().ok()?;
    let ip = header.parse().ok()?;
    Some(ip)
}

fn request_identities_from_headers(headers: &HeaderMap) -> Option<MononokeIdentitySet> {
    let encoded_identities = headers.get(ENCODED_CLIENT_IDENTITY)?;
    let json_identities = percent_decode(encoded_identities.as_bytes())
        .decode_utf8()
        .ok()?;
    MononokeIdentity::try_from_json_encoded(&json_identities).ok()
}

#[async_trait::async_trait]
impl Middleware for MetadataMiddleware {
    async fn inbound(&self, state: &mut State) -> Option<Response<Body>> {
        let cert_idents = TlsCertificateIdentities::try_take_from(state);
        let mut metadata = Metadata::default();

        if let Some(headers) = HeaderMap::try_borrow_from(state) {
            metadata = metadata.set_client_ip(request_ip_from_headers(headers));

            let maybe_identities = {
                let maybe_cat_idents =
                    match try_get_cats_idents(self.fb, headers, &self.internal_identity) {
                        Err(e) => {
                            let msg = format!("Error extracting CATs identities: {}.", &e,);
                            error!(self.logger, "{}", &msg,);
                            let response = Response::builder()
                                .status(StatusCode::UNAUTHORIZED)
                                .body(
                                    format!(
                                        "{{\"message:\"{}\", \"request_id\":\"{}\"}}",
                                        msg,
                                        state.short_request_id()
                                    )
                                    .into(),
                                )
                                .expect("Couldn't build http response");

                            return Some(response);
                        }
                        Ok(maybe_cats) => maybe_cats,
                    };

                let maybe_tls_or_proxied_idents: Option<MononokeIdentitySet> =
                    cert_idents.and_then(|x| self.extract_client_identities(x, headers));

                match (maybe_cat_idents, maybe_tls_or_proxied_idents) {
                    (None, None) => None,
                    (Some(cat_idents), Some(tls_or_proxied_idents)) => {
                        Some(cat_idents.union(&tls_or_proxied_idents).cloned().collect())
                    }
                    (Some(cat_idents), None) => Some(cat_idents),
                    (None, Some(tls_or_proxied_idents)) => Some(tls_or_proxied_idents),
                }
            };
            if let Some(identities) = maybe_identities {
                metadata = metadata.set_identities(identities)
            }

            metadata.update_client_untrusted(
                metadata::security::is_client_untrusted(|h| {
                    Ok(headers
                        .get(h)
                        .map(|h| h.to_str().map(|s| s.to_owned()))
                        .transpose()?)
                })
                .unwrap_or_default(),
            );
        }

        // For the IP, we can fallback to the peer IP
        if metadata.client_ip().is_none() {
            metadata = metadata.set_client_ip(client_addr(state).as_ref().map(SocketAddr::ip));
        }

        state.put(MetadataState(metadata));

        None
    }
}


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

use std::panic::RefUnwindSafe;

use gotham::state::State;
use hyper::Body;
use hyper::Response;

pub mod load;
pub mod log;
pub mod metadata;
pub mod post_request;
pub mod scuba;
pub mod server_identity;
pub mod timer;
pub mod tls_session_data;

pub use self::load::LoadMiddleware;
pub use self::load::RequestLoad;
pub use self::log::LogMiddleware;
pub use self::metadata::MetadataMiddleware;
pub use self::metadata::MetadataState;
pub use self::post_request::PostResponseCallbacks;
pub use self::post_request::PostResponseConfig;
pub use self::post_request::PostResponseInfo;
pub use self::post_request::PostResponseMiddleware;
pub use self::scuba::HttpScubaKey;
pub use self::scuba::ScubaHandler;
pub use self::scuba::ScubaMiddleware;
pub use self::scuba::ScubaMiddlewareState;
pub use self::server_identity::ServerIdentityMiddleware;
pub use self::timer::HeadersDuration;
pub use self::timer::RequestStartTime;
pub use self::timer::TimerMiddleware;
pub use self::tls_session_data::TlsSessionDataMiddleware;

#[async_trait::async_trait]
pub trait Middleware: 'static + RefUnwindSafe + Send + Sync {
    async fn inbound(&self, _state: &mut State) -> Option<Response<Body>> {
        // Implement inbound to perform pre-request actions, such as putting something in the
        // state.

        None
    }

    async fn outbound(&self, _state: &mut State, _response: &mut Response<Body>) {
        // Implement outbound to perform post-request actions, such as logging the response status
        // code.
    }
}


/*
 * 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::panic::RefUnwindSafe;
use std::time::Duration;
use std::time::Instant;

use anyhow::Error;
use async_trait::async_trait;
use cached_config::ConfigHandle;
use futures::future;
use futures::prelude::*;
use futures::Future;
use futures_stats::StreamStats;
use gotham::state::FromState;
use gotham::state::State;
use gotham_derive::StateData;
use hyper::body::Body;
use hyper::Response;
use permission_checker::MononokeIdentitySetExt;
use tokio::task;
use trust_dns_resolver::TokioAsyncResolver;

use super::MetadataState;
use super::Middleware;
use super::RequestStartTime;
use crate::response::PendingResponseMeta;
use crate::response::PendingStreamStats;
use crate::response::ResponseMeta;

type Callback = Box<dyn FnOnce(&PostResponseInfo) + Send + 'static>;

/// Information passed to each post-request callback.
pub struct PostResponseInfo {
    pub duration: Option<Duration>,
    pub client_hostname: Option<String>,
    pub meta: Option<ResponseMeta>,
    pub stream_stats: Option<StreamStats>,
}

impl PostResponseInfo {
    pub fn first_error(&self) -> Option<&Error> {
        self.meta.as_ref()?.body().error_meta.errors.first()
    }

    pub fn error_count(&self) -> u64 {
        self.meta
            .as_ref()
            .map_or(0, |m| m.body().error_meta.error_count())
    }
}

/// Trait allowing post-request callbacks to be configured dynamically.
pub trait PostResponseConfig: Clone + Send + Sync + RefUnwindSafe + 'static {
    /// Specify whether the middleware should perform a potentially
    /// expensive reverse DNS lookup of the client's hostname.
    fn resolve_hostname(&self) -> bool {
        true
    }
}

#[derive(Clone)]
pub struct DefaultConfig;

impl PostResponseConfig for DefaultConfig {}

impl<C: PostResponseConfig> PostResponseConfig for ConfigHandle<C> {
    fn resolve_hostname(&self) -> bool {
        self.get().resolve_hostname()
    }
}

/// Middleware that allows the application to register callbacks which will
/// be run upon request completion.
pub struct PostResponseMiddleware<C> {
    config: C,
}

impl<C> PostResponseMiddleware<C> {
    pub fn with_config(config: C) -> Self {
        Self { config }
    }
}

impl Default for PostResponseMiddleware<DefaultConfig> {
    fn default() -> Self {
        PostResponseMiddleware::with_config(DefaultConfig)
    }
}

#[async_trait]
impl<C: PostResponseConfig> Middleware for PostResponseMiddleware<C> {
    async fn inbound(&self, state: &mut State) -> Option<Response<Body>> {
        state.put(PostResponseCallbacks::new());
        None
    }

    async fn outbound(&self, state: &mut State, _response: &mut Response<Body>) {
        let config = self.config.clone();
        let start_time = RequestStartTime::try_borrow_from(state).map(|t| t.0);
        let hostname_future = MetadataState::try_borrow_from(state).map(resolve_hostname);
        let meta = PendingResponseMeta::try_take_from(state);
        let stream_stats = PendingStreamStats::try_take_from(state);

        if let Some(callbacks) = state.try_take::<PostResponseCallbacks>() {
            task::spawn(callbacks.run(config, start_time, meta, stream_stats, hostname_future));
        }
    }
}

/// A collection of callbacks that will run once the request has completed.
#[derive(StateData)]
pub struct PostResponseCallbacks {
    callbacks: Vec<Callback>,
}

impl PostResponseCallbacks {
    fn new() -> Self {
        Self {
            callbacks: Vec::new(),
        }
    }

    /// Add a callback that will be run once the request has completed. This is
    /// primarily useful for things like logging.
    ///
    /// Note that the callbacks are run serially in a task on the Tokio runtime.
    /// Although these callbacks are not asynchronous, they SHOULD NOT BLOCK as
    /// this could block the server from handling other requests.
    pub fn add<F>(&mut self, callback: F)
    where
        F: FnOnce(&PostResponseInfo) + Send + 'static,
    {
        self.callbacks.push(Box::new(callback));
    }

    async fn run<C, H>(
        self,
        config: C,
        start_time: Option<Instant>,
        meta: Option<PendingResponseMeta>,
        stream_stats: Option<PendingStreamStats>,
        hostname_future: Option<H>,
    ) where
        C: PostResponseConfig,
        H: Future<Output = Option<String>> + Send + 'static,
    {
        let Self { callbacks } = self;

        let meta = match meta {
            Some(meta) => Some(meta.finish().await),
            None => None,
        };

        let stream_stats = match stream_stats {
            Some(stream_stats) => stream_stats.finish().await,
            None => None,
        };

        // Capture elapsed time before waiting for the client hostname to resolve.
        let duration = start_time.map(|start| start.elapsed());

        // Resolve client hostname if enabled.
        let client_hostname = match hostname_future {
            Some(hostname) if config.resolve_hostname() => hostname.await,
            _ => None,
        };

        let info = PostResponseInfo {
            duration,
            client_hostname,
            meta,
            stream_stats,
        };

        for callback in callbacks {
            callback(&info);
        }
    }
}

// Hostname of the client is for non-critical use only (best-effort lookup):
pub fn resolve_hostname(
    metadata_state: &MetadataState,
) -> impl Future<Output = Option<String>> + 'static {
    // XXX: Can't make this an async fn because the resulting Future would
    // have a non-'static lifetime (due to the &ClientIdentity argument).

    let metadata = metadata_state.metadata();
    // 1) We're extracting it from identities (which requires no remote calls)
    if let Some(client_hostname) = metadata.identities().hostname().map(|h| h.to_string()) {
        return future::ready(Some(client_hostname)).left_future();
    }
    // 2) Perform a reverse DNS lookup of the client's IP address to determine
    // its hostname.
    let address = metadata.client_ip().cloned();
    (async move {
        let resolver = TokioAsyncResolver::tokio_from_system_conf().ok()?;
        let hosts = resolver.reverse_lookup(address?).await.ok()?;
        let host = hosts.iter().next()?;
        Some(host.to_string().trim_end_matches('.').to_string())
    })
    .right_future()
}


/*
 * 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::marker::PhantomData;
use std::num::NonZeroU64;
use std::panic::RefUnwindSafe;

use futures_stats::FutureStats;
use gotham::state::FromState;
use gotham::state::State;
use gotham_derive::StateData;
use hyper::header;
use hyper::header::AsHeaderName;
use hyper::header::HeaderMap;
use hyper::Body;
use hyper::Method;
use hyper::Response;
use hyper::StatusCode;
use hyper::Uri;
use scuba_ext::MononokeScubaSampleBuilder;
use scuba_ext::ScubaValue;
use time_ext::DurationExt;

use super::HeadersDuration;
use super::RequestLoad;
use crate::middleware::MetadataState;
use crate::middleware::Middleware;
use crate::middleware::PostResponseCallbacks;
use crate::middleware::PostResponseInfo;
use crate::response::HeadersMeta;
use crate::state_ext::StateExt;

/// HTTP header used to correlate the request with the client-side logging
const CLIENT_CORRELATOR: &str = "x-client-correlator";

/// Common HTTP-related Scuba columns that the middlware will set automatically.
/// Applications using the middleware are encouraged to follow a similar pattern
/// when adding application-specific columns to the `ScubaMiddlewareState`.
#[derive(Copy, Clone, Debug)]
pub enum HttpScubaKey {
    /// The status code for this response
    HttpStatus,
    /// The HTTP Path requested by the client.
    HttpPath,
    /// The HTTP Query string provided by the client.
    HttpQuery,
    /// The HTTP Method requested by the client.
    HttpMethod,
    /// The Http "Host" header sent by the client.
    HttpHost,
    /// The HTTP User Agent provided by the client.
    HttpUserAgent,
    /// The "Content-Length" advertised by the client in their request.
    RequestContentLength,
    /// The "Content-Length" we returned in our response.
    ResponseContentLength,
    /// The Content-Encoding we used for our response.
    ResponseContentEncoding,
    /// The IP of the connecting client.
    ClientIp,
    /// The client correlator submitted by the client, if any.
    ClientCorrelator,
    /// The client identities received for the client, if any.
    ClientIdentities,
    /// The request load when this request was admitted.
    RequestLoad,
    /// A unique ID identifying this request.
    RequestId,
    /// How long it took to send headers.
    HeadersDurationMs,
    /// How long it took to finish sending the response.
    DurationMs,
    /// The hostname of the connecting client.
    ClientHostname,
    /// How many bytes were sent to the client (should normally equal the content length)
    ResponseBytesSent,
    /// How many bytes were received from the client (should normally equal the content length)
    RequestBytesReceived,
}

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

        match self {
            HttpStatus => "http_status",
            HttpPath => "http_path",
            HttpQuery => "http_query",
            HttpMethod => "http_method",
            HttpHost => "http_host",
            HttpUserAgent => "http_user_agent",
            RequestContentLength => "request_content_length",
            ResponseContentLength => "response_content_length",
            ResponseContentEncoding => "response_content_encoding",
            ClientIp => "client_ip",
            ClientCorrelator => "client_correlator",
            ClientIdentities => "client_identities",
            RequestLoad => "request_load",
            RequestId => "request_id",
            HeadersDurationMs => "headers_duration_ms",
            DurationMs => "duration_ms",
            ClientHostname => "client_hostname",
            ResponseBytesSent => "response_bytes_sent",
            RequestBytesReceived => "request_bytes_received",
        }
    }
}

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

pub trait ScubaHandler: Send + 'static {
    fn from_state(state: &State) -> Self;

    fn populate_scuba(self, info: &PostResponseInfo, scuba: &mut MononokeScubaSampleBuilder);
}

#[derive(Clone)]
pub struct ScubaMiddleware<H> {
    scuba: MononokeScubaSampleBuilder,
    _phantom: PhantomHandler<H>,
}

impl<H> ScubaMiddleware<H> {
    pub fn new(scuba: MononokeScubaSampleBuilder) -> Self {
        Self {
            scuba,
            _phantom: PhantomHandler(PhantomData),
        }
    }
}

/// Phantom type that ensures that `ScubaMiddleware` can be `RefUnwindSafe` and
/// `Sync` without imposing those constraints on its type parameter.
///
/// Since `ScubaMiddleware` is generic over its handler type, in order for it
/// to automatically implement `Sync` and `RefUnwindSafe` (which are required
/// by the `Middleware` trait), the handler would ordinarily need to also
/// be subject to those constraints.
///
/// This isn't actually necessary since the middleware itself does not contain
/// an instance of the handler. (The handler is instantiated shortly before it
/// is used in a post-request callback.) Therefore, it is safe to manually mark
/// `PhantomData<H>` with these traits via a wrapper struct, ensuring that
/// the middleware automatically implements the required marker traits.
#[derive(Clone)]
struct PhantomHandler<H>(PhantomData<H>);

impl<H> RefUnwindSafe for PhantomHandler<H> {}

unsafe impl<H> Sync for PhantomHandler<H> {}

fn add_header<'a, Header, Converter, Value>(
    scuba: &mut MononokeScubaSampleBuilder,
    headers: &'a HeaderMap,
    scuba_key: HttpScubaKey,
    header: Header,
    convert: Converter,
) -> Option<&'a str>
where
    Header: AsHeaderName,
    Converter: FnOnce(&str) -> Value,
    Value: Into<ScubaValue>,
{
    if let Some(header_val) = headers.get(header) {
        if let Ok(header_val) = header_val.to_str() {
            scuba
                .entry(scuba_key)
                .or_insert_with(|| convert(header_val).into());
            return Some(header_val);
        }
    }

    None
}

fn log_stats<H: ScubaHandler>(state: &mut State, status_code: &StatusCode) -> Option<()> {
    let mut scuba = state.try_take::<ScubaMiddlewareState>()?.0;

    scuba.add(HttpScubaKey::HttpStatus, status_code.as_u16());

    if let Some(uri) = Uri::try_borrow_from(state) {
        scuba.add(HttpScubaKey::HttpPath, uri.path());
        if let Some(query) = uri.query() {
            scuba.add(HttpScubaKey::HttpQuery, query);
        }
    }

    if let Some(method) = Method::try_borrow_from(state) {
        scuba.add(HttpScubaKey::HttpMethod, method.to_string());
    }

    if let Some(headers) = HeaderMap::try_borrow_from(state) {
        add_header(
            &mut scuba,
            headers,
            HttpScubaKey::HttpHost,
            header::HOST,
            |header| header.to_string(),
        );

        add_header(
            &mut scuba,
            headers,
            HttpScubaKey::RequestContentLength,
            header::CONTENT_LENGTH,
            |header| header.parse::<u64>().unwrap_or(0),
        );

        add_header(
            &mut scuba,
            headers,
            HttpScubaKey::HttpUserAgent,
            header::USER_AGENT,
            |header| header.to_string(),
        );

        add_header(
            &mut scuba,
            headers,
            HttpScubaKey::ClientCorrelator,
            CLIENT_CORRELATOR,
            |header| header.to_string(),
        );
    }

    if let Some(metadata_state) = MetadataState::try_borrow_from(state) {
        let metadata = metadata_state.metadata();
        if let Some(ref address) = metadata.client_ip() {
            scuba.add(HttpScubaKey::ClientIp, address.to_string());
        }

        let identities = metadata.identities();
        scuba.sample_for_identities(identities);
        let identities: Vec<_> = identities.iter().map(|i| i.to_string()).collect();
        scuba.add(HttpScubaKey::ClientIdentities, identities);
    }

    if let Some(request_load) = RequestLoad::try_borrow_from(state) {
        scuba.add(HttpScubaKey::RequestLoad, request_load.0);
    }

    scuba.add(HttpScubaKey::RequestId, state.short_request_id());

    if let Some(HeadersDuration(duration)) = HeadersDuration::try_borrow_from(state) {
        scuba.add(
            HttpScubaKey::HeadersDurationMs,
            duration.as_millis_unchecked(),
        );
    }

    let handler = H::from_state(state);

    let callbacks = state.try_borrow_mut::<PostResponseCallbacks>()?;
    callbacks.add(move |info| {
        if let Some(duration) = info.duration {
            // If tunables say we should, log high values unsampled
            if let Ok(threshold) = tunables::tunables()
                .edenapi_unsampled_duration_threshold_ms()
                .unwrap_or_default()
                .try_into()
            {
                if duration.as_millis_unchecked() > threshold {
                    scuba.unsampled();
                }
            }
            scuba.add(HttpScubaKey::DurationMs, duration.as_millis_unchecked());
        }

        if let Some(client_hostname) = info.client_hostname.as_deref() {
            scuba.add(HttpScubaKey::ClientHostname, client_hostname);
        }

        if let Some(meta) = info.meta.as_ref() {
            match *meta.headers() {
                HeadersMeta::Sized(content_length) => {
                    scuba.add(HttpScubaKey::ResponseContentLength, content_length);
                }
                HeadersMeta::Compressed(compression) => {
                    scuba.add(HttpScubaKey::ResponseContentEncoding, compression.as_str());
                }
                HeadersMeta::Chunked => {}
            }

            scuba.add(HttpScubaKey::ResponseBytesSent, meta.body().bytes_sent);
        }

        if let Some(stats) = info.stream_stats.as_ref() {
            scuba.add_prefixed_stream_stats(stats);
        }

        handler.populate_scuba(info, &mut scuba);

        scuba.log();
    });

    Some(())
}

#[derive(StateData)]
pub struct ScubaMiddlewareState(MononokeScubaSampleBuilder);

impl ScubaMiddlewareState {
    pub fn add<K, V>(&mut self, key: K, value: V) -> &mut Self
    where
        K: Into<String>,
        V: Into<ScubaValue>,
    {
        self.0.add(key, value);
        self
    }

    /// Borrow the ScubaMiddlewareState, if any, and add a key-value pair to it.
    pub fn try_borrow_add<K, V>(state: &mut State, key: K, value: V)
    where
        K: Into<String>,
        V: Into<ScubaValue>,
    {
        let mut scuba = state.try_borrow_mut::<Self>();
        if let Some(ref mut scuba) = scuba {
            scuba.add(key, value);
        }
    }

    pub fn try_set_sampling_rate(state: &mut State, rate: NonZeroU64) {
        let mut scuba = state.try_borrow_mut::<Self>();
        if let Some(ref mut scuba) = scuba {
            scuba.0.sampled_unless_verbose(rate);
        }
    }

    pub fn try_set_future_stats(state: &mut State, future_stats: &FutureStats) {
        let mut scuba = state.try_borrow_mut::<Self>();
        if let Some(ref mut scuba) = scuba {
            scuba.0.add_future_stats(future_stats);
        }
    }

    pub fn maybe_add<K, V>(scuba: &mut Option<&mut ScubaMiddlewareState>, key: K, value: V)
    where
        K: Into<String>,
        V: Into<ScubaValue>,
    {
        if let Some(ref mut scuba) = scuba {
            scuba.add(key, value);
        }
    }
}

#[async_trait::async_trait]
impl<H: ScubaHandler> Middleware for ScubaMiddleware<H> {
    async fn inbound(&self, state: &mut State) -> Option<Response<Body>> {
        // Reset the scuba sequence counter for each request.
        state.put(ScubaMiddlewareState(self.scuba.clone().with_seq("seq")));
        None
    }

    async fn outbound(&self, state: &mut State, response: &mut Response<Body>) {
        if let Some(uri) = Uri::try_borrow_from(state) {
            if uri.path() == "/health_check" {
                return;
            }
        }

        log_stats::<H>(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::collections::HashMap;
use std::env;

use anyhow::Error;
use gotham::helpers::http::header::X_REQUEST_ID;
use gotham::state::State;
use hyper::header::HeaderValue;
use hyper::Body;
use hyper::Response;

use super::Middleware;
use crate::state_ext::StateExt;

pub struct ServerIdentityMiddleware {
    headers: HashMap<&'static str, HeaderValue>,
}

impl ServerIdentityMiddleware {
    pub fn new(server_name: HeaderValue) -> Self {
        let mut headers = HashMap::new();

        headers.insert("Server", server_name);

        // NOTE: We ignore errors here â€” those will happen if environment variables are missing,
        // which is fine.
        let _ = Self::add_tw_task(&mut headers);
        let _ = Self::add_tw_task_version(&mut headers);
        let _ = Self::add_tw_canary_id(&mut headers);

        Self { headers }
    }

    fn add_tw_task(headers: &mut HashMap<&'static str, HeaderValue>) -> Result<(), Error> {
        let tw_job_cluster = env::var("TW_JOB_CLUSTER")?;
        let tw_job_user = env::var("TW_JOB_USER")?;
        let tw_job_name = env::var("TW_JOB_NAME")?;
        let tw_task_id = env::var("TW_TASK_ID")?;
        let task = format!(
            "{}/{}/{}/{}",
            tw_job_cluster, tw_job_user, tw_job_name, tw_task_id
        );
        let header = HeaderValue::from_str(&task)?;
        headers.insert("X-TW-Task", header);
        Ok(())
    }

    fn add_tw_task_version(headers: &mut HashMap<&'static str, HeaderValue>) -> Result<(), Error> {
        let tw_task_version = env::var("TW_TASK_VERSION")?;
        let header = HeaderValue::from_str(&tw_task_version)?;
        headers.insert("X-TW-Task-Version", header);
        Ok(())
    }

    fn add_tw_canary_id(headers: &mut HashMap<&'static str, HeaderValue>) -> Result<(), Error> {
        let tw_canary_id = env::var("TW_CANARY_ID")?;
        let header = HeaderValue::from_str(&tw_canary_id)?;
        headers.insert("X-TW-Canary-Id", header);
        Ok(())
    }
}

#[async_trait::async_trait]
impl Middleware for ServerIdentityMiddleware {
    async fn outbound(&self, state: &mut State, response: &mut Response<Body>) {
        let headers = response.headers_mut();

        for (header, value) in self.headers.iter() {
            headers.insert(*header, value.clone());
        }

        if let Ok(id) = HeaderValue::from_str(state.short_request_id()) {
            headers.insert(X_REQUEST_ID, id);
        }
    }
}


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

use gotham::state::State;
use gotham_derive::StateData;
use hyper::Body;
use hyper::Response;

use super::Middleware;

#[derive(StateData, Debug, Copy, Clone)]
pub struct RequestStartTime(pub Instant);

#[derive(StateData, Debug, Copy, Clone)]
pub struct HeadersDuration(pub Duration);

#[derive(Clone)]
pub struct TimerMiddleware;

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

#[async_trait::async_trait]
impl Middleware for TimerMiddleware {
    async fn inbound(&self, state: &mut State) -> Option<Response<Body>> {
        state.put(RequestStartTime(Instant::now()));
        None
    }

    async fn outbound(&self, state: &mut State, _response: &mut Response<Body>) {
        if let Some(RequestStartTime(start)) = state.try_borrow() {
            let headers_duration = start.elapsed();
            state.put(HeadersDuration(headers_duration));
        }
    }
}


/*
 * 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::fs::File;
use std::fs::OpenOptions;
use std::io::Write;
use std::path::Path;
use std::sync::Arc;
use std::sync::Mutex;

use anyhow::Error;
use gotham::state::FromState;
use gotham::state::State;
use hyper::Body;
use hyper::Response;

use super::Middleware;
use crate::socket_data::TlsSessionData;

pub struct TlsSessionDataMiddleware {
    log_file: Option<Arc<Mutex<File>>>,
}

impl TlsSessionDataMiddleware {
    pub fn new<L: AsRef<Path>>(log_file: Option<L>) -> Result<Self, Error> {
        let log_file: Result<_, Error> = log_file
            .map(|log_file| {
                let log_file = OpenOptions::new()
                    .create(true)
                    .append(true)
                    .open(log_file)?;
                Ok(Arc::new(Mutex::new(log_file)))
            })
            .transpose();

        Ok(Self {
            log_file: log_file?,
        })
    }

    fn log(&self, state: &mut State) -> Option<()> {
        let mut log_file = self.log_file.as_ref()?.lock().expect("Poisoned lock");
        let session_data = TlsSessionData::try_take_from(state)?;

        // See https://developer.mozilla.org/en-US/docs/Mozilla/Projects/NSS/Key_Log_Format for
        // formatting.
        let payload = format!(
            "CLIENT_RANDOM {} {}\n",
            hex::encode(session_data.client_random.as_ref()),
            hex::encode(session_data.master_key.as_ref())
        );

        log_file.write_all(payload.as_bytes()).ok()?;

        Some(())
    }
}

#[async_trait::async_trait]
impl Middleware for TlsSessionDataMiddleware {
    async fn inbound(&self, state: &mut State) -> Option<Response<Body>> {
        self.log(state);
        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 super::stream::CompressedResponseStream;
use super::stream::ResponseStream;
use super::stream_ext::CaptureFirstErr;
use super::stream_ext::EndOnErr;
use crate::content_encoding::ContentEncoding;

pub trait ContentMetaProvider {
    /// Provide the content (i.e. Content-Encoding) for the underlying content. This will be sent
    /// to the client.
    fn content_encoding(&self) -> ContentEncoding;

    /// Provide the length of the content in this stream, if available (i.e. Content-Length). If
    /// provided, this must be the actual length of the stream. If missing, the transfer will be
    /// chunked.
    fn content_length(&self) -> Option<u64>;
}

impl<S> ContentMetaProvider for ResponseStream<S> {
    fn content_length(&self) -> Option<u64> {
        ResponseStream::content_length(self)
    }

    fn content_encoding(&self) -> ContentEncoding {
        ContentEncoding::Identity
    }
}

impl ContentMetaProvider for CompressedResponseStream<'_> {
    fn content_length(&self) -> Option<u64> {
        None
    }

    fn content_encoding(&self) -> ContentEncoding {
        ContentEncoding::Compressed(self.content_compression())
    }
}

/// Provide an implementation of ContentMetaProvider that propagates through Either (i.e. left_stream(),
/// right_stream()).
impl<A, B> ContentMetaProvider for futures::future::Either<A, B>
where
    A: ContentMetaProvider,
    B: ContentMetaProvider,
{
    fn content_length(&self) -> Option<u64> {
        // left_stream(), right_stream() doesn't change the stream data.
        match self {
            Self::Left(a) => a.content_length(),
            Self::Right(b) => b.content_length(),
        }
    }

    fn content_encoding(&self) -> ContentEncoding {
        // left_stream(), right_stream() doesn't change the stream data.
        match self {
            Self::Left(a) => a.content_encoding(),
            Self::Right(b) => b.content_encoding(),
        }
    }
}

impl<S, F> ContentMetaProvider for futures::stream::InspectOk<S, F>
where
    S: ContentMetaProvider,
{
    fn content_length(&self) -> Option<u64> {
        // inspect_ok doesn't change the stream data.
        self.get_ref().content_length()
    }

    fn content_encoding(&self) -> ContentEncoding {
        // inspect_ok doesn't change the stream data.
        self.get_ref().content_encoding()
    }
}

impl<S, E> ContentMetaProvider for EndOnErr<S, E>
where
    S: ContentMetaProvider,
{
    fn content_length(&self) -> Option<u64> {
        // If an error occurs, the stream will end prematurely, so the content
        // length here may not be correct. However, this should be OK as it
        // would allow the client to detect an incomplete response and raise
        // an appropriate error.
        self.get_ref().content_length()
    }

    fn content_encoding(&self) -> ContentEncoding {
        // end_on_err doesn't change the data's encoding.
        self.get_ref().content_encoding()
    }
}

impl<S, E> ContentMetaProvider for CaptureFirstErr<S, E>
where
    S: ContentMetaProvider,
{
    fn content_length(&self) -> Option<u64> {
        // Since errors may be filtered out without intending to fail the response, we cannot
        // provide a content length.
        None
    }

    fn content_encoding(&self) -> ContentEncoding {
        // CaptureFirstErr does not change the stream's encoding.
        self.get_ref().content_encoding()
    }
}


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

pub struct ErrorMeta<E> {
    /// Errors that were captured
    pub errors: Vec<E>,

    /// Extra erorrs that were observed but not captured.
    pub extra_error_count: u64,
}

impl<E> ErrorMeta<E> {
    pub fn new() -> Self {
        Self {
            errors: Vec::new(),
            extra_error_count: 0,
        }
    }

    pub fn one_error(e: E) -> Self {
        Self {
            errors: vec![e],
            extra_error_count: 0,
        }
    }

    pub fn error_count(&self) -> u64 {
        // NOTE: unwrap is safe since we if we have all those elements in a Vec, we're going to be
        // able to fit that into a u64.
        let n: u64 = self.errors.len().try_into().unwrap();
        n + self.extra_error_count
    }
}

pub trait ErrorMetaProvider<E> {
    fn report_errors(self: Pin<&mut Self>, error_meta: &mut ErrorMeta<E>);
}


/*
 * 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 content_meta;
mod error_meta;
mod response;
mod response_meta;
mod signal_stream;
mod stream;
mod stream_ext;
mod stream_stats;

pub use content_meta::ContentMetaProvider;
pub use error_meta::ErrorMeta;
pub use error_meta::ErrorMetaProvider;
pub use response::build_error_response;
pub use response::build_response;
pub use response::BytesBody;
pub use response::EmptyBody;
pub use response::StreamBody;
pub use response::TryIntoResponse;
pub use response_meta::BodyMeta;
pub use response_meta::HeadersMeta;
pub use response_meta::PendingResponseMeta;
pub use response_meta::ResponseMeta;
pub use stream::encode_stream;
pub use stream::CompressedResponseStream;
pub use stream::ResponseStream;
pub use stream_ext::ResponseTryStreamExt;
pub use stream_stats::PendingStreamStats;


/*
 * 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 bytes::Bytes;
use futures::channel::oneshot;
use futures::stream::Stream;
use futures::stream::StreamExt;
use futures_stats::TimedStreamExt;
use gotham::handler::HandlerError;
use gotham::helpers::http::response::create_response;
use gotham::state::State;
use hyper::header::HeaderValue;
use hyper::header::CONTENT_ENCODING;
use hyper::header::CONTENT_LENGTH;
use hyper::header::CONTENT_TYPE;
use hyper::Body;
use hyper::Response;
use hyper::StatusCode;
use mime::Mime;

use super::content_meta::ContentMetaProvider;
use super::error_meta::ErrorMetaProvider;
use super::response_meta::HeadersMeta;
use super::response_meta::PendingResponseMeta;
use super::signal_stream::SignalStream;
use super::stream_stats::PendingStreamStats;
use crate::content_encoding::ContentEncoding;
use crate::error::ErrorFormatter;
use crate::error::HttpError;

pub trait TryIntoResponse {
    fn try_into_response(self, state: &mut State) -> Result<Response<Body>, Error>;
}

pub fn build_response<IR: TryIntoResponse, F: ErrorFormatter>(
    res: Result<IR, HttpError>,
    mut state: State,
    formatter: &F,
) -> Result<(State, Response<Body>), (State, HandlerError)> {
    let res = res.and_then(|c| {
        c.try_into_response(&mut state)
            .context("try_into_response failed!")
            .map_err(HttpError::e500)
    });

    match res {
        Ok(res) => Ok((state, res)),
        Err(err) => build_error_response(err, state, formatter),
    }
}

pub fn build_error_response<F: ErrorFormatter>(
    err: HttpError,
    mut state: State,
    formatter: &F,
) -> Result<(State, Response<Body>), (State, HandlerError)> {
    let formatted = formatter.format(&err.error, &state);

    state.put(PendingResponseMeta::error(err.error));

    match formatted {
        Ok((body, mime)) => {
            let res = create_response(&state, err.status_code, mime, body);
            Ok((state, res))
        }
        Err(error) => Err((state, error.into())),
    }
}

pub struct EmptyBody;

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

impl TryIntoResponse for EmptyBody {
    fn try_into_response(self, state: &mut State) -> Result<Response<Body>, Error> {
        state.put(PendingStreamStats::none());
        state.put(PendingResponseMeta::immediate(0));

        Response::builder()
            .status(StatusCode::OK)
            .header(CONTENT_LENGTH, 0)
            .body(Body::empty())
            .map_err(Error::from)
    }
}

pub struct BytesBody<B> {
    bytes: B,
    mime: Mime,
}

impl<B> BytesBody<B> {
    pub fn new(bytes: B, mime: Mime) -> Self {
        Self { bytes, mime }
    }
}

impl<B> TryIntoResponse for BytesBody<B>
where
    B: Into<Bytes>,
{
    fn try_into_response(self, state: &mut State) -> Result<Response<Body>, Error> {
        let bytes = self.bytes.into();
        let mime_header: HeaderValue = self.mime.as_ref().parse()?;

        let size = bytes.len().try_into()?;
        state.put(PendingResponseMeta::immediate(size));
        state.put(PendingStreamStats::none());

        Response::builder()
            .header(CONTENT_TYPE, mime_header)
            .status(StatusCode::OK)
            .body(bytes.into())
            .map_err(Error::from)
    }
}

pub struct StreamBody<S> {
    stream: S,
    mime: Mime,
    pub partial: bool,
}

impl<S> StreamBody<S> {
    pub fn new(stream: S, mime: Mime) -> Self {
        Self {
            stream,
            mime,
            partial: false,
        }
    }
}

impl<S> TryIntoResponse for StreamBody<S>
where
    S: Stream<Item = Bytes> + ContentMetaProvider + ErrorMetaProvider<Error> + Send + 'static,
{
    fn try_into_response(self, state: &mut State) -> Result<Response<Body>, Error> {
        let Self {
            stream,
            mime,
            partial,
        } = self;

        let status = if partial {
            StatusCode::PARTIAL_CONTENT
        } else {
            StatusCode::OK
        };

        let mime_header: HeaderValue = mime.as_ref().parse()?;

        let content_encoding = stream.content_encoding();
        let content_length = stream.content_length();

        let res = Response::builder()
            .header(CONTENT_TYPE, mime_header)
            .header(CONTENT_ENCODING, content_encoding)
            .status(status);

        let (res, headers_meta) = match content_encoding {
            ContentEncoding::Compressed(compression) => (res, HeadersMeta::Compressed(compression)),
            ContentEncoding::Identity => match content_length {
                Some(content_length) => (
                    res.header(CONTENT_LENGTH, content_length),
                    HeadersMeta::Sized(content_length),
                ),
                None => (res, HeadersMeta::Chunked),
            },
        };

        let (meta_tx, meta_rx) = oneshot::channel();
        state.put(PendingResponseMeta::deferred(headers_meta, meta_rx));

        // Set up a SignalStream to send the PostSendMeta.
        let stream = SignalStream::new(stream, meta_tx);

        let (stats_tx, stats_rx) = oneshot::channel();
        state.put(PendingStreamStats::deferred(stats_rx));

        let stream = stream.timed(move |stats| {
            let _ = stats_tx.send(stats);
        });

        // Turn the stream into a TryStream, as expected by hyper::Body.
        let stream = stream.map(<Result<_, Error>>::Ok);

        Ok(res.body(Body::wrap_stream(stream))?)
    }
}


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

use anyhow::Error;
use futures::channel::oneshot::Receiver;
use gotham_derive::StateData;

use super::error_meta::ErrorMeta;
use crate::content_encoding::ContentCompression;

#[derive(Debug, Copy, Clone)]
pub enum HeadersMeta {
    Sized(u64),
    Chunked,
    Compressed(ContentCompression),
}

impl HeadersMeta {
    pub fn content_length(&self) -> Option<u64> {
        match self {
            Self::Sized(s) => Some(*s),
            Self::Compressed(..) => None,
            Self::Chunked => None,
        }
    }
}

pub struct BodyMeta {
    pub bytes_sent: u64,
    pub error_meta: ErrorMeta<Error>,
}

pub struct ResponseMeta {
    headers: HeadersMeta,
    body: BodyMeta,
}

impl ResponseMeta {
    pub fn headers(&self) -> &HeadersMeta {
        &self.headers
    }

    pub fn body(&self) -> &BodyMeta {
        &self.body
    }
}

enum PendingBodyMeta {
    Immediate(u64),
    Deferred(Receiver<BodyMeta>),
    Error(Error),
}

#[derive(StateData)]
pub struct PendingResponseMeta {
    headers: HeadersMeta,
    body: PendingBodyMeta,
}

impl PendingResponseMeta {
    /// Instantiate PendingResponseMeta for a response that will be sent synchronously (i.e. a
    /// set of bytes).
    pub fn immediate(size: u64) -> Self {
        Self {
            headers: HeadersMeta::Sized(size),
            body: PendingBodyMeta::Immediate(size),
        }
    }

    /// Instantiate PendingResponseMeta for a response that will not be sent synchronously (i.e. a
    /// stream).
    pub fn deferred(headers: HeadersMeta, receiver: Receiver<BodyMeta>) -> Self {
        Self {
            headers,
            body: PendingBodyMeta::Deferred(receiver),
        }
    }

    pub fn error(error: Error) -> Self {
        Self {
            headers: HeadersMeta::Sized(0),
            body: PendingBodyMeta::Error(error),
        }
    }

    /// Wait for this response to finish, and return the associated ResponseMeta. For a
    pub async fn finish(self) -> ResponseMeta {
        let Self { headers, body } = self;

        let body = match body {
            PendingBodyMeta::Immediate(bytes_sent) => BodyMeta {
                bytes_sent,
                error_meta: ErrorMeta::new(),
            },
            PendingBodyMeta::Deferred(receiver) => match receiver.await {
                Ok(body) => body,
                Err(_) => BodyMeta {
                    bytes_sent: 0,
                    error_meta: ErrorMeta::one_error(Error::msg("Deferred body meta was not sent")),
                },
            },
            PendingBodyMeta::Error(e) => BodyMeta {
                bytes_sent: 0,
                error_meta: ErrorMeta::one_error(e),
            },
        };

        ResponseMeta { headers, body }
    }
}