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

//! Construct and serialize headers for bundle2 parts.

use std::collections::HashMap;

use anyhow::bail;
use anyhow::Context;
use anyhow::Result;
use bytes::Bytes as BytesNew;
use bytes_old::BufMut;
use bytes_old::Bytes;
use quickcheck::Arbitrary;
use quickcheck::Gen;

use crate::chunk::Chunk;
use crate::errors::ErrorKind;
use crate::utils::BytesExt;

pub type PartId = u32;

#[derive(Clone, Copy, Debug, Eq, PartialEq, Hash)]
pub enum PartHeaderType {
    /// Responsible for sending Changesets and Filelogs during a push. In Mercurial it also sends
    /// flat Manifests, but with Mononoke we support only TreeManifests.
    Changegroup,
    /// When responding for bundle2 this part contains the response for the corresponding
    /// Changegroup.
    ReplyChangegroup,
    /// The bundle2 sender describes his own capabilities to handle the bundle2 that will be send
    /// in response for this one.
    Replycaps,
    /// Contains keys that are being used in this bundle2, f.e. bookmarks
    Listkeys,
    /// Contains wirepacks that are encoded TreeManifests required in the push.
    B2xTreegroup2,
    /// In Mercurial this is used to verify that during the push the heads did not change. In
    /// Mononoke this parameter will be ignored, because it does not provide transaction during
    /// push
    CheckHeads,
    /// Contains list of heads that are present client-side and server-side.
    /// Used in pushrebase to find out what to send to the client.
    B2xCommonHeads,
    /// Contains changegroup for infinitepush commits
    B2xInfinitepush,
    /// Contains bookmarks for infinitepush backups (won't be used in Mononoke,
    /// but they needs to be parsed).
    B2xInfinitepushBookmarks,
    /// Contains mutation info for infinitepush commits
    B2xInfinitepushMutation,
    /// Pushrebase part with changegroup
    B2xRebase,
    /// Pushrebase part that contains packs
    B2xRebasePack,
    /// Pushkey part is used to update different namespaces: phases, bookmarks, etc.
    /// In Mononoke it's used to update bookmarks.
    Pushkey,
    /// Respond to a corresponding pushkey part
    ReplyPushkey,
    /// Contains parameters that can be used by hooks
    Pushvars,
    /// Contains phase heads part
    /// Used in communicating phases between Mononoke and clients
    /// Pushkey / Listkeys are not used to communicate phases
    PhaseHeads,
    // RemoteChangegroup,       // We don't wish to support this functionality
    // CheckBookmarks,          // TODO Do we want to support this?
    // CheckHeads,              // TODO Do we want to support this?
    // CheckUpdatedHeads,       // TODO Do we want to support this?
    // CheckPhases,             // TODO Do we want to support this?
    // Output,                  // TODO Do we want to support this?
    // ErrorAbort,              // TODO Do we want to support this?
    // ErrorPushkey,            // TODO Do we want to support this?
    // ErrorUnsupportedContent, // TODO Do we want to support this?
    // ErrorPushRaced,          // TODO Do we want to support this?
    // Pushkey,                 // TODO Do we want to support this?
    // Bookmarks,               // TODO Do we want to support this?
    // ReplyPushkey,            // TODO Do we want to support this?
    Obsmarkers,
    // ReplyObsmarkers,         // TODO Do we want to support this?
    // HgtagsFnodes,            // TODO Do we want to support this?
}

impl PartHeaderType {
    fn decode(data: &str) -> Result<Self> {
        use self::PartHeaderType::*;
        match data.to_ascii_lowercase().as_str() {
            "changegroup" => Ok(Changegroup),
            "reply:changegroup" => Ok(ReplyChangegroup),
            "replycaps" => Ok(Replycaps),
            "listkeys" => Ok(Listkeys),
            "b2x:treegroup2" => Ok(B2xTreegroup2),
            "b2x:infinitepush" => Ok(B2xInfinitepush),
            "b2x:infinitepushscratchbookmarks" => Ok(B2xInfinitepushBookmarks),
            "b2x:infinitepushmutation" => Ok(B2xInfinitepushMutation),
            "b2x:commonheads" => Ok(B2xCommonHeads),
            "b2x:rebase" => Ok(B2xRebase),
            "b2x:rebasepackpart" => Ok(B2xRebasePack),
            "check:heads" => Ok(CheckHeads),
            "pushkey" => Ok(Pushkey),
            "reply:pushkey" => Ok(ReplyPushkey),
            "pushvars" => Ok(Pushvars),
            "phase-heads" => Ok(PhaseHeads),
            "obsmarkers" => Ok(Obsmarkers),
            bad => bail!("unknown header type {}", bad),
        }
    }

    fn as_str(&self) -> &str {
        use self::PartHeaderType::*;
        match *self {
            Changegroup => "changegroup",
            ReplyChangegroup => "reply:changegroup",
            Replycaps => "replycaps",
            Listkeys => "listkeys",
            B2xTreegroup2 => "b2x:treegroup2",
            B2xCommonHeads => "b2x:commonheads",
            B2xInfinitepush => "b2x:infinitepush",
            B2xInfinitepushBookmarks => "b2x:infinitepushscratchbookmarks",
            B2xInfinitepushMutation => "b2x:infinitepushmutation",
            B2xRebase => "b2x:rebase",
            B2xRebasePack => "b2x:rebasepackpart",
            CheckHeads => "check:heads",
            Pushkey => "pushkey",
            Pushvars => "pushvars",
            ReplyPushkey => "reply:pushkey",
            PhaseHeads => "phase-heads",
            Obsmarkers => "obsmarkers",
        }
    }
}

#[derive(Clone, Debug, Eq, PartialEq)]
pub struct PartHeaderInner {
    pub part_type: PartHeaderType,
    pub mandatory: bool,
    pub part_id: PartId,
    // Part parameter keys are strings and values are arbitrary bytestrings
    // (which can even include null characters).
    pub mparams: HashMap<String, BytesNew>,
    pub aparams: HashMap<String, BytesNew>,
}

/// A bundle2 part header.
#[derive(Clone, Debug, Eq, PartialEq)]
pub struct PartHeader(PartHeaderInner);

impl PartHeader {
    #[inline]
    pub fn part_type(&self) -> &PartHeaderType {
        &self.0.part_type
    }

    #[inline]
    pub fn part_id(&self) -> PartId {
        self.0.part_id
    }

    #[inline]
    pub fn mparams(&self) -> &HashMap<String, BytesNew> {
        &self.0.mparams
    }

    #[inline]
    pub fn aparams(&self) -> &HashMap<String, BytesNew> {
        &self.0.aparams
    }

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

    pub fn into_inner(self) -> PartHeaderInner {
        self.0
    }

    pub fn encode(self) -> Chunk {
        let mut out_buf: Vec<u8> = Vec::new();

        // part type
        let part_type = self.0.part_type.as_str();
        let part_type = if self.0.mandatory {
            part_type.to_ascii_uppercase()
        } else {
            part_type.to_owned()
        };
        let part_type = part_type.as_bytes();
        out_buf.put_u8(part_type.len() as u8);
        out_buf.put_slice(part_type);

        // part id
        out_buf.put_u32_be(self.0.part_id);

        // mandatory/advisory params
        let num_mparams = self.0.mparams.len() as u8;
        let num_aparams = self.0.aparams.len() as u8;

        out_buf.put_u8(num_mparams);
        out_buf.put_u8(num_aparams);

        // sort the params to ensure determinism
        let mut mparams: Vec<(String, Bytes)> = self
            .0
            .mparams
            .into_iter()
            .map(|(k, v)| (k, bytes_ext::copy_from_new(v)))
            .collect();
        mparams.sort();

        let mut aparams: Vec<(String, Bytes)> = self
            .0
            .aparams
            .into_iter()
            .map(|(k, v)| (k, bytes_ext::copy_from_new(v)))
            .collect();
        aparams.sort();

        // param sizes
        for (key, val) in &mparams {
            out_buf.put_u8(key.len() as u8);
            out_buf.put_u8(val.len() as u8);
        }
        for (key, val) in &aparams {
            out_buf.put_u8(key.len() as u8);
            out_buf.put_u8(val.len() as u8);
        }

        // the actual params themselves
        for (key, val) in &mparams {
            out_buf.put_slice(key.as_bytes());
            out_buf.put_slice(val);
        }
        for (key, val) in &aparams {
            out_buf.put_slice(key.as_bytes());
            out_buf.put_slice(val);
        }

        // This can only fail because the chunk is too big, but the restrictions
        // on sizes above put a cap on the size of the chunk that's much
        // smaller.
        Chunk::new(out_buf).expect("chunk cannot be too big")
    }
}

pub fn decode(mut header_bytes: Bytes) -> Result<PartHeader> {
    // Header internals:
    // ---
    // type_size: u8
    // part_type: str (type_size bytes)
    // part_id: u32
    // number of mandatory params: u8
    // number of advisory params: u8
    // param_sizes: (key: u8, val: u8) * (number of mandatory + advisory params)
    // parameters: key, val are both strings of lengths corresponding to index in param_sizes
    // ---
    // This function assumes that the full header is available.
    let type_size = header_bytes.drain_u8() as usize;
    let part_type_encoded = header_bytes
        .drain_str(type_size)
        .with_context(|| ErrorKind::Bundle2Decode("invalid part type".into()))?;
    let part_type = PartHeaderType::decode(&part_type_encoded)?;

    let mandatory = part_type_encoded.chars().any(|c| c.is_ascii_uppercase());

    let part_id = header_bytes.drain_u32();

    let nmparams = header_bytes.drain_u8() as usize;
    let naparams = header_bytes.drain_u8() as usize;

    let mut param_sizes = Vec::with_capacity(nmparams + naparams);
    let mut header = PartHeaderBuilder::with_capacity(part_type, mandatory, nmparams, naparams)
        .with_context(|| ErrorKind::Bundle2Decode("invalid part header".into()))?;

    for _ in 0..(nmparams + naparams) {
        // TODO: ensure none of the params is empty
        param_sizes.push((
            header_bytes.drain_u8() as usize,
            header_bytes.drain_u8() as usize,
        ));
    }

    for (cur, (ksize, vsize)) in param_sizes.iter().cloned().enumerate().take(nmparams) {
        let (key, val) =
            decode_header_param(&mut header_bytes, ksize, vsize).with_context(|| {
                let err_msg = format!(
                    "part '{:?}' (id {}): invalid param {}",
                    header.part_type(),
                    part_id,
                    cur
                );
                ErrorKind::Bundle2Decode(err_msg)
            })?;
        header
            .add_mparam(key, val)
            .with_context(|| ErrorKind::Bundle2Decode("invalid part header".into()))?;
    }

    for (cur, (ksize, vsize)) in param_sizes
        .iter()
        .cloned()
        .enumerate()
        .skip(nmparams)
        .take(naparams)
    {
        let (key, val) =
            decode_header_param(&mut header_bytes, ksize, vsize).with_context(|| {
                let err_msg = format!(
                    "part '{:?}' (id {}): invalid param {}",
                    header.part_type(),
                    part_id,
                    cur
                );
                ErrorKind::Bundle2Decode(err_msg)
            })?;
        header
            .add_aparam(key, val)
            .with_context(|| ErrorKind::Bundle2Decode("invalid part header".into()))?;
    }

    Ok(header.build(part_id))
}

