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

use std::pin::Pin;

use anyhow::Error;
use bytes::Bytes;
use futures::channel::oneshot::Sender;
use futures::ready;
use futures::stream::Stream;
use futures::task::Context;
use futures::task::Poll;
use pin_project::pin_project;
use pin_project::pinned_drop;

use super::error_meta::ErrorMeta;
use super::error_meta::ErrorMetaProvider;
use super::response_meta::BodyMeta;

pub trait Sizeable {
    fn size(&self) -> u64;
}

impl Sizeable for Bytes {
    fn size(&self) -> u64 {
        // NOTE: It is reasonable to unwrap here because we're not going to have buffers of bytes
        // that are larger than a u64.
        self.len().try_into().unwrap()
    }
}

/// A stream that will fire to the sender associated upon completing or being dropped. The Sender
/// will receive the amount of data that passed through the stream.
#[pin_project(PinnedDrop, project = SignalStreamProjection)]
pub struct SignalStream<S: ErrorMetaProvider<Error>> {
    #[pin]
    stream: S,
    sender: Option<Sender<BodyMeta>>,
    size_sent: u64,
}

impl<S> SignalStream<S>
where
    S: ErrorMetaProvider<Error>,
{
    pub fn new(stream: S, sender: Sender<BodyMeta>) -> Self {
        Self {
            stream,
            sender: Some(sender),
            size_sent: 0,
        }
    }
}

impl<S, T> Stream for SignalStream<S>
where
    S: Stream<Item = T> + ErrorMetaProvider<Error>,
    T: Sizeable,
{
    type Item = T;

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

        let poll = ready!(this.stream.as_mut().poll_next(ctx));

        match poll {
            // We have an item: increment the amount of data we sent.
            Some(ref item) => *this.size_sent += item.size(),
            // No items left: signal our receiver.
            None => send_body_meta(this),
        }

        Poll::Ready(poll)
    }
}

#[pinned_drop]
impl<S> PinnedDrop for SignalStream<S>
where
    S: ErrorMetaProvider<Error>,
{
    fn drop(self: Pin<&mut Self>) {
        let this = self.project();
        send_body_meta(this);
    }
}

fn send_body_meta<S>(this: SignalStreamProjection<S>)
where
    S: ErrorMetaProvider<Error>,
{
    if let Some(sender) = this.sender.take() {
        let bytes_sent = *this.size_sent;
        let mut error_meta = ErrorMeta::new();

        this.stream.report_errors(&mut error_meta);

        let _ = sender.send(BodyMeta {
            bytes_sent,
            error_meta,
        });
    }
}


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

use std::pin::Pin;

use anyhow::Error;
use async_compression::tokio::bufread::GzipEncoder;
use async_compression::tokio::bufread::ZstdEncoder;
use bytes::Bytes;
use futures::future::Either;
use futures::stream::BoxStream;
use futures::stream::Stream;
use futures::stream::StreamExt;
use futures::stream::TryStreamExt;
use futures::task::Context;
use futures::task::Poll;
use pin_project::pin_project;
use tokio_util::io::ReaderStream;
use tokio_util::io::StreamReader;
use yield_stream::YieldStream;

use crate::content_encoding::ContentCompression;
use crate::content_encoding::ContentEncoding;

/// Create a response stream using the specified Content-Encoding.
///
/// The resulting stream may or may not be compressed depending on the chosen encoding. Optionally,
/// the caller can specify the value for the `Content-Length` header. This is only useful in cases
/// where the response isn't compressed (i.e., the encoding is set to `ContentEncoding::Identity`)
/// because otherwise, we would need to send the post-compression size of the content, which cannot
/// be known in advance.
pub fn encode_stream<S>(
    stream: S,
    encoding: ContentEncoding,
    length: Option<u64>,
) -> Either<ResponseStream<S>, CompressedResponseStream<'static>>
where
    S: Stream<Item = Result<Bytes, Error>> + Send + 'static,
{
    match (encoding, length) {
        (ContentEncoding::Identity, Some(size)) => ResponseStream::new(stream)
            .set_content_length(size)
            .left_stream(),
        (ContentEncoding::Identity, None) => ResponseStream::new(stream).left_stream(),
        (ContentEncoding::Compressed(c), _) => {
            CompressedResponseStream::new(stream, c).right_stream()
        }
    }
}

#[pin_project]
pub struct CompressedResponseStream<'a> {
    inner: BoxStream<'a, Result<Bytes, Error>>,
    content_compression: ContentCompression,
}

impl<'a> CompressedResponseStream<'a> {
    pub fn new<S>(inner: S, content_compression: ContentCompression) -> Self
    where
        S: Stream<Item = Result<Bytes, Error>> + Send + 'a,
    {
        use std::io;

        // 2MiB, for LFS that's at least once every content chunk.
        const YIELD_EVERY: usize = 2 * 1024 * 1024;

        let inner = inner.map_err(|e| io::Error::new(io::ErrorKind::Other, e));
        let inner = YieldStream::new(inner, YIELD_EVERY, |data| {
            data.as_ref().map_or(0, |b| b.len())
        });
        let inner = StreamReader::new(inner);

        let inner = match content_compression {
            ContentCompression::Zstd => ReaderStream::new(ZstdEncoder::new(inner))
                .map_err(Error::from)
                .boxed(),
            ContentCompression::Gzip => ReaderStream::new(GzipEncoder::new(inner))
                .map_err(Error::from)
                .boxed(),
        };

        Self {
            inner,
            content_compression,
        }
    }

    pub fn content_compression(&self) -> ContentCompression {
        self.content_compression
    }
}

impl Stream for CompressedResponseStream<'_> {
    type Item = Result<Bytes, Error>;

    fn poll_next(self: Pin<&mut Self>, ctx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
        self.project().inner.poll_next_unpin(ctx)
    }
}

#[pin_project]
pub struct ResponseStream<S> {
    #[pin]
    inner: S,
    content_length: Option<u64>,
}

impl<S> ResponseStream<S> {
    pub fn new(inner: S) -> Self {
        Self {
            inner,
            content_length: None,
        }
    }

