/*
 * 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::io;
use std::io::Write;
use std::net::SocketAddr;
use std::path::PathBuf;
use std::sync::atomic::AtomicBool;
use std::sync::atomic::AtomicUsize;
use std::sync::atomic::Ordering;
use std::sync::Arc;
use std::time::Duration;

use anyhow::Context;
use anyhow::Error;
use anyhow::Result;
use bytes::Bytes;
use cached_config::ConfigStore;
use connection_security_checker::ConnectionSecurityChecker;
use edenapi_service::EdenApi;
use failure_ext::SlogKVError;
use fbinit::FacebookInit;
use futures::channel::oneshot;
use futures::future::Future;
use futures::select_biased;
use futures_01_ext::BoxStream;
use futures_ext::FbFutureExt;
use futures_old::stream;
use futures_old::sync::mpsc;
use futures_old::Stream;
use futures_util::compat::Stream01CompatExt;
use futures_util::future::AbortHandle;
use futures_util::future::FutureExt;
use futures_util::stream::StreamExt;
use futures_util::stream::TryStreamExt;
use hostname::get_hostname;
use hyper::server::conn::Http;
use lazy_static::lazy_static;
use metaconfig_types::CommonConfig;
use metadata::Metadata;
use mononoke_api::Mononoke;
use mononoke_app::fb303::ReadyFlagService;
use openssl::ssl::Ssl;
use openssl::ssl::SslAcceptor;
use permission_checker::AclProvider;
use permission_checker::MononokeIdentity;
use permission_checker::MononokeIdentitySet;
use qps::Qps;
use quiet_stream::QuietShutdownStream;
use rate_limiting::RateLimitEnvironment;
use scribe_ext::Scribe;
use scuba_ext::MononokeScubaSampleBuilder;
use slog::debug;
use slog::error;
use slog::info;
use slog::warn;
use slog::Logger;
use sshrelay::IoStream;
use sshrelay::SshDecoder;
use sshrelay::SshEncoder;
use sshrelay::SshMsg;
use sshrelay::Stdio;
use stats::prelude::*;
use tokio::io::AsyncRead;
use tokio::io::AsyncWrite;
use tokio::net::TcpListener;
use tokio::net::TcpStream;
use tokio::task::JoinHandle;
use tokio_openssl::SslStream;
use tokio_util::codec::FramedRead;
use tokio_util::codec::FramedWrite;

use crate::errors::ErrorKind;
use crate::http_service::MononokeHttpService;
use crate::request_handler::create_conn_logger;
use crate::request_handler::request_handler;
use crate::wireproto_sink::WireprotoSink;

define_stats! {
    prefix = "mononoke.connection_acceptor";
    http_accepted: timeseries(Sum),
}

pub trait MononokeStream: AsyncRead + AsyncWrite + Unpin + Send + Sync + 'static {}

impl<T> MononokeStream for T where T: AsyncRead + AsyncWrite + Unpin + Send + Sync + 'static {}

const KEEP_ALIVE_INTERVAL: Duration = Duration::from_millis(5000);
const CHUNK_SIZE: usize = 10000;
lazy_static! {
    static ref OPEN_CONNECTIONS: AtomicUsize = AtomicUsize::new(0);
}

pub async fn wait_for_connections_closed(logger: &Logger) {
    loop {
        let conns = OPEN_CONNECTIONS.load(Ordering::Relaxed);
        if conns == 0 {
            break;
        }

        slog::info!(logger, "Waiting for {} connections to close", conns);
        tokio::time::sleep(Duration::new(1, 0)).await;
    }
}

pub async fn connection_acceptor(
    fb: FacebookInit,
    common_config: CommonConfig,
    sockname: String,
    service: ReadyFlagService,
    root_log: Logger,
    mononoke: Arc<Mononoke>,
    tls_acceptor: SslAcceptor,
    terminate_process: oneshot::Receiver<()>,
    rate_limiter: Option<RateLimitEnvironment>,
    scribe: Scribe,
    edenapi: EdenApi,
    will_exit: Arc<AtomicBool>,
    config_store: &ConfigStore,
    cslb_config: Option<String>,
    wireproto_scuba: MononokeScubaSampleBuilder,
    bound_addr_path: Option<PathBuf>,
    acl_provider: &dyn AclProvider,
    readonly: bool,
) -> Result<()> {
    let enable_http_control_api = common_config.enable_http_control_api;

    let security_checker = ConnectionSecurityChecker::new(acl_provider, &common_config).await?;
    let addr: SocketAddr = sockname
        .parse()
        .with_context(|| format!("could not parse '{}'", sockname))?;
    let listener = TcpListener::bind(&addr)
        .await
        .with_context(|| format!("could not bind mononoke on '{}'", sockname))?;

    let mut terminate_process = terminate_process.fuse();

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

    // Now that we are listening and ready to accept connections, report that we are alive.
    service.set_ready();

    let bound_addr = listener.local_addr()?.to_string();
    debug!(root_log, "server is listening on {}", bound_addr);

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

    let acceptor = Arc::new(Acceptor {
        fb,
        tls_acceptor,
        mononoke,
        security_checker,
        rate_limiter,
        scribe,
        logger: root_log.clone(),
        edenapi,
        enable_http_control_api,
        server_hostname: get_hostname().unwrap_or_else(|_| "unknown_hostname".to_string()),
        will_exit,
        config_store: config_store.clone(),
        qps,
        wireproto_scuba,
        common_config,
        readonly,
    });

    loop {
        select_biased! {
            _ = terminate_process => {
                debug!(root_log, "Received shutdown handler, stop accepting connections...");
                return Ok(());
            },
            sock_tuple = listener.accept().fuse() => match sock_tuple {
                Ok((stream, addr)) => {
                    let conn = PendingConnection { acceptor: acceptor.clone(), addr };
                    let task = handle_connection(conn.clone(), stream);
                    conn.spawn_task(task, "Failed to handle_connection");
                }
                Err(err) => {
                    error!(root_log, "{}", err.to_string(); SlogKVError(Error::from(err)));
                }
            },
        };
    }
}

/// Our environment for accepting connections.
pub struct Acceptor {
    pub fb: FacebookInit,
    pub tls_acceptor: SslAcceptor,
    pub mononoke: Arc<Mononoke>,
    pub security_checker: ConnectionSecurityChecker,
    pub rate_limiter: Option<RateLimitEnvironment>,
    pub scribe: Scribe,
    pub logger: Logger,
    pub edenapi: EdenApi,
    pub enable_http_control_api: bool,
    pub server_hostname: String,
    pub will_exit: Arc<AtomicBool>,
    pub config_store: ConfigStore,
    pub qps: Option<Arc<Qps>>,
    pub wireproto_scuba: MononokeScubaSampleBuilder,
    pub common_config: CommonConfig,
    pub readonly: bool,
}

/// Details for a socket we've just opened.
#[derive(Clone)]
pub struct PendingConnection {
    pub acceptor: Arc<Acceptor>,
    pub addr: SocketAddr,
}

/// A connection where we completed the initial TLS handshake.
#[derive(Clone)]
pub struct AcceptedConnection {
    pub pending: PendingConnection,
    pub is_trusted: bool,
    pub identities: Arc<MononokeIdentitySet>,
}

impl PendingConnection {
    /// Spawn a task that is dedicated to this connection. This will block server shutdown, and
    /// also log on error or cancellation.
    pub fn spawn_task(
        &self,
        task: impl Future<Output = Result<()>> + Send + 'static,
        label: &'static str,
    ) {
        let this = self.clone();

        OPEN_CONNECTIONS.fetch_add(1, Ordering::Relaxed);

        tokio::task::spawn(async move {
            let logger = &this.acceptor.logger;
            let res = task
                .on_cancel(|| warn!(logger, "connection to {} was cancelled", this.addr))
                .await
                .context(label)
                .with_context(|| format!("Failed to handle connection to {}", this.addr));

            if let Err(e) = res {
                error!(logger, "connection_acceptor error: {:#}", e);
            }

            OPEN_CONNECTIONS.fetch_sub(1, Ordering::Relaxed);
        });
    }
}

async fn handle_connection(conn: PendingConnection, sock: TcpStream) -> Result<()> {
    let ssl = Ssl::new(conn.acceptor.tls_acceptor.context()).context("Error creating Ssl")?;
    let ssl_socket = SslStream::new(ssl, sock).context("Error creating SslStream")?;
    let mut ssl_socket = Box::pin(ssl_socket);

    ssl_socket
        .as_mut()
        .accept()
        .await
        .context("Failed to perform tls handshake")?;

    let identities = match ssl_socket.ssl().peer_certificate() {
        Some(cert) => MononokeIdentity::try_from_x509(&cert),
        None => Err(ErrorKind::ConnectionNoClientCertificate.into()),
    }?;

    let is_trusted = conn
        .acceptor
        .security_checker
        .check_if_trusted(&identities)
        .await;

    let conn = AcceptedConnection {
        pending: conn,
        is_trusted,
        identities: Arc::new(identities),
    };

    let ssl_socket = QuietShutdownStream::new(ssl_socket);

    handle_http(conn, ssl_socket)
        .await
        .context("Failed to handle_http")?;

    Ok(())
}

async fn handle_http<S: MononokeStream>(conn: AcceptedConnection, stream: S) -> Result<()> {
    STATS::http_accepted.add_value(1);
    let svc = MononokeHttpService::<S>::new(conn);
    Http::new()
        .serve_connection(stream, svc)
        .with_upgrades()
        .await
        .context("Failed to serve_connection")?;

    Ok(())
}

pub async fn handle_wireproto<R, W>(
    conn: AcceptedConnection,
    framed: FramedConn<R, W>,
    reponame: String,
    metadata: Metadata,
) -> Result<()>
where
    R: AsyncRead + Send + std::marker::Unpin + 'static,
    W: AsyncWrite + Send + std::marker::Unpin + 'static,
{
    let metadata = Arc::new(metadata);

    let ChannelConn {
        stdin,
        stdout,
        stderr,
        logger,
        keep_alive,
        join_handle,
    } = ChannelConn::setup(framed, conn.clone(), metadata.clone());

    if metadata.client_debug() {
        info!(&logger, "{:#?}", metadata; "remote" => "true");
    }

    // Don't let the logger hold onto the channel. This is a bit fragile (but at least it breaks
    // tests deterministically).
    drop(logger);

    let stdio = Stdio {
        metadata,
        stdin,
        stdout,
        stderr,
    };

    // Don't immediately return error here, we need to cleanup our
    // handlers like keep alive, otherwise they will run forever.
    let result = request_handler(
        conn.pending.acceptor.fb,
        reponame,
        Arc::clone(&conn.pending.acceptor.mononoke),
        &conn.pending.acceptor.security_checker,
        stdio,
        conn.pending.acceptor.rate_limiter.clone(),
        conn.pending.acceptor.scribe.clone(),
        conn.pending.acceptor.qps.clone(),
        conn.pending.acceptor.readonly,
    )
    .await
    .context("Failed to execute request_handler");

    // Shutdown our keepalive handler
    keep_alive.abort();

    join_handle
        .await
        .context("Failed to join ChannelConn")?
        .context("Failed to close ChannelConn")?;

    result
}

pub struct FramedConn<R, W> {
    rd: FramedRead<R, SshDecoder>,
    wr: FramedWrite<W, SshEncoder>,
}

impl<R, W> FramedConn<R, W>
where
    R: AsyncRead + Send + std::marker::Unpin + 'static,
    W: AsyncWrite + Send + std::marker::Unpin + 'static,
{
    pub fn setup(rd: R, wr: W, compression_writes: Option<i32>) -> Result<Self> {
        // NOTE: FramedRead does buffering, so no need to wrap with a BufReader here.
        let rd = FramedRead::new(rd, SshDecoder::new());
        let wr = FramedWrite::new(wr, SshEncoder::new(compression_writes)?);
        Ok(Self { rd, wr })
    }
}

pub struct ChannelConn {
    stdin: BoxStream<Bytes, io::Error>,
    stdout: mpsc::Sender<Bytes>,
    stderr: mpsc::UnboundedSender<Bytes>,
    logger: Logger,
    keep_alive: AbortHandle,
    join_handle: JoinHandle<Result<(), io::Error>>,
}

impl ChannelConn {
    pub fn setup<R, W>(
        framed: FramedConn<R, W>,
        conn: AcceptedConnection,
        metadata: Arc<Metadata>,
    ) -> Self
    where
        R: AsyncRead + Send + std::marker::Unpin + 'static,
        W: AsyncWrite + Send + std::marker::Unpin + 'static,
    {
        let FramedConn { rd, wr } = framed;

        let stdin = Box::new(rd.compat().filter_map(|s| {
            if s.stream() == IoStream::Stdin {
                Some(s.data())
            } else {
                None
            }
        }));

        let (stdout, stderr, keep_alive, join_handle) = {
            let (otx, orx) = mpsc::channel(1);
            let (etx, erx) = mpsc::unbounded();
            let (ktx, krx) = mpsc::unbounded();

            let orx = orx
                .map(|blob| split_bytes_in_chunk(blob, CHUNK_SIZE))
                .flatten()
                .map(|v| SshMsg::new(IoStream::Stdout, v));
            let erx = erx
                .map(|blob| split_bytes_in_chunk(blob, CHUNK_SIZE))
                .flatten()
                .map(|v| SshMsg::new(IoStream::Stderr, v));
            let krx = krx.map(|v| SshMsg::new(IoStream::Stderr, v));

            // Glue them together
            let fwd = async move {
                let wr = WireprotoSink::new(wr);

                futures::pin_mut!(wr);

                let res = orx
                    .select(erx)
                    .select(krx)
                    .compat()
                    .map_err(|()| io::Error::new(io::ErrorKind::Other, "huh?"))
                    .forward(wr.as_mut())
                    .await;

                if let Err(e) = res.as_ref() {
                    let projected_wr = wr.as_mut().project();
                    let data = projected_wr.data;

                    let mut scuba = conn.pending.acceptor.wireproto_scuba.clone();
                    scuba.add_metadata(&metadata);
                    scuba.add_opt(
                        "last_successful_flush",
                        data.last_successful_flush.map(|dt| dt.timestamp()),
                    );
                    scuba.add_opt(
                        "last_successful_io",
                        data.last_successful_io.map(|dt| dt.timestamp()),
                    );
                    scuba.add_opt(
                        "last_failed_io",
                        data.last_failed_io.map(|dt| dt.timestamp()),
                    );
                    scuba.add("stdout_bytes", data.stdout.bytes);
                    scuba.add("stdout_messages", data.stdout.messages);
                    scuba.add("stderr_bytes", data.stderr.bytes);
                    scuba.add("stderr_messages", data.stderr.messages);
                    scuba.log_with_msg("Forwarding failed", format!("{:#}", e));
                }

                Ok(())
            };

            let keep_alive_sender = async move {
                loop {
                    tokio::time::sleep(KEEP_ALIVE_INTERVAL).await;
                    if ktx.unbounded_send(Bytes::new()).is_err() {
                        break;
                    }
                }
            };
            let (keep_alive_sender, keep_alive_abort) =
                futures::future::abortable(keep_alive_sender);

            // spawn a task for sending keepalive messages
            tokio::spawn(keep_alive_sender);

            // spawn a task for forwarding stdout/err into stream
            let join_handle = tokio::spawn(fwd);

            // NOTE: This might seem useless, but it's not. When you spawn a task, Tokio puts it on
            // a "LIFO slot" associated with the current thread. While the task is in the LIFO
            // slot, it is not eligible to be run by other threads. If the thread that just spawned
            // `fwd` above goes and does some expensive synchronous work in `request_handler` (we'd
            // like to avoid that but sometimes that happens), then that will delay `fwd`. This
            // means that notably keepalives will not be sent (you can repro by putting a
            // `std::thread::sleep` right after we spawn `fwd`). To mitigate this, we spawn another
            // dummy taks here. This task will take `fwd`'s place in the LIFO slot, thus pushing
            // `fwd` onto a task queue where other runtime threads can claim it. This way, even if
            // this thread goes do some expensive CPU-bound work, we won't delay keepalives.
            tokio::spawn(async {});

            (otx, etx, keep_alive_abort, join_handle)
        };

        let logger = create_conn_logger(stderr.clone(), None, None);

        ChannelConn {
            stdin,
            stdout,
            stderr,
            logger,
            keep_alive,
            join_handle,
        }
    }
}

// TODO(stash): T33775046 we had to chunk responses because hgcli
// can't cope with big chunks
fn split_bytes_in_chunk<E>(blob: Bytes, chunksize: usize) -> impl Stream<Item = Bytes, Error = E> {
    stream::unfold(blob, move |mut remain| {
        let len = remain.len();
        if len > 0 {
            let ret = remain.split_to(::std::cmp::min(chunksize, len));
            Some(Ok((ret, remain)))
        } else {
            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 mononoke_types::RepositoryId;
use thiserror::Error;

#[derive(Debug, Error)]
pub enum ErrorKind {
    #[error("connection does not have a client certificate")]
    ConnectionNoClientCertificate,
    #[error("Unauthorized access, permission denied")]
    AuthorizationFailed,
    #[error("Large repo not found: {0}")]
    LargeRepoNotFound(RepositoryId),
}


/*
 * 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::io::Cursor;
use std::marker::PhantomData;
use std::str::FromStr;
use std::sync::atomic::Ordering;
use std::task;

use anyhow::anyhow;
use anyhow::Context;
use anyhow::Error;
use anyhow::Result;
use bookmarks::BookmarksRef;
#[cfg(fbcode_build)]
use clientinfo::ClientInfo;
#[cfg(fbcode_build)]
use clientinfo::CLIENT_INFO_HEADER;
use futures::future::BoxFuture;
use futures::future::FutureExt;
use gotham_ext::socket_data::TlsSocketData;
use http::HeaderMap;
use http::HeaderValue;
use http::Method;
use http::Request;
use http::Response;
use http::Uri;
use hyper::service::Service;
use hyper::Body;
use metadata::Metadata;
use percent_encoding::percent_decode;
use qps::Qps;
use session_id::generate_session_id;
use sha1::Digest;
use sha1::Sha1;
use slog::debug;
use slog::error;
use slog::trace;
use slog::Logger;
use thiserror::Error;
use tokio::io::AsyncReadExt;
use tunables::force_update_tunables;
use tunables::tunables;

use crate::connection_acceptor;
use crate::connection_acceptor::AcceptedConnection;
use crate::connection_acceptor::Acceptor;
use crate::connection_acceptor::FramedConn;
use crate::connection_acceptor::MononokeStream;

const HEADER_CLIENT_COMPRESSION: &str = "x-client-compression";
const HEADER_CLIENT_DEBUG: &str = "x-client-debug";
const HEADER_WEBSOCKET_KEY: &str = "sec-websocket-key";
const HEADER_WEBSOCKET_ACCEPT: &str = "sec-websocket-accept";
const HEADER_MONONOKE_ENCODING: &str = "x-mononoke-encoding";
const HEADER_MONONOKE_HOST: &str = "x-mononoke-host";
const HEADER_REVPROXY_REGION: &str = "x-fb-revproxy-region";

// See https://tools.ietf.org/html/rfc6455#section-1.3
const WEBSOCKET_MAGIC_KEY: &str = "258EAFA5-E914-47DA-95CA-C5AB0DC85B11";

#[derive(Error, Debug)]
pub enum HttpError {
    #[error("Bad request")]
    BadRequest(#[source] Error),

    #[error("Forbidden")]
    Forbidden,

    #[error("Not found")]
    NotFound,

    #[error("Method not allowed")]
    MethodNotAllowed,

    #[error("Internal server error")]
    InternalServerError(#[source] Error),
}

impl HttpError {
    pub fn internal(e: impl Into<Error>) -> Self {
        Self::InternalServerError(e.into())
    }

    pub fn http_response(&self) -> http::Result<Response<Body>> {
        let status = match self {
            Self::BadRequest(..) => http::StatusCode::BAD_REQUEST,
            Self::Forbidden => http::StatusCode::FORBIDDEN,
            Self::NotFound => http::StatusCode::NOT_FOUND,
            Self::MethodNotAllowed => http::StatusCode::METHOD_NOT_ALLOWED,
            Self::InternalServerError(..) => http::StatusCode::INTERNAL_SERVER_ERROR,
        };

        let body = match self {
            Self::BadRequest(ref e) => Body::from(format!("{:#}", e)),
            Self::Forbidden => Body::empty(),
            Self::NotFound => Body::empty(),
            Self::MethodNotAllowed => Body::empty(),
            Self::InternalServerError(ref e) => Body::from(format!("{:#}", e)),
        };

        Response::builder().status(status).body(body)
    }
}

pub struct MononokeHttpService<S> {
    pub conn: AcceptedConnection,
    sock: PhantomData<S>,
}

impl<S> MononokeHttpService<S> {
    pub fn new(conn: AcceptedConnection) -> Self {
        Self {
            conn,
            sock: PhantomData,
        }
    }
}

impl<S> Clone for MononokeHttpService<S> {
    fn clone(&self) -> Self {
        Self {
            conn: self.conn.clone(),
            sock: PhantomData,
        }
    }
}

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

impl<S> MononokeHttpService<S>
where
    S: MononokeStream,
{
    async fn handle(&self, req: Request<Body>) -> Result<Response<Body>, HttpError> {
        if req.method() == Method::GET
            && (req.uri().path() == "/" || req.uri().path() == "/health_check")
        {
            let res = if self.acceptor().will_exit.load(Ordering::Relaxed) {
                "EXITING"
            } else {
                "I_AM_ALIVE"
            };

            let res = Response::builder()
                .status(http::StatusCode::OK)
                .body(res.into())
                .map_err(HttpError::internal)?;

            return Ok(res);
        }

        let upgrade = req
            .headers()
            .get(http::header::UPGRADE)
            .as_ref()
            .map(|h| h.to_str())
            .transpose()
            .with_context(|| {
                // NOTE: We're just stringifying here: the borrow is fine.
                #[allow(clippy::borrow_interior_mutable_const)]
                let header = &http::header::UPGRADE;
                format!("Invalid header: {}", header)
            })
            .map_err(HttpError::BadRequest)?;

        if upgrade == Some("websocket") {
            return self.handle_websocket_request(req).await;
        }

        let (req, body) = req.into_parts();

        if req.uri.path() == "/netspeedtest" {
            return crate::netspeedtest::handle(req.method, &req.headers, body).await;
        }

        if let Some(path) = req.uri.path().strip_prefix("/control") {
            return self.handle_control_request(req.method, path).await;
        }

        let edenapi_path_and_query = req
            .uri
            .path_and_query()
            .as_ref()
            .and_then(|pq| pq.as_str().strip_prefix("/edenapi"));

        if let Some(edenapi_path_and_query) = edenapi_path_and_query {
            let pq = http::uri::PathAndQuery::from_str(edenapi_path_and_query)
                .context("Error translating EdenAPI request path")
                .map_err(HttpError::internal)?;
            return self.handle_eden_api_request(req, pq, body).await;
        }

        Err(HttpError::NotFound)
    }

    async fn handle_websocket_request(
        &self,
        mut req: Request<Body>,
    ) -> Result<Response<Body>, HttpError> {
        let reponame_urlencoded = req.uri().path().trim_matches('/').to_string();

        let reponame = percent_decode(reponame_urlencoded.as_bytes())
            .decode_utf8()
            .context("reponame must be url-encoded utf-8")
            .map_err(HttpError::BadRequest)?
            .into_owned();

        let websocket_key = calculate_websocket_accept(req.headers());

        let mut builder = Response::builder()
            .status(http::StatusCode::SWITCHING_PROTOCOLS)
            .header(http::header::CONNECTION, "upgrade")
            .header(http::header::UPGRADE, "websocket")
            .header(HEADER_WEBSOCKET_ACCEPT, websocket_key);

        let metadata = h2m::try_convert_headers_to_metadata(&self.conn, req.headers())
            .await
            .context("Invalid metadata")
            .map_err(HttpError::BadRequest)?;

        let zstd_level: i32 = tunables::tunables()
            .zstd_compression_level()
            .unwrap_or_default()
            .try_into()
            .unwrap_or_default();
        let compression = match req.headers().get(HEADER_CLIENT_COMPRESSION) {
            Some(header_value) => match header_value.as_bytes() {
                b"zstd=stdin" if zstd_level > 0 => Ok(Some(zstd_level)),
                header_value_bytes => Err(anyhow!(
                    "'{}' is not a recognized compression value",
                    String::from_utf8_lossy(header_value_bytes),
                )),
            },
            None => Ok(None),
        }
        .map_err(HttpError::BadRequest)?;

        match compression {
            Some(zstd_level) => {
                builder = builder.header(HEADER_MONONOKE_ENCODING, format!("zstd={}", zstd_level));
            }
            _ => {}
        };

        let res = builder.body(Body::empty()).map_err(HttpError::internal)?;

        let this = self.clone();

        let fut = async move {
            let io = hyper::upgrade::on(&mut req)
                .await
                .context("Failed to upgrade connection")?;

            // NOTE: We unwrap() here because we explicitly parameterize the MononokeHttpService
            // over its socket type. If we get it wrong then that'd be a deterministic failure that
            // would show up in tests.
            let hyper::upgrade::Parts { io, read_buf, .. } = io.downcast::<S>().unwrap();

            let (rx, tx) = tokio::io::split(io);
            let mut rx = AsyncReadExt::chain(Cursor::new(read_buf), rx);

            // Sometimes server rejects client's request quickly. So quickly,
            // that right after sending 101 Switching Protocols, it immediately
            // sends wireproto error message to the client on the stderr channel,
            // and terminates the connection right away.
            //
            // Client sees HTTP 101 and assumes it can speak wireproto to the server.
            // It tries to send wireproto's "hello" command, but fails miserably, because
            // the connection has already been closed.
            //
            // Example:
            // https://pxl.cl/1XcnH
            //
            // Lines bellow make server wait for the client to send any data before
            // any wireproto handling can take place. Normally server that speaks just
            // wireproto shouldn't send anything to the client until it sees "hello".
            // Here we try to replicate that behavior by making sure the client
            // sent something. We assume it's wireproto's "hello".
            let mut buffer = [0; 1];
            rx.read_exact(&mut buffer).await?;
            let rx = AsyncReadExt::chain(Cursor::new(buffer), rx);

            let framed = FramedConn::setup(rx, tx, compression)?;

            connection_acceptor::handle_wireproto(this.conn, framed, reponame, metadata)
                .await
                .context("Failed to handle_wireproto")?;

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

        // Spawning concurrent task handling wireproto
        self.conn
            .pending
            .spawn_task(fut, "Failed to handle websocket channel");

        // Returning with HTTP 101. The task spawned above will handle
        // upgraded connection.
        Ok(res)
    }

    async fn handle_control_request(
        &self,
        method: Method,
        path: &str,
    ) -> Result<Response<Body>, HttpError> {
        if method != Method::POST {
            return Err(HttpError::MethodNotAllowed);
        }

        if !self.acceptor().enable_http_control_api {
            return Err(HttpError::Forbidden);
        }

        let ok = Response::builder()
            .status(http::StatusCode::OK)
            .body(Body::empty())
            .map_err(HttpError::internal)?;

        if path == "/drop_bookmarks_cache" {
            for repo in self.acceptor().mononoke.repos() {
                repo.blob_repo().bookmarks().drop_caches();
            }

            return Ok(ok);
        }

        if path == "/force_update_configerator" {
            self.acceptor().config_store.force_update_configs();
            force_update_tunables();
            return Ok(ok);
        }

        Err(HttpError::NotFound)
    }

    async fn handle_eden_api_request(
        &self,
        mut req: http::request::Parts,
        pq: http::uri::PathAndQuery,
        body: Body,
    ) -> Result<Response<Body>, HttpError> {
        if tunables()
            .disable_http_service_edenapi()
            .unwrap_or_default()
        {
            let res = Response::builder()
                .status(http::StatusCode::SERVICE_UNAVAILABLE)
                .body("EdenAPI service is killswitched".into())
                .map_err(HttpError::internal)?;
            return Ok(res);
        }

        let mut uri_parts = req.uri.into_parts();

        uri_parts.path_and_query = Some(pq);

        req.uri = Uri::from_parts(uri_parts)
            .context("Error translating EdenAPI request")
            .map_err(HttpError::internal)?;

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

        let tls_socket_data = if self.conn.is_trusted {
            TlsSocketData::trusted_proxy((*self.conn.identities).clone())
        } else {
            TlsSocketData::authenticated_identities((*self.conn.identities).clone())
        };

        let req = Request::from_parts(req, body);

        let res = self
            .acceptor()
            .edenapi
            .clone()
            .into_service(self.conn.pending.addr, Some(tls_socket_data))
            .call_gotham(req)
            .await;

        Ok(res)
    }

    fn acceptor(&self) -> &Acceptor {
        &self.conn.pending.acceptor
    }

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

impl<S> Service<Request<Body>> for MononokeHttpService<S>
where
    S: MononokeStream,
{
    type Response = Response<Body>;
    type Error = http::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 method = req.method().clone();
            let uri = req.uri().clone();
            debug!(this.logger(), "{} {}", method, uri);

            this.handle(req)
                .await
                .map(|mut res| {
                    match HeaderValue::from_str(this.conn.pending.acceptor.server_hostname.as_str())
                    {
                        Ok(header) => {
                            res.headers_mut().insert(HEADER_MONONOKE_HOST, header);
                        }
                        Err(e) => {
                            error!(
                                this.logger(),
                                "http service error: can't set {} header: {}",
                                HEADER_MONONOKE_HOST,
                                anyhow::Error::from(e),
                            );
                        }
                    };
                    res
                })
                .or_else(|e| {
                    let res = e.http_response();

                    error!(
                        this.logger(),
                        "http service error: {} {}: {:#}",
                        method,
                        uri,
                        anyhow::Error::from(e)
                    );

                    res
                })
            // NOTE: If we fail to even generate the response here, this will crash
            // serve_connection in Hyper, so we don't actually need to log this here.
        }
        .boxed()
    }
}

// See https://tools.ietf.org/html/rfc6455#section-1.3
fn calculate_websocket_accept(headers: &HeaderMap<HeaderValue>) -> String {
    let mut sha1 = Sha1::new();

    // This is OK to fall back to empty, because we only need to give
    // this header, if it's asked for. In case of hg<->mononoke with
    // no Proxygen in between, this header will be missing and the result
    // ignored.
    if let Some(header) = headers.get(HEADER_WEBSOCKET_KEY) {
        sha1.update(header.as_ref());
    }
    sha1.update(WEBSOCKET_MAGIC_KEY.as_bytes());
    let hash: [u8; 20] = sha1.finalize().into();
    base64::encode(&hash)
}

#[cfg(not(fbcode_build))]
mod h2m {
    use super::*;

    pub async fn try_convert_headers_to_metadata(
        conn: &AcceptedConnection,
        headers: &HeaderMap<HeaderValue>,
    ) -> Result<Metadata> {
        let debug = headers.contains_key(HEADER_CLIENT_DEBUG);

        Ok(Metadata::new(
            Some(&generate_session_id().to_string()),
            (*conn.identities).clone(),
            debug,
            metadata::security::is_client_untrusted(|h| {
                Ok(headers
                    .get(h)
                    .map(|h| h.to_str().map(|s| s.to_owned()))
                    .transpose()?)
            })?,
            Some(conn.pending.addr.ip()),
        )
        .await)
    }
}

#[cfg(fbcode_build)]
mod h2m {
    use std::net::IpAddr;

    use cats::try_get_cats_idents;
    use percent_encoding::percent_decode;
    use permission_checker::MononokeIdentity;

    use super::*;

    const HEADER_ENCODED_CLIENT_IDENTITY: &str = "x-fb-validated-client-encoded-identity";
    const HEADER_CLIENT_IP: &str = "tfb-orig-client-ip";
    const HEADER_FORWARDED_CATS: &str = "x-forwarded-cats";

    fn metadata_populate_trusted(
        metadata: &mut Metadata,
        headers: &HeaderMap<HeaderValue>,
    ) -> Result<()> {
        if let Some(cats) = headers.get(HEADER_FORWARDED_CATS) {
            metadata
                .add_raw_encoded_cats(cats.to_str().context("Invalid encoded cats")?.to_string());
        }

        let src_region = headers
            .get(HEADER_REVPROXY_REGION)
            .and_then(|r| r.to_str().ok().map(|r| r.to_string()));

        if let Some(src_region) = src_region {
            metadata.add_revproxy_region(src_region);
        }

        let client_info: Option<ClientInfo> = headers
            .get(CLIENT_INFO_HEADER)
            .and_then(|h| h.to_str().ok())
            .and_then(|ci| serde_json::from_str(ci).ok());

        if let Some(client_info) = client_info {
            metadata.add_client_info(client_info);
        }

        Ok(())
    }

    /// Used only for wireproto handling.
    pub async fn try_convert_headers_to_metadata(
        conn: &AcceptedConnection,
        headers: &HeaderMap<HeaderValue>,
    ) -> Result<Metadata> {
        let debug = headers.contains_key(HEADER_CLIENT_DEBUG);
        let internal_identity = &conn.pending.acceptor.common_config.internal_identity;
        let is_trusted = conn.is_trusted;

        // CATs are verifiable - we know that only the signer could have
        // generated them. We extract the signer's identity. The connecting
        // party doesn't have to be trusted.
        //
        // This correctly returns error if cats are present but are invalid.
        let cats_identities =
            try_get_cats_idents(conn.pending.acceptor.fb.clone(), headers, internal_identity)?;

        if is_trusted {
            if let (Some(encoded_identities), Some(client_address)) = (
                headers.get(HEADER_ENCODED_CLIENT_IDENTITY),
                headers.get(HEADER_CLIENT_IP),
            ) {
                let json_identities = percent_decode(encoded_identities.as_ref())
                    .decode_utf8()
                    .context("Invalid encoded identities")?;

                let mut identities = MononokeIdentity::try_from_json_encoded(&json_identities)
                    .context("Invalid identities")?;
                let ip_addr = client_address
                    .to_str()?
                    .parse::<IpAddr>()
                    .context("Invalid IP Address")?;

                identities.extend(cats_identities.unwrap_or_default().into_iter());

                let mut metadata = Metadata::new(
                    Some(&generate_session_id().to_string()),
                    identities,
                    debug,
                    metadata::security::is_client_untrusted(|h| {
                        Ok(headers
                            .get(h)
                            .map(|h| h.to_str().map(|s| s.to_owned()))
                            .transpose()?)
                    })?,
                    Some(ip_addr),
                )
                .await;

                metadata_populate_trusted(&mut metadata, headers)?;

                return Ok(metadata);
            }
        }

        let mut identities = cats_identities.unwrap_or_default();
        identities.extend(conn.identities.iter().cloned());

        // Generic fallback
        Ok(Metadata::new(
            Some(&generate_session_id().to_string()),
            identities,
            debug,
            metadata::security::is_client_untrusted(|h| {
                Ok(headers
                    .get(h)
                    .map(|h| h.to_str().map(|s| s.to_owned()))
                    .transpose()?)
            })?,
            Some(conn.pending.addr.ip()),
        )
        .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.
 */