fn decode_header_param(buf: &mut Bytes, ksize: usize, vsize: usize) -> Result<(String, BytesNew)> {
    let key = buf.drain_str(ksize).context("invalid key")?;
    let val = buf.split_to(vsize);
    let val = bytes_ext::copy_from_old(val);
    Ok((key, val))
}

/// Builder for a bundle2 part header.
#[derive(Debug, Eq, PartialEq)]
pub struct PartHeaderBuilder {
    part_type: PartHeaderType,
    mandatory: bool,
    mparams: HashMap<String, BytesNew>,
    aparams: HashMap<String, BytesNew>,
}

impl PartHeaderBuilder {
    pub fn new(part_type: PartHeaderType, mandatory: bool) -> Result<Self> {
        Self::with_capacity(part_type, mandatory, 0, 0)
    }

    pub fn with_capacity(
        part_type: PartHeaderType,
        mandatory: bool,
        mparam_capacity: usize,
        aparam_capacity: usize,
    ) -> Result<Self> {
        Ok(PartHeaderBuilder {
            part_type,
            mandatory,
            mparams: HashMap::with_capacity(mparam_capacity),
            aparams: HashMap::with_capacity(aparam_capacity),
        })
    }

    pub fn add_mparam<S, B>(&mut self, key: S, val: B) -> Result<&mut Self>
    where
        S: Into<String>,
        B: Into<BytesNew>,
    {
        let key = key.into();
        let val = val.into();
        self.check_param(&key, &val)?;
        if self.mparams.len() >= u8::max_value() as usize {
            bail!(
                "number of mandatory params exceeds maximum {}",
                u8::max_value()
            );
        }
        self.mparams.insert(key, val);
        Ok(self)
    }

    pub fn add_aparam<S, B>(&mut self, key: S, val: B) -> Result<&mut Self>
    where
        S: Into<String>,
        B: Into<BytesNew>,
    {
        let key = key.into();
        let val = val.into();
        self.check_param(&key, &val)?;
        if self.aparams.len() >= u8::max_value() as usize {
            bail!(
                "number of advisory params exceeds maximum {}",
                u8::max_value()
            );
        }
        self.aparams.insert(key, val);
        Ok(self)
    }

    pub fn part_type(&self) -> &PartHeaderType {
        &self.part_type
    }

    /// Turn this `PartHeaderBuilder` into a `PartHeader`.
    ///
    /// We only accept part_id at this point because in the serialization use
    /// case, a part id is only assigned when the header is finalized.
    pub fn build(self, part_id: PartId) -> PartHeader {
        PartHeader(PartHeaderInner {
            part_type: self.part_type,
            mandatory: self.mandatory,
            part_id,
            mparams: self.mparams,
            aparams: self.aparams,
        })
    }

    fn check_param(&self, key: &str, val: &[u8]) -> Result<()> {
        if self.mparams.contains_key(key) || self.aparams.contains_key(key) {
            bail!(
                "part '{:?}': key '{}' already present in this part",
                self.part_type,
                key
            );
        }
        if key.is_empty() {
            bail!("part '{:?}': empty key", self.part_type);
        }
        if key.len() > u8::max_value() as usize {
            bail!(
                "part '{:?}': key '{}' exceeds max length {}",
                self.part_type,
                key,
                u8::max_value()
            );
        }
        if val.len() > u8::max_value() as usize {
            bail!(
                "part '{:?}': value for key '{}' exceeds max length {}",
                self.part_type,
                key,
                u8::max_value()
            );
        }
        Ok(())
    }
}

impl Arbitrary for PartHeaderType {
    fn arbitrary(g: &mut Gen) -> Self {
        use self::PartHeaderType::*;
        g.choose(&[
            Changegroup,
            ReplyChangegroup,
            Replycaps,
            Listkeys,
            B2xTreegroup2,
            CheckHeads,
        ])
        .expect("empty choice provided")
        .clone()
    }
}

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

    use super::*;
    use crate::quickcheck_types::QCBytes;

    const MAX_LEN: usize = ::std::u8::MAX as usize;

    #[test]
    fn test_check_params() {
        let mut header = PartHeaderBuilder::new(PartHeaderType::Changegroup, false).unwrap();

        assert_param(&mut header, "", &b"val"[..], false);
        assert_param(&mut header, "key", &b"val"[..], true);
        // if a key was already stored, reject it the second time
        assert_param(&mut header, "key", &b"val"[..], false);

        assert_param(&mut header, "k".repeat(MAX_LEN), &b"val"[..], true);
        assert_param(&mut header, "k".repeat(MAX_LEN + 1), &b"val"[..], false);
        assert_param(&mut header, "key2", "v".repeat(MAX_LEN), true);
        assert_param(&mut header, "key3", "v".repeat(MAX_LEN + 1), false);
    }

    #[test]
    fn test_roundtrip() {
        quickcheck(
            roundtrip
                as fn(
                    PartHeaderType,
                    bool,
                    PartId,
                    HashMap<String, QCBytes>,
                    HashMap<String, QCBytes>,
                ) -> TestResult,
        );
    }

    fn roundtrip(
        part_type: PartHeaderType,
        mandatory: bool,
        part_id: PartId,
        mparams: HashMap<String, QCBytes>,
        aparams: HashMap<String, QCBytes>,
    ) -> TestResult {
        match roundtrip_inner(part_type, mandatory, part_id, mparams, aparams) {
            Ok(test_result) => test_result,
            Err(_err) => TestResult::discard(),
        }
    }

    /// Test that roundtrip encoding -> decoding works.
    ///
    /// For convenience, errors here are treated as skipped tests. Panics are
    /// test failures.
    fn roundtrip_inner(
        part_type: PartHeaderType,
        mandatory: bool,
        part_id: PartId,
        mparams: HashMap<String, QCBytes>,
        aparams: HashMap<String, QCBytes>,
    ) -> Result<TestResult> {
        let mut builder = PartHeaderBuilder::new(part_type, mandatory)?;
        for (k, v) in mparams {
            builder.add_mparam(k, v)?;
        }
        for (k, v) in aparams {
            builder.add_aparam(k, v)?;
        }

        let header = builder.build(part_id);

        let header_chunk = header.clone().encode();
        let header_bytes = header_chunk.into_bytes().unwrap();
        let decoded_header = decode(header_bytes).unwrap();

        assert_eq!(header, decoded_header);

        Ok(TestResult::passed())
    }

    fn assert_param<S: Into<String>, B: Into<BytesNew>>(
        header: &mut PartHeaderBuilder,
        key: S,
        val: B,
        valid: bool,
    ) {
        let res = header.add_aparam(key, val);
        if valid {
            res.unwrap();
        } else {
            res.unwrap_err();
        }
    }
}


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

//! Type definitions for inner streams.

use std::collections::HashMap;
use std::collections::HashSet;
use std::io::BufRead;
use std::str;

use anyhow::bail;
use anyhow::ensure;
use anyhow::Error;
use anyhow::Result;
use bytes_old::Bytes;
use bytes_old::BytesMut;
use futures_ext::BoxFuture;
use futures_ext::FutureExt;
use futures_ext::StreamExt;
use futures_ext::StreamLayeredExt;
use futures_old::future;
use futures_old::Future;
use futures_old::Stream;
use lazy_static::lazy_static;
use maplit::hashset;
use slog::o;
use slog::warn;
use slog::Logger;
use tokio_io::codec::Decoder;
use tokio_io::AsyncRead;

use crate::capabilities;
use crate::changegroup;
use crate::errors::ErrorKind;
use crate::infinitepush;
use crate::part_header::PartHeader;
use crate::part_header::PartHeaderType;
use crate::part_outer::OuterFrame;
use crate::part_outer::OuterStream;
use crate::pushrebase;
use crate::wirepack;
use crate::OldBundle2Item;

// --- Part parameters

lazy_static! {
    static ref KNOWN_PARAMS: HashMap<PartHeaderType, HashSet<&'static str>> = {
        let mut m: HashMap<PartHeaderType, HashSet<&'static str>> = HashMap::new();
        m.insert(PartHeaderType::Changegroup, hashset!{"version", "nbchanges", "treemanifest"});
        // TODO(stash): currently ignore all the parameters. Later we'll
        // support 'bookmark' parameter, and maybe 'create' and 'force' (although 'force' will
        // probably) be renamed T26385545. 'bookprevnode' and 'pushbackbookmarks' will be
        // removed T26384190.
        m.insert(PartHeaderType::B2xInfinitepush, hashset!{
            "pushbackbookmarks", "cgversion", "bookmark", "bookprevnode", "create", "force"});
        m.insert(PartHeaderType::B2xInfinitepushBookmarks, hashset!{});
        m.insert(PartHeaderType::B2xInfinitepushMutation, hashset!{});
        m.insert(PartHeaderType::B2xCommonHeads, hashset!{});
        m.insert(PartHeaderType::B2xRebase, hashset!{"onto", "newhead", "cgversion", "obsmarkerversions"});
        m.insert(PartHeaderType::B2xRebasePack, hashset!{"version", "cache", "category"});
        m.insert(PartHeaderType::B2xTreegroup2, hashset!{"version", "cache", "category"});
        m.insert(PartHeaderType::Replycaps, hashset!{});
        m.insert(PartHeaderType::Pushkey, hashset!{ "namespace", "key", "old", "new" });
        m.insert(PartHeaderType::Pushvars, hashset!{});
        m
    };
}

pub fn validate_header(header: PartHeader) -> Result<Option<PartHeader>> {
    match KNOWN_PARAMS.get(header.part_type()) {
        Some(known_params) => {
            // Make sure all the mandatory params are recognized.
            let unknown_params: Vec<_> = header
                .mparams()
                .keys()
                .filter(|param| !known_params.contains(param.as_str()))
                .cloned()
                .collect();
            if !unknown_params.is_empty() {
                bail!(ErrorKind::BundleUnknownPartParams(
                    *header.part_type(),
                    unknown_params,
                ));
            }
            Ok(Some(header))
        }
        None => {
            if header.mandatory() {
                bail!(ErrorKind::BundleUnknownPart(header));
            }
            Ok(None)
        }
    }
}

pub fn get_cg_version(header: PartHeader, field: &str) -> Result<changegroup::unpacker::CgVersion> {
    let version = header
        .mparams()
        .get(field)
        .or_else(|| header.aparams().get(field));
    let err = ErrorKind::CgDecode(format!(
        "No changegroup version in Part Header in field {}",
        field
    ))
    .into();

    version
        .ok_or(err)
        .and_then(|version_bytes| {
            str::from_utf8(version_bytes)
                .map_err(|e| ErrorKind::CgDecode(format!("{:?}", e)).into())
        })
        .and_then(|version_str| {
            version_str
                .parse::<changegroup::unpacker::CgVersion>()
                .map_err(|e| ErrorKind::CgDecode(format!("{:?}", e)).into())
        })
}

pub fn get_cg_unpacker(
    logger: Logger,
    header: PartHeader,
    field: &str,
) -> changegroup::unpacker::CgUnpacker {
    // TODO(anastasiyaz): T34812941 return Result here, no default packer (version should be specified)
    let _logger = logger.clone();
    match get_cg_version(header, field) {
        Ok(version) => changegroup::unpacker::CgUnpacker::new(logger, version),
        Err(e) => {
            // ChangeGroup2 by default
            warn!(_logger.clone(), "{:?}", e);
            let default_version = changegroup::unpacker::CgVersion::Cg2Version;
            changegroup::unpacker::CgUnpacker::new(_logger, default_version)
        }
    }
}