    /// Set a Content-Length for this stream. This *must* match the exact size of the uncompressed
    /// content that will be sent, since that is what the client will expect.
    pub fn set_content_length(self, content_length: u64) -> Self {
        Self {
            content_length: Some(content_length),
            ..self
        }
    }

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

impl<S> Stream for ResponseStream<S>
where
    S: Stream,
{
    type Item = S::Item;

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


/*
 * 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 futures::prelude::*;
use futures::ready;
use futures::stream::Stream;
use futures::task::Context;
use futures::task::Poll;
use pin_project::pin_project;

use super::error_meta::ErrorMeta;
use super::error_meta::ErrorMetaProvider;

pub trait ResponseTryStreamExt: TryStream {
    /// Filter out errors from a `TryStream` and capture the first one (then count further errors),
    /// transforming the `TryStream` into a `Stream<Item=Self::Ok>`.
    fn capture_first_err<E>(self) -> CaptureFirstErr<Self, E>
    where
        Self: Sized,
    {
        CaptureFirstErr::new(self)
    }

    /// Immediately end the `TryStream` upon encountering an error.
    ///
    /// The error will be passed to the given callback, and the stream will be
    /// fused to prevent the underlying `TryStream` from being polled again.
    fn end_on_err<E>(self) -> EndOnErr<Self, E>
    where
        Self: Sized,
    {
        EndOnErr::new(self)
    }
}

impl<S: TryStream + ?Sized> ResponseTryStreamExt for S {}

/// A stream that ignores errors raised by the internal stream and captures the first one.
#[pin_project]
pub struct CaptureFirstErr<S, E> {
    #[pin]
    stream: S,
    /// First error this stream encountered.
    first_error: Option<E>,
    /// Count of errors that were observed but not captured.
    extra_error_count: u64,
}

impl<S, E> CaptureFirstErr<S, E> {
    pub fn new(stream: S) -> Self {
        Self {
            stream,
            first_error: None,
            extra_error_count: 0,
        }
    }

    pub fn get_ref(&self) -> &S {
        &self.stream
    }
}

impl<S, E> Stream for CaptureFirstErr<S, E>
where
    S: TryStream<Error = E>,
{
    type Item = S::Ok;

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

        loop {
            match ready!(this.stream.as_mut().try_poll_next(ctx)) {
                Some(Ok(item)) => return Poll::Ready(Some(item)),
                Some(Err(e)) => {
                    if this.first_error.is_some() {
                        *this.extra_error_count += 1;
                    } else {
                        this.first_error.replace(e);
                    }
                }
                None => return Poll::Ready(None),
            }
        }
    }
}

impl<S, E> ErrorMetaProvider<E> for CaptureFirstErr<S, E>
where
    S: TryStream<Error = E>,
{
    fn report_errors(self: Pin<&mut Self>, errors: &mut ErrorMeta<E>) {
        let this = self.project();
        errors.errors.extend(this.first_error.take());
        errors.extra_error_count += *this.extra_error_count;
    }
}

#[pin_project]
pub struct EndOnErr<S, E> {
    #[pin]
    stream: S,
    errored: bool,
    error: Option<E>,
}

impl<S, E> EndOnErr<S, E> {
    pub fn new(stream: S) -> Self {
        Self {
            stream,
            errored: false,
            error: None,
        }
    }

    pub fn get_ref(&self) -> &S {
        &self.stream
    }
}

impl<S, E> Stream for EndOnErr<S, E>
where
    S: TryStream<Error = E>,
{
    type Item = S::Ok;

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

        // Fuse the stream once the error callback has fired.
        if *this.errored {
            return Poll::Ready(None);
        }

        match ready!(this.stream.try_poll_next(ctx)) {
            Some(Ok(item)) => Poll::Ready(Some(item)),
            Some(Err(e)) => {
                this.error.replace(e);
                *this.errored = true;
                Poll::Ready(None)
            }
            None => Poll::Ready(None),
        }
    }
}

impl<S, E> ErrorMetaProvider<E> for EndOnErr<S, E>
where
    S: TryStream<Error = E>,
{
    fn report_errors(self: Pin<&mut Self>, error_meta: &mut ErrorMeta<E>) {
        let this = self.project();
        error_meta.errors.extend(this.error.take());
    }
}

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

    #[tokio::test]
    async fn test_end_on_err() {
        let s = stream::iter(vec![
            Ok("hello"),
            Ok("world"),
            Err("error"),
            Ok("foo"),
            Err("bar"),
        ])
        .end_on_err();

        futures::pin_mut!(s);

        assert_eq!(s.next().await, Some("hello"));
        assert_eq!(s.next().await, Some("world"));
        assert_eq!(s.next().await, None);

        let mut errors = ErrorMeta::new();
        s.report_errors(&mut errors);
        assert_eq!(&errors.errors, &["error"]);
    }

    #[tokio::test]
    async fn test_capture_first_err() {
        let s = stream::iter(vec![
            Ok("hello"),
            Ok("world"),
            Err("error"),
            Ok("foo"),
            Err("bar"),
        ])
        .capture_first_err();

        futures::pin_mut!(s);

        assert_eq!(s.next().await, Some("hello"));
        assert_eq!(s.next().await, Some("world"));
        assert_eq!(s.next().await, Some("foo"));
        assert_eq!(s.next().await, None);

        let mut errors = ErrorMeta::new();
        s.report_errors(&mut errors);
        assert_eq!(&errors.errors, &["error"]);
        assert_eq!(errors.extra_error_count, 1);
    }
}


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

use futures::channel::oneshot::Receiver;
use futures_stats::StreamStats;
use gotham_derive::StateData;

#[derive(StateData)]
pub struct PendingStreamStats {
    stats: Option<Receiver<StreamStats>>,
}

impl PendingStreamStats {
    pub fn none() -> Self {
        Self { stats: None }
    }

    pub fn deferred(receiver: Receiver<StreamStats>) -> Self {
        Self {
            stats: Some(receiver),
        }
    }

    pub async fn finish(self) -> Option<StreamStats> {
        match self.stats {
            Some(receiver) => receiver.await.ok(),
            None => 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.
 */

//! Operations on the "batch" command.
//!
//! Based on the Mercurial wire protocol documentation. See
//! <https://www.mercurial-scm.org/repo/hg/file/@/mercurial/help/internals/wireprotocol.txt>

use anyhow::bail;
use anyhow::Result;

use crate::errors::ErrorKind;

/// Unescape a batch-escaped argument key or value.
pub fn unescape(bs: &[u8]) -> Result<Vec<u8>> {
    let mut out = Vec::with_capacity(bs.len());
    for (idx, slice) in bs.split(|b| b == &b':').enumerate() {
        if idx > 0 {
            // "::" or ":<end of string>" are both illegal.
            if slice.is_empty() {
                bail!(ErrorKind::BatchInvalid(
                    String::from_utf8_lossy(bs).into_owned(),
                ));
            }
            out.push(match slice[0] {
                b'c' => b':',
                b'o' => b',',
                b's' => b';',
                b'e' => b'=',
                ch => bail!(ErrorKind::BatchEscape(ch)),
            });
            out.extend_from_slice(&slice[1..]);
        } else {
            out.extend_from_slice(slice);
        }
    }
    Ok(out)
}

/// Escape a batch result.
pub fn escape<T: AsRef<[u8]>>(res: T) -> Vec<u8> {
    let res = res.as_ref();
    let mut out = Vec::with_capacity(res.len());
    for b in res {
        match b {
            b':' => out.extend_from_slice(b":c"),
            b',' => out.extend_from_slice(b":o"),
            b';' => out.extend_from_slice(b":s"),
            b'=' => out.extend_from_slice(b":e"),
            ch => out.push(*ch),
        }
    }
    out
}

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

    use super::*;

    const BAD_BYTES: [u8; 3] = [b',', b';', b'='];
    const BYTES_TO_ESCAPE: [u8; 4] = [b':', b',', b';', b'='];

    quickcheck! {
        fn test_roundtrip(input: Vec<u8>) -> bool {
            let bytes = Bytes::from(input);
            let escaped = escape(&bytes);
            let unescaped = unescape(&escaped).unwrap();
            unescaped == bytes
        }

        fn test_bad_bytes(input: Vec<u8>) -> bool {
            let bytes = Bytes::from(input);
            let escaped = escape(&bytes);
            escaped.iter().all(|&b| !BAD_BYTES.contains(&b))
        }

        fn test_escaped_length(input: Vec<u8>) -> bool {
            // Each character in BAD_BYTES should cause the length to be
            // extended by 1.
            let nbad = input.iter().fold(0, |acc, &b| {
                if BYTES_TO_ESCAPE.contains(&b) {
                    acc + 1
                } else {
                    acc
                }
            });
            let bytes = Bytes::from(input);
            let escaped = escape(&bytes);
            bytes.len() + nbad == escaped.len()
        }

        fn test_reverse_roundtrip(input: Vec<u8>) -> TestResult {
            // In this case we're not guaranteed that the string is well-formed.
            if input.iter().any(|&b| BAD_BYTES.contains(&b)) {
                return TestResult::discard();
            }
            let unescaped = unescape(&input);
            let unescaped_bytes = match unescaped {
                Err(_) => return TestResult::discard(),
                Ok(v) => Bytes::from(v),
            };
            let escaped = escape(&unescaped_bytes);
            TestResult::from_bool(escaped == input)
        }
    }
}


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

//! Mercurial protocol service framework
//!
//! To implement a Mercurial service, implement `HgCommands` and then use it to handle incominng
//! connections.
use std::collections::BTreeMap;
use std::collections::HashMap;
use std::collections::HashSet;
use std::io;
use std::io::BufRead;
use std::io::Cursor;
use std::sync::Arc;

use anyhow::bail;
use anyhow::Error;
use anyhow::Result;
use bytes_old::Buf;
use bytes_old::Bytes;
use bytes_old::BytesMut;
use failure_ext::FutureErrorContext;
use futures::future;
use futures::future::err;
use futures::future::ok;
use futures::future::Either;
use futures::future::Future;
use futures::stream;
use futures::stream::futures_ordered;
use futures::stream::once;
use futures::stream::Stream;
use futures::sync::oneshot;
use futures::IntoFuture;
use futures_ext::BoxFuture;
use futures_ext::BoxStream;
use futures_ext::BytesStream;
use futures_ext::FutureExt;
use futures_ext::StreamExt;
use limited_async_read::LimitedAsyncRead;
use mercurial_bundles::bundle2;
use mercurial_bundles::bundle2::Bundle2Stream;
use mercurial_bundles::bundle2::StreamEvent;
use mercurial_bundles::Bundle2Item;
use mercurial_types::HgChangesetId;
use mercurial_types::HgFileNodeId;
use mercurial_types::MPath;
use qps::Qps;
use slog::Logger;
use tokio_io::codec::Decoder;
use tokio_io::AsyncRead;

use crate::dechunker::Dechunker;
use crate::errors::*;
use crate::GetbundleArgs;
use crate::GettreepackArgs;
use crate::SingleRequest;
use crate::SingleResponse;

pub struct HgCommandHandler<H> {
    logger: Logger,
    commands: H,
    qps: Option<Arc<Qps>>,
    src_region: Option<String>,
}

impl<H: HgCommands + Send + 'static> HgCommandHandler<H> {
    pub fn new(
        logger: Logger,
        commands: H,
        qps: Option<Arc<Qps>>,
        src_region: Option<String>,
    ) -> Self {
        HgCommandHandler {
            logger,
            commands,
            qps,
            src_region,
        }
    }