#![cfg_attr(not(fbcode_build), allow(unused_crate_dependencies))]
#![feature(never_type)]
#![recursion_limit = "256"]

mod connection_acceptor;
mod errors;
mod http_service;
mod netspeedtest;
mod repo_handlers;
mod request_handler;
mod wireproto_sink;

use std::path::PathBuf;
use std::sync::atomic::AtomicBool;
use std::sync::Arc;

use anyhow::Context as _;
use anyhow::Result;
use cached_config::ConfigStore;
use fbinit::FacebookInit;
use futures::channel::oneshot;
use metaconfig_types::CommonConfig;
use mononoke_api::Mononoke;
use mononoke_app::fb303::ReadyFlagService;
use openssl::ssl::SslAcceptor;
use permission_checker::AclProvider;
use rate_limiting::RateLimitEnvironment;
use scribe_ext::Scribe;
use scuba_ext::MononokeScubaSampleBuilder;
use slog::o;
use slog::Logger;

use crate::connection_acceptor::connection_acceptor;
pub use crate::connection_acceptor::wait_for_connections_closed;

const CONFIGERATOR_RATE_LIMITING_CONFIG: &str = "scm/mononoke/ratelimiting/ratelimits";

pub async fn create_repo_listeners<'a>(
    fb: FacebookInit,
    common_config: CommonConfig,
    mononoke: Arc<Mononoke>,
    root_log: Logger,
    sockname: String,
    tls_acceptor: SslAcceptor,
    service: ReadyFlagService,
    terminate_process: oneshot::Receiver<()>,
    config_store: &'a ConfigStore,
    scribe: Scribe,
    scuba: &'a MononokeScubaSampleBuilder,
    will_exit: Arc<AtomicBool>,
    cslb_config: Option<String>,
    bound_addr_file: Option<PathBuf>,
    acl_provider: &dyn AclProvider,
    readonly: bool,
) -> Result<()> {
    let rate_limiter = {
        let handle = config_store
            .get_config_handle_DEPRECATED(CONFIGERATOR_RATE_LIMITING_CONFIG.to_string())
            .ok();

        handle.and_then(|handle| {
            common_config
                .loadlimiter_category
                .clone()
                .map(|category| RateLimitEnvironment::new(fb, category, handle))
        })
    };

    let edenapi = {
        let mut scuba = scuba.clone();
        scuba.add("service", "edenapi");

        edenapi_service::build(
            fb,
            root_log.new(o!("service" => "edenapi")),
            scuba,
            Arc::clone(&mononoke),
            will_exit.clone(),
            false,
            None,
            rate_limiter.clone(),
            &common_config,
            readonly,
        )
        .context("Error instantiating EdenAPI")?
    };

    connection_acceptor(
        fb,
        common_config,
        sockname,
        service,
        root_log,
        mononoke,
        tls_acceptor,
        terminate_process,
        rate_limiter,
        scribe,
        edenapi,
        will_exit,
        config_store,
        cslb_config,
        {
            let mut scuba = scuba.clone();
            scuba.add("service", "wireproto");
            scuba
        },
        bound_addr_file,
        acl_provider,
        readonly,
    )
    .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::time::Duration;

use anyhow::anyhow;
use anyhow::Context;
use anyhow::Error;
use anyhow::Result;
use futures::stream::TryStreamExt;
use futures_ext::stream::StreamTimeoutError;
use futures_ext::FbStreamExt;
use futures_ext::FbTryStreamExt;
use http::HeaderMap;
use http::HeaderValue;
use http::Method;
use http::Response;
use hyper::Body;
use thiserror::Error;
use tokio::io::AsyncReadExt;
use tokio_util::codec::BytesCodec;
use tokio_util::codec::FramedRead;

use crate::http_service::HttpError;

const NETSPEEDTEST_MAX_NBYTES: u64 = 100 * 1024 * 1024;
const NETSPEEDTEST_TIMEOUT: Duration = Duration::from_secs(300);

const HEADER_DOWNLOAD_NBYTES: &str = "x-netspeedtest-nbytes";

#[derive(Error, Debug)]
pub enum RequestError {
    #[error("Request is invalid")]
    Invalid(#[source] Error),

    #[error("Client hung up")]
    Hangup(#[source] hyper::Error),

    #[error("Request is too large (only up to {} bytes are allowed)", .0)]
    TooLarge(u64),

    #[error("Request timed out")]
    Timeout,
}

impl From<StreamTimeoutError> for RequestError {
    fn from(_: StreamTimeoutError) -> Self {
        Self::Timeout
    }
}

impl From<RequestError> for HttpError {
    fn from(r: RequestError) -> Self {
        Self::BadRequest(Error::from(r).context("Invalid NetSpeedTest request"))
    }
}

pub async fn handle(
    method: Method,
    headers: &HeaderMap<HeaderValue>,
    body: Body,
) -> Result<Response<Body>, HttpError> {
    if method == Method::GET {
        return download(headers);
    }

    if method == Method::POST {
        return upload(body).await;
    }

    Err(HttpError::MethodNotAllowed)
}

fn download(headers: &HeaderMap<HeaderValue>) -> Result<Response<Body>, HttpError> {
    fn read_byte_count(headers: &HeaderMap<HeaderValue>) -> Result<u64, Error> {
        headers
            .get(HEADER_DOWNLOAD_NBYTES)
            .ok_or_else(|| anyhow!("Missing {} header", HEADER_DOWNLOAD_NBYTES))?
            .to_str()
            .with_context(|| format!("Invalid {} header (not UTF-8)", HEADER_DOWNLOAD_NBYTES))?
            .parse()
            .with_context(|| format!("Invalid {} header (invalid usize)", HEADER_DOWNLOAD_NBYTES))
    }

    let byte_count = read_byte_count(headers).map_err(RequestError::Invalid)?;
    let byte_count = std::cmp::min(byte_count, NETSPEEDTEST_MAX_NBYTES);

    let repeat = tokio::io::repeat(0x42).take(byte_count);
    let stream = FramedRead::new(repeat, BytesCodec::new());
    let stream = stream
        .map_err(Error::from)
        .whole_stream_timeout(NETSPEEDTEST_TIMEOUT)
        .flatten_err();

    let res = Response::builder()
        .status(http::StatusCode::OK)
        .header(http::header::CONTENT_LENGTH, byte_count.to_string())
        .body(Body::wrap_stream(stream))
        .map_err(HttpError::internal)?;

    Ok(res)
}

async fn upload(body: Body) -> Result<Response<Body>, HttpError> {
    let mut size = 0;
    let body = body
        .map_err(RequestError::Hangup)
        .whole_stream_timeout(NETSPEEDTEST_TIMEOUT)
        .flatten_err();

    futures::pin_mut!(body);

    while let Some(chunk) = body
        .try_next()
        .await
        .context("Error reading body")
        .map_err(HttpError::internal)?
    {
        let chunk_size: u64 = chunk
            .len()
            .try_into()
            .context("Chunk too large")
            .map_err(HttpError::internal)?;

        size += chunk_size;

        if size > NETSPEEDTEST_MAX_NBYTES {
            return Err(RequestError::TooLarge(NETSPEEDTEST_MAX_NBYTES).into());
        }
    }

    let res = Response::builder()
        .status(http::StatusCode::OK)
        .body(Body::empty())
        .map_err(HttpError::internal)?;

    Ok(res)
}


/*
 * 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::sync::Arc;

use anyhow::anyhow;
use metaconfig_types::RepoClientKnobs;
use mononoke_api::Mononoke;
use mononoke_api::Repo;
use repo_client::PushRedirectorArgs;
use scuba_ext::MononokeScubaSampleBuilder;
use slog::Logger;
use wireproto_handler::BackupSourceRepo;

use crate::errors::ErrorKind;

#[derive(Clone)]
pub struct RepoHandler {
    pub logger: Logger,
    pub scuba: MononokeScubaSampleBuilder,
    pub repo: Arc<Repo>,
    pub maybe_push_redirector_args: Option<PushRedirectorArgs<Repo>>,
    pub repo_client_knobs: RepoClientKnobs,
    pub maybe_backup_repo_source: Option<BackupSourceRepo>,
}

pub fn repo_handler(mononoke: Arc<Mononoke>, repo_name: &str) -> anyhow::Result<RepoHandler> {
    let source_repo = mononoke.raw_repo(repo_name).ok_or_else(|| {
        anyhow!(
            "Requested repo {} is not being served by this server",
            &repo_name
        )
    })?;
    let base = source_repo.repo_handler_base.clone();
    let maybe_push_redirector_args = match &base.maybe_push_redirector_base {
        Some(push_redirector_base) => {
            let large_repo_id = push_redirector_base.common_commit_sync_config.large_repo_id;
            let target_repo = mononoke
                .raw_repo_by_id(large_repo_id.id())
                .ok_or(ErrorKind::LargeRepoNotFound(large_repo_id))?;
            Some(PushRedirectorArgs::new(
                target_repo,
                Arc::clone(&source_repo),
                push_redirector_base.synced_commit_mapping.clone(),
                Arc::clone(&push_redirector_base.target_repo_dbs),
            ))
        }
        None => None,
    };
    let maybe_backup_repo_source = match &base.backup_repo_config {
        None => None,
        Some(ref backup_repo_config) => {
            let (orig_repo_name, source_repo_name) =
                (source_repo.name(), &backup_repo_config.source_repo_name);
            // If the repo itself serves as its backup source, then it's not a backup repo.
            // Hence, no need to setup backup_repo_source
            if *orig_repo_name == *source_repo_name {
                None
            } else {
                let backup_repo_source = mononoke.raw_repo(source_repo_name).ok_or_else(|| {
                    anyhow!(
                        "Backup source repo {} for core repo {} is not being served by this server",
                        source_repo_name,
                        orig_repo_name,
                    )
                })?;
                Some(BackupSourceRepo::from_blob_repo(
                    backup_repo_source.blob_repo(),
                ))
            }
        }
    };

    Ok(RepoHandler {
        logger: base.logger.clone(),
        scuba: base.scuba.clone(),
        repo_client_knobs: base.repo_client_knobs.clone(),
        repo: source_repo,
        maybe_push_redirector_args,
        maybe_backup_repo_source,
    })
}


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

use anyhow::anyhow;
use anyhow::Context;
use anyhow::Error;
use anyhow::Result;
use bytes::Bytes;
use connection_security_checker::ConnectionSecurityChecker;
use context::LoggingContainer;
use context::SessionContainer;
use context::SessionId;
use failure_ext::SlogKVError;
use fbinit::FacebookInit;
use futures::compat::Future01CompatExt;
use futures_old::sync::mpsc;
use futures_old::Future;
use futures_old::Stream;
use futures_stats::TimedFutureExt;
use hgproto::sshproto;
use hgproto::HgProtoHandler;
use maplit::hashmap;
use maplit::hashset;
use mononoke_api::Mononoke;
use qps::Qps;
use rate_limiting::Metric;
use rate_limiting::RateLimitEnvironment;
use repo_client::RepoClient;
use repo_permission_checker::RepoPermissionCheckerRef;
use scribe_ext::Scribe;
use slog::error;
use slog::o;
use slog::Drain;
use slog::Level;
use slog::Logger;
use slog_ext::SimpleFormatWithError;
use slog_kvfilter::KVFilter;
use sshrelay::SenderBytesWrite;
use sshrelay::Stdio;
use stats::prelude::*;
use time_ext::DurationExt;

use crate::errors::ErrorKind;
use crate::repo_handlers::repo_handler;
use crate::repo_handlers::RepoHandler;

define_stats! {
    prefix = "mononoke.request_handler";
    wireproto_ms:
        histogram(500, 0, 100_000, Average, Sum, Count; P 5; P 25; P 50; P 75; P 95; P 97; P 99),
    request_success: timeseries(Rate, Sum),
    request_failure: timeseries(Rate, Sum),
    request_outcome_permille: timeseries(Average),
}

pub async fn request_handler(
    fb: FacebookInit,
    reponame: String,
    mononoke: Arc<Mononoke>,
    _security_checker: &ConnectionSecurityChecker,
    stdio: Stdio,
    rate_limiter: Option<RateLimitEnvironment>,
    scribe: Scribe,
    qps: Option<Arc<Qps>>,
    readonly: bool,
) -> Result<()> {
    let Stdio {
        stdin,
        stdout,
        stderr,
        metadata,
    } = stdio;

    let session_id = metadata.session_id();

    // We don't have a repository yet, so create without server drain
    let conn_log = create_conn_logger(stderr.clone(), None, Some(session_id));

    let handler = repo_handler(mononoke, &reponame).with_context(|| {
        error!(
            conn_log,
            "Requested repo \"{}\" does not exist or is disabled", &reponame;
            "remote" => "true"
        );

        anyhow!("Unknown Repo: {}", &reponame)
    })?;

    let RepoHandler {
        logger,
        mut scuba,
        repo,
        maybe_push_redirector_args,
        repo_client_knobs,
        maybe_backup_repo_source,
    } = handler;

    // Upgrade log to include server drain
    let conn_log = create_conn_logger(stderr.clone(), Some(logger), Some(session_id));

    scuba = scuba.with_seq("seq");
    scuba.add("repo", reponame);
    scuba.add_metadata(&metadata);
    scuba.sample_for_identities(metadata.identities());

    let rate_limiter = rate_limiter.map(|r| r.get_rate_limiter());
    if let Some(ref rate_limiter) = rate_limiter {
        if let Err(err) = rate_limiter.check_load_shed(metadata.identities()) {
            scuba.log_with_msg("Request rejected due to load shedding", format!("{}", err));
            error!(conn_log, "Request rejected due to load shedding: {}", err; "remote" => "true");

            return Err(err.into());
        }
    }

    let is_allowed_to_repo = repo
        .repo_permission_checker()
        .check_if_read_access_allowed(metadata.identities())
        .await;

    if !is_allowed_to_repo {
        let err: Error = ErrorKind::AuthorizationFailed.into();
        scuba.log_with_msg("Authorization failed", format!("{}", err));
        error!(conn_log, "Authorization failed: {}", err; "remote" => "true");

        return Err(err);
    }

    // Info per wireproto command within this session
    let wireproto_calls = Arc::new(Mutex::new(Vec::new()));

    scuba.log_with_msg("Connection established", None);

    let session_builder = SessionContainer::builder(fb)
        .metadata(metadata.clone())
        .readonly(readonly)
        .rate_limiter(rate_limiter);

    let session = session_builder.build();

    let mut logging = LoggingContainer::new(fb, conn_log.clone(), scuba.clone());
    logging.with_scribe(scribe);

    let repo_client = RepoClient::new(
        repo,
        session.clone(),
        logging,
        maybe_push_redirector_args,
        repo_client_knobs,
        maybe_backup_repo_source,
    );
    let request_perf_counters = repo_client.request_perf_counters();

    // Construct a hg protocol handler
    let proto_handler = HgProtoHandler::new(
        conn_log.clone(),
        stdin.map(|b| bytes_old::Bytes::from(b.as_ref())),
        repo_client,
        sshproto::HgSshCommandDecode,
        sshproto::HgSshCommandEncode,
        wireproto_calls.clone(),
        qps.clone(),
        metadata.revproxy_region().clone(),
    );

    // send responses back
    let endres = proto_handler
        .inspect(move |bytes| session.bump_load(Metric::EgressBytes, bytes.len() as f64))
        .map_err(Error::from)
        .map(|b| Bytes::copy_from_slice(b.as_ref()))
        .forward(stdout)
        .map(|_| ());

    // If we got an error at this point, then catch it and print a message
    let (stats, result) = endres.compat().timed().await;

    let wireproto_calls = {
        let mut wireproto_calls = wireproto_calls.lock().expect("lock poisoned");
        std::mem::take(&mut *wireproto_calls)
    };

    STATS::wireproto_ms.add_value(stats.completion_time.as_millis_unchecked() as i64);

    let mut scuba = scuba.clone();

    scuba
        .add_future_stats(&stats)
        .add("wireproto_commands", wireproto_calls);

    // Populate stats no matter what to avoid dead detectors firing.
    STATS::request_success.add_value(0);
    STATS::request_failure.add_value(0);

    // Log request level perf counters
    request_perf_counters.insert_perf_counters(&mut scuba);

    match &result {
        Ok(_) => {
            STATS::request_success.add_value(1);
            STATS::request_outcome_permille.add_value(1000);
            scuba.log_with_msg("Request finished - Success", None)
        }
        Err(err) => {
            if err.is::<mpsc::SendError<Bytes>>() {
                STATS::request_outcome_permille.add_value(0);
                scuba.log_with_msg("Request finished - Client Disconnected", format!("{}", err));
            } else {
                STATS::request_failure.add_value(1);
                STATS::request_outcome_permille.add_value(0);
                scuba.log_with_msg("Request finished - Failure", format!("{:#?}", err));
            }
        }
    }

    if let Err(err) = result {
        error!(&conn_log, "Command failed";
            SlogKVError(err),
            "remote" => "true"
        );
    }

    Ok(())
}

pub fn create_conn_logger(
    stderr: mpsc::UnboundedSender<Bytes>,
    server_logger: Option<Logger>,
    session_id: Option<&SessionId>,
) -> Logger {
    let session_id = match session_id {
        Some(session_id) => session_id.to_string(),
        None => "".to_string(),
    };
    let decorator = o!("session_uuid" => format!("{}", session_id));

    let stderr_write = SenderBytesWrite { chan: stderr };
    let client_drain = slog_term::PlainSyncDecorator::new(stderr_write);
    let client_drain = SimpleFormatWithError::new(client_drain);
    let client_drain = KVFilter::new(client_drain, Level::Critical).only_pass_any_on_all_keys(
        (hashmap! {
            "remote".into() => hashset!["true".into(), "remote_only".into()],
        })
        .into(),
    );

    if let Some(logger) = server_logger {
        let server_drain = KVFilter::new(logger, Level::Critical).always_suppress_any(
            (hashmap! {
                "remote".into() => hashset!["remote_only".into()],
            })
            .into(),
        );

        // Don't fail logging if the client goes away
        let drain = slog::Duplicate::new(client_drain, server_drain).ignore_res();
        Logger::root(drain, decorator)
    } else {
        Logger::root(client_drain.ignore_res(), decorator)
    }
}


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

use std::pin::Pin;

use chrono::DateTime;
use chrono::Utc;
use futures::sink::Sink;
use futures::task::Context;
use futures::task::Poll;
use pin_project::pin_project;
use sshrelay::IoStream;
use sshrelay::SshMsg;

#[pin_project]
pub struct WireprotoSink<T> {
    #[pin]
    inner: T,
    pub data: WireprotoSinkData,
}

impl<T> WireprotoSink<T> {
    pub fn new(inner: T) -> Self {
        Self {
            inner,
            data: WireprotoSinkData::new(),
        }
    }
}

impl<T> Sink<SshMsg> for WireprotoSink<T>
where
    T: Sink<SshMsg>,
{
    type Error = <T as Sink<SshMsg>>::Error;

    fn poll_ready(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
        let this = self.project();
        let ret = this.inner.poll_ready(cx);
        this.data.peek_io(&ret);
        ret
    }

    fn start_send(self: Pin<&mut Self>, item: SshMsg) -> Result<(), Self::Error> {
        let this = self.project();
        this.data.peek_message(&item);
        this.inner.start_send(item)
    }

    fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
        let this = self.project();
        let ret = this.inner.poll_flush(cx);
        this.data.peek_io(&ret);
        this.data.peek_flush(&ret);
        ret
    }

    fn poll_close(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
        let this = self.project();
        let ret = this.inner.poll_close(cx);
        this.data.peek_io(&ret);
        ret
    }
}

pub struct WireprotoSinkData {
    pub last_successful_flush: Option<DateTime<Utc>>,
    pub last_successful_io: Option<DateTime<Utc>>,
    pub last_failed_io: Option<DateTime<Utc>>,
    pub stdout: ChannelData,
    pub stderr: ChannelData,
}

impl WireprotoSinkData {
    fn new() -> Self {
        Self {
            last_successful_flush: None,
            last_successful_io: None,
            last_failed_io: None,
            stdout: ChannelData::default(),
            stderr: ChannelData::default(),
        }
    }

    fn peek_message(&mut self, item: &SshMsg) {
        match item.stream_ref() {
            IoStream::Stdout => self.stdout.peek(item.as_ref()),
            IoStream::Stderr => self.stderr.peek(item.as_ref()),
            IoStream::Stdin => {}
        }
    }

    fn peek_io<E>(&mut self, res: &Poll<Result<(), E>>) {
        match res {
            Poll::Pending => {}
            Poll::Ready(Ok(())) => {
                self.last_successful_io = Some(Utc::now());
            }
            Poll::Ready(Err(..)) => {
                self.last_failed_io = Some(Utc::now());
            }
        }
    }

    fn peek_flush<E>(&mut self, res: &Poll<Result<(), E>>) {
        match res {
            Poll::Pending => {}
            Poll::Ready(Ok(())) => {
                self.last_successful_flush = Some(Utc::now());
            }
            Poll::Ready(Err(..)) => {
                // No need for this it's already tracked in peek_io.
            }
        }
    }
}

#[derive(Default)]
pub struct ChannelData {
    pub messages: u64,
    pub bytes: u64,
}

impl ChannelData {
    pub fn peek(&mut self, data: &[u8]) {
        let len: u64 = data
            .len()
            .try_into()
            .expect("The length of a buffer that exists will fit in a u64");

        self.messages += 1;
        self.bytes += 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 rand::distributions::Alphanumeric;
use rand::thread_rng;
use rand::Rng;

#[derive(Debug, Clone, Hash, Eq, PartialEq)]
pub struct SessionId(String);

pub fn generate_session_id() -> SessionId {
    let s: String = thread_rng()
        .sample_iter(&Alphanumeric)
        .map(char::from)
        .take(16)
        .collect();
    SessionId::from_string(s)
}

impl SessionId {
    pub fn from_string<T: ToString>(s: T) -> Self {
        Self(s.to_string())
    }

    pub fn into_string(self) -> String {
        self.0
    }

    pub fn as_str(&self) -> &str {
        self.0.as_str()
    }
}

impl Default for SessionId {
    fn default() -> Self {
        generate_session_id()
    }
}

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


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

#![feature(never_type)]

use std::path::PathBuf;
use std::sync::atomic::AtomicBool;
use std::sync::atomic::Ordering;
use std::sync::Arc;

use anyhow::Context;
use anyhow::Result;
use async_trait::async_trait;
use cache_warmup::cache_warmup;
use clap::Parser;
use cloned::cloned;
use cmdlib_logging::ScribeLoggingArgs;
use environment::WarmBookmarksCacheDerivedData;
use executor_lib::args::ShardedExecutorArgs;
use executor_lib::RepoShardedProcess;
use executor_lib::RepoShardedProcessExecutor;
use fbinit::FacebookInit;
use futures::channel::oneshot;
use futures::stream;
use futures::stream::StreamExt;
use futures::stream::TryStreamExt;
use metaconfig_types::ShardedService;
use mononoke_api::CoreContext;
use mononoke_api::Repo;
use mononoke_app::args::HooksAppExtension;
use mononoke_app::args::McrouterAppExtension;
use mononoke_app::args::ReadonlyArgs;
use mononoke_app::args::RepoFilterAppExtension;
use mononoke_app::args::ShutdownTimeoutArgs;
use mononoke_app::args::WarmBookmarksCacheExtension;
use mononoke_app::fb303::Fb303AppExtension;
use mononoke_app::fb303::ReadyFlagService;
use mononoke_app::MononokeApp;
use mononoke_app::MononokeAppBuilder;
use mononoke_app::MononokeReposManager;
use openssl::ssl::AlpnError;
use sharding_ext::RepoShard;
use slog::error;
use slog::info;
use slog::o;
use slog::Logger;

const SM_CLEANUP_TIMEOUT_SECS: u64 = 120;

// We will select the first protocol supported by the server which is also supported by the client.
// Order of preferences: hgcli, h2, http/1.1.
pub const ALPN_MONONOKE_PROTOS_OFFERS: &[u8] = b"\x05hgcli\x02h2\x08http/1.1";

/// Mononoke Server
#[derive(Parser)]
struct MononokeServerArgs {
    #[clap(flatten)]
    shutdown_timeout_args: ShutdownTimeoutArgs,
    #[clap(flatten)]
    scribe_logging_args: ScribeLoggingArgs,
    /// TCP address to listen to in format `host:port
    #[clap(long)]
    listening_host_port: String,
    /// Path for file in which to write the bound tcp address in rust std::net::SocketAddr format
    #[clap(long)]
    bound_address_file: Option<PathBuf>,
    /// If provided the thrift server will start on this port
    #[clap(long, short = 'p')]
    thrift_port: Option<String>,
    /// Path to a file with server certificate
    #[clap(long)]
    cert: String,
    /// Path to a file with server private key
    #[clap(long)]
    private_key: String,
    /// Path to a file with CA certificate
    #[clap(long)]
    ca_pem: String,
    /// Path to a file with encryption keys for SSL tickets
    #[clap(long)]
    ssl_ticket_seeds: Option<String>,
    /// Top level Mononoke tier where CSLB publishes routing table
    #[clap(long)]
    cslb_config: Option<String>,
    #[clap(flatten)]
    readonly: ReadonlyArgs,
    #[clap(flatten)]
    sharded_executor_args: ShardedExecutorArgs,
    /// Path to a file with land service client certificate
    #[clap(long)]
    land_service_client_cert: Option<String>,
    /// Path to a file with land service client private key
    #[clap(long, requires = "land_service_client_cert")]
    land_service_client_private_key: Option<String>,
}

/// Struct representing the Mononoke server process when sharding by repo.
pub struct MononokeServerProcess {
    fb: FacebookInit,
    repos_mgr: Arc<MononokeReposManager<Repo>>,
}

impl MononokeServerProcess {
    fn new(fb: FacebookInit, repos_mgr: MononokeReposManager<Repo>) -> Self {
        let repos_mgr = Arc::new(repos_mgr);
        Self { fb, repos_mgr }
    }

    async fn add_repo(&self, repo_name: &str, logger: &Logger) -> Result<()> {
        // Check if the input repo is already initialized. This can happen if the repo is a
        // shallow-sharded repo, in which case it would already be initialized during service startup.
        if self.repos_mgr.repos().get_by_name(repo_name).is_none() {
            // The input repo is a deep-sharded repo, so it needs to be added now.
            let repo = self.repos_mgr.add_repo(repo_name).await?;
            let blob_repo = repo.blob_repo().clone();
            let cache_warmup_params = repo.config().cache_warmup.clone();
            let ctx = CoreContext::new_with_logger(self.fb, logger.clone());
            cache_warmup(&ctx, &blob_repo, cache_warmup_params)
                .await
                .with_context(|| format!("Error while warming up cache for repo {}", repo_name))?;
            info!(
                &logger,
                "Completed repo {} setup in Mononoke service", repo_name
            );
        } else {
            info!(
                &logger,
                "Repo {} is already setup in Mononoke service", repo_name
            );
        }
        Ok(())
    }
}

#[async_trait]
impl RepoShardedProcess for MononokeServerProcess {
    async fn setup(&self, repo: &RepoShard) -> anyhow::Result<Arc<dyn RepoShardedProcessExecutor>> {
        let repo_name = repo.repo_name.as_str();
        let logger = self.repos_mgr.repo_logger(repo_name);
        info!(&logger, "Setting up repo {} in Mononoke service", repo_name);
        self.add_repo(repo_name, &logger).await.with_context(|| {
            format!(
                "Failure in setting up repo {} in Mononoke service",
                repo_name
            )
        })?;
        Ok(Arc::new(MononokeServerProcessExecutor {
            repo_name: repo_name.to_string(),
            repos_mgr: self.repos_mgr.clone(),
        }))
    }
}

/// Struct representing the execution of the Mononoke server for a particular
/// repo when sharding by repo.
pub struct MononokeServerProcessExecutor {
    repo_name: String,
    repos_mgr: Arc<MononokeReposManager<Repo>>,
}

impl MononokeServerProcessExecutor {
    fn remove_repo(&self, repo_name: &str) -> Result<()> {
        let config = self.repos_mgr.repo_config(repo_name).with_context(|| {
            format!(
                "Failure in remove repo {}. The config for repo doesn't exist",
                repo_name
            )
        })?;
        // Check if the current repo is a deep-sharded or shallow-sharded repo. If the
        // repo is deep-sharded, then remove it since SM wants some other host to serve it.
        // If repo is shallow-sharded, then keep it since regardless of SM sharding, shallow
        // sharded repos need to be present on each host.
        let is_deep_sharded = config
            .deep_sharding_config
            .and_then(|c| c.status.get(&ShardedService::EdenApi).copied())
            .unwrap_or(false);
        if is_deep_sharded {
            self.repos_mgr.remove_repo(repo_name);
            info!(
                self.repos_mgr.logger(),
                "No longer serving repo {} in Mononoke service.", repo_name,
            );
        } else {
            info!(
                self.repos_mgr.logger(),
                "Continuing serving repo {} in Mononoke service because it's shallow-sharded.",
                repo_name,
            );
        }
        Ok(())
    }
}

#[async_trait]
impl RepoShardedProcessExecutor for MononokeServerProcessExecutor {
    async fn execute(&self) -> anyhow::Result<()> {
        info!(
            self.repos_mgr.logger(),
            "Serving repo {} in Mononoke service", &self.repo_name,
        );
        Ok(())
    }

    async fn stop(&self) -> anyhow::Result<()> {
        self.remove_repo(&self.repo_name)
            .with_context(|| format!("Failure in stopping repo {}", &self.repo_name))
    }
}

#[fbinit::main]
fn main(fb: FacebookInit) -> Result<()> {
    let app = MononokeAppBuilder::new(fb)
        .with_default_scuba_dataset("mononoke_test_perf")
        .with_warm_bookmarks_cache(WarmBookmarksCacheDerivedData::HgOnly)
        .with_app_extension(WarmBookmarksCacheExtension {})
        .with_app_extension(McrouterAppExtension {})
        .with_app_extension(Fb303AppExtension {})
        .with_app_extension(HooksAppExtension {})
        .with_app_extension(RepoFilterAppExtension {})
        .build::<MononokeServerArgs>()?;
    let args: MononokeServerArgs = app.args()?;

    let root_log = app.logger().clone();
    let runtime = app.runtime().clone();

    let cslb_config = args.cslb_config.clone();
    info!(root_log, "Starting up");

    #[cfg(fbcode_build)]
    if let (Some(land_service_cert_path), Some(land_service_key_path)) = (
        &args.land_service_client_cert,
        &args.land_service_client_private_key,
    ) {
        pushrebase_client::land_service_override_certificate_paths(
            land_service_cert_path,
            land_service_key_path,
            &args.ca_pem,
        );
    }

    let configs = app.repo_configs();

    let acceptor = {
        let mut builder = secure_utils::SslConfig::new(
            args.ca_pem,
            args.cert,
            args.private_key,
            args.ssl_ticket_seeds,
        )
        .tls_acceptor_builder(root_log.clone())
        .context("Failed to instantiate TLS Acceptor builder")?;

        builder.set_alpn_protos(ALPN_MONONOKE_PROTOS_OFFERS)?;
        builder.set_alpn_select_callback(|_, list| {
            openssl::ssl::select_next_proto(ALPN_MONONOKE_PROTOS_OFFERS, list)
                .ok_or(AlpnError::NOACK)
        });

        builder.build()
    };

    info!(root_log, "Creating repo listeners");

    let scribe = args.scribe_logging_args.get_scribe(fb)?;
    let host_port = args.listening_host_port;
    let bound_addr_file = args.bound_address_file;

    let service = ReadyFlagService::new();
    let (terminate_sender, terminate_receiver) = oneshot::channel::<()>();
    let will_exit = Arc::new(AtomicBool::new(false));

    let env = app.environment();
    let scuba = env.scuba_sample_builder.clone();
    // Service name is used for shallow or deep sharding. If sharding itself is disabled, provide
    // service name as None while opening repos.
    let service_name = args
        .sharded_executor_args
        .sharded_service_name
        .as_ref()
        .map(|_| ShardedService::EdenApi);
    app.start_monitoring("mononoke_server", service.clone())?;
    app.start_stats_aggregation()?;

    let repo_listeners = {
        cloned!(root_log, will_exit, env, runtime, service_name);
        move |app: MononokeApp| async move {
            let common = configs.common.clone();
            let repos_mgr = app.open_managed_repos(service_name).await?;
            let mononoke = Arc::new(repos_mgr.make_mononoke_api()?);
            info!(&root_log, "Built Mononoke");

            info!(&root_log, "Warming up cache");
            stream::iter(mononoke.repos())
                .map(|repo| {
                    let repo_name = repo.name().to_string();
                    let blob_repo = repo.blob_repo().clone();
                    let root_log = root_log.clone();
                    let cache_warmup_params = repo.config().cache_warmup.clone();
                    async move {
                        let logger = root_log.new(o!("repo" => repo_name.clone()));
                        let ctx = CoreContext::new_with_logger(fb, logger);
                        cache_warmup(&ctx, &blob_repo, cache_warmup_params)
                            .await
                            .with_context(|| {
                                format!("Error while warming up cache for repo {}", repo_name)
                            })
                    }
                })
                // Repo cache warmup can be quite expensive, let's limit to 40
                // at a time.
                .buffer_unordered(40)
                .try_collect()
                .await?;
            info!(&root_log, "Cache warmup completed");
            if let Some(mut executor) = args.sharded_executor_args.build_executor(
                app.fb,
                runtime.clone(),
                app.logger(),
                || Arc::new(MononokeServerProcess::new(app.fb, repos_mgr)),
                false, // disable shard (repo) level healing
                SM_CLEANUP_TIMEOUT_SECS,
            )? {
                // The Sharded Process Executor needs to branch off and execute
                // on its own dedicated task spawned off the common tokio runtime.
                runtime.spawn({
                    let logger = app.logger().clone();
                    {
                        cloned!(will_exit);
                        async move { executor.block_and_execute(&logger, will_exit).await }
                    }
                });
            }
            repo_listener::create_repo_listeners(
                fb,
                common,
                mononoke.clone(),
                root_log,
                host_port,
                acceptor,
                service,
                terminate_receiver,
                &env.config_store,
                scribe,
                &scuba,
                will_exit,
                cslb_config,
                bound_addr_file,
                env.acl_provider.as_ref(),
                args.readonly.readonly,
            )
            .await
        }
    };

    app.run_until_terminated(
        repo_listeners,
        move || will_exit.store(true, Ordering::Relaxed),
        args.shutdown_timeout_args.shutdown_grace_period,
        async {
            if let Err(err) = terminate_sender.send(()) {
                error!(root_log, "could not send termination signal: {:?}", err);
            }
            repo_listener::wait_for_connections_closed(&root_log).await;
        },
        args.shutdown_timeout_args.shutdown_timeout,
    )
}


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

use anyhow::anyhow;
use anyhow::Result;
use bytes::BufMut;
use bytes::Bytes;
use bytes::BytesMut;
use futures::sync::mpsc;
use futures_ext::BoxStream;
use metadata::Metadata;
use netstring::NetstringDecoder;
use netstring::NetstringEncoder;
use tokio_util::codec::Decoder;
use tokio_util::codec::Encoder;
use zstd::stream::raw::Encoder as ZstdEncoder;
use zstd::stream::raw::InBuffer;
use zstd::stream::raw::Operation;
use zstd::stream::raw::OutBuffer;

// Multiplex stdin/out/err over a single stream using netstring as framing
#[derive(Debug)]
pub struct SshDecoder(NetstringDecoder);

pub struct SshEncoder {
    netstring: NetstringEncoder<Bytes>,
    compressor: Option<ZstdEncoder<'static>>,
}

pub struct Stdio {
    pub metadata: Arc<Metadata>,
    pub stdin: BoxStream<Bytes, io::Error>,
    pub stdout: mpsc::Sender<Bytes>,
    pub stderr: mpsc::UnboundedSender<Bytes>,
}

pub struct SenderBytesWrite {
    pub chan: mpsc::UnboundedSender<Bytes>,
}

impl io::Write for SenderBytesWrite {
    fn flush(&mut self) -> io::Result<()> {
        // Nothing to do. We don't control what happens in our receiver.
        Ok(())
    }

    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
        self.chan
            .unbounded_send(Bytes::copy_from_slice(buf))
            .map(|_| buf.len())
            .map_err(|e| io::Error::new(io::ErrorKind::BrokenPipe, e))
    }
}

// Matches Iostream in Mercurial mononokepeer.py
#[derive(Clone, Debug, Eq, PartialEq)]
pub enum IoStream {
    Stdin,
    Stdout,
    Stderr,
}

#[derive(Debug, Clone, Eq, PartialEq)]
pub struct SshMsg(IoStream, Bytes);

impl SshMsg {
    pub fn new(stream: IoStream, data: Bytes) -> Self {
        SshMsg(stream, data)
    }

    pub fn from_slice<T>(stream: IoStream, t: T) -> Self
    where
        T: AsRef<[u8]>,
    {
        Self::new(stream, Bytes::copy_from_slice(t.as_ref()))
    }

    pub fn stream(&self) -> IoStream {
        self.0.clone()
    }

    pub fn stream_ref(&self) -> &IoStream {
        &self.0
    }

    pub fn data(self) -> Bytes {
        self.1
    }
}

impl AsRef<[u8]> for SshMsg {
    fn as_ref(&self) -> &[u8] {
        self.1.as_ref()
    }
}

impl SshDecoder {
    pub fn new() -> Self {
        SshDecoder(NetstringDecoder::default())
    }
}

fn ioerr_cvt(err: anyhow::Error) -> io::Error {
    io::Error::new(io::ErrorKind::InvalidData, format!("{}", err))
}

impl Decoder for SshDecoder {
    type Item = SshMsg;
    type Error = io::Error;

    fn decode(&mut self, buf: &mut BytesMut) -> io::Result<Option<SshMsg>> {
        if let Some(mut data) = self.0.decode(buf).map_err(ioerr_cvt)? {
            if data.is_empty() {
                return Ok(None);
            }
            match data.split_to(1)[0] {
                0 => Ok(Some(SshMsg(IoStream::Stdin, data.freeze()))),
                1 => Ok(Some(SshMsg(IoStream::Stdout, data.freeze()))),
                2 => Ok(Some(SshMsg(IoStream::Stderr, data.freeze()))),
                _ => Err(io::Error::new(
                    io::ErrorKind::InvalidInput,
                    "bad ssh stream",
                )),
            }
        } else {
            Ok(None)
        }
    }
}

impl SshEncoder {
    pub fn new(compression_level: Option<i32>) -> Result<Self> {
        match compression_level {
            Some(level) => Ok(SshEncoder {
                netstring: NetstringEncoder::default(),
                compressor: Some(ZstdEncoder::new(level)?),
            }),
            _ => Ok(SshEncoder {
                netstring: NetstringEncoder::default(),
                compressor: None,
            }),
        }
    }

    fn compress_into<'a>(&mut self, out: &mut BytesMut, input: &'a [u8]) -> Result<()> {
        match &mut self.compressor {
            Some(compressor) => {
                let buflen = zstd_safe::compress_bound(input.len());
                if buflen >= zstd_safe::dstream_out_size() {
                    return Err(anyhow!(
                        "block is too big to compress in to a single zstd block"
                    ));
                }

                let mut src = InBuffer::around(input);
                let mut dst = vec![0u8; buflen];
                let mut dst = OutBuffer::around(&mut dst);

                while src.pos < src.src.len() {
                    compressor.run(&mut src, &mut dst)?;
                }
                loop {
                    let remaining = compressor.flush(&mut dst)?;
                    if remaining == 0 {
                        break;
                    }
                }
                out.put_slice(dst.as_slice());
            }
            None => out.put_slice(input),
        };

        Ok(())
    }
}

impl Encoder<SshMsg> for SshEncoder {
    type Error = io::Error;

    fn encode(&mut self, msg: SshMsg, buf: &mut BytesMut) -> io::Result<()> {
        let mut v = BytesMut::with_capacity(1 + msg.1.len());

        match msg.0 {
            IoStream::Stdin => {
                v.put_u8(0);
                self.compress_into(&mut v, &msg.1).map_err(ioerr_cvt)?;
                Ok(self.netstring.encode(v.freeze(), buf).map_err(ioerr_cvt)?)
            }
            IoStream::Stdout => {
                v.put_u8(1);
                self.compress_into(&mut v, &msg.1).map_err(ioerr_cvt)?;
                Ok(self.netstring.encode(v.freeze(), buf).map_err(ioerr_cvt)?)
            }
            IoStream::Stderr => {
                v.put_u8(2);
                self.compress_into(&mut v, &msg.1).map_err(ioerr_cvt)?;
                Ok(self.netstring.encode(v.freeze(), buf).map_err(ioerr_cvt)?)
            }
        }
    }
}

#[cfg(test)]
mod test {
    use bytes::BufMut;
    use bytes::BytesMut;
    use tokio_util::codec::Decoder;
    use tokio_util::codec::Encoder;

    use super::IoStream::*;
    use super::*;

    trait ToBytes: AsRef<[u8]> {
        fn bytes(&self) -> Bytes {
            Bytes::copy_from_slice(self.as_ref())
        }
    }

    impl<T> ToBytes for T where T: AsRef<[u8]> {}

    #[test]
    fn encode_simple() {
        let mut buf = BytesMut::with_capacity(1024);
        let mut encoder = SshEncoder::new(None).unwrap();

        encoder
            .encode(SshMsg::new(Stdin, b"ls -l".bytes()), &mut buf)
            .expect("encode failed");

        assert_eq!(buf.as_ref(), b"6:\x00ls -l,");
    }

    #[test]
    fn encode_zero() {
        let mut buf = BytesMut::with_capacity(1024);
        let mut encoder = SshEncoder::new(None).unwrap();

        encoder
            .encode(SshMsg::new(Stdin, b"".bytes()), &mut buf)
            .expect("encode failed");

        assert_eq!(buf.as_ref(), b"1:\x00,");
    }

    #[test]
    fn encode_one() {
        let mut buf = BytesMut::with_capacity(1024);
        let mut encoder = SshEncoder::new(None).unwrap();

        encoder
            .encode(SshMsg::new(Stdin, b"X".bytes()), &mut buf)
            .expect("encode failed");

        assert_eq!(buf.as_ref(), b"2:\x00X,");
    }

    #[test]
    fn encode_multi() {
        let mut buf = BytesMut::with_capacity(1024);
        let mut encoder = SshEncoder::new(None).unwrap();

        encoder
            .encode(SshMsg::new(Stdin, b"X".bytes()), &mut buf)
            .expect("encode failed");
        encoder
            .encode(SshMsg::new(Stdout, b"Y".bytes()), &mut buf)
            .expect("encode failed");
        encoder
            .encode(SshMsg::new(Stderr, b"Z".bytes()), &mut buf)
            .expect("encode failed");

        assert_eq!(buf.as_ref(), b"2:\x00X,2:\x01Y,2:\x02Z,");
    }

    #[test]
    fn encode_compressed() {
        let mut buf = BytesMut::with_capacity(1024);
        let mut encoder = SshEncoder::new(Some(3)).unwrap();

        encoder
            .encode(
                SshMsg::new(
                    Stdin,
                    b"hello hello hello hello hello hello hello hello hello".bytes(),
                ),
                &mut buf,
            )
            .expect("encode failed");
        assert_eq!(buf.as_ref(), b"22:\x00\x28\xb5\x2f\xfd\x00\x58\x64\x00\x00\x30\x68\x65\x6c\x6c\x6f\x20\x01\x00\x24\x2a\x45\x2c");
    }

    #[test]
    fn encode_compressed_too_big() {
        let mut buf = BytesMut::with_capacity(1024);
        let mut encoder = SshEncoder::new(Some(3)).unwrap();

        // 1MB, which is larger then 128KB zstd streaming buffer
        let message = vec![0u8; 1048576];
        let result = encoder.encode(SshMsg::new(Stdin, message.as_slice().bytes()), &mut buf);
        assert!(result.is_err());
    }

    #[test]
    fn decode_simple() {
        let mut buf = BytesMut::with_capacity(1024);
        buf.put_slice(b"6:\x00ls -l,");

        let mut decoder = SshDecoder::new();

        match decoder.decode(&mut buf) {
            Ok(Some(ref res)) if res == &SshMsg::new(Stdin, b"ls -l".bytes()) => {}
            bad => panic!("decode failed: {:?}", bad.as_ref()),
        }
    }

    #[test]
    fn decode_zero() {
        let mut buf = BytesMut::with_capacity(1024);
        buf.put_slice(b"1:\x00,");

        let mut decoder = SshDecoder::new();

        match decoder.decode(&mut buf) {
            Ok(Some(ref res)) if res == &SshMsg::new(Stdin, b"".bytes()) => {}
            bad => panic!("decode failed: {:?}", bad.as_ref()),
        }
    }

    #[test]
    fn decode_one() {
        let mut buf = BytesMut::with_capacity(1024);
        buf.put_slice(b"2:\x00X,");

        let mut decoder = SshDecoder::new();

        match decoder.decode(&mut buf) {
            Ok(Some(ref res)) if res == &SshMsg::new(Stdin, b"X".bytes()) => {}
            bad => panic!("decode failed: {:?}", bad.as_ref()),
        }
    }

    #[test]
    fn decode_multi() {
        let mut buf = BytesMut::with_capacity(1024);
        buf.put_slice(b"2:\x00X,2:\x01Y,2:\x02Z,");

        let mut decoder = SshDecoder::new();

        match decoder.decode(&mut buf) {
            Ok(Some(ref res)) if res == &SshMsg::new(Stdin, b"X".bytes()) => {}
            bad => panic!("decode failed: {:?}", bad.as_ref()),
        }
        match decoder.decode(&mut buf) {
            Ok(Some(ref res)) if res == &SshMsg::new(Stdout, b"Y".bytes()) => {}
            bad => panic!("decode failed: {:?}", bad.as_ref()),
        }
        match decoder.decode(&mut buf) {
            Ok(Some(ref res)) if res == &SshMsg::new(Stderr, b"Z".bytes()) => {}
            bad => panic!("decode failed: {:?}", bad.as_ref()),
        }
    }

    #[test]
    fn decode_bad() {
        let mut buf = BytesMut::with_capacity(1024);
        buf.put_slice(b"2:\x03X,");

        let mut decoder = SshDecoder::new();

        match decoder.decode(&mut buf) {
            Ok(bad) => panic!("unexpected success: {:?}", bad),
            Err(_err) => {}
        }
    }

    #[test]
    fn decode_short_framing() {
        let mut buf = BytesMut::with_capacity(1024);
        buf.put_slice(b"3:\x02X,");

        let mut decoder = SshDecoder::new();

        match decoder.decode(&mut buf) {
            Ok(None) => {}
            bad => panic!("bad framing: {:?}", bad),
        }
    }

    #[test]
    fn decode_broken_framing() {
        let mut buf = BytesMut::with_capacity(1024);
        buf.put_slice(b"1:\x02X,");

        let mut decoder = SshDecoder::new();

        match decoder.decode(&mut buf) {
            Ok(bad) => panic!("unexpected success: {:?}", bad),
            Err(_err) => {}
        }
    }
}