/// Convert an OuterStream into an InnerStream using the part header.
pub(crate) fn inner_stream<R: AsyncRead + BufRead + 'static + Send>(
    logger: Logger,
    header: PartHeader,
    stream: OuterStream<R>,
) -> (OldBundle2Item, BoxFuture<OuterStream<R>, Error>) {
    let wrapped_stream = stream
        .take_while(|frame| future::ok(frame.is_payload()))
        .map(OuterFrame::get_payload as fn(OuterFrame) -> Bytes);
    let (wrapped_stream, remainder) = wrapped_stream.return_remainder();

    let bundle2item = match *header.part_type() {
        PartHeaderType::Changegroup => {
            let cg2_stream = wrapped_stream.decode(get_cg_unpacker(
                logger.new(o!("stream" => "changegroup")),
                header.clone(),
                "version",
            ));
            OldBundle2Item::Changegroup(header, cg2_stream.boxify())
        }
        PartHeaderType::B2xCommonHeads => {
            let heads_stream = wrapped_stream.decode(pushrebase::CommonHeadsUnpacker::new());
            OldBundle2Item::B2xCommonHeads(header, heads_stream.boxify())
        }
        PartHeaderType::B2xInfinitepush => {
            let cg2_stream = wrapped_stream.decode(get_cg_unpacker(
                logger.new(o!("stream" => "b2xinfinitepush")),
                header.clone(),
                "cgversion",
            ));
            OldBundle2Item::B2xInfinitepush(header, cg2_stream.boxify())
        }
        PartHeaderType::B2xInfinitepushBookmarks => {
            let bookmarks_stream =
                wrapped_stream.decode(infinitepush::InfinitepushBookmarksUnpacker::new());
            OldBundle2Item::B2xInfinitepushBookmarks(header, bookmarks_stream.boxify())
        }
        PartHeaderType::B2xInfinitepushMutation => {
            let mutation_stream =
                wrapped_stream.decode(infinitepush::InfinitepushMutationUnpacker::new());
            OldBundle2Item::B2xInfinitepushMutation(header, mutation_stream.boxify())
        }
        PartHeaderType::B2xTreegroup2 => {
            let wirepack_stream = wrapped_stream.decode(wirepack::unpacker::new(
                logger.new(o!("stream" => "wirepack")),
                // Mercurial only knows how to send trees at the moment.
                // TODO: add support for file wirepacks once that's a thing
                wirepack::Kind::Tree,
            ));
            OldBundle2Item::B2xTreegroup2(header, wirepack_stream.boxify())
        }
        PartHeaderType::Replycaps => {
            let caps = wrapped_stream
                .decode(capabilities::CapabilitiesUnpacker)
                .collect()
                .and_then(|caps| {
                    ensure!(caps.len() == 1, "Unexpected Replycaps payload: {:?}", caps);
                    Ok(caps.into_iter().next().unwrap())
                });
            OldBundle2Item::Replycaps(header, caps.boxify())
        }
        PartHeaderType::B2xRebasePack => {
            let wirepack_stream = wrapped_stream.decode(wirepack::unpacker::new(
                logger.new(o!("stream" => "wirepack")),
                // Mercurial only knows how to send trees at the moment.
                // TODO: add support for file wirepacks once that's a thing
                wirepack::Kind::Tree,
            ));
            OldBundle2Item::B2xRebasePack(header, wirepack_stream.boxify())
        }
        PartHeaderType::B2xRebase => {
            let cg2_stream = wrapped_stream.decode(get_cg_unpacker(
                logger.new(o!("stream" => "bx2rebase")),
                header.clone(),
                "cgversion",
            ));
            OldBundle2Item::B2xRebase(header, cg2_stream.boxify())
        }
        PartHeaderType::Pushkey => {
            // Pushkey part has an empty part payload, but we still need to "parse" it
            // Otherwise polling remainder stream may fail.
            let empty = wrapped_stream.decode(EmptyUnpacker).for_each(|_| Ok(()));
            OldBundle2Item::Pushkey(header, Box::new(empty))
        }
        PartHeaderType::Pushvars => {
            // Pushvars part has an empty part payload, but we still need to "parse" it
            // Otherwise polling remainder stream may fail.
            let empty = wrapped_stream.decode(EmptyUnpacker).for_each(|_| Ok(()));
            OldBundle2Item::Pushvars(header, Box::new(empty))
        }
        _ => panic!("TODO: make this an error"),
    };

    (
        bundle2item,
        remainder
            .map(|s| s.into_inner().into_inner())
            .from_err()
            .boxify(),
    )
}

// Decoder for an empty part (for example, pushkey)
pub struct EmptyUnpacker;

impl Decoder for EmptyUnpacker {
    type Item = ();
    type Error = Error;

    fn decode(&mut self, _buf: &mut BytesMut) -> Result<Option<Self::Item>> {
        Ok(None)
    }
}

#[cfg(test)]
mod test {
    use crate::changegroup::unpacker::CgVersion;
    use crate::part_header::PartHeaderBuilder;
    use crate::part_header::PartHeaderType;
    use crate::part_inner::*;

    #[test]
    fn test_cg_unpacker_v3() {
        let mut header_builder =
            PartHeaderBuilder::new(PartHeaderType::Changegroup, false).unwrap();
        header_builder.add_aparam("version", "03").unwrap();
        let header = header_builder.build(1);

        assert_eq!(
            get_cg_version(header, "version").unwrap(),
            CgVersion::Cg3Version
        );
    }

    #[test]
    fn test_cg_unpacker_v2() {
        let mut header_builder =
            PartHeaderBuilder::new(PartHeaderType::Changegroup, false).unwrap();
        header_builder.add_aparam("version", "02").unwrap();
        let header = header_builder.build(1);

        assert_eq!(
            get_cg_version(header, "version").unwrap(),
            CgVersion::Cg2Version
        );
    }

    #[test]
    fn test_cg_unpacker_default_v2() {
        let header_builder = PartHeaderBuilder::new(PartHeaderType::Changegroup, false).unwrap();
        let h = header_builder.build(1);

        assert!(get_cg_version(h, "version").is_err());
    }
}


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

//! Codec to parse the bits that are the same for every bundle2, except for
//! stream-level parameters (see `stream_start` for those). This parses bundle2
//! part headers and puts together chunks for inner codecs to parse.

use std::io::BufRead;
use std::mem;

use anyhow::Error;
use anyhow::Result;
use async_compression::Decompressor;
use bytes_old::Bytes;
use bytes_old::BytesMut;
use futures_ext::io::Either;
use futures_ext::io::Either::A as UncompressedRead;
use futures_ext::io::Either::B as CompressedRead;
use limited_async_read::LimitedAsyncRead;
use slog::debug;
use slog::o;
use slog::Logger;
use tokio_codec::Decoder;
use tokio_codec::Framed;
use tokio_codec::FramedParts;
use tokio_io::AsyncRead;

use crate::errors::ErrorKind;
use crate::part_header;
use crate::part_header::PartHeader;
use crate::part_header::PartHeaderType;
use crate::part_header::PartId;
use crate::part_inner::validate_header;
use crate::types::StreamHeader;
use crate::utils::get_decompressor_type;
use crate::utils::BytesExt;

pub fn outer_stream<R: AsyncRead + BufRead + Send + 'static>(
    logger: Logger,
    stream_header: &StreamHeader,
    r: R,
) -> Result<OuterStream<R>> {
    let decompressor_type = get_decompressor_type(
        stream_header
            .m_stream_params
            .get("compression")
            .map(String::as_ref),
    )?;

    Ok(Framed::from_parts(FramedParts::new(
        match decompressor_type {
            None => UncompressedRead(r),
            Some(decompressor_type) => CompressedRead(LimitedAsyncRead::new(Decompressor::new(
                r,
                decompressor_type,
            ))),
        },
        OuterDecoder::new(logger.new(o!("stream" => "outer"))),
    )))
}

pub type OuterStream<R> =
    Framed<Either<R, LimitedAsyncRead<Decompressor<'static, R>>>, OuterDecoder>;

#[derive(Debug)]
enum OuterState {
    Header,
    Payload {
        part_type: PartHeaderType,
        part_id: PartId,
    },
    DiscardPayload,
    StreamEnd,
    Invalid,
}

impl OuterState {
    pub fn take(&mut self) -> Self {
        mem::replace(self, OuterState::Invalid)
    }

    pub fn payload_frame(&self, data: BytesMut) -> OuterFrame {
        match *self {
            OuterState::Payload {
                ref part_type,
                ref part_id,
            } => OuterFrame::Payload {
                part_type: part_type.clone(),
                part_id: *part_id,
                payload: data.freeze(),
            },
            OuterState::DiscardPayload => OuterFrame::Discard,
            _ => panic!("payload_frame called for state without payloads"),
        }
    }

    pub fn part_end_frame(self) -> OuterFrame {
        match self {
            OuterState::Payload { part_type, part_id } => {
                OuterFrame::PartEnd { part_type, part_id }
            }
            OuterState::DiscardPayload => OuterFrame::Discard,
            _ => panic!("part_end_frame called for state without payloads"),
        }
    }
}

#[derive(Debug)]
pub struct OuterDecoder {
    logger: Logger,
    state: OuterState,
}

impl Decoder for OuterDecoder {
    type Item = OuterFrame;
    type Error = Error;

    fn decode(&mut self, buf: &mut BytesMut) -> Result<Option<Self::Item>> {
        let (ret, next_state) = Self::decode_next(&self.logger, buf, self.state.take());
        self.state = next_state;
        ret
    }
}

impl OuterDecoder {
    pub fn new(logger: Logger) -> Self {
        OuterDecoder {
            logger,
            state: OuterState::Header,
        }
    }

    fn decode_next(
        logger: &Logger,
        buf: &mut BytesMut,
        mut state: OuterState,
    ) -> (Result<Option<OuterFrame>>, OuterState) {
        // TODO: the only state valid when the stream terminates is
        // StreamEnd. Communicate that to callers.
        match state.take() {
            OuterState::Header => {
                // The header is structured as:
                // ---
                // header_len: u32
                // header: header_len bytes
                // ---
                // See part_header::decode for information about the internal structure.
                if buf.len() < 4 {
                    return (Ok(None), OuterState::Header);
                }

                let header_len = buf.peek_u32() as usize;
                if buf.len() < 4 + header_len {
                    return (Ok(None), OuterState::Header);
                }

                let _ = buf.split_to(4);
                if header_len == 0 {
                    // A zero-length header indicates that the stream has ended.
                    return (Ok(Some(OuterFrame::StreamEnd)), OuterState::StreamEnd);
                }

                let part_header = Self::decode_header(logger, buf.split_to(header_len).freeze());
                if let Err(e) = part_header {
                    let next = match e.downcast::<ErrorKind>() {
                        Ok(ek) => {
                            if ek.is_app_error() {
                                (Err(ek.into()), OuterState::DiscardPayload)
                            } else {
                                (Err(ek.into()), OuterState::Invalid)
                            }
                        }
                        Err(e) => (Err(e), OuterState::Invalid),
                    };
                    return next;
                };
                let part_header = part_header.unwrap();
                // If no part header was returned, this part wasn't
                // recognized. Throw it away.
                match part_header {
                    None => (Ok(Some(OuterFrame::Discard)), OuterState::DiscardPayload),
                    Some(header) => {
                        let part_type = *header.part_type();
                        let part_id = header.part_id();
                        (
                            Ok(Some(OuterFrame::Header(header))),
                            OuterState::Payload { part_type, part_id },
                        )
                    }
                }
            }

            cur_state @ OuterState::Payload { .. } | cur_state @ OuterState::DiscardPayload => {
                let (payload, next_state) = Self::decode_payload(buf, cur_state);
                (payload, next_state)
            }

            OuterState::StreamEnd => (Ok(Some(OuterFrame::StreamEnd)), OuterState::StreamEnd),

            OuterState::Invalid => (
                Err(ErrorKind::Bundle2Decode("byte stream corrupt".into()).into()),
                OuterState::Invalid,
            ),
        }
    }