    /// Handles a single command (not batched) by returning a stream of responses and a future
    /// resolving to the remainder unused input available only after the entire stream of responses
    /// have been consumed.
    pub fn handle<S>(
        &self,
        req: SingleRequest,
        instream: BytesStream<S>,
    ) -> (
        BoxStream<SingleResponse, Error>,
        BoxFuture<BytesStream<S>, Error>,
    )
    where
        S: Stream<Item = Bytes, Error = io::Error> + Send + 'static,
    {
        let hgcmds = &self.commands;

        if let (Some(qps), Some(src_region)) = (self.qps.as_ref(), self.src_region.as_ref()) {
            let _res = qps.bump(src_region);
        }

        match req {
            SingleRequest::Between { pairs } => (
                hgcmds
                    .between(pairs)
                    .map(SingleResponse::Between)
                    .into_stream()
                    .boxify(),
                ok(instream).boxify(),
            ),
            SingleRequest::Branchmap => (
                hgcmds
                    .branchmap()
                    .map(SingleResponse::Branchmap)
                    .into_stream()
                    .boxify(),
                ok(instream).boxify(),
            ),
            SingleRequest::Capabilities => (
                hgcmds
                    .capabilities()
                    .map(SingleResponse::Capabilities)
                    .into_stream()
                    .boxify(),
                ok(instream).boxify(),
            ),
            SingleRequest::ClientTelemetry { args } => (
                hgcmds
                    .clienttelemetry(args)
                    .map(SingleResponse::ClientTelemetry)
                    .into_stream()
                    .boxify(),
                ok(instream).boxify(),
            ),
            SingleRequest::Debugwireargs { one, two, all_args } => (
                self.debugwireargs(one, two, all_args)
                    .map(SingleResponse::Debugwireargs)
                    .into_stream()
                    .boxify(),
                ok(instream).boxify(),
            ),
            SingleRequest::Getbundle(args) => (
                hgcmds
                    .getbundle(args)
                    .map(SingleResponse::Getbundle)
                    .boxify(),
                ok(instream).boxify(),
            ),
            SingleRequest::Heads => (
                hgcmds
                    .heads()
                    .map(SingleResponse::Heads)
                    .into_stream()
                    .boxify(),
                ok(instream).boxify(),
            ),
            SingleRequest::Hello => (
                hgcmds
                    .hello()
                    .map(SingleResponse::Hello)
                    .into_stream()
                    .boxify(),
                ok(instream).boxify(),
            ),
            SingleRequest::Listkeys { namespace } => (
                hgcmds
                    .listkeys(namespace)
                    .map(SingleResponse::Listkeys)
                    .into_stream()
                    .boxify(),
                ok(instream).boxify(),
            ),
            SingleRequest::ListKeysPatterns {
                namespace,
                patterns,
            } => (
                hgcmds
                    .listkeyspatterns(namespace, patterns)
                    .map(SingleResponse::ListKeysPatterns)
                    .into_stream()
                    .boxify(),
                ok(instream).boxify(),
            ),
            SingleRequest::Lookup { key } => (
                hgcmds
                    .lookup(key)
                    .map(SingleResponse::Lookup)
                    .into_stream()
                    .boxify(),
                ok(instream).boxify(),
            ),
            SingleRequest::Known { nodes } => (
                hgcmds
                    .known(nodes)
                    .map(SingleResponse::Known)
                    .into_stream()
                    .boxify(),
                ok(instream).boxify(),
            ),
            SingleRequest::Knownnodes { nodes } => (
                hgcmds
                    .knownnodes(nodes)
                    .map(SingleResponse::Known)
                    .into_stream()
                    .boxify(),
                ok(instream).boxify(),
            ),
            SingleRequest::Unbundle { heads } => self.handle_unbundle(instream, heads, None, None),
            SingleRequest::UnbundleReplay {
                heads,
                replaydata,
                respondlightly,
            } => {
                let (resp, instream) =
                    self.handle_unbundle(instream, heads, Some(respondlightly), Some(replaydata));
                (resp, instream)
            }
            SingleRequest::Gettreepack(args) => (
                hgcmds
                    .gettreepack(args)
                    .map(SingleResponse::Gettreepack)
                    .boxify(),
                ok(instream).boxify(),
            ),
            SingleRequest::StreamOutShallow { tag } => (
                hgcmds
                    .stream_out_shallow(tag)
                    .map(SingleResponse::StreamOutShallow)
                    .boxify(),
                ok(instream).boxify(),
            ),
            SingleRequest::GetpackV1 => {
                let (reqs, instream) =
                    decode_getpack_arg_stream(instream, Getpackv1ArgDecoder::new);
                (
                    hgcmds
                        .getpackv1(reqs)
                        .map(SingleResponse::Getpackv1)
                        .boxify(),
                    instream,
                )
            }
            SingleRequest::GetpackV2 => {
                let (reqs, instream) =
                    decode_getpack_arg_stream(instream, Getpackv1ArgDecoder::new);
                (
                    hgcmds
                        .getpackv2(reqs)
                        .map(SingleResponse::Getpackv2)
                        .boxify(),
                    instream,
                )
            }
            SingleRequest::GetCommitData { nodes } => (
                hgcmds
                    .getcommitdata(nodes)
                    .map(SingleResponse::GetCommitData)
                    .boxify(),
                ok(instream).boxify(),
            ),
        }
    }

    fn handle_unbundle<S>(
        &self,
        instream: BytesStream<S>,
        heads: Vec<String>,
        respondlightly: Option<bool>,
        replaydata: Option<String>,
    ) -> (
        BoxStream<SingleResponse, Error>,
        BoxFuture<BytesStream<S>, Error>,
    )
    where
        S: Stream<Item = Bytes, Error = io::Error> + Send + 'static,
    {
        let hgcmds = &self.commands;
        let dechunker = Dechunker::new(instream);

        let bundle2stream =
            Bundle2Stream::new(self.logger.clone(), LimitedAsyncRead::new(dechunker));
        let (bundle2stream, remainder) = extract_remainder_from_bundle2(bundle2stream);

        let remainder = remainder
            .then(|rest| {
                let (bytes, remainder) = match rest {
                    Err(e) => return Either::A(err(e)),
                    Ok(rest) => rest,
                };
                if !bytes.is_empty() {
                    Either::A(err(ErrorKind::UnconsumedData(
                        String::from_utf8_lossy(bytes.as_ref()).into_owned(),
                    )
                    .into()))
                } else {
                    Either::B(remainder.into_inner().check_is_done().from_err())
                }
            })
            .then(
                |check_is_done: Result<(bool, Dechunker<_>)>| match check_is_done {
                    Ok((true, remainder)) => ok(remainder.into_inner()),
                    Ok((false, mut remainder)) => match remainder.fill_buf() {
                        Err(e) => err(e.into()),
                        Ok(buf) => err(ErrorKind::UnconsumedData(
                            String::from_utf8_lossy(buf).into_owned(),
                        )
                        .into()),
                    },
                    Err(e) => err(e),
                },
            )
            .boxify();

        let resps = futures_ordered(vec![
            Either::A(ok(SingleResponse::ReadyForStream)),
            Either::B({
                hgcmds
                    .unbundle(heads, bundle2stream, respondlightly, replaydata)
                    .map(SingleResponse::Unbundle)
            }),
        ]);
        (resps.boxify(), remainder)
    }

    // @wireprotocommand('debugwireargs', 'one two *')
    // Handled here because this is a meta debugging command that isn't
    // specific to a repo.
    fn debugwireargs(
        &self,
        one: Vec<u8>,
        two: Vec<u8>,
        all_args: HashMap<Vec<u8>, Vec<u8>>,
    ) -> HgCommandRes<Bytes> {
        let mut out = Vec::<u8>::new();
        out.extend_from_slice(&one[..]);
        out.push(b' ');
        out.extend_from_slice(&two[..]);
        out.push(b' ');
        out.extend_from_slice(get_or_none(&all_args, &b"three"[..]));
        out.push(b' ');
        out.extend_from_slice(get_or_none(&all_args, &b"four"[..]));
        out.push(b' ');
        // Note that "five" isn't actually read off the wire -- instead, the
        // default value "None" is used.
        out.extend_from_slice(NONE);

        future::ok(out.into()).boxify()
    }
}