    fn decode_header(logger: &Logger, header_bytes: Bytes) -> Result<Option<PartHeader>> {
        let header = part_header::decode(header_bytes)?;
        debug!(logger, "Decoded header: {:?}", header);
        match validate_header(header)? {
            Some(header) => Ok(Some(header)),
            None => {
                // The part couldn't be recognized but wasn't important anyway.
                // Throw it away (the state machine will throw away any associated
                // chunks it finds).
                Ok(None)
            }
        }
    }

    fn decode_payload(
        buf: &mut BytesMut,
        state: OuterState,
    ) -> (Result<Option<OuterFrame>>, OuterState) {
        if buf.len() < 4 {
            return (Ok(None), state);
        }

        // Payloads are in the format:
        // ---
        // total_len: i32
        // payload: Vec<u8>, total_len bytes
        // ---
        // A payload is guaranteed to be < 2**31 bytes, so buffer up
        // until the whole payload is available.
        //
        // TODO: -1 means this part has been interrupted. Handle that
        // case.

        let total_len = buf.peek_i32();
        if total_len == 0 {
            let _ = buf.drain_i32();
            // A zero-size chunk indicates that this part has
            // ended. More parts might be coming up, so go back to the
            // header state.
            (Ok(Some(state.part_end_frame())), OuterState::Header)
        } else {
            let payload = Self::decode_payload_chunk(buf, &state, total_len as usize);
            (Ok(payload), state)
        }
    }

    fn decode_payload_chunk(
        buf: &mut BytesMut,
        state: &OuterState,
        total_len: usize,
    ) -> Option<OuterFrame> {
        // + 4 bytes for the header
        if buf.len() < total_len + 4 {
            return None;
        }

        let _ = buf.drain_i32();
        let chunk = buf.split_to(total_len);

        Some(state.payload_frame(chunk))
    }
}

#[derive(Debug, Eq, PartialEq)]
pub enum OuterFrame {
    Header(PartHeader),
    Payload {
        part_type: PartHeaderType,
        part_id: PartId,
        payload: Bytes,
    },
    PartEnd {
        part_type: PartHeaderType,
        part_id: PartId,
    },
    Discard,
    StreamEnd,
}

impl OuterFrame {
    pub fn is_payload(&self) -> bool {
        match self {
            &OuterFrame::Payload { .. } => true,
            _ => false,
        }
    }

    pub fn get_payload(self) -> Bytes {
        match self {
            OuterFrame::Payload { payload, .. } => payload,
            _ => panic!("get_payload called on an OuterFrame without a payload!"),
        }
    }
}


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

//! Quickcheck support for a few types that don't have support upstream,
//! and for a few other test types.

use std::iter;
#[cfg(test)]
use std::vec::IntoIter;

#[cfg(test)]
use anyhow::Error;
#[cfg(test)]
use anyhow::Result;
use bytes::Bytes;
#[cfg(test)]
use futures_old::stream;
use mercurial_types::Delta;
use mercurial_types::HgNodeHash;
use mercurial_types::MPath;
use mercurial_types::RevFlags;
use quickcheck::empty_shrinker;
use quickcheck::Arbitrary;
use quickcheck::Gen;

use crate::changegroup;

#[derive(Clone, Debug)]
pub struct QCBytes(Bytes);

impl From<QCBytes> for Bytes {
    fn from(qcbytes: QCBytes) -> Bytes {
        qcbytes.0
    }
}

impl Arbitrary for QCBytes {
    fn arbitrary(g: &mut Gen) -> Self {
        // Just use the Vec<u8> quickcheck underneath.
        let v: Vec<u8> = Vec::arbitrary(g);
        QCBytes(v.into())
    }

    fn shrink(&self) -> Box<dyn Iterator<Item = Self>> {
        Box::new(self.0.to_vec().shrink().map(|v| QCBytes(v.into())))
    }
}

#[derive(Clone, Debug)]
pub struct CgPartSequence {
    // Storing the ends in here bypasses a number of lifetime issues.
    changesets: Vec<changegroup::Part>,
    changesets_end: changegroup::Part,
    manifests: Vec<changegroup::Part>,
    manifests_end: changegroup::Part,
    treemanifest_end: Option<changegroup::Part>,
    filelogs: Vec<(Vec<changegroup::Part>, changegroup::Part)>,
    end: changegroup::Part,

    version: changegroup::unpacker::CgVersion,
}

impl CgPartSequence {
    /// Combine all the changesets, manifests and filelogs into a single iterator.
    pub fn as_iter<'a>(&'a self) -> Box<dyn Iterator<Item = &'a changegroup::Part> + 'a> {
        // Trying to describe the type here is madness. Just box it.
        Box::new(
            self.changesets
                .iter()
                .chain(iter::once(&self.changesets_end))
                .chain(self.manifests.iter())
                .chain(iter::once(&self.manifests_end))
                .chain(self.treemanifest_end.iter())
                .chain(
                    self.filelogs
                        .iter()
                        .filter(|&(parts, _)| {
                            // If there are no filelog parts, it isn't valid to return a
                            // SectionEnd since that won't be referring to anything. So
                            // just skip the whole filelog.
                            !parts.is_empty()
                        })
                        .flat_map(|(parts, end)| parts.iter().chain(iter::once(end))),
                )
                .chain(iter::once(&self.end)),
        )
    }

    /// Combine all the changesets, manifests and filelogs into a single stream.
    ///
    /// This returns a clone of everything because streams can't really return
    /// references at the moment.
    #[cfg(test)]
    pub fn to_stream(&self) -> stream::IterOk<IntoIter<Result<changegroup::Part>>, Error> {
        let part_results: Vec<_> = self.as_iter().cloned().map(Ok).collect();
        stream::iter_ok(part_results.into_iter())
    }

    #[cfg(test)]
    pub fn version(&self) -> &changegroup::unpacker::CgVersion {
        &self.version
    }
}

impl PartialEq<[changegroup::Part]> for CgPartSequence {
    fn eq(&self, other: &[changegroup::Part]) -> bool {
        self.as_iter().eq(other.iter())
    }
}

impl Arbitrary for CgPartSequence {
    fn arbitrary(g: &mut Gen) -> Self {
        use crate::changegroup::unpacker::*;
        use crate::changegroup::*;

        let version_ind = bool::arbitrary(g);
        let version = match version_ind {
            true => CgVersion::Cg2Version,
            false => CgVersion::Cg3Version,
        };
        let gen_parts = match version {
            CgVersion::Cg2Version => gen_parts_v2,
            CgVersion::Cg3Version => gen_parts_v3,
        };
        let gen_sequence = match version {
            CgVersion::Cg2Version => gen_sequence_v2,
            CgVersion::Cg3Version => gen_sequence_v3,
        };

        // Generate a valid part sequence (changegroup, then manifest, then filelogs).
        let size = g.size();

        let changesets = gen_parts(Section::Changeset, g);
        let manifests = gen_parts(Section::Manifest, g);

        let nfilelogs = usize::arbitrary(g) % size;
        let mut filelogs = Vec::with_capacity(nfilelogs);

        for _ in 0..nfilelogs {
            // Changegroups can't support empty paths, so skip over those.
            let path = MPath::arbitrary(g);
            let section_end = Part::SectionEnd(Section::Filelog(path.clone()));
            filelogs.push((gen_parts(Section::Filelog(path), g), section_end));
        }

        gen_sequence(changesets, manifests, filelogs)
    }

    fn shrink(&self) -> Box<dyn Iterator<Item = Self>> {
        use crate::changegroup::unpacker::*;

        let gen_sequence = match self.version {
            CgVersion::Cg2Version => gen_sequence_v2,
            CgVersion::Cg3Version => gen_sequence_v3,
        };

        // All the parts can be shrinked independently as long as the section
        // remains the same (ensured in the impl of Arbitrary for
        // changegroup::Part).
        Box::new(
            (
                self.changesets.clone(),
                self.manifests.clone(),
                self.filelogs.clone(),
            )
                .shrink()
                .map(move |(c, m, f)| gen_sequence(c, m, f)),
        )
    }
}

fn gen_parts_v2(section: changegroup::Section, g: &mut Gen) -> Vec<changegroup::Part> {
    let size = g.size();
    (0..usize::arbitrary(g) % size)
        .map(|_| {
            changegroup::Part::CgChunk(section.clone(), changegroup::CgDeltaChunk::arbitrary(g))
        })
        .collect()
}

fn gen_parts_v3(section: changegroup::Section, g: &mut Gen) -> Vec<changegroup::Part> {
    let size = g.size();
    (0..usize::arbitrary(g) % size)
        .map(|_| {
            changegroup::Part::CgChunk(
                section.clone(),
                changegroup::CgDeltaChunk::arbitrary_with_flags(g),
            )
        })
        .collect()
}

fn gen_sequence_v2(
    changesets: Vec<changegroup::Part>,
    manifests: Vec<changegroup::Part>,
    filelogs: Vec<(Vec<changegroup::Part>, changegroup::Part)>,
) -> CgPartSequence {
    use crate::changegroup::*;
    CgPartSequence {
        changesets,
        changesets_end: Part::SectionEnd(Section::Changeset),
        manifests,
        manifests_end: Part::SectionEnd(Section::Manifest),
        treemanifest_end: None,
        filelogs,
        end: Part::End,

        version: changegroup::unpacker::CgVersion::Cg2Version,
    }
}
fn gen_sequence_v3(
    changesets: Vec<changegroup::Part>,
    manifests: Vec<changegroup::Part>,
    filelogs: Vec<(Vec<changegroup::Part>, changegroup::Part)>,
) -> CgPartSequence {
    use crate::changegroup::*;
    CgPartSequence {
        changesets,
        changesets_end: Part::SectionEnd(Section::Changeset),
        manifests,
        manifests_end: Part::SectionEnd(Section::Manifest),
        treemanifest_end: Some(Part::SectionEnd(Section::Treemanifest)),
        filelogs,
        end: Part::End,

        version: changegroup::unpacker::CgVersion::Cg3Version,
    }
}

impl Arbitrary for changegroup::Part {
    fn arbitrary(_g: &mut Gen) -> Self {
        unimplemented!()
    }

    fn shrink(&self) -> Box<dyn Iterator<Item = Self>> {
        use crate::changegroup::Part::CgChunk;

        match self {
            CgChunk(section, delta_chunk) => {
                // Keep the section the same, but shrink the delta chunk.
                let section = section.clone();
                Box::new(
                    delta_chunk
                        .shrink()
                        .map(move |chunk| CgChunk(section.clone(), chunk)),
                )
            }
            _ => empty_shrinker(),
        }
    }
}