const NONE: &[u8] = b"None";

fn decode_getpack_arg_stream<D, RType, S>(
    input: BytesStream<S>,
    create_decoder: impl Fn() -> D + Send + 'static,
) -> (BoxStream<RType, Error>, BoxFuture<BytesStream<S>, Error>)
where
    D: Decoder<Item = Option<RType>, Error = Error> + Send + 'static,
    RType: Send + 'static,
    S: Stream<Item = Bytes, Error = io::Error> + Send + 'static,
{
    let (send, recv) = oneshot::channel();

    // stream::unfold() requires us to to return None if it's finished, or Some(Future) if not.
    // We can't say if node file stream is finished before we parse the entry, that means that
    // we can't stop unfolding by returning None. Instead we return a "fake" error. This fake
    // error is a Result. If this fake error is Ok(...) then no real error happened.
    // Note that fake error also contains input stream that will be send back to the future that
    // waits for it.
    let entry_stream: BoxStream<_, Result<BytesStream<S>, (_, BytesStream<S>)>> =
        stream::unfold(input, move |input| {
            let fut_decode = input.into_future_decode(create_decoder());
            let fut = fut_decode
                .map_err(Err) // Real error happened, wrap it in result
                .and_then(|(maybe_item, instream)| match maybe_item {
                    None => {
                        // None here means we hit EOF, but that shouldn't happen
                        Err(Err((Error::msg("unexpected EOF"), instream)))
                            .into_future()
                            .boxify()
                    }
                    Some(maybe_nodehash) => {
                        match maybe_nodehash {
                            None => {
                                // None here means that we've read all the node-file pairs
                                // that client has sent us. Return fake error that means that
                                // we've successfully parsed the stream.
                                Err(Ok(instream)).into_future().boxify()
                            }
                            Some(nodehash) => {
                                // Parsed one more entry - continue
                                Ok((nodehash, instream)).into_future().boxify()
                            }
                        }
                    }
                });

            Some(fut)
        })
        .boxify();

    let try_send_instream =
        |wrapped_send: &mut Option<oneshot::Sender<_>>, instream: BytesStream<S>| -> Result<()> {
            let send = wrapped_send.take();
            let send =
                send.ok_or_else(|| Error::msg("internal error: tried to send input stream twice"))?;
            match send.send(instream) {
                Ok(_) => Ok(()), // Finished
                Err(_) => bail!("internal error while sending input stream back"),
            }
        };

    // We are parsing errors (both fake and real), and sending instream to the future
    // that awaits it. Note: instream should be send only once!
    let entry_stream = entry_stream.then({
        let mut wrapped_send = Some(send);
        move |val| {
            match val {
                Ok(nodefile) => Ok(Some(nodefile)),
                Err(Ok(instream)) => try_send_instream(&mut wrapped_send, instream).map(|_| None),
                Err(Err((err, instream))) => {
                    match try_send_instream(&mut wrapped_send, instream) {
                        // TODO(stash): if send fails, then Mononoke is deadlocked
                        // ignore send errors
                        Ok(_) => Err(err),
                        Err(_) => Err(err),
                    }
                }
            }
        }
    });

    // Finally, filter out last None value
    let entry_stream = entry_stream.filter_map(|val| val);
    (entry_stream.boxify(), recv.map_err(Error::from).boxify())
}

#[derive(Clone)]
enum GetPackv1ParsingState {
    Start,
    ParsingFilename(u16),
    ParsedFilename(MPath),
    ParsingFileNodes(MPath, u32, Vec<HgFileNodeId>),
}

// Request format:
//
// [<filerequest>,...]\0\0
// filerequest = <filename len: 2 byte><filename><count: 4 byte>
//               [<node: 20 byte>,...]
//
// Getpackv1ArgDecoder parses one `filerequest` entry i.e. one filename and a few filenodes
struct Getpackv1ArgDecoder {
    state: GetPackv1ParsingState,
}

impl Getpackv1ArgDecoder {
    #[allow(unused)]
    pub fn new() -> Self {
        Self {
            state: GetPackv1ParsingState::Start,
        }
    }
}

impl Decoder for Getpackv1ArgDecoder {
    // If None has been decoded, then that means that client has sent all the data
    type Item = Option<(MPath, Vec<HgFileNodeId>)>;
    type Error = Error;

    fn decode(&mut self, src: &mut BytesMut) -> Result<Option<Self::Item>> {
        use self::GetPackv1ParsingState::*;

        let mut state = self.state.clone();
        let (result, state) = loop {
            let new_state = match state {
                Start => {
                    let prefix_len = 2;
                    if src.len() < prefix_len {
                        break (Ok(None), Start);
                    }
                    let len_bytes = src.split_to(prefix_len);
                    let len = Cursor::new(len_bytes.freeze()).get_u16_be();
                    if len == 0 {
                        // Finished parsing the stream
                        // 'Ok' means no error, 'Some' means that no more bytes needed,
                        // None means that getfiles stream has finished
                        return Ok(Some(None));
                    }
                    ParsingFilename(len)
                }
                ParsingFilename(filelen) => {
                    let filelen = filelen as usize;
                    if src.len() < filelen {
                        break (Ok(None), ParsingFilename(filelen as u16));
                    }

                    let filename_bytes = src.split_to(filelen);
                    ParsedFilename(MPath::new(&filename_bytes)?)
                }
                ParsedFilename(file) => {
                    let prefix_len = 4;
                    if src.len() < prefix_len {
                        break (Ok(None), ParsedFilename(file));
                    }

                    let len_bytes = src.split_to(prefix_len);
                    let nodes_count = Cursor::new(len_bytes.freeze()).get_u32_be();

                    ParsingFileNodes(file, nodes_count, vec![])
                }
                ParsingFileNodes(file, file_nodes_count, mut file_nodes) => {
                    if file_nodes_count as usize == file_nodes.len() {
                        return Ok(Some(Some((file, file_nodes))));
                    }
                    let node_size = 20;
                    if src.len() < node_size {
                        break (
                            Ok(None),
                            ParsingFileNodes(file, file_nodes_count, file_nodes),
                        );
                    }

                    let node = src.split_to(node_size);
                    let node = HgFileNodeId::from_bytes(&node)?;
                    file_nodes.push(node);
                    ParsingFileNodes(file, file_nodes_count, file_nodes)
                }
            };

            state = new_state;
        };

        self.state = state;
        result
    }
}

fn extract_remainder_from_bundle2<R>(
    bundle2: Bundle2Stream<R>,
) -> (
    BoxStream<Bundle2Item<'static>, Error>,
    BoxFuture<bundle2::Remainder<R>, Error>,
)
where
    R: AsyncRead + BufRead + 'static + Send,
{
    let (send, recv) = oneshot::channel();
    let mut send = Some(send);

    let bundle2items = bundle2
        .then(move |res_stream_event| {
            match res_stream_event {
                Ok(StreamEvent::Next(bundle2item)) => Ok(Some(bundle2item)),
                Ok(StreamEvent::Done(remainder)) => {
                    let send = send.take().ok_or_else(|| {
                        ErrorKind::Bundle2Invalid("stream remainder was sent twice".into())
                    })?;
                    // Receiving end will deal with failures
                    let _ = send.send(remainder);
                    Ok(None)
                }
                Err(err) => Err(err),
            }
        })
        .filter_map(|val| val)
        .boxify();

    (
        bundle2items,
        recv.context("While extracting bundle2 remainder")
            .from_err()
            .boxify(),
    )
}

#[inline]
fn get_or_none<'a>(map: &'a HashMap<Vec<u8>, Vec<u8>>, key: &'a [u8]) -> &'a [u8] {
    match map.get(key) {
        Some(val) => val,
        None => NONE,
    }
}

#[inline]
fn unimplemented<S, T>(op: S) -> HgCommandRes<T>
where
    S: Into<String>,
    T: Send + 'static,
{
    future::err(ErrorKind::Unimplemented(op.into()).into()).boxify()
}

// Async response from an Hg command
pub type HgCommandRes<T> = BoxFuture<T, Error>;

// Trait representing Mercurial protocol operations, generic across protocols
// Derived from hg/mercurial/wireprotocol.py, functions with the `@wireprotocommand`
// decorator.
//
// XXX Do we need to do all of these? Are some historical/obsolete and can be ignored?
//
// TODO: placeholder types are generally `()`
pub trait HgCommands {
    // @wireprotocommand('between', 'pairs')
    fn between(
        &self,
        _pairs: Vec<(HgChangesetId, HgChangesetId)>,
    ) -> HgCommandRes<Vec<Vec<HgChangesetId>>> {
        unimplemented("between")
    }