impl Arbitrary for changegroup::CgDeltaChunk {
    fn arbitrary(g: &mut Gen) -> Self {
        // TODO: should these be more structured? e.g. base = p1 some of the time
        changegroup::CgDeltaChunk {
            node: HgNodeHash::arbitrary(g),
            p1: HgNodeHash::arbitrary(g),
            p2: HgNodeHash::arbitrary(g),
            base: HgNodeHash::arbitrary(g),
            linknode: HgNodeHash::arbitrary(g),
            delta: Delta::arbitrary(g),
            flags: None,
        }
    }

    fn shrink(&self) -> Box<dyn Iterator<Item = Self>> {
        // Don't bother trying to shrink node hashes -- the meat is in the delta.
        let clone = self.clone();
        Box::new(
            self.delta
                .shrink()
                .map(move |delta| changegroup::CgDeltaChunk {
                    node: clone.node.clone(),
                    p1: clone.p1.clone(),
                    p2: clone.p2.clone(),
                    base: clone.base.clone(),
                    linknode: clone.linknode.clone(),
                    delta,
                    flags: clone.flags.clone(),
                }),
        )
    }
}

impl changegroup::CgDeltaChunk {
    fn arbitrary_with_flags(g: &mut Gen) -> Self {
        let flags = u8::arbitrary(g) % 3;
        let flags = if flags == 0 {
            RevFlags::REVIDX_DEFAULT_FLAGS
        } else if flags == 1 {
            RevFlags::REVIDX_EXTSTORED
        } else {
            RevFlags::REVIDX_ELLIPSIS
        };
        changegroup::CgDeltaChunk {
            node: HgNodeHash::arbitrary(g),
            p1: HgNodeHash::arbitrary(g),
            p2: HgNodeHash::arbitrary(g),
            base: HgNodeHash::arbitrary(g),
            linknode: HgNodeHash::arbitrary(g),
            delta: Delta::arbitrary(g),
            flags: Some(flags),
        }
    }
}


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

use std::collections::HashMap;

use anyhow::bail;
use anyhow::Context;
use anyhow::Error;
use anyhow::Result;
use byteorder::BigEndian;
use byteorder::ByteOrder;
use bytes_old::BytesMut;
use tokio_io::codec::Decoder;

use crate::errors::ErrorKind;
use crate::types::StreamHeader;
use crate::utils::is_mandatory_param;

#[derive(Debug)]
pub struct StartDecoder;

impl Decoder for StartDecoder {
    type Item = StreamHeader;
    type Error = Error;

    fn decode(&mut self, buf: &mut BytesMut) -> Result<Option<StreamHeader>> {
        // bundle2 spec: "HG20" + u32 length of header + header + payload
        if buf.len() <= 8 {
            return Ok(None);
        }

        let header_len = {
            if &buf[..4] != b"HG20" {
                bail!(ErrorKind::Bundle2Decode(
                    "invalid bundle magic string".into(),
                ));
            }
            BigEndian::read_u32(&buf[4..8]) as usize
        };

        if buf.len() <= 8 + header_len {
            // Still more data to read.
            return Ok(None);
        }

        let _ = buf.split_to(8);

        let (m_stream_params, a_stream_params) = decode_stream_params(buf, header_len)?;

        Ok(Some(StreamHeader {
            m_stream_params,
            a_stream_params,
        }))
    }
}

fn decode_stream_params(
    buf: &mut BytesMut,
    header_len: usize,
) -> Result<(HashMap<String, String>, HashMap<String, String>)> {
    let mut m_stream_params = HashMap::new();
    let mut a_stream_params = HashMap::new();

    if header_len == 0 {
        return Ok((m_stream_params, a_stream_params));
    }

    // mutate the buffer to get the headers out of the way
    let header_buf = buf.split_to(header_len);
    let buf_slice = header_buf.as_ref();

    let headers = buf_slice.split(|c| *c == b' ');
    for header in headers {
        let mut key_val = header.splitn(2, |c| *c == b'=');
        let key = key_val
            .next()
            .ok_or_else(|| ErrorKind::Bundle2Decode("bad stream level key".into()))?;
        let val = key_val
            .next()
            .ok_or_else(|| ErrorKind::Bundle2Decode("bad stream level val".into()))?;
        let key_decoded = percent_encoding::percent_decode(key);
        let val_decoded = percent_encoding::percent_decode(val);
        let key_str = key_decoded.decode_utf8().with_context(|| {
            ErrorKind::Bundle2Decode("stream level key is invalid UTF-8".into())
        })?;
        let val_str = val_decoded.decode_utf8().with_context(|| {
            ErrorKind::Bundle2Decode("stream level val is invalid UTF-8".into())
        })?;
        if is_mandatory_param(&key_str)
            .with_context(|| ErrorKind::Bundle2Decode("stream key is invalid".to_string()))?
        {
            m_stream_params.insert(key_str.to_lowercase(), val_str.into_owned());
        } else {
            a_stream_params.insert(key_str.to_lowercase(), val_str.into_owned());
        }
    }

    Ok((m_stream_params, a_stream_params))
}


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

use std::collections::HashMap;
use std::fmt::Debug;
use std::io::BufRead;
use std::io::BufReader;
use std::io::Cursor;
use std::str::FromStr;

use anyhow::Result;
use assert_matches::assert_matches;
use async_compression::membuf::MemBuf;
use async_compression::Bzip2Compression;
use async_compression::CompressorType;
use async_compression::FlateCompression;
use context::CoreContext;
use fbinit::FacebookInit;
use futures::compat::Future01CompatExt;
use futures::stream::BoxStream;
use futures::TryStreamExt;
use futures_old::stream;
use futures_old::stream::Stream;
use futures_old::stream::Stream as OldStream;
use mercurial_types::HgChangesetId;
use mercurial_types::HgNodeHash;
use mercurial_types::MPath;
use mercurial_types::RepoPath;
use mercurial_types::NULL_HASH;
use partial_io::GenWouldBlock;
use partial_io::PartialAsyncRead;
use partial_io::PartialWithErrors;
use phases::Phase;
use quickcheck::Gen;
use quickcheck::QuickCheck;
use slog::o;
use slog::Discard;
use slog::Logger;
use tokio::runtime::Runtime;
use tokio_io::AsyncRead;

use crate::bundle2::Bundle2Stream;
use crate::bundle2::StreamEvent;
use crate::bundle2_encode::Bundle2EncodeBuilder;
use crate::changegroup;
use crate::errors::ErrorKind;
use crate::part_encode::PartEncodeBuilder;
use crate::part_header::PartHeaderBuilder;
use crate::part_header::PartHeaderType;
use crate::parts::phases_part;
use crate::types::StreamHeader;
use crate::utils::get_compression_param;
use crate::wirepack;
use crate::Bundle2Item;

const BZIP2_BUNDLE2: &[u8] = include_bytes!("fixtures/bzip2.bin");
const UNCOMP_BUNDLE2: &[u8] = include_bytes!("fixtures/uncompressed.bin");
const UNKNOWN_COMPRESSION_BUNDLE2: &[u8] = include_bytes!("fixtures/unknown-compression.bin");
const WIREPACK_BUNDLE2: &[u8] = include_bytes!("fixtures/wirepack.bin");

const CHANGESET1_HASH_STR: &str = "b2040b24fd5cdfaf36e3164ddc357e834167b14a";
const CHANGESET2_HASH_STR: &str = "415ab71954c98ea93dab4b8f61f04ca57bc5c33c";
const MANIFEST1_HASH_STR: &str = "afcff2144f55cfa5d9b04ac4ed6598f26035aa77";
const MANIFEST2_HASH_STR: &str = "aa93dc3435cbfecd0c4c245b80b2a0b9ed35a015";
const ABC_HASH_STR: &str = "b80de5d138758541c5f05265ad144ab9fa86d1db";
const DEF_HASH_STR: &str = "bb969a19e8853962b4347bea4c24796324f10d8b";

#[derive(PartialEq)]
struct ByteBuf<'a>(&'a [u8]);

impl<'a> Debug for ByteBuf<'a> {
    fn fmt(&self, fmt: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        for byte in self.0 {
            fmt.write_fmt(format_args!(r"\x{:02x}", byte))?;
        }
        Ok(())
    }
}

#[test]
fn test_parse_bzip2() {
    let gen = Gen::new(20);
    let mut quickcheck = QuickCheck::new().gen(gen);
    quickcheck.quickcheck(parse_bzip2 as fn(PartialWithErrors<GenWouldBlock>) -> ());
}

fn parse_bzip2(read_ops: PartialWithErrors<GenWouldBlock>) {
    parse_bundle(BZIP2_BUNDLE2, Some("BZ"), read_ops);
}

#[test]
fn test_parse_uncompressed() {
    let gen = Gen::new(20);
    let mut quickcheck = QuickCheck::new().gen(gen);
    quickcheck.quickcheck(parse_uncompressed as fn(PartialWithErrors<GenWouldBlock>) -> ());
}

fn parse_uncompressed(read_ops: PartialWithErrors<GenWouldBlock>) {
    parse_bundle(UNCOMP_BUNDLE2, None, read_ops);
}

#[test]
fn test_parse_unknown_compression() {
    let runtime = Runtime::new().unwrap();
    let bundle2_buf = BufReader::new(MemBuf::from(Vec::from(UNKNOWN_COMPRESSION_BUNDLE2)));
    let outer_stream_err = parse_stream_start(&runtime, bundle2_buf, Some("IL")).unwrap_err();
    assert_matches!(outer_stream_err.downcast::<ErrorKind>().unwrap(),
                    ErrorKind::Bundle2Decode(ref msg) if msg == "unknown compression 'IL'");
}

#[test]
fn test_empty_bundle_roundtrip_bzip() {
    empty_bundle_roundtrip(Some(CompressorType::Bzip2(Bzip2Compression::Default)));
}

#[test]
fn test_empty_bundle_roundtrip_gzip() {
    empty_bundle_roundtrip(Some(CompressorType::Gzip(FlateCompression::best())));
}

#[test]
fn test_empty_bundle_roundtrip_uncompressed() {
    empty_bundle_roundtrip(None);
}

fn empty_bundle_roundtrip(ct: Option<CompressorType>) {
    // Encode an empty bundle.
    let cursor = Cursor::new(Vec::with_capacity(32 * 1024));
    let mut builder = Bundle2EncodeBuilder::new(cursor);
    builder.set_compressor_type(ct);
    builder
        .add_stream_param("Foo".into(), "123".into())
        .unwrap();
    builder
        .add_stream_param("bar".into(), "456".into())
        .unwrap();
    let encode_fut = builder.build();

    let runtime = Runtime::new().unwrap();
    let mut buf = runtime.block_on(encode_fut.compat()).unwrap();
    buf.set_position(0);

    // Now decode it.
    let logger = Logger::root(Discard, o!());
    let stream = Bundle2Stream::new(logger, buf);
    let (item, stream) = runtime.block_on(stream.into_future().compat()).unwrap();

    let mut mparams = HashMap::new();
    let mut aparams = HashMap::new();
    mparams.insert("foo".into(), "123".into());
    mparams.insert("compression".into(), get_compression_param(&ct).into());
    aparams.insert("bar".into(), "456".into());
    let expected_header = StreamHeader {
        m_stream_params: mparams,
        a_stream_params: aparams,
    };

    assert_matches!(
        item,
        Some(StreamEvent::Next(Bundle2Item::Start(ref header))) if header == &expected_header
    );

    let (item, stream) = runtime.block_on(stream.into_future().compat()).unwrap();
    assert_matches!(item, Some(StreamEvent::Done(_)));

    let (item, _stream) = runtime.block_on(stream.into_future().compat()).unwrap();
    assert!(item.is_none());
}

#[fbinit::test]
fn test_phases_part_encording(fb: FacebookInit) {
    let phases_entries = stream::iter_ok(vec![
        (
            HgChangesetId::from_bytes(b"bbbbbbbbbbbbbbbbbbbb").unwrap(),
            Phase::Public,
        ),
        (
            HgChangesetId::from_bytes(b"cccccccccccccccccccc").unwrap(),
            Phase::Public,
        ),
        (
            HgChangesetId::from_bytes(b"aaaaaaaaaaaaaaaaaaaa").unwrap(),
            Phase::Draft,
        ),
    ]);

    let cursor = Cursor::new(Vec::new());
    let mut builder = Bundle2EncodeBuilder::new(cursor);
    builder.set_compressor_type(None);

    let ctx = CoreContext::test_mock(fb);
    let part = phases_part(ctx, phases_entries).unwrap();
    builder.add_part(part);

    let mut cursor = Runtime::new()
        .unwrap()
        .block_on(builder.build().compat())
        .unwrap();
    cursor.set_position(0);
    let buf = cursor.fill_buf().unwrap();

    let res = b"HG20\x00\x00\x00\x0eCompression\x3dUN\x00\x00\x00\x12\x0bPHASE-HEADS\x00\x00\x00\x00\x00\x00\x00\x00\x00H\x00\x00\x00\x00bbbbbbbbbbbbbbbbbbbb\x00\x00\x00\x00cccccccccccccccccccc\x00\x00\x00\x01aaaaaaaaaaaaaaaaaaaa\x00\x00\x00\x00\x00\x00\x00\x00";
    assert_eq!(
        ByteBuf(buf),
        ByteBuf(res),
        "Compare phase-heads bundle2 part encoding against binary representation from mercurial"
    );
}

#[test]
fn test_unknown_part_bzip() {
    unknown_part(Some(CompressorType::Bzip2(Bzip2Compression::Default)));
}

#[test]
fn test_unknown_part_gzip() {
    unknown_part(Some(CompressorType::Gzip(FlateCompression::best())));
}

#[test]
fn test_unknown_part_uncompressed() {
    unknown_part(None);
}

fn unknown_part(ct: Option<CompressorType>) {
    let cursor = Cursor::new(Vec::with_capacity(32 * 1024));
    let mut builder = Bundle2EncodeBuilder::new(cursor);

    builder.set_compressor_type(ct);

    let unknown_part = PartEncodeBuilder::mandatory(PartHeaderType::Listkeys).unwrap();

    builder.add_part(unknown_part);
    let encode_fut = builder.build();

    let runtime = Runtime::new().unwrap();
    let mut buf = runtime.block_on(encode_fut.compat()).unwrap();
    buf.set_position(0);

    let logger = Logger::root(Discard, o!());
    let stream = Bundle2Stream::new(logger, buf);
    let parts = Vec::new();

    let decode_fut = stream
        .map_err(|e| panic!("unexpected error: {:?}", e))
        .forward(parts);
    let (stream, parts) = runtime.block_on(decode_fut.compat()).unwrap();

    // Only the stream header should have been returned.
    let mut m_stream_params = HashMap::new();
    m_stream_params.insert("compression".into(), get_compression_param(&ct).into());
    let expected = StreamHeader {
        m_stream_params,
        a_stream_params: HashMap::new(),
    };

    let mut parts = parts.into_iter();
    assert_matches!(
        parts.next().unwrap().into_next().unwrap(),
        Bundle2Item::Start(ref header) if header == &expected
    );
    assert_matches!(parts.next(), Some(StreamEvent::Done(_)));
    assert!(parts.next().is_none());

    // Make sure the error was accumulated.
    let stream = stream.into_inner();
    let app_errors = stream.app_errors();
    assert_eq!(app_errors.len(), 1);
    assert_matches!(&app_errors[0],
                    ErrorKind::BundleUnknownPart(header)
                    if header.part_type() == &PartHeaderType::Listkeys && header.mandatory());
}

fn parse_bundle(
    input: &[u8],
    compression: Option<&str>,
    read_ops: PartialWithErrors<GenWouldBlock>,
) {
    let runtime = Runtime::new().unwrap();

    let bundle2_buf = MemBuf::from(Vec::from(input));
    let partial_read = BufReader::new(PartialAsyncRead::new(bundle2_buf, read_ops));
    let stream = parse_stream_start(&runtime, partial_read, compression).unwrap();

    let (stream, cg2s) = {
        let (res, stream) = runtime.old_next_stream(stream);
        let mut header = PartHeaderBuilder::new(PartHeaderType::Changegroup, true).unwrap();
        header.add_mparam("version", "02").unwrap();
        header.add_aparam("nbchanges", "2").unwrap();
        let header = header.build(0);
        let cg2s = match res.unwrap().into_next().unwrap() {
            Bundle2Item::Changegroup(h, cg2s) => {
                assert_eq!(h, header);
                cg2s
            }
            bad => panic!("Unexpected bundle2 item: {:?}", bad),
        };
        (stream, cg2s)
    };

    verify_cg2(&runtime, cg2s);

    let (res, stream) = runtime.old_next_stream(stream);
    assert_matches!(res, Some(StreamEvent::Done(_)));

    let (res, stream) = runtime.old_next_stream(stream);
    assert!(res.is_none());

    // Make sure the stream is fused.
    let (res, stream) = runtime.old_next_stream(stream);
    assert!(res.is_none());

    assert!(stream.app_errors().is_empty());
}

fn verify_cg2(runtime: &Runtime, stream: BoxStream<'static, Result<changegroup::Part>>) {
    let (res, stream) = runtime.next_stream(stream);
    let res = res.expect("expected part");

    assert_eq!(*res.section(), changegroup::Section::Changeset);
    let chunk = res.chunk();

    // Verify that changesets parsed correctly.
    let changeset1_hash = HgNodeHash::from_str(CHANGESET1_HASH_STR).unwrap();
    assert_eq!(chunk.node, changeset1_hash);
    assert_eq!(chunk.p1, NULL_HASH);
    assert_eq!(chunk.p2, NULL_HASH);
    assert_eq!(chunk.base, NULL_HASH);
    assert_eq!(chunk.linknode, changeset1_hash);
    let frags = chunk.delta.fragments();
    assert_eq!(frags.len(), 1);
    assert_eq!(frags[0].start, 0);
    assert_eq!(frags[0].end, 0);
    assert_eq!(frags[0].content.len(), 98);

    let (res, stream) = runtime.next_stream(stream);
    let res = res.expect("expected part");

    assert_eq!(*res.section(), changegroup::Section::Changeset);
    let chunk = res.chunk();

    let changeset2_hash = HgNodeHash::from_str(CHANGESET2_HASH_STR).unwrap();
    assert_eq!(chunk.node, changeset2_hash);
    assert_eq!(chunk.p1, changeset1_hash);
    assert_eq!(chunk.p2, NULL_HASH);
    assert_eq!(chunk.base, NULL_HASH);
    assert_eq!(chunk.linknode, changeset2_hash);
    let frags = chunk.delta.fragments();
    assert_eq!(frags.len(), 1);
    assert_eq!(frags[0].start, 0);
    assert_eq!(frags[0].end, 0);
    assert_eq!(frags[0].content.len(), 102);

    let (res, stream) = runtime.next_stream(stream);
    let res = res.expect("expected part");

    assert_matches!(
        res,
        changegroup::Part::SectionEnd(changegroup::Section::Changeset)
    );

    // Verify basic properties of manifests.
    let (res, stream) = runtime.next_stream(stream);
    let res = res.expect("expected part");

    assert_eq!(*res.section(), changegroup::Section::Manifest);
    let chunk = res.chunk();

    let manifest1_hash = HgNodeHash::from_str(MANIFEST1_HASH_STR).unwrap();
    assert_eq!(chunk.node, manifest1_hash);
    assert_eq!(chunk.p1, NULL_HASH);
    assert_eq!(chunk.p2, NULL_HASH);
    assert_eq!(chunk.base, NULL_HASH);
    assert_eq!(chunk.linknode, changeset1_hash);

    let (res, stream) = runtime.next_stream(stream);
    let res = res.expect("expected part");

    assert_eq!(*res.section(), changegroup::Section::Manifest);
    let chunk = res.chunk();

    let manifest2_hash = HgNodeHash::from_str(MANIFEST2_HASH_STR).unwrap();
    assert_eq!(chunk.node, manifest2_hash);
    assert_eq!(chunk.p1, manifest1_hash);
    assert_eq!(chunk.p2, NULL_HASH);
    assert_eq!(chunk.base, manifest1_hash); // In this case there's a delta.
    assert_eq!(chunk.linknode, changeset2_hash);

    let (res, stream) = runtime.next_stream(stream);
    let res = res.expect("expected part");

    assert_matches!(
        res,
        changegroup::Part::SectionEnd(changegroup::Section::Manifest)
    );

    // Filelog section
    let (res, stream) = runtime.next_stream(stream);
    let res = res.expect("expected part");

    assert_eq!(*res.section(), changegroup::Section::Filelog(path(b"abc")));
    let chunk = res.chunk();

    let abch = HgNodeHash::from_str(ABC_HASH_STR).unwrap();
    assert_eq!(chunk.node, abch);
    assert_eq!(chunk.p1, NULL_HASH);
    assert_eq!(chunk.p2, NULL_HASH);
    assert_eq!(chunk.base, NULL_HASH);
    assert_eq!(chunk.linknode, changeset1_hash);
    assert_eq!(chunk.delta.fragments().len(), 0); // empty file

    let (res, stream) = runtime.next_stream(stream);
    let res = res.expect("expected part");

    assert_matches!(res,
                    changegroup::Part::SectionEnd(ref section)
                    if *section == changegroup::Section::Filelog(path(b"abc")));

    let (res, stream) = runtime.next_stream(stream);
    let res = res.expect("expected part");

    assert_eq!(*res.section(), changegroup::Section::Filelog(path(b"def")));
    let chunk = res.chunk();

    let defh = HgNodeHash::from_str(DEF_HASH_STR).unwrap();
    assert_eq!(chunk.node, defh);
    assert_eq!(chunk.p1, NULL_HASH);
    assert_eq!(chunk.p2, NULL_HASH);
    assert_eq!(chunk.base, NULL_HASH);
    assert_eq!(chunk.linknode, changeset2_hash);
    assert_eq!(chunk.delta.fragments().len(), 1);

    // That's it, wrap this up.
    let (res, stream) = runtime.next_stream(stream);
    let res = res.expect("expected part");

    assert_matches!(res,
                    changegroup::Part::SectionEnd(ref section)
                    if *section == changegroup::Section::Filelog(path(b"def")));

    let (res, stream) = runtime.next_stream(stream);
    let res = res.expect("expected part");

    assert_matches!(res, changegroup::Part::End);

    let (res, _) = runtime.next_stream(stream);
    assert!(
        res.is_none(),
        "after the End part this stream should be empty"
    );
}