    // @wireprotocommand('branchmap')
    fn branchmap(&self) -> HgCommandRes<HashMap<String, HashSet<HgChangesetId>>> {
        // We have no plans to support mercurial branches and hence no plans for branchmap,
        // so just return fake response.
        future::ok(HashMap::new()).boxify()
    }

    // @wireprotocommand('capabilities')
    fn capabilities(&self) -> HgCommandRes<Vec<String>> {
        unimplemented("capabilities")
    }

    // @wireprotocommand('clienttelemetry')
    fn clienttelemetry(&self, _args: HashMap<Vec<u8>, Vec<u8>>) -> HgCommandRes<String> {
        unimplemented("clienttelemetry")
    }

    // @wireprotocommand('getbundle', '*')
    // TODO: make this streaming
    fn getbundle(&self, _args: GetbundleArgs) -> BoxStream<Bytes, Error> {
        once(Err(ErrorKind::Unimplemented("getbundle".into()).into())).boxify()
    }

    // @wireprotocommand('heads')
    fn heads(&self) -> HgCommandRes<HashSet<HgChangesetId>> {
        unimplemented("heads")
    }

    // @wireprotocommand('hello')
    fn hello(&self) -> HgCommandRes<HashMap<String, Vec<String>>> {
        unimplemented("hello")
    }

    // @wireprotocommand('listkeys', 'namespace')
    fn listkeys(&self, _namespace: String) -> HgCommandRes<HashMap<Vec<u8>, Vec<u8>>> {
        unimplemented("listkeys")
    }

    // @wireprotocommand('listkeyspatterns', 'namespace', 'patterns *')
    fn listkeyspatterns(
        &self,
        _namespace: String,
        _patterns: Vec<String>,
    ) -> HgCommandRes<BTreeMap<String, HgChangesetId>> {
        unimplemented("listkeyspatterns")
    }

    // @wireprotocommand('lookup', 'key')
    fn lookup(&self, _key: String) -> HgCommandRes<Bytes> {
        unimplemented("lookup")
    }

    // @wireprotocommand('known', 'nodes *')
    fn known(&self, _nodes: Vec<HgChangesetId>) -> HgCommandRes<Vec<bool>> {
        unimplemented("known")
    }

    // @wireprotocommand('known', 'nodes *')
    fn knownnodes(&self, _nodes: Vec<HgChangesetId>) -> HgCommandRes<Vec<bool>> {
        unimplemented("knownnodes")
    }

    // @wireprotocommand('unbundle', 'heads')
    fn unbundle(
        &self,
        _heads: Vec<String>,
        _stream: BoxStream<Bundle2Item<'static>, Error>,
        _respondlightly: Option<bool>,
        _replaydata: Option<String>,
    ) -> HgCommandRes<Bytes> {
        unimplemented("unbundle")
    }

    // @wireprotocommand('gettreepack', 'rootdir mfnodes basemfnodes directories')
    fn gettreepack(&self, _params: GettreepackArgs) -> BoxStream<Bytes, Error> {
        once(Err(ErrorKind::Unimplemented("gettreepack".into()).into())).boxify()
    }

    // @wireprotocommand('stream_out_shallow', '*')
    fn stream_out_shallow(&self, _tag: Option<String>) -> BoxStream<Bytes, Error> {
        once(Err(
            ErrorKind::Unimplemented("stream_out_shallow".into()).into()
        ))
        .boxify()
    }

    // @wireprotocommand()
    fn getpackv1(
        &self,
        _params: BoxStream<(MPath, Vec<HgFileNodeId>), Error>,
    ) -> BoxStream<Bytes, Error> {
        once(Err(ErrorKind::Unimplemented("getpackv1".into()).into())).boxify()
    }

    fn getpackv2(
        &self,
        _params: BoxStream<(MPath, Vec<HgFileNodeId>), Error>,
    ) -> BoxStream<Bytes, Error> {
        once(Err(ErrorKind::Unimplemented("getpackv2".into()).into())).boxify()
    }

    // @wireprotocommand('getcommitdata', 'nodes *')
    fn getcommitdata(&self, _nodes: Vec<HgChangesetId>) -> BoxStream<Bytes, Error> {
        once(Err(ErrorKind::Unimplemented("getcommitdata".into()).into())).boxify()
    }
}

#[cfg(test)]
mod test {
    use bytes_old::BufMut;
    use bytes_old::BytesMut;
    use futures::future;
    use futures::stream;
    use slog::o;
    use slog::Discard;

    use super::*;

    struct Dummy;
    impl HgCommands for Dummy {
        fn hello(&self) -> HgCommandRes<HashMap<String, Vec<String>>> {
            let mut res = HashMap::new();
            res.insert("capabilities".into(), vec!["something".into()]);

            future::ok(res).boxify()
        }
    }

    fn assert_one<T>(vs: Vec<T>) -> T {
        assert_eq!(vs.len(), 1);
        vs.into_iter().next().unwrap()
    }

    fn hash_ones() -> HgFileNodeId {
        HgFileNodeId::new("1111111111111111111111111111111111111111".parse().unwrap())
    }

    #[test]
    fn hello() {
        let logger = Logger::root(Discard, o!());
        let handler = HgCommandHandler::new(logger, Dummy, None, None);

        let (r, _) = handler.handle(SingleRequest::Hello, BytesStream::new(stream::empty()));
        let r = assert_one(r.wait().collect::<Vec<_>>());
        println!("hello r = {:?}", r);

        let mut res: HashMap<String, Vec<String>> = HashMap::new();
        res.insert("capabilities".into(), vec!["something".into()]);

        match r {
            Ok(SingleResponse::Hello(ref r)) if r == &res => {}
            bad => panic!("Bad result {:?}", bad),
        }
    }

    #[test]
    fn unimpl() {
        let logger = Logger::root(Discard, o!());
        let handler = HgCommandHandler::new(logger, Dummy, None, None);

        let (r, _) = handler.handle(SingleRequest::Heads, BytesStream::new(stream::empty()));
        let r = assert_one(r.wait().collect::<Vec<_>>());
        println!("heads r = {:?}", r);

        match r {
            Err(ref err) => println!("got expected error {:?}", err),
            bad => panic!("Bad result {:?}", bad),
        }
    }

    #[test]
    fn getpackv1decoder() {
        let mut decoder = Getpackv1ArgDecoder::new();
        let mut buf = vec![];
        buf.put_u16_be(0);
        assert_eq!(
            decoder
                .decode(&mut BytesMut::from(buf))
                .expect("unexpected error"),
            Some(None)
        );

        let mut buf = vec![];
        let path = MPath::new("file".as_bytes()).unwrap();
        buf.put_u16_be(4);
        buf.put_slice(&path.to_vec());
        buf.put_u32_be(1);
        buf.put_slice(hash_ones().as_bytes());
        assert_eq!(
            decoder
                .decode(&mut BytesMut::from(buf))
                .expect("unexpected error"),
            Some(Some((path, vec![hash_ones()])))
        );
    }

    #[test]
    fn getpackv1() {
        let input = "\u{0}\u{4}path\u{0}\u{0}\u{0}\u{0}\u{0}\u{0}\u{0}\u{0}";
        let (paramstream, _input) = decode_getpack_arg_stream(
            BytesStream::new(stream::once(Ok(Bytes::from(input)))),
            Getpackv1ArgDecoder::new,
        );
        let res = paramstream.collect().wait().unwrap();
        assert_eq!(res, vec![(MPath::new("path").unwrap(), vec![])]);
    }
}


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

//! Decodes stream of data that is "chunked" in the following format:
//!
//! ```text
//! stream := <chunk>
//! chunk := <numbytes> '\n' <byte>{numbytes} <chunk> | '0\n'
//! ```
//!
//! 0-sized chunk is the indication of end of stream, so a proper stream of data should not
//! contain empty chunks inside.

use std::io;
use std::io::BufRead;
use std::io::Read;

use futures::future::poll_fn;
use futures::Async;
use futures::Future;
use tokio_io::try_nb;
use tokio_io::AsyncRead;

/// Structure that wraps around a `AsyncRead + BufRead` object to provide chunked-based encoding
/// over it. See this module's doc for the description of the format.
///
/// This structure ensures that you don't override the underlying BufRead, so if you read EOF while
/// reading from `Dechunker` and call `Dechunker::into_inner` on it then you will get the original
/// `BufRead` object that contains data following the encoded chunks.
pub struct Dechunker<R> {
    bufread: R,
    state: DechunkerState,
}