#[test]
fn test_parse_wirepack() {
    let gen = Gen::new(20);
    let mut quickcheck = QuickCheck::new().gen(gen);
    quickcheck.quickcheck(parse_wirepack as fn(PartialWithErrors<GenWouldBlock>) -> ());
}

fn parse_wirepack(read_ops: PartialWithErrors<GenWouldBlock>) {
    let runtime = Runtime::new().unwrap();

    let cursor = Cursor::new(WIREPACK_BUNDLE2);
    let partial_read = BufReader::new(PartialAsyncRead::new(cursor, read_ops));

    let stream = parse_stream_start(&runtime, partial_read, None).unwrap();

    let stream = {
        let (res, stream) = runtime.old_next_stream(stream);
        match res {
            Some(StreamEvent::Next(Bundle2Item::Changegroup(_, cg2s))) => {
                runtime
                    .block_on(cg2s.compat().for_each(|_| Ok(())).compat())
                    .unwrap();
            }
            bad => panic!("Unexpected Bundle2Item: {:?}", bad),
        }
        stream
    };

    let (stream, wirepacks) = {
        let (res, stream) = runtime.old_next_stream(stream);
        // Header
        let mut header = PartHeaderBuilder::new(PartHeaderType::B2xTreegroup2, true).unwrap();
        header.add_mparam("version", "1").unwrap();
        header.add_mparam("cache", "False").unwrap();
        header.add_mparam("category", "manifests").unwrap();
        let header = header.build(1);
        let wirepacks = match res.unwrap().into_next().unwrap() {
            Bundle2Item::B2xTreegroup2(h, wirepacks) => {
                assert_eq!(h, header);
                wirepacks
            }
            bad => panic!("Unexpected bundle2 item: {:?}", bad),
        };
        (stream, wirepacks)
    };

    // These are a few identifiers present in the bundle.
    let baz_dir = RepoPath::dir("baz").unwrap();
    let baz_hash = HgNodeHash::from_str("dcb9fa4bb7cdb673cd5752088b48d4c3f9c1fc23").unwrap();
    let root_hash = HgNodeHash::from_str("7d315c7a04cce5404f7ef16bf55eb7f4e90d159f").unwrap();
    let root_p1 = HgNodeHash::from_str("e313fc172615835d205f5881f8f34dd9bb0f0092").unwrap();

    let (res, wirepacks) = runtime.next_stream(wirepacks);
    let res = res.expect("expected part");

    // First entries received are for the directory "baz".
    let (path, entry_count) = res.unwrap_history_meta();
    assert_eq!(path, baz_dir);
    assert_eq!(entry_count, 1);

    let (res, wirepacks) = runtime.next_stream(wirepacks);
    let res = res.expect("expected part");

    let history_entry = res.unwrap_history();
    assert_eq!(history_entry.node, baz_hash);
    assert_eq!(history_entry.p1, NULL_HASH);
    assert_eq!(history_entry.p2, NULL_HASH);
    assert_eq!(history_entry.linknode, NULL_HASH);
    assert_eq!(history_entry.copy_from, None);

    let (res, wirepacks) = runtime.next_stream(wirepacks);
    let res = res.expect("expected part");

    let (path, entry_count) = res.unwrap_data_meta();
    assert_eq!(path, baz_dir);
    assert_eq!(entry_count, 1);

    let (res, wirepacks) = runtime.next_stream(wirepacks);
    let res = res.expect("expected part");

    let data_entry = res.unwrap_data();
    assert_eq!(path, baz_dir);
    assert_eq!(data_entry.node, baz_hash);
    assert_eq!(data_entry.delta_base, NULL_HASH);
    let fragments = data_entry.delta.fragments();
    assert_eq!(fragments.len(), 1);
    assert_eq!(fragments[0].start, 0);
    assert_eq!(fragments[0].end, 0);
    assert_eq!(fragments[0].content.len(), 46);

    let (res, wirepacks) = runtime.next_stream(wirepacks);
    let res = res.expect("expected part");

    // Next entries received are for the root manifest.
    let (path, entry_count) = res.unwrap_history_meta();
    assert_eq!(path, RepoPath::root());
    assert_eq!(entry_count, 1);

    let (res, wirepacks) = runtime.next_stream(wirepacks);
    let res = res.expect("expected part");

    let history_entry = res.unwrap_history();
    assert_eq!(history_entry.node, root_hash);
    assert_eq!(history_entry.p1, root_p1);
    assert_eq!(history_entry.p2, NULL_HASH);
    assert_eq!(history_entry.linknode, NULL_HASH);
    assert_eq!(history_entry.copy_from, None);

    let (res, wirepacks) = runtime.next_stream(wirepacks);
    let res = res.expect("expected part");

    let (path, entry_count) = res.unwrap_data_meta();
    assert_eq!(path, RepoPath::root());
    assert_eq!(entry_count, 1);

    let (res, wirepacks) = runtime.next_stream(wirepacks);
    let res = res.expect("expected part");

    let data_entry = res.unwrap_data();
    assert_eq!(data_entry.node, root_hash);
    assert_eq!(data_entry.delta_base, NULL_HASH);
    let fragments = data_entry.delta.fragments();
    assert_eq!(fragments.len(), 1);
    assert_eq!(fragments[0].start, 0);
    assert_eq!(fragments[0].end, 0);
    assert_eq!(fragments[0].content.len(), 136);

    let (res, wirepacks) = runtime.next_stream(wirepacks);
    let res = res.expect("expected part");

    // Finally the end.
    assert_eq!(res, wirepack::Part::End);
    let (res, _) = runtime.next_stream(wirepacks);
    assert!(
        res.is_none(),
        "after the End part this stream should be empty"
    );

    let (res, stream) = runtime.old_next_stream(stream);
    assert_matches!(res, Some(StreamEvent::Done(_)));
    assert!(stream.app_errors().is_empty());
}

fn path(bytes: &[u8]) -> MPath {
    MPath::new(bytes).unwrap()
}

fn parse_stream_start<R: AsyncRead + BufRead + 'static + Send>(
    runtime: &Runtime,
    reader: R,
    compression: Option<&str>,
) -> Result<Bundle2Stream<R>> {
    let mut m_stream_params = HashMap::new();
    let a_stream_params = HashMap::new();
    if let Some(compression) = compression {
        m_stream_params.insert("compression".into(), compression.into());
    }
    let expected = StreamHeader {
        m_stream_params,
        a_stream_params,
    };

    let logger = Logger::root(Discard, o!());
    let stream = Bundle2Stream::new(logger, reader);
    match runtime.block_on(stream.into_future().compat()) {
        Ok((item, stream)) => {
            let stream_start = item.unwrap();
            assert_eq!(stream_start.into_next().unwrap().unwrap_start(), expected);
            Ok(stream)
        }
        Err((e, _)) => Err(e),
    }
}

trait RuntimeExt {
    fn next_stream<T>(
        &self,
        stream: BoxStream<'static, Result<T>>,
    ) -> (Option<T>, BoxStream<'static, Result<T>>)
    where
        T: Send + 'static;

    fn old_next_stream<S>(&self, stream: S) -> (Option<S::Item>, S)
    where
        S: OldStream + Send + 'static,
        <S as OldStream>::Item: Send,
        <S as OldStream>::Error: Debug + Send;
}

impl RuntimeExt for Runtime {
    fn next_stream<T>(
        &self,
        stream: BoxStream<'static, Result<T>>,
    ) -> (Option<T>, BoxStream<'static, Result<T>>)
    where
        T: Send + 'static,
    {
        let (item, stream) = self.old_next_stream(stream.compat());
        (item, stream.into_inner())
    }

    fn old_next_stream<S>(&self, stream: S) -> (Option<S::Item>, S)
    where
        S: OldStream + Send + 'static,
        <S as OldStream>::Item: Send,
        <S as OldStream>::Error: Debug + Send,
    {
        match self.block_on(stream.into_future().compat()) {
            Ok((res, stream)) => (res, stream),
            Err((e, _)) => {
                panic!("stream failed to produce the next value! {:?}", e);
            }
        }
    }
}


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

//! Common types used by both decode and encode code.

use std::collections::HashMap;

#[derive(Debug, Eq, PartialEq)]
pub struct StreamHeader {
    // Stream parameters are specified as using a "simple textual format", which we
    // take to be valid UTF-8-encoded strings.
    pub m_stream_params: HashMap<String, String>,
    pub a_stream_params: HashMap<String, String>,
}


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

//! Utilities for decoding bundles.

use std::ops::Deref;
use std::str;

use anyhow::bail;
use anyhow::Context;
use anyhow::Error;
use anyhow::Result;
use async_compression::CompressorType;
use async_compression::DecompressorType;
use byteorder::BigEndian;
use byteorder::ByteOrder;
use bytes_old::Bytes;
use bytes_old::BytesMut;
use mercurial_types::HgNodeHash;
use mercurial_types::MPath;

use crate::errors::ErrorKind;

pub trait SplitTo {
    fn split_to(&mut self, at: usize) -> Self;
}

impl SplitTo for Bytes {
    #[inline]
    fn split_to(&mut self, at: usize) -> Self {
        Bytes::split_to(self, at)
    }
}

impl SplitTo for BytesMut {
    #[inline]
    fn split_to(&mut self, at: usize) -> Self {
        BytesMut::split_to(self, at)
    }
}

pub trait BytesExt {
    fn drain_u8(&mut self) -> u8;
    fn drain_u16(&mut self) -> u16;
    fn drain_u32(&mut self) -> u32;
    fn drain_u64(&mut self) -> u64;
    fn drain_i32(&mut self) -> i32;
    fn drain_str(&mut self, len: usize) -> Result<String>;
    fn drain_path(&mut self, len: usize) -> Result<MPath>;
    fn drain_node(&mut self) -> HgNodeHash;
    fn peek_u16(&self) -> u16;
    fn peek_u32(&self) -> u32;
    fn peek_i32(&self) -> i32;
}

impl<T> BytesExt for T
where
    T: SplitTo + AsRef<[u8]> + Deref<Target = [u8]>,
{
    #[inline]
    fn drain_u8(&mut self) -> u8 {
        self.split_to(1)[0]
    }

    #[inline]
    fn drain_u16(&mut self) -> u16 {
        BigEndian::read_u16(self.split_to(2).as_ref())
    }

    #[inline]
    fn drain_u32(&mut self) -> u32 {
        BigEndian::read_u32(self.split_to(4).as_ref())
    }

    #[inline]
    fn drain_u64(&mut self) -> u64 {
        BigEndian::read_u64(self.split_to(8).as_ref())
    }

    #[inline]
    fn drain_i32(&mut self) -> i32 {
        BigEndian::read_i32(self.split_to(4).as_ref())
    }

    #[inline]
    fn drain_str(&mut self, len: usize) -> Result<String> {
        Ok(str::from_utf8(self.split_to(len).as_ref())
            .context("invalid UTF-8")?
            .into())
    }

    #[inline]
    fn drain_path(&mut self, len: usize) -> Result<MPath> {
        MPath::new(self.split_to(len))
            .context("invalid path")
            .map_err(Error::from)
    }

    #[inline]
    fn drain_node(&mut self) -> HgNodeHash {
        // This only fails if the size of input passed in isn't 20
        // bytes. drain_to would have panicked in that case anyway.
        HgNodeHash::from_bytes(self.split_to(20).as_ref()).unwrap()
    }

    #[inline]
    fn peek_u16(&self) -> u16 {
        BigEndian::read_u16(&self[..2])
    }

    #[inline]
    fn peek_u32(&self) -> u32 {
        BigEndian::read_u32(&self[..4])
    }

    #[inline]
    fn peek_i32(&self) -> i32 {
        BigEndian::read_i32(&self[..4])
    }
}

pub fn is_mandatory_param(s: &str) -> Result<bool> {
    match s.chars().next() {
        Some(ch) => {
            if !ch.is_alphabetic() {
                bail!("'{}': first char '{}' is not alphabetic", s, ch);
            }
            Ok(ch.is_uppercase())
        }
        None => bail!("string is empty"),
    }
}

pub fn get_decompressor_type(compression: Option<&str>) -> Result<Option<DecompressorType>> {
    match compression {
        Some("BZ") => Ok(Some(DecompressorType::Bzip2)),
        Some("GZ") => Ok(Some(DecompressorType::Gzip)),
        Some("ZS") => Ok(Some(DecompressorType::OverreadingZstd)),
        Some("UN") => Ok(None),
        Some(s) => bail!(ErrorKind::Bundle2Decode(format!(
            "unknown compression '{}'",
            s
        ),)),
        None => Ok(None),
    }
}

pub fn get_compression_param(ct: &Option<CompressorType>) -> &'static str {
    match *ct {
        Some(CompressorType::Bzip2(_)) => "BZ",
        Some(CompressorType::Gzip(_)) => "GZ",
        Some(CompressorType::Zstd { .. }) => "ZS",
        None => "UN",
    }
}

pub fn capitalize_first(s: String) -> String {
    // Capitalize Unicode style, since capitalizing a single code point can
    // produce multiple code points.
    // TODO: just enforce ASCII-only and make this simpler.
    let mut chars = s.chars();
    match chars.next() {
        None => String::new(),
        Some(ch) => ch.to_uppercase().collect::<String>() + chars.as_str(),
    }
}

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