enum DechunkerState {
    ParsingInt(Vec<u8>),
    ReadingChunk(usize),
    Done,
}

use self::DechunkerState::*;

impl<R> Dechunker<R>
where
    R: AsyncRead + BufRead,
{
    pub fn new(bufread: R) -> Self {
        Self {
            bufread,
            state: ParsingInt(Vec::new()),
        }
    }

    pub fn check_is_done(self) -> impl Future<Item = (bool, Self), Error = io::Error> {
        let mut this = Some(self);
        poll_fn(move || {
            let is_done = match this {
                None => panic!("called poll after completed"),
                Some(ref mut this) => try_nb!(this.fill_buf()).is_empty(),
            };

            let this = this.take().expect("This was Some few lines above");
            Ok(Async::Ready((is_done, this)))
        })
    }

    pub fn into_inner(self) -> R {
        self.bufread
    }

    /// If the self.state is ParsingInt then we try to parse the following content of the buffer
    /// as:
    ///
    /// ```text
    /// <integer> '\n'
    /// integer := [digit]*
    /// digit := '0' | '1' | ... | '9'
    /// ```
    ///
    /// If the integer parsed was of value `0` then the parsing is done, otherwise we continue by
    /// reading a chunk of the size equal to the provided integer.
    fn advance_parsing_int(&mut self) -> io::Result<()> {
        let chunk_size = match &mut self.state {
            &mut ParsingInt(ref mut buf) => {
                self.bufread.read_until(b'\n', buf)?;

                let mut size = 0;
                for inp in &*buf {
                    match *inp {
                        digit @ b'0'..=b'9' => size = size * 10 + ((digit - b'0') as usize),
                        b'\n' => break,
                        _ => {
                            return Err(io::Error::new(
                                io::ErrorKind::InvalidData,
                                format!("Failed to parse int for Dechunker from '{:?}'", buf),
                            ));
                        }
                    }
                }

                size
            }
            _ => return Ok(()),
        };

        if chunk_size == 0 {
            self.state = Done;
        } else {
            self.state = ReadingChunk(chunk_size);
        }
        Ok(())
    }

    fn consume_chunk(&mut self, amt: usize) {
        if amt > 0 {
            let chunk_size = match &self.state {
                ReadingChunk(chunk_size) => *chunk_size,
                _ => panic!("Trying to consume bytes while internally not reading chunk yet"),
            };

            if amt == chunk_size {
                self.state = ParsingInt(Vec::new());
            } else {
                assert!(
                    chunk_size > amt,
                    "Trying to consume more bytes than the size of chunk"
                );
                self.state = ReadingChunk(chunk_size - amt);
            }
        }
    }
}

impl<R> Read for Dechunker<R>
where
    R: AsyncRead + BufRead,
{
    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
        self.advance_parsing_int()?;
        let chunk_size = match self.state {
            ParsingInt(_) => panic!("expected to get pass parsing int state"),
            ReadingChunk(ref chunk_size) => *chunk_size,
            Done => return Ok(0),
        };

        let buf_size = if buf.len() > chunk_size {
            chunk_size
        } else {
            buf.len()
        };

        let buf_size = self.bufread.read(&mut buf[0..buf_size])?;
        self.consume_chunk(buf_size);
        Ok(buf_size)
    }
}

impl<R> AsyncRead for Dechunker<R> where R: AsyncRead + BufRead {}

impl<R> BufRead for Dechunker<R>
where
    R: AsyncRead + BufRead,
{
    fn fill_buf(&mut self) -> io::Result<&[u8]> {
        self.advance_parsing_int()?;
        let chunk_size = match self.state {
            ParsingInt(_) => panic!("expected to get pass parsing int state"),
            ReadingChunk(ref chunk_size) => *chunk_size,
            Done => return Ok(&[]),
        };

        let buf = self.bufread.fill_buf()?;
        let buf_size = if buf.len() > chunk_size {
            chunk_size
        } else {
            buf.len()
        };
        Ok(&buf[0..buf_size])
    }

    fn consume(&mut self, amt: usize) {
        self.consume_chunk(amt);
        self.bufread.consume(amt);
    }
}

#[cfg(test)]
mod tests {
    use std::io::Cursor;

    use anyhow::ensure;
    use anyhow::Result;
    use quickcheck::quickcheck;
    use quickcheck::Arbitrary;
    use quickcheck::Gen;
    use quickcheck::TestResult;

    use super::*;

    // Vec of non empty Vec<u8> for quickcheck::Arbitrary
    #[derive(Clone, Debug)]
    struct Chunks(Vec<Vec<u8>>);
    impl Arbitrary for Chunks {
        fn arbitrary(g: &mut Gen) -> Self {
            Chunks(
                (0..g.size())
                    .map(|_| Arbitrary::arbitrary(g))
                    .filter_map(|v: Vec<u8>| if v.is_empty() { None } else { Some(v) })
                    .collect(),
            )
        }

        fn shrink(&self) -> Box<dyn Iterator<Item = Self>> {
            Box::new(
                Arbitrary::shrink(&self.0)
                    .map(|v| {
                        v.into_iter()
                            .filter_map(|v| if v.is_empty() { None } else { Some(v) })
                            .collect()
                    })
                    .map(Chunks),
            )
        }
    }

    quickcheck! {
        fn test_bufread_api(chunks: Chunks, remainder: Vec<u8>) -> TestResult {
            let chunks = &chunks;
            let remainder = remainder.as_slice();
            let concat_chunks = concat_chunks(chunks, remainder);

            let dechunker = Dechunker::new(Cursor::new(&concat_chunks));
            match check_bufread_api(
                dechunker,
                chunks,
                remainder,
            ) {
                Ok(()) => TestResult::passed(),
                Err(e) =>TestResult::error(format!("{}", e)),
            }
        }

        fn test_read_api(chunks: Chunks, remainder: Vec<u8>) -> TestResult {
            let chunks = &chunks;
            let remainder = remainder.as_slice();
            let concat_chunks = concat_chunks(chunks, remainder);

            let dechunker = Dechunker::new(Cursor::new(&concat_chunks));

            match check_read_api(
                dechunker,
                chunks,
                remainder,
            ) {
                Ok(()) => TestResult::passed(),
                Err(e) => TestResult::error(format!("{}", e)),
            }
        }
    }

    fn concat_chunks(chunks: &Chunks, remainder: &[u8]) -> Vec<u8> {
        let mut buf = Vec::new();
        for chunk in &chunks.0 {
            buf.extend_from_slice(format!("{}\n", chunk.len()).as_bytes());
            buf.extend_from_slice(chunk.as_slice());
        }
        buf.extend_from_slice(b"0\n");
        buf.extend_from_slice(remainder);
        buf
    }

    fn check_bufread_api<R: AsyncRead + BufRead>(
        mut d: Dechunker<R>,
        chunks: &Chunks,
        remainder: &[u8],
    ) -> Result<()> {
        for chunk in &chunks.0 {
            let buf_len = {
                let buf = d.fill_buf()?;
                ensure!(
                    buf == chunk.as_slice(),
                    "expected {:?} found {:?} in bufread api check",
                    chunk,
                    buf
                );
                buf.len()
            };
            d.consume(buf_len);
        }

        check_remainder(d, remainder)
    }

    fn check_read_api<R: AsyncRead + BufRead>(
        mut d: Dechunker<R>,
        chunks: &Chunks,
        remainder: &[u8],
    ) -> Result<()> {
        let concat_chunks = {
            let mut buf = Vec::new();
            for chunk in &chunks.0 {
                buf.extend_from_slice(chunk.as_slice());
            }
            buf
        };

        let mut buf = Vec::new();
        let buf_len = d.read_to_end(&mut buf)?;
        ensure!(
            buf_len == concat_chunks.len(),
            "expected read_to_end {:?} bytes, but read {:?}",
            concat_chunks.len(),
            buf_len
        );
        ensure!(
            buf == concat_chunks,
            "expected read_to_end {:?}, but read {:?}",
            concat_chunks,
            buf
        );

        check_remainder(d, remainder)
    }

    fn check_remainder<R: AsyncRead + BufRead>(d: Dechunker<R>, remainder: &[u8]) -> Result<()> {
        let (is_done, d) = d.check_is_done().wait()?;
        ensure!(is_done, "expected the dechunker to be done");

        let mut inner = d.into_inner();
        let mut buf = Vec::new();
        let buf_len = inner.read_to_end(&mut buf)?;
        ensure!(
            buf_len == remainder.len(),
            "expected read_to_end {:?} bytes from inner reader, but read {:?}",
            remainder.len(),
            buf_len
        );
        ensure!(
            buf.as_slice() == remainder,
            "expected read_to_end {:?} from inner reader, but read {:?}",
            remainder,
            buf
        );
        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 thiserror::Error;

#[derive(Debug, Error)]
pub enum ErrorKind {
    #[error("Unimplemented operation '{0}'")]
    Unimplemented(String),
    #[error("command parse failed for '{0}'")]
    CommandParse(String),
    #[error("unconsumed data left after parsing '{0}'")]
    UnconsumedData(String),
    #[error("malformed batch with command '{0}'")]
    BatchInvalid(String),
    #[error("malformed bundle2 '{0}'")]
    Bundle2Invalid(String),
    #[error("unknown escape character in batch command '{0}'")]
    BatchEscape(u8),
    #[error("Repo error")]
    RepoError,
    #[error("cannot serve revlog repos")]
    CantServeRevlogRepo,
}


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

use anyhow::Error;
use bytes_old::Bytes;
use failure_ext::FutureErrorContext;
use futures::future::err;
use futures::future::ok;
use futures::future::Either;
use futures::stream;
use futures::sync::oneshot;
use futures::Future;
use futures::Poll;
use futures::Stream;
use futures_ext::BoxFuture;
use futures_ext::BoxStream;
use futures_ext::BytesStream;
use futures_ext::FutureExt;
use futures_ext::StreamExt;
use qps::Qps;
use slog::Logger;
use tokio_io::codec::Decoder;

use crate::commands::HgCommandHandler;
use crate::errors::*;
use crate::HgCommands;
use crate::Request;
use crate::Response;

pub type OutputStream = BoxStream<Bytes, Error>;

pub trait ResponseEncoder {
    fn encode(&self, response: Response) -> OutputStream;
}

pub struct HgProtoHandler {
    outstream: OutputStream,
}

struct HgProtoHandlerInner<H, Dec, Enc> {
    commands_handler: HgCommandHandler<H>,
    reqdec: Dec,
    respenc: Enc,
    wireproto_calls: Arc<Mutex<Vec<String>>>,
}

impl HgProtoHandler {
    pub fn new<In, H, Dec, Enc>(
        logger: Logger,
        input: In,
        commands: H,
        reqdec: Dec,
        respenc: Enc,
        wireproto_calls: Arc<Mutex<Vec<String>>>,
        qps: Option<Arc<Qps>>,
        src_region: Option<String>,
    ) -> Self
    where
        In: Stream<Item = Bytes, Error = io::Error> + Send + 'static,
        H: HgCommands + Send + Sync + 'static,
        Dec: Decoder<Item = Request> + Clone + Send + Sync + 'static,
        Dec::Error: From<io::Error> + Send + 'static,
        Enc: ResponseEncoder + Clone + Send + Sync + 'static,
        Error: From<Dec::Error>,
    {
        let inner = Arc::new(HgProtoHandlerInner {
            commands_handler: HgCommandHandler::new(logger, commands, qps, src_region),
            reqdec,
            respenc,
            wireproto_calls,
        });

        HgProtoHandler {
            outstream: handle(input, inner),
        }
    }
}

impl Stream for HgProtoHandler {
    type Item = Bytes;
    type Error = Error;

    fn poll(&mut self) -> Poll<Option<Self::Item>, Self::Error> {
        self.outstream.poll()
    }
}

fn handle<In, H, Dec, Enc>(
    input: In,
    handler: Arc<HgProtoHandlerInner<H, Dec, Enc>>,
) -> OutputStream
where
    In: Stream<Item = Bytes, Error = io::Error> + Send + 'static,
    H: HgCommands + Send + Sync + 'static,
    Dec: Decoder<Item = Request> + Clone + Send + Sync + 'static,
    Dec::Error: From<io::Error> + Send + 'static,
    Enc: ResponseEncoder + Clone + Send + Sync + 'static,
    Error: From<Dec::Error>,
{
    let input = BytesStream::new(input);

    stream::unfold(Some(ok(input).boxify()), move |input| {
        let input = match input {
            None => return None,
            Some(input) => input,
        };

        let future = input.and_then({
            let handler = handler.clone();
            move |input| {
                if input.is_empty() {
                    return Either::A(ok((None, None)));
                }

                let future = input
                    .into_future_decode(handler.reqdec.clone())
                    .map_err(|(err, _)| -> Error { err.into() })
                    .and_then({
                        let handler = handler.clone();
                        move |(req, remainder)| match req {
                            None => Either::A(if remainder.is_empty() {
                                ok((None, None))
                            } else {
                                let (bytes, _) = remainder.into_parts();
                                err(ErrorKind::UnconsumedData(
                                    String::from_utf8_lossy(bytes.as_ref()).into_owned(),
                                )
                                .into())
                            }),
                            Some(req) => {
                                let (resps, remainder) =
                                    handle_request(req, remainder, handler.clone());
                                Either::B(ok((
                                    Some(
                                        resps
                                            .map(move |resp| handler.respenc.encode(resp))
                                            .flatten()
                                            .boxify(),
                                    ),
                                    Some(remainder),
                                )))
                            }
                        }
                    });
                Either::B(future)
            }
        });

        Some(future)
    })
    .filter_map(|x| x)
    .flatten()
    .boxify()
}

/// Handles a singular request regardless if it contains multiple batched commands or a single one
/// It returns stream of responses that should be send to the client as soon as they are produced
/// and a future containing the remainder of the input that might contain more requests and that
/// will become available once the stream of responses is consumed.
fn handle_request<In, H, Dec, Enc>(
    req: Request,
    input: BytesStream<In>,
    handler: Arc<HgProtoHandlerInner<H, Dec, Enc>>,
) -> (
    BoxStream<Response, Error>,
    BoxFuture<BytesStream<In>, Error>,
)
where
    In: Stream<Item = Bytes, Error = io::Error> + Send + 'static,
    H: HgCommands + Send + Sync + 'static,
    Dec: Decoder<Item = Request> + Clone + Send + Sync + 'static,
    Dec::Error: From<io::Error> + Send + 'static,
    Enc: ResponseEncoder + Clone + Send + Sync + 'static,
    Error: From<Dec::Error>,
{
    req.record_request(&handler.wireproto_calls);
    match req {
        Request::Batch(reqs) => {
            let (send, recv) = oneshot::channel();
            let responses = stream::unfold(
                (reqs.into_iter(), ok(input).boxify(), send),
                move |(mut reqs, input, send)| match reqs.next() {
                    None => {
                        let _ = send.send(input);
                        None
                    }
                    Some(req) => Some(input.map({
                        let handler = handler.clone();
                        move |input| {
                            let (resps, remainder) = handler.commands_handler.handle(req, input);
                            (resps, (reqs, remainder, send))
                        }
                    })),
                },
            )
            .flatten()
            .collect()
            .map(Response::Batch)
            .into_stream();
            (
                responses.boxify(),
                recv.context("While handling batch command")
                    .from_err()
                    .and_then(|input| input)
                    .boxify(),
            )
        }
        Request::Single(req) => {
            let (resps, remainder) = handler.commands_handler.handle(req, input);
            (resps.map(Response::Single).boxify(), remainder)
        }
    }
}


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

//! Mercurial command protocol
//!
//! Mercurial has a set of commands which are implemented across at least two protocols:
//! SSH and HTTP. This module defines enums representing requests and responses for those
//! protocols, and a Tokio Service framework for them via a trait.

use std::collections::BTreeMap;
use std::collections::BTreeSet;
use std::collections::HashMap;
use std::collections::HashSet;
use std::fmt;
use std::fmt::Debug;
use std::sync::Mutex;

use bytes_old::Bytes;
use mercurial_types::HgChangesetId;
use mercurial_types::HgManifestId;
use mononoke_types::MPath;

pub mod batch;
mod commands;
mod dechunker;
mod errors;
mod handler;
pub mod sshproto;

const MAX_NODES_TO_LOG: usize = 5;

#[derive(Debug, Eq, PartialEq)]
pub enum Request {
    Batch(Vec<SingleRequest>),
    Single(SingleRequest),
}

impl Request {
    pub fn record_request(&self, record: &Mutex<Vec<String>>) {
        let mut record = record.lock().expect("lock poisoned");
        match *self {
            Request::Batch(ref batch) => record.extend(batch.iter().map(|s| s.name().into())),
            Request::Single(ref req) => record.push(req.name().into()),
        }
    }
}

#[derive(Debug, Eq, PartialEq)]
pub enum SingleRequest {
    Between {
        pairs: Vec<(HgChangesetId, HgChangesetId)>,
    },
    Branchmap,
    Capabilities,
    ClientTelemetry {
        args: HashMap<Vec<u8>, Vec<u8>>,
    },
    Debugwireargs {
        one: Vec<u8>,
        two: Vec<u8>,
        all_args: HashMap<Vec<u8>, Vec<u8>>,
    },
    Getbundle(GetbundleArgs),
    Heads,
    Hello,
    Listkeys {
        namespace: String,
    },
    ListKeysPatterns {
        namespace: String,
        patterns: Vec<String>,
    },
    Lookup {
        key: String,
    },
    Known {
        nodes: Vec<HgChangesetId>,
    },
    Knownnodes {
        nodes: Vec<HgChangesetId>,
    },
    Unbundle {
        heads: Vec<String>,
    },
    UnbundleReplay {
        heads: Vec<String>,
        replaydata: String,
        respondlightly: bool,
    },
    Gettreepack(GettreepackArgs),
    StreamOutShallow {
        tag: Option<String>,
    },
    GetpackV1,
    GetpackV2,
    GetCommitData {
        nodes: Vec<HgChangesetId>,
    },
}

impl SingleRequest {
    pub fn name(&self) -> &'static str {
        match *self {
            SingleRequest::Between { .. } => "between",
            SingleRequest::Branchmap => "branchmap",
            SingleRequest::Capabilities => "capabilities",
            SingleRequest::ClientTelemetry { .. } => "clienttelemetry",
            SingleRequest::Debugwireargs { .. } => "debugwireargs",
            SingleRequest::Getbundle(_) => "getbundle",
            SingleRequest::Heads => "heads",
            SingleRequest::Hello => "hello",
            SingleRequest::Listkeys { .. } => "listkeys",
            SingleRequest::Lookup { .. } => "lookup",
            SingleRequest::Known { .. } => "known",
            SingleRequest::Knownnodes { .. } => "knownnodes",
            SingleRequest::Unbundle { .. } => "unbundle",
            SingleRequest::UnbundleReplay { .. } => "unbundlereplay",
            SingleRequest::Gettreepack(_) => "gettreepack",
            SingleRequest::StreamOutShallow { .. } => "stream_out_shallow",
            SingleRequest::GetpackV1 => "getpackv1",
            SingleRequest::GetpackV2 => "getpackv2",
            SingleRequest::ListKeysPatterns { .. } => "listkeyspatterns",
            SingleRequest::GetCommitData { .. } => "getcommitdata",
        }
    }
}

/// The arguments that `getbundle` accepts, in a separate struct for
/// the convenience of callers.
#[derive(Eq, PartialEq)]
pub struct GetbundleArgs {
    /// List of space-delimited hex nodes of heads to retrieve
    pub heads: Vec<HgChangesetId>,
    /// List of space-delimited hex nodes that the client has in common with the server
    pub common: Vec<HgChangesetId>,
    /// Comma-delimited set of strings defining client bundle capabilities.
    pub bundlecaps: HashSet<Vec<u8>>,
    /// Comma-delimited list of strings of ``pushkey`` namespaces. For each namespace listed, a bundle2 part will be included with the content of that namespace.
    pub listkeys: Vec<Vec<u8>>,
    /// phases: Boolean indicating whether phases data is requested
    pub phases: bool,
}

impl Debug for GetbundleArgs {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        let bcaps: HashSet<_> = self
            .bundlecaps
            .iter()
            .map(|s| String::from_utf8_lossy(s))
            .collect();
        let listkeys: Vec<_> = self
            .listkeys
            .iter()
            .map(|s| String::from_utf8_lossy(s))
            .collect();
        let heads: Vec<_> = self.heads.iter().take(MAX_NODES_TO_LOG).collect();
        let common: Vec<_> = self.common.iter().take(MAX_NODES_TO_LOG).collect();
        fmt.debug_struct("GetbundleArgs")
            .field("heads_len", &self.heads.len())
            .field("heads", &heads)
            .field("common_len", &self.common.len())
            .field("common", &common)
            .field("bundlecaps", &bcaps)
            .field("listkeys", &listkeys)
            .field("phases", &self.phases)
            .finish()
    }
}