    #[test]
    fn test_is_mandatory_param() {
        assert!(is_mandatory_param("Foo").unwrap());
        assert!(!is_mandatory_param("bar").unwrap());
        assert_eq!(
            format!("{}", is_mandatory_param("").unwrap_err()),
            "string is empty"
        );
        assert_eq!(
            format!("{}", is_mandatory_param("123").unwrap_err()),
            "'123': first char '1' is not alphabetic"
        );
    }
}


/*
 * 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 mercurial_types::Delta;
use mercurial_types::HgNodeHash;
use mercurial_types::MPath;
use mercurial_types::RevFlags;

pub mod packer;
pub mod unpacker;
pub use unpacker::CgVersion;

#[derive(Clone, Debug, Eq, PartialEq)]
pub enum Section {
    Changeset,
    Manifest,
    Treemanifest,
    Filelog(MPath),
}

#[derive(Clone, Debug, Eq, PartialEq)]
pub enum Part {
    CgChunk(Section, CgDeltaChunk),
    SectionEnd(Section),
    End,
}

impl Part {
    pub fn is_section_end(&self) -> bool {
        match self {
            &Part::SectionEnd(_) => true,
            _ => false,
        }
    }
}

#[derive(Clone, Debug, Eq, PartialEq)]
pub struct CgDeltaChunk {
    pub node: HgNodeHash,
    pub p1: HgNodeHash,
    pub p2: HgNodeHash,
    pub base: HgNodeHash,
    pub linknode: HgNodeHash,
    pub delta: Delta,
    pub flags: Option<RevFlags>,
}

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

    use futures::compat::Future01CompatExt;
    use futures_ext::StreamLayeredExt;
    use futures_old::Future;
    use futures_old::Stream;
    use partial_io::GenWouldBlock;
    use partial_io::PartialAsyncRead;
    use partial_io::PartialAsyncWrite;
    use partial_io::PartialWithErrors;
    use quickcheck::Gen;
    use quickcheck::QuickCheck;
    use quickcheck::TestResult;
    use slog::o;
    use slog::Discard;
    use slog::Logger;
    use tokio_codec::FramedRead;
    use tokio_codec::FramedWrite;

    use super::*;
    use crate::chunk::ChunkDecoder;
    use crate::chunk::ChunkEncoder;
    use crate::quickcheck_types::CgPartSequence;

    #[test]
    fn test_roundtrip() {
        // Each test case gets pretty big (O(size**2) parts (because of
        // filelogs), each part with O(size) deltas), so reduce the size a bit
        // and generate a smaller number of test cases.
        let gen = Gen::new(50);
        let mut quickcheck = QuickCheck::new().gen(gen).tests(50);
        // Use NoErrors here because:
        // - AsyncWrite impls aren't supposed to return Interrupted errors.
        // - WouldBlock would require parking and unparking the task, which
        //   isn't yet supported in partial-io.
        quickcheck.quickcheck(
            roundtrip
                as fn(
                    CgPartSequence,
                    PartialWithErrors<GenWouldBlock>,
                    PartialWithErrors<GenWouldBlock>,
                ) -> TestResult,
        );
    }

    #[test]
    fn test_roundtrip_giant() {
        // Test a smaller number of cases with much larger inputs.
        let gen = Gen::new(200);
        let mut quickcheck = QuickCheck::new().gen(gen).tests(1);
        quickcheck.quickcheck(
            roundtrip
                as fn(
                    CgPartSequence,
                    PartialWithErrors<GenWouldBlock>,
                    PartialWithErrors<GenWouldBlock>,
                ) -> TestResult,
        );
    }

    fn roundtrip(
        seq: CgPartSequence,
        write_ops: PartialWithErrors<GenWouldBlock>,
        read_ops: PartialWithErrors<GenWouldBlock>,
    ) -> TestResult {
        // Encode this sequence.
        let cursor = Cursor::new(Vec::with_capacity(32 * 1024));
        let partial_write = PartialAsyncWrite::new(cursor, write_ops);
        let packer = packer::CgPacker::new(seq.to_stream().and_then(|x| x));
        let sink = FramedWrite::new(partial_write, ChunkEncoder);
        let unpacker_version = seq.version().clone();

        let fut = packer
            .forward(sink)
            .map_err(|e| panic!("unexpected error: {}", e))
            .and_then(move |(_, sink)| {
                let mut cursor = sink.into_inner().into_inner();

                // Decode it.
                cursor.set_position(0);

                let partial_read = PartialAsyncRead::new(cursor, read_ops);
                let chunks = FramedRead::new(partial_read, ChunkDecoder)
                    .map(|chunk| chunk.into_bytes().expect("expected normal chunk"));

                let logger = Logger::root(Discard, o!());
                let unpacker = unpacker::CgUnpacker::new(logger, unpacker_version);
                let part_stream = chunks.decode(unpacker);

                let parts = Vec::new();
                part_stream
                    .map_err(|e| panic!("unexpected error: {}", e))
                    .forward(parts)
            })
            .map(move |(_, parts)| {
                if seq != parts[..] {
                    TestResult::failed()
                } else {
                    TestResult::passed()
                }
            });

        let runtime = tokio::runtime::Runtime::new().unwrap();
        let result = runtime.block_on(fut.compat());
        result.unwrap()
    }
}


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

use anyhow::Error;
use anyhow::Result;
use byteorder::BigEndian;
use byteorder::ByteOrder;
use bytes_old::BufMut;
use futures_old::try_ready;
use futures_old::Async;
use futures_old::Poll;
use futures_old::Stream;

use super::CgDeltaChunk;
use super::Part;
use super::Section;
use crate::chunk::Chunk;
use crate::delta;

pub struct CgPacker<S> {
    delta_stream: S,
    last_seen: Section,
}

impl<S> CgPacker<S> {
    pub fn new(delta_stream: S) -> Self {
        CgPacker {
            delta_stream,
            last_seen: Section::Changeset,
        }
    }
}

impl<S> Stream for CgPacker<S>
where
    S: Stream<Item = Part>,
    Error: From<S::Error>,
{
    type Item = Chunk;
    type Error = Error;

    fn poll(&mut self) -> Poll<Option<Chunk>, Error> {
        use self::Part::*;

        match try_ready!(self.delta_stream.poll()) {
            None => Ok(Async::Ready(None)),
            Some(CgChunk(section, delta_chunk)) => {
                let mut builder = ChunkBuilder::new();
                if self.last_seen != section {
                    builder.encode_section(&section)?;
                    self.last_seen = section;
                }
                builder.encode_delta_chunk(delta_chunk);
                Ok(Async::Ready(Some(builder.build()?)))
            }
            Some(SectionEnd(_section)) => Ok(Async::Ready(Some(empty_cg_chunk()))),
            Some(End) => Ok(Async::Ready(Some(empty_cg_chunk()))),
        }
    }
}

/// Produce an empty changegroup chunk.
///
/// Note that this is distinct from Chunk::empty() -- this is an actual chunk
/// with a 4-byte payload.
fn empty_cg_chunk() -> Chunk {
    Chunk::new(vec![0, 0, 0, 0]).expect("Chunk::new should not fail for a 4-byte chunk")
}

#[derive(Debug)]
struct ChunkBuilder {
    inner: Vec<u8>,
    len_offset: usize,
}

impl ChunkBuilder {
    pub fn new() -> Self {
        ChunkBuilder {
            // Reserve four bytes in the beginning for the length.
            inner: vec![0, 0, 0, 0],
            len_offset: 0,
        }
    }

    /// Encode the beginning of a section. This should always happen before any
    /// delta chunks are encoded.
    pub fn encode_section(&mut self, section: &Section) -> Result<&mut Self> {
        assert_eq!(
            self.inner.len(),
            4,
            "encode_section must only be called once at the start"
        );
        // Changeset and manifest sections are implicitly encoded, so we don't
        // need to do anything there.
        // TODO: will need to encode tree manifests here as well
        if let Section::Filelog(f) = section {
            let f_vec = f.to_vec();
            // Note that the filename length must include the four bytes for itself.
            BigEndian::write_i32(&mut self.inner[0..], (f_vec.len() + 4) as i32);
            self.inner.put_slice(f_vec.as_slice());
            // Add four more bytes for the start of the section.
            self.len_offset = self.inner.len();
            self.inner.put_slice(&[0, 0, 0, 0]);
        }
        Ok(self)
    }

    pub fn encode_delta_chunk(&mut self, chunk: CgDeltaChunk) -> &mut Self {
        self.inner.put_slice(chunk.node.as_ref());
        self.inner.put_slice(chunk.p1.as_ref());
        self.inner.put_slice(chunk.p2.as_ref());
        self.inner.put_slice(chunk.base.as_ref());
        self.inner.put_slice(chunk.linknode.as_ref());

        if let Some(flags) = chunk.flags {
            let flags = flags.bits();
            self.inner
                .put_slice(&[(flags >> 8) as u8, (flags & 0xff) as u8]);
        }

        delta::encode_delta(&chunk.delta, &mut self.inner);

        self
    }

    pub fn build(self) -> Result<Chunk> {
        let len = self.inner.len() - self.len_offset;
        let mut inner = self.inner;
        BigEndian::write_i32(&mut inner[self.len_offset..], len as i32);
        Chunk::new(inner)
    }
}