/// The arguments that `gettreepack` accepts, in a separate struct for
/// the convenience of callers.
#[derive(Debug, Clone, Eq, PartialEq)]
pub struct GettreepackArgs {
    /// The directory of the tree to send (including its subdirectories).
    /// If set to `None`, that means "root of the repo".
    pub rootdir: Option<MPath>,
    /// The manifest nodes of the specified root directory to send.
    pub mfnodes: Vec<HgManifestId>,
    /// The manifest nodes of the rootdir that are already on the client.
    pub basemfnodes: BTreeSet<HgManifestId>,
    /// The fullpath (not relative path) of directories underneath
    /// the rootdir that should be sent.
    pub directories: Vec<Bytes>,
    /// The depth from the root that should be sent.
    pub depth: Option<usize>,
}

#[derive(Debug)]
pub enum Response {
    Batch(Vec<SingleResponse>),
    Single(SingleResponse),
}

#[derive(Debug)]
pub enum SingleResponse {
    Between(Vec<Vec<HgChangesetId>>),
    Branchmap(HashMap<String, HashSet<HgChangesetId>>),
    Capabilities(Vec<String>),
    ClientTelemetry(String),
    Debugwireargs(Bytes),
    Getbundle(Bytes),
    Heads(HashSet<HgChangesetId>),
    Hello(HashMap<String, Vec<String>>),
    Listkeys(HashMap<Vec<u8>, Vec<u8>>),
    ListKeysPatterns(BTreeMap<String, HgChangesetId>),
    Lookup(Bytes),
    Known(Vec<bool>),
    Knownnodes(Vec<bool>),
    ReadyForStream,
    Unbundle(Bytes),
    Gettreepack(Bytes),
    StreamOutShallow(Bytes),
    Getpackv1(Bytes),
    Getpackv2(Bytes),
    GetCommitData(Bytes),
}

impl SingleResponse {
    /// Whether this represents a streaming response. Streaming responses don't have any framing.
    pub fn is_stream(&self) -> bool {
        use SingleResponse::*;

        match self {
            &Getbundle(_) | &ReadyForStream | &Unbundle(_) | &Gettreepack(_)
            | &StreamOutShallow(_) | &Getpackv1(_) | &Getpackv2(_) => true,
            _ => false,
        }
    }
}

pub use commands::HgCommandRes;
pub use commands::HgCommands;
pub use errors::ErrorKind;
pub use handler::HgProtoHandler;


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

//! SSH/stdio line-oriented protocol
//!
//! References are <https://www.mercurial-scm.org/wiki/SshCommandProtocol> and
//! <https://www.mercurial-scm.org/wiki/WireProtocol> though they're scant on detail.
//!
//! The encoding is:
//! ```text
//! command := <command> '\n' <key-value>{N}
//! key-value := star | kv
//! star := '*' ' ' <count> kv{count}
//! kv := <name> ' ' <numbytes> '\n' <byte>{numbytes}
//! ```
//!
//! Where `N` is the specific number of `<key-value>` pairs expected by the command.
//!
//! The types of the values are implied by the command rather than explicitly encoded.
//!
//! pair := <hash> '-' <hash>
//! list := (<hash> ' ')* <item>
//!
//! Responses to commands are almost always:
//! ```text
//! <numbytes> '\n'
//! <byte>{numbytes}
//! ```
//!
//! The expections are requests that pass streaming arguments (f.e. unbundle). After such a
//! requests the responder should respond with
//! ```text
//! '0\n'
//! ```
//! to acknowledge readiness for processing the stream. After the stream is fully read the
//! responder should respond with a stream (no acknowledgment is required).
//!
//! Each command has its own encoding of the regular or streaming responses, although by
//! convention the streaming responses are chunked. See `hgproto/dechunker.rs` for the format of
//! chunking.

use anyhow::Error;
use anyhow::Result;
use bytes_old::BytesMut;
use tokio_io::codec::Decoder;

use crate::handler::OutputStream;
use crate::handler::ResponseEncoder;
use crate::Request;
use crate::Response;

pub mod request;
pub mod response;

#[derive(Clone)]
pub struct HgSshCommandEncode;
#[derive(Clone)]
pub struct HgSshCommandDecode;

impl ResponseEncoder for HgSshCommandEncode {
    fn encode(&self, response: Response) -> OutputStream {
        response::encode(response)
    }
}

impl Decoder for HgSshCommandDecode {
    type Item = Request;
    type Error = Error;

    fn decode(&mut self, buf: &mut BytesMut) -> Result<Option<Request>> {
        request::parse_request(buf)
    }
}