/*
 * 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::borrow::Cow;
use std::collections::HashMap;
use std::fmt;
use std::fs::File;
use std::io::prelude::*;
use std::path::Path;
use std::path::PathBuf;
use std::sync::atomic::AtomicUsize;
use std::sync::atomic::Ordering::AcqRel;
use std::time::Duration;
use std::time::SystemTime;

use curl::easy::Easy2;
use curl::easy::HttpVersion;
use curl::easy::List;
use http::header;
use lru_cache::LruCache;
use maplit::hashmap;
use once_cell::sync::Lazy;
use openssl::pkcs12::Pkcs12;
use openssl::pkey::PKey;
use openssl::x509::X509;
use parking_lot::Mutex;
use parking_lot::RwLock;
use serde::Serialize;
use url::Url;

use crate::errors::HttpClientError;
use crate::event_listeners::RequestCreationEventListeners;
use crate::event_listeners::RequestEventListeners;
use crate::handler::Buffered;
use crate::handler::HandlerExt;
use crate::handler::Streaming;
use crate::receiver::ChannelReceiver;
use crate::receiver::Receiver;
use crate::response::AsyncResponse;
use crate::response::Response;

#[derive(Copy, Clone, Debug, PartialEq)]
pub enum Method {
    Get,
    Head,
    Post,
    Put,
}

#[derive(Copy, Clone, Debug, PartialEq)]
pub struct MinTransferSpeed {
    pub min_bytes_per_second: u32,
    pub grace_period: Duration,
}

impl fmt::Display for Method {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(
            f,
            "{}",
            match self {
                Method::Get => "GET",
                Method::Head => "HEAD",
                Method::Post => "POST",
                Method::Put => "PUT",
            }
        )
    }
}

#[derive(Clone, Debug, PartialEq)]
pub enum Encoding {
    Identity,
    Brotli,
    Deflate,
    Gzip,
    Zstd,
    Other(String),
}

impl Encoding {
    pub fn all() -> Vec<Self> {
        use Encoding::*;
        vec![Zstd, Brotli, Gzip, Deflate]
    }
}

impl<'a> From<&'a str> for Encoding {
    fn from(encoding: &'a str) -> Self {
        use Encoding::*;
        match encoding {
            "identity" => Identity,
            "br" => Brotli,
            "deflate" => Deflate,
            "gzip" => Gzip,
            "zstd" => Zstd,
            other => Other(other.into()),
        }
    }
}

impl AsRef<str> for Encoding {
    fn as_ref(&self) -> &str {
        use Encoding::*;
        match self {
            Identity => "identity",
            Brotli => "br",
            Deflate => "deflate",
            Gzip => "gzip",
            Zstd => "zstd",
            Other(s) => s,
        }
    }
}

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

/// Metadata about this request.
#[derive(Debug, Clone)]
pub struct RequestInfo {
    id: RequestId,
    url: Url,
    method: Method,
}

impl RequestInfo {
    /// Obtain the URL of the request.
    pub fn url(&self) -> &Url {
        &self.url
    }

    /// Obtain the HTTP method of the request.
    pub fn method(&self) -> &Method {
        &self.method
    }

    /// Obtain the request ID
    ///
    /// The ID is automatically assigned and uniquely identifies the requests
    /// in this process.
    pub fn id(&self) -> RequestId {
        self.id
    }
}

/// A subset of the `Request` builder. Preserved in curl types.
/// Expose the request in curl handler callback context.
#[derive(Clone, Debug)]
pub struct RequestContext {
    pub(crate) info: RequestInfo,
    pub(crate) body: Option<Vec<u8>>,
    pub(crate) event_listeners: RequestEventListeners,
}

/// Identity of a request.
#[derive(Copy, Clone, Debug, Eq, PartialEq, Ord, PartialOrd, Hash)]
pub struct RequestId(usize);

/// A builder struct for HTTP requests, designed to be
/// a more egonomic API for setting up a curl handle.
#[derive(Clone, Debug)]
pub struct Request {
    ctx: RequestContext,
    headers: HashMap<String, String>,
    cert: Option<PathBuf>,
    key: Option<PathBuf>,
    cainfo: Option<PathBuf>,
    timeout: Option<Duration>,
    http_version: HttpVersion,
    accept_encoding: Vec<Encoding>,
    min_transfer_speed: Option<MinTransferSpeed>,
    verify_tls_host: bool,
    verify_tls_cert: bool,
    verbose: bool,
    convert_cert: bool,
    auth_proxy_socket_path: Option<String>,
}

static REQUEST_CREATION_LISTENERS: Lazy<RwLock<RequestCreationEventListeners>> =
    Lazy::new(Default::default);

// Attempt to use HTTP/2 by default. Will fall back to HTTP/1.1
// if version negotiation with the server fails or if the binary
// is built against some Curl version that doesn't have HTTP/2 available
static DEFAULT_HTTP_VERSION: Lazy<HttpVersion> = Lazy::new(|| {
    if curl::Version::get().feature_http2() {
        HttpVersion::V2
    } else {
        HttpVersion::Any
    }
});

impl RequestContext {
    /// Create a [`RequestContext`].
    pub fn new(url: Url, method: Method) -> Self {
        static ID: AtomicUsize = AtomicUsize::new(0);
        let id = RequestId(ID.fetch_add(1, AcqRel));
        Self {
            info: RequestInfo { id, url, method },
            body: None,
            event_listeners: Default::default(),
        }
    }

    /// Obtain the HTTP url of the request.
    pub fn url(&self) -> &Url {
        &self.info.url()
    }

    /// Obtain the HTTP method of the request.
    pub fn method(&self) -> &Method {
        &self.info.method()
    }

    /// Obtain the request Id.
    ///
    /// The Id is automatically assigned and uniquely identifies the requests
    /// in this process.
    pub fn id(&self) -> RequestId {
        self.info.id()
    }

    /// Obtain the request metadata.
    pub fn info(&self) -> &RequestInfo {
        &self.info
    }

    /// Set the data to be uploaded in the request body.
    pub fn body<B: Into<Vec<u8>>>(mut self, data: B) -> Self {
        self.body = Some(data.into());
        self
    }

    /// Set the data to be uploaded in the request body.
    pub fn set_body<B: Into<Vec<u8>>>(&mut self, data: B) {
        self.body = Some(data.into());
    }

    /// Provide a way to register event callbacks.
    pub fn event_listeners(&mut self) -> &mut RequestEventListeners {
        &mut self.event_listeners
    }
}

impl Request {
    pub fn new(url: Url, method: Method) -> Self {
        let ctx = RequestContext::new(url, method);
        Self {
            ctx,
            // Always set Expect so we can disable curl automatically expecting "100-continue".
            // That would require two response reads, which breaks the http_client model.
            headers: hashmap! {
                "Expect".to_string() => "".to_string(),
            },
            cert: None,
            key: None,
            cainfo: None,
            timeout: None,
            http_version: DEFAULT_HTTP_VERSION.clone(),
            accept_encoding: Vec::new(),
            min_transfer_speed: None,
            verify_tls_host: true,
            verify_tls_cert: true,
            verbose: false,
            convert_cert: false,
            auth_proxy_socket_path: None,
        }
    }

    /// Obtain the request Id.
    ///
    /// The Id is automatically assigned and uniquely identifies the requests
    /// in this process.
    pub fn id(&self) -> RequestId {
        self.ctx.id()
    }

    /// Get a reference to this request's context.
    ///
    /// The request context contains most of the information about the request,
    /// such as URL, method, etc.
    pub fn ctx(&self) -> &RequestContext {
        &self.ctx
    }

    /// Get a mutable reference to this request's context.
    pub fn ctx_mut(&mut self) -> &mut RequestContext {
        &mut self.ctx
    }

    /// Create a GET request.
    pub(crate) fn get(url: Url) -> Self {
        Self::new(url, Method::Get)
    }

    /// Create a HEAD request.
    pub(crate) fn head(url: Url) -> Self {
        Self::new(url, Method::Head)
    }

    /// Create a POST request.
    pub(crate) fn post(url: Url) -> Self {
        Self::new(url, Method::Post)
    }

    /// Create a PUT request.
    pub(crate) fn put(url: Url) -> Self {
        Self::new(url, Method::Put)
    }

    /// Set the data to be uploaded in the request body.
    pub fn body<B: Into<Vec<u8>>>(mut self, data: B) -> Self {
        self.set_body(data);
        self
    }

    /// Set the data to be uploaded in the request body.
    pub fn set_body<B: Into<Vec<u8>>>(&mut self, data: B) -> &mut Self {
        self.ctx.set_body(data);
        self
    }

    /// Set the http version for this request. Defaults to HTTP/2.
    pub fn http_version(mut self, version: HttpVersion) -> Self {
        self.set_http_version(version);
        self
    }

    /// Set the http version for this request. Defaults to HTTP/2.
    pub fn set_http_version(&mut self, version: HttpVersion) -> &mut Self {
        self.http_version = version;
        self
    }

    /// Specify the content compression formats the client should advertise to
    /// the server. By default, this will be every format supported by libcurl.
    pub fn accept_encoding(mut self, formats: impl IntoIterator<Item = Encoding>) -> Self {
        self.set_accept_encoding(formats);
        self
    }

    /// Specify the content compression formats the client should advertise to
    /// the server. By default, this will be every format supported by libcurl.
    pub fn set_accept_encoding(
        &mut self,
        formats: impl IntoIterator<Item = Encoding>,
    ) -> &mut Self {
        self.accept_encoding = formats.into_iter().collect();
        self
    }

    /// Set transfer speed options for this request.
    pub fn min_transfer_speed(mut self, min_transfer_speed: Option<MinTransferSpeed>) -> Self {
        self.min_transfer_speed = min_transfer_speed;
        self
    }

    /// Set transfer speed options for this request.
    pub fn set_min_transfer_speed(&mut self, min_transfer_speed: MinTransferSpeed) -> &mut Self {
        self.min_transfer_speed = Some(min_transfer_speed);
        self
    }

    /// Serialize the given value as JSON and use it as the request body.
    pub fn json<S: Serialize>(mut self, value: &S) -> Result<Self, serde_json::Error> {
        self.set_json_body(value)?;
        Ok(self)
    }

    /// Serialize the given value as JSON and use it as the request body.
    pub fn set_json_body<S: Serialize>(
        &mut self,
        value: &S,
    ) -> Result<&mut Self, serde_json::Error> {
        self.set_header("Content-Type", "application/json")
            .set_body(serde_json::to_vec(value)?);
        Ok(self)
    }

    /// Serialize the given value as CBOR and use it as the request body.
    pub fn cbor<S: Serialize>(mut self, value: &S) -> Result<Self, serde_cbor::Error> {
        self.set_cbor_body(value)?;
        Ok(self)
    }

    /// Serialize the given value as CBOR and use it as the request body.
    pub fn set_cbor_body<S: Serialize>(
        &mut self,
        value: &S,
    ) -> Result<&mut Self, serde_cbor::Error> {
        self.set_header("Content-Type", "application/cbor")
            .set_body(serde_cbor::to_vec(value)?);
        Ok(self)
    }

    /// Set a request header.
    pub fn header(mut self, name: impl ToString, value: impl ToString) -> Self {
        self.set_header(name, value);
        self
    }

    /// Set a request header.
    pub fn set_header(&mut self, name: impl ToString, value: impl ToString) -> &mut Self {
        self.headers
            .insert(name.to_string().to_lowercase(), value.to_string());
        self
    }

    pub fn get_header_mut<'a>(&'a mut self, name: impl ToString) -> Option<&'a mut String> {
        self.headers.get_mut(&name.to_string().to_lowercase())
    }

    /// Specify a client certificate for TLS mutual authentiation.
    ///
    /// This should be a path to a base64-encoded PEM file containing the
    /// client's X.509 certificate. When using a client certificate, the client
    /// must also provide the corresponding private key; this can either be
    /// concatenated to the certificate in the PEM file (in which case it will
    /// be used automatically), or specified separately via the `key` method.
    pub fn cert(mut self, cert: impl AsRef<Path>) -> Self {
        self.set_cert(cert);
        self
    }

    /// Specify a client certificate for TLS mutual authentiation.
    ///
    /// This should be a path to a base64-encoded PEM file containing the
    /// client's X.509 certificate. When using a client certificate, the client
    /// must also provide the corresponding private key; this can either be
    /// concatenated to the certificate in the PEM file (in which case it will
    /// be used automatically), or specified separately via the `key` method.
    pub fn set_cert(&mut self, cert: impl AsRef<Path>) -> &mut Self {
        self.cert = Some(cert.as_ref().into());
        self
    }

    /// Specify a client private key for TLS mutual authentiation.
    ///
    /// This method can be used to specify the path to the client's private
    /// key if this key was not included in the certificate file specified via
    /// the `cert` method.
    pub fn key(mut self, key: impl AsRef<Path>) -> Self {
        self.set_key(key);
        self
    }

    /// Specify a client private key for TLS mutual authentiation.
    ///
    /// This method can be used to specify the path to the client's private
    /// key if this key was not included in the certificate file specified via
    /// the `cert` method.
    pub fn set_key(&mut self, key: impl AsRef<Path>) -> &mut Self {
        self.key = Some(key.as_ref().into());
        self
    }

    /// Specify a CA certificate bundle to be used to verify the
    /// server's certificate. If not specified, the client will
    /// use the system default CA certificate bundle.
    pub fn cainfo(mut self, cainfo: impl AsRef<Path>) -> Self {
        self.set_cainfo(cainfo);
        self
    }

    /// Specify a CA certificate bundle to be used to verify the
    /// server's certificate. If not specified, the client will
    /// use the system default CA certificate bundle.
    pub fn set_cainfo(&mut self, cainfo: impl AsRef<Path>) -> &mut Self {
        self.cainfo = Some(cainfo.as_ref().into());
        self
    }

    /// Set the maximum time this request is allowed to take.
    pub fn timeout(mut self, timeout: Duration) -> Self {
        self.set_timeout(timeout);
        self
    }

    /// Set the maximum time this request is allowed to take.
    pub fn set_timeout(&mut self, timeout: Duration) -> &mut Self {
        self.timeout = Some(timeout);
        self
    }

    /// Configure whether the client should verify that the server's hostname
    /// matches either the common name (CN) or a subject alternate name (SAN)
    /// present in the server's TLS certificate. Disabling this option will make
    /// the connection insecure. This is primarily useful for testing.
    pub fn verify_tls_host(mut self, verify: bool) -> Self {
        self.set_verify_tls_host(verify);
        self
    }

    /// Configure whether the client should verify that the server's hostname
    /// matches either the common name (CN) or a subject alternate name (SAN)
    /// present in the server's TLS certificate. Disabling this option will make
    /// the connection insecure. This is primarily useful for testing.
    pub fn set_verify_tls_host(&mut self, verify: bool) -> &mut Self {
        self.verify_tls_host = verify;
        self
    }

    /// Configure whether the client should verify the authenticity of the
    /// server's TLS certificate using the CA certificate bundle specified
    /// via `cainfo` (or the default CA bundle if not set). This option is
    /// enabled by default; disabling it will make the connection insecure.
    /// This is primarily useful for testing.
    pub fn verify_tls_cert(mut self, verify: bool) -> Self {
        self.set_verify_tls_cert(verify);
        self
    }

    /// Configure whether the client should verify the authenticity of the
    /// server's TLS certificate using the CA certificate bundle specified
    /// via `cainfo` (or the default CA bundle if not set). This option is
    /// enabled by default; disabling it will make the connection insecure.
    /// This is primarily useful for testing.
    pub fn set_verify_tls_cert(&mut self, verify: bool) -> &mut Self {
        self.verify_tls_cert = verify;
        self
    }

    /// Turn on libcurl's verbose output. This will cause libcurl to print lots
    /// of verbose debug messages to stderr. This can be useful when trying to
    /// understand exactly what libcurl is doing under the hood, which can help
    /// to debug low-level protocol issues.
    pub fn verbose(mut self, verbose: bool) -> Self {
        self.set_verbose(verbose);
        self
    }

    pub fn set_client_info(&mut self, client_info: &Option<String>) -> &mut Self {
        if let Some(info) = client_info {
            self.set_header("X-Client-Info", info);
        }
        self
    }

    /// Turn on libcurl's verbose output. This will cause libcurl to print lots
    /// of verbose debug messages to stderr. This can be useful when trying to
    /// understand exactly what libcurl is doing under the hood, which can help
    /// to debug low-level protocol issues.
    pub fn set_verbose(&mut self, verbose: bool) -> &mut Self {
        self.verbose = verbose;
        self
    }

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

    /// Convert the client's X.509 certificate from a PEM file into an in-memory
    /// PKCS#12 archive before passing it to libcurl. This is necessary on some
    /// platforms (most notably Windows) where the system crypto APIs (SChannel
    /// in the case of Windows) do not support loading PEM files.
    pub fn convert_cert(mut self, convert: bool) -> Self {
        self.set_convert_cert(convert);
        self
    }

    /// Convert the client's X.509 certificate from a PEM file into an in-memory
    /// PKCS#12 archive before passing it to libcurl. This is necessary on some
    /// platforms (most notably Windows) where the system crypto APIs (SChannel
    /// in the case of Windows) do not support loading PEM files.
    pub fn set_convert_cert(&mut self, convert: bool) -> &mut Self {
        self.convert_cert = convert;
        self
    }

    /// Execute the request, blocking until completion.
    ///
    /// This method is intended as a simple way to perform
    /// one-off HTTP requests. `HttpClient::send` should be
    /// used instead of this method when working with many
    /// concurrent requests or large requests that require
    /// progress reporting.
    pub fn send(self) -> Result<Response, HttpClientError> {
        let mut easy: Easy2<Buffered> = self.try_into()?;
        let res = easy.perform();
        let ctx = easy.get_mut().request_context_mut();
        let info = ctx.info().clone();

        match res {
            Ok(()) => {
                ctx.event_listeners().trigger_success(&info);
            }
            Err(e) => {
                ctx.event_listeners().trigger_failure(&info);
                return Err(e.into());
            }
        }

        Response::try_from(easy.get_mut())
    }

    /// Execute this request asynchronously.
    pub async fn send_async(self) -> Result<AsyncResponse, HttpClientError> {
        let request_info = self.ctx().info().clone();
        let (receiver, streams) = ChannelReceiver::new();
        let request = self.into_streaming(receiver);

        // Spawn the request as another task, which will block
        // the worker it is scheduled on until completion.
        let io_task = tokio::task::spawn_blocking(move || request.send());

        match AsyncResponse::new(streams, request_info).await {
            Ok(res) => Ok(res),
            // If the request was dropped before completion, this likely means
            // that configuring or sending the request failed. The IO task will
            // likely return a more meaningful error message, so return that
            // instead of a generic "this request was dropped" error.
            e @ Err(HttpClientError::RequestDropped(_)) => io_task.await?.and(e),
            Err(e) => Err(e),
        }
    }

    /// Turn this `Request` into a streaming request. The
    /// received data for this request will be passed as
    /// it arrives to the given `Receiver`.
    pub fn into_streaming<R>(self, receiver: R) -> StreamRequest<R> {
        StreamRequest {
            request: self,
            receiver,
        }
    }

    /// Turn this `Request` into a `curl::Easy2` handle using the given
    /// `Handler` to process the response.
    pub(crate) fn into_handle<H: HandlerExt>(
        mut self,
        create_handler: impl FnOnce(RequestContext) -> H,
    ) -> Result<Easy2<H>, HttpClientError> {
        // Allow request creation listeners to configure the Request before we
        // use it, potentially overriding settings explicitly configured via
        // the methods on Request.
        REQUEST_CREATION_LISTENERS
            .read()
            .trigger_new_request(&mut self);

        let body_size = self.ctx.body.as_ref().map(|body| body.len() as u64);
        let mut url = self.ctx.url().clone();
        if self.auth_proxy_socket_path.is_some() {
            url.set_scheme("http")
                .expect("Failed setting url scheme to http");
            self.set_verify_tls_cert(false)
                .set_verify_tls_host(false)
                .set_convert_cert(false);

            if let Some(user_agent) = self.get_header_mut("user-agent") {
                user_agent.push_str("+x2pagentd");
            }
        }
        let handler = create_handler(self.ctx);

        let mut easy = Easy2::new(handler);

        easy.url(url.as_str())?;
        easy.verbose(self.verbose)?;
        easy.unix_socket_path(self.auth_proxy_socket_path)?;

        // Configure the handle for the desired HTTP method.
        match easy.get_ref().request_context().method() {
            Method::Get => {}
            Method::Head => {
                easy.nobody(true)?;
            }
            Method::Post => {
                easy.post(true)?;
                if let Some(size) = body_size {
                    easy.post_field_size(size)?;
                }
            }
            Method::Put => {
                easy.upload(true)?;
                if let Some(size) = body_size {
                    easy.in_filesize(size)?;
                }
            }
        }

        if !self.accept_encoding.is_empty() {
            // To maintain compatibility with libcurl, if the Accept-Encoding is explicitly set to
            // the empty string, advertise all formats the client supports.
            if self.accept_encoding.len() == 1
                && self.accept_encoding[0] == Encoding::Other("".into())
            {
                self.accept_encoding = Encoding::all()
            }

            let encoding = self
                .accept_encoding
                .iter()
                .map(|s| s.as_ref())
                .collect::<Vec<_>>()
                .join(", ");

            // XXX: Ideally, we should set the Accept-Encoding via the accept_encoding() method
            // (which corresponds to CURLOPT_ACCEPT_ENCODING). This will cause libcurl to decode
            // the response body automatically if the received Content-Encoding matches one of the
            // requested formats.
            //
            // Unfortunately, although libcurl can be built to support many compression formats,
            // the Rust bindings configure it so that only a few formats (e.g., gzip and deflate)
            // are supported. To work around this, right now we just set the Accept-Encoding header
            // as a regular header (without setting CURLOPT_ACCEPT_ENCODING) and decode the response
            // manually. This allows us to ensure support for formats we care about (e.g., zstd).
            self.headers
                .insert(header::ACCEPT_ENCODING.as_str().into(), encoding);
        }

        // Add headers.
        let mut headers = List::new();
        for (name, value) in self.headers.iter() {
            let header = format!("{}: {}", name, value);
            headers.append(&header)?;
        }
        easy.http_headers(headers)?;

        // Configure TLS verification.
        easy.ssl_verify_host(self.verify_tls_host)?;
        easy.ssl_verify_peer(self.verify_tls_cert)?;

        match &self.cert {
            Some(cert) if self.convert_cert => {
                // Convert certificate to PKCS#12 format for platforms that do
                // not support loading PEM files (notably Windows).
                tracing::debug!("Converting certificate {:?} to PKCS#12 format", cert);
                let blob = pem_to_pkcs12(cert, self.key)?;
                easy.ssl_cert_type("P12")?;
                easy.ssl_cert_blob(&blob)?;
            }
            Some(cert) => {
                easy.ssl_cert(cert)?;
                if let Some(key) = &self.key {
                    easy.ssl_key(key)?;
                }
            }
            None => {}
        }

        // Windows enables ssl revocation checking by default, which doesn't work inside the
        // datacenter.
        #[cfg(windows)]
        {
            use curl::easy::SslOpt;
            let mut ssl_opts = SslOpt::new();
            ssl_opts.no_revoke(true);
            easy.ssl_options(&ssl_opts)?;
        }

        if let Some(cainfo) = self.cainfo {
            easy.cainfo(cainfo)?;
        }

        if let Some(timeout) = self.timeout {
            easy.timeout(timeout)?;
        }

        easy.http_version(self.http_version)?;

        if let Some(mts) = self.min_transfer_speed {
            easy.low_speed_limit(mts.min_bytes_per_second)?;
            easy.low_speed_time(mts.grace_period)?;
        }

        // Tell libcurl to report progress to the handler.
        easy.progress(true)?;

        Ok(easy)
    }

    /// Register a callback function that is called on new requests.
    pub fn on_new_request(f: impl Fn(&mut Self) + Send + Sync + 'static) {
        REQUEST_CREATION_LISTENERS.write().on_new_request(f);
    }
}

impl TryFrom<Request> for Easy2<Buffered> {
    type Error = HttpClientError;

    fn try_from(req: Request) -> Result<Self, Self::Error> {
        req.into_handle(Buffered::new)
    }
}

pub struct StreamRequest<R> {
    pub(crate) request: Request,
    pub(crate) receiver: R,
}

impl<R: Receiver> StreamRequest<R> {
    pub fn send(self) -> Result<(), HttpClientError> {
        let mut easy: Easy2<Streaming<R>> = self.try_into()?;
        let res = easy.perform().map_err(Into::into);
        let _ = easy
            .get_mut()
            .take_receiver()
            .expect("Receiver is gone; this should never happen")
            .done(res);
        Ok(())
    }
}

impl<R: Receiver> TryFrom<StreamRequest<R>> for Easy2<Streaming<R>> {
    type Error = HttpClientError;

    fn try_from(req: StreamRequest<R>) -> Result<Self, Self::Error> {
        let StreamRequest { request, receiver } = req;
        request.into_handle(|ctx| Streaming::new(receiver, ctx))
    }
}

fn read_file(path: impl AsRef<Path>) -> Result<Vec<u8>, anyhow::Error> {
    let mut f = File::open(path)?;
    let mut buf = Vec::new();
    f.read_to_end(&mut buf)?;
    Ok(buf)
}

#[derive(Eq, Hash, PartialEq)]
struct PemCacheKey {
    pub cert: PathBuf,
    pub key: Option<PathBuf>,
    pub cert_mtime: SystemTime,
    pub key_mtime: Option<SystemTime>,
}

static PEM_CONVERT_CACHE: Lazy<Mutex<LruCache<PemCacheKey, Vec<u8>>>> =
    Lazy::new(|| Mutex::new(LruCache::new(10)));

/// Convert a PEM-formatted X.509 certificate chain and private key into a
/// PKCS#12 archive, which can then be directly passed to libcurl using
/// `CURLOPT_SSLCERT_BLOB`. This is useful because not all TLS engines (notably
/// SChannel (WinSSL) on Windows) support loading PEM files, but all major TLS
/// engines support PKCS#12. Returns a DER-encoded binary representation of
/// the combined certificate chain and private key.
fn pem_to_pkcs12(
    cert: impl AsRef<Path>,
    key: Option<impl AsRef<Path>>,
) -> Result<Vec<u8>, anyhow::Error> {
    let mut cache = PEM_CONVERT_CACHE.lock();
    let cert_mtime = cert.as_ref().metadata()?.modified()?;
    let key_mtime = match &key {
        Some(key) => Some(key.as_ref().metadata()?.modified()?),
        None => None,
    };
    let cache_key = PemCacheKey {
        cert: cert.as_ref().to_owned(),
        key: key.as_ref().map(|k| k.as_ref().to_owned()),
        cert_mtime,
        key_mtime,
    };
    if let Some(data) = cache.get_mut(&cache_key) {
        return Ok(data.clone());
    }

    // It's common for the certificate and private key to be concatenated
    // together in the same PEM file. If a key path isn't specified, assume
    // this is the case and use the certificate PEM for the key as well.
    let cert_bytes = read_file(cert)?;
    let key_bytes = match key {
        Some(key) => Cow::Owned(read_file(key)?),
        None => Cow::Borrowed(&cert_bytes),
    };

    let cert = X509::from_pem(&cert_bytes)?;
    let key = PKey::private_key_from_pem(&key_bytes)?;

    // PKCS#12 archives are encrypted, so we need to specify a password when
    // creating one. Here we just use an empty password since it seems like most
    // TLS engines will attempt to decrypt using the empty string if no password
    // is specified.
    let pkcs12 = Pkcs12::builder()
        .name("")
        .pkey(&key)
        .cert(&cert)
        .build2("")?;

    let result = pkcs12.to_der()?;
    cache.insert(cache_key, result.clone());

    Ok(result)
}

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

    use anyhow::Result;
    use futures::TryStreamExt;
    use http::header;
    use http::header::HeaderName;
    use http::header::HeaderValue;
    use http::StatusCode;
    use mockito::mock;
    use mockito::Matcher;
    use serde_json::json;

    use super::*;
    use crate::Config;
    use crate::HttpClient;

    #[test]
    fn test_get() -> Result<()> {
        let mock = mock("GET", "/test")
            .with_status(200)
            .match_header("X-Api-Key", "1234")
            .with_header("Content-Type", "text/plain")
            .with_header("X-Served-By", "mock")
            .with_body("Hello, world!")
            .create();

        let url = Url::parse(&mockito::server_url())?.join("test")?;
        let res = Request::get(url).header("X-Api-Key", "1234").send()?;

        mock.assert();

        assert_eq!(res.head.status, StatusCode::OK);
        assert_eq!(&*res.body, &b"Hello, world!"[..]);
        assert_eq!(
            res.head.headers.get(header::CONTENT_TYPE).unwrap(),
            HeaderValue::from_static("text/plain")
        );
        assert_eq!(
            res.head
                .headers
                .get(HeaderName::from_bytes(b"X-Served-By")?)
                .unwrap(),
            HeaderValue::from_static("mock")
        );

        Ok(())
    }

    #[tokio::test]
    async fn test_async_get() -> Result<()> {
        let mock = mock("GET", "/test")
            .with_status(200)
            .match_header("X-Api-Key", "1234")
            .with_header("Content-Type", "text/plain")
            .with_header("X-Served-By", "mock")
            .with_body("Hello, world!")
            .create();

        let url = Url::parse(&mockito::server_url())?.join("test")?;
        let res = Request::get(url)
            .header("X-Api-Key", "1234")
            .send_async()
            .await?;

        mock.assert();

        assert_eq!(res.head.status, StatusCode::OK);
        assert_eq!(
            res.head.headers.get(header::CONTENT_TYPE).unwrap(),
            HeaderValue::from_static("text/plain")
        );
        assert_eq!(
            res.head
                .headers
                .get(HeaderName::from_bytes(b"X-Served-By")?)
                .unwrap(),
            HeaderValue::from_static("mock")
        );

        let body = res.into_body().raw().try_concat().await?;
        assert_eq!(&*body, &b"Hello, world!"[..]);

        Ok(())
    }

    #[test]
    fn test_head() -> Result<()> {
        let mock = mock("HEAD", "/test")
            .with_status(200)
            .match_header("X-Api-Key", "1234")
            .with_header("Content-Type", "text/plain")
            .with_header("X-Served-By", "mock")
            .create();

        let url = Url::parse(&mockito::server_url())?.join("test")?;
        let res = Request::head(url).header("X-Api-Key", "1234").send()?;

        mock.assert();

        assert_eq!(res.head.status, StatusCode::OK);
        assert!(res.body.is_empty());
        assert_eq!(
            res.head.headers.get(header::CONTENT_TYPE).unwrap(),
            HeaderValue::from_static("text/plain")
        );
        assert_eq!(
            res.head
                .headers
                .get(HeaderName::from_bytes(b"X-Served-By")?)
                .unwrap(),
            HeaderValue::from_static("mock")
        );

        Ok(())
    }

    #[test]
    fn test_post() -> Result<()> {
        let body = "foo=hello&bar=world";

        let mock = mock("POST", "/test")
            .with_status(201)
            .match_header("Content-Type", "application/x-www-form-urlencoded")
            .match_body(Matcher::Exact(body.into()))
            .create();

        let url = Url::parse(&mockito::server_url())?.join("test")?;
        let res = Request::post(url).body(body.as_bytes()).send()?;

        mock.assert();
        assert_eq!(res.head.status, StatusCode::CREATED);

        Ok(())
    }

    #[test]
    fn test_post_large() -> Result<()> {
        let body_bytes = vec![65; 1024 * 1024];
        let body = String::from_utf8_lossy(body_bytes.as_ref());

        let mock = mock("POST", "/test")
            .with_status(201)
            .match_header("Expect", Matcher::Missing)
            .match_body(Matcher::Exact(body.into()))
            .create();

        let url = Url::parse(&mockito::server_url())?.join("test")?;
        let res = Request::post(url).body(body_bytes).send()?;

        mock.assert();
        assert_eq!(res.head.status, StatusCode::CREATED);

        Ok(())
    }

    #[test]
    fn test_put() -> Result<()> {
        let body = "Hello, world!";

        let mock = mock("PUT", "/test")
            .with_status(201)
            .match_header("Content-Type", "text/plain")
            .match_body(Matcher::Exact(body.into()))
            .create();

        let url = Url::parse(&mockito::server_url())?.join("test")?;
        let res = Request::put(url)
            .header("Content-Type", "text/plain")
            .body(body.as_bytes())
            .send()?;

        mock.assert();
        assert_eq!(res.head.status, StatusCode::CREATED);

        Ok(())
    }

    #[test]
    fn test_json() -> Result<()> {
        let body = json!({
            "foo": "bar",
            "hello": "world"
        });

        let mock = mock("POST", "/test")
            .with_status(201)
            .match_header("Content-Type", "application/json")
            .match_body(Matcher::Json(body.clone()))
            .create();

        let url = Url::parse(&mockito::server_url())?.join("test")?;
        let res = Request::post(url).json(&body)?.send()?;

        mock.assert();
        assert_eq!(res.head.status, StatusCode::CREATED);

        Ok(())
    }

    #[test]
    fn test_cbor() -> Result<()> {
        #[derive(Serialize)]
        struct Body<'a> {
            foo: &'a str,
            hello: &'a str,
        }

        let body = Body {
            foo: "bar",
            hello: "world",
        };

        let mock = mock("POST", "/test")
            .with_status(201)
            .match_header("Content-Type", "application/cbor")
            // As of v0.25, mockito doesn't support matching binary bodies.
            .match_body(Matcher::Any)
            .create();

        let url = Url::parse(&mockito::server_url())?.join("test")?;
        let res = Request::post(url).cbor(&body)?.send()?;

        mock.assert();
        assert_eq!(res.head.status, StatusCode::CREATED);

        Ok(())
    }

    #[test]
    fn test_accept_encoding() -> Result<()> {
        let mock = mock("GET", "/test")
            .with_status(200)
            .match_header("Accept-Encoding", "zstd, gzip, foobar")
            .create();

        let encodings = vec![
            Encoding::Zstd,
            Encoding::Gzip,
            Encoding::Other("foobar".into()),
        ];

        let url = Url::parse(&mockito::server_url())?.join("test")?;
        let _ = Request::get(url).accept_encoding(encodings).send()?;

        mock.assert();
        Ok(())
    }

    const DUMMY_URL_STR: &str = "https://a.example.com/b";
    const DUMMY_METOD: Method = Method::Get;

    impl RequestContext {
        /// Dummy RequestContext for testing.
        pub(crate) fn dummy() -> Self {
            Self::new(Url::parse(DUMMY_URL_STR).unwrap(), DUMMY_METOD)
        }
    }

    #[test]
    fn test_request_context() {
        let req = RequestContext::dummy();
        assert_eq!(req.url().as_str(), DUMMY_URL_STR);
        assert_eq!(req.method(), &DUMMY_METOD);

        let req2 = RequestContext::dummy();
        assert_ne!(req.id(), req2.id());
    }

    #[test]
    fn test_request_callback() -> Result<()> {
        let called = Arc::new(AtomicUsize::new(0));
        Request::on_new_request({
            let called = called.clone();
            move |req| {
                // The callback can receive requests in other tests.
                // So we need to check the request is sent by this test.
                if req.ctx().url().path() == "/test_callback" {
                    called.fetch_add(1, AcqRel);
                }
            }
        });

        let mock = mock("HEAD", "/test_callback").with_status(200).create();
        let url = Url::parse(&mockito::server_url())?.join("test_callback")?;
        let _res = Request::head(url).send()?;

        mock.assert();
        assert_eq!(called.load(Acquire), 1);

        Ok(())
    }

    #[test]
    fn test_convert_cert_flag() -> Result<()> {
        let client = HttpClient::new();
        let url: Url = "https://example.com".parse()?;

        // Make sure convert_cert defaults to cfg!(windows) and gets
        // passed along to request.
        assert_eq!(cfg!(windows), client.get(url.clone()).convert_cert);

        let client = HttpClient::from_config(Config {
            convert_cert: true,
            ..Default::default()
        });
        assert!(client.get(url).convert_cert);

        Ok(())
    }
}


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

use std::io::Cursor;
use std::pin::Pin;

use anyhow::anyhow;
use async_compression::tokio::bufread::BrotliDecoder;
use async_compression::tokio::bufread::DeflateDecoder;
use async_compression::tokio::bufread::GzipDecoder;
use async_compression::tokio::bufread::ZstdDecoder;
use futures::prelude::*;
use http::header;
use http::header::HeaderMap;
use http::status::StatusCode;
use http::version::Version;
use serde::de::DeserializeOwned;
use tokio::io::BufReader;
use tokio_util::io::ReaderStream;
use tokio_util::io::StreamReader;
use url::Url;

use crate::errors::HttpClientError;
use crate::handler::Buffered;
use crate::handler::HandlerExt;
use crate::header::Header;
use crate::receiver::ResponseStreams;
use crate::request::Encoding;
use crate::request::RequestInfo;
use crate::stream::BufferedStream;
use crate::stream::CborStream;

#[derive(Debug)]
pub struct Head {
    pub(crate) version: Version,
    pub(crate) status: StatusCode,
    pub(crate) headers: HeaderMap,
    pub(crate) request_info: RequestInfo,
}

impl Head {
    /// Get the HTTP version of the response.
    pub fn version(&self) -> Version {
        self.version
    }

    /// Get the HTTP status code of the response.
    pub fn status(&self) -> StatusCode {
        self.status
    }

    /// Get the response's headers.
    pub fn headers(&self) -> &HeaderMap {
        &self.headers
    }

    /// Get metadata about the response's corresponding HTTP request.
    pub fn request_info(&self) -> &RequestInfo {
        &self.request_info
    }

    /// Get the response's encoding from the Content-Encoding header.
    pub fn encoding(&self) -> Result<Encoding, HttpClientError> {
        self.headers
            .get(header::CONTENT_ENCODING)
            .map(|encoding| Ok(encoding.to_str()?.into()))
            .unwrap_or(Ok(Encoding::Identity))
            .map_err(|_: header::ToStrError| {
                HttpClientError::BadResponse(anyhow!("Invalid Content-Encoding"))
            })
    }
}

#[derive(Debug)]
pub struct Response {
    pub(crate) head: Head,
    pub(crate) body: Vec<u8>,
}

impl Response {
    /// Get the HTTP version of the response.
    pub fn version(&self) -> Version {
        self.head.version
    }

    /// Get the HTTP status code of the response.
    pub fn status(&self) -> StatusCode {
        self.head.status
    }

    /// Get the response's headers.
    pub fn headers(&self) -> &HeaderMap {
        &self.head.headers
    }

    /// Get the response's body.
    pub fn body(&self) -> &[u8] {
        &self.body
    }

    /// Split the
    pub fn into_parts(self) -> (Head, Vec<u8>) {
        (self.head, self.body)
    }

    /// Deserialize the response body from JSON.
    pub fn json<T: DeserializeOwned>(&self) -> Result<T, serde_json::Error> {
        serde_json::from_slice(&self.body)
    }

    /// Deserialize the response body from CBOR.
    pub fn cbor<T: DeserializeOwned>(&self) -> Result<T, serde_cbor::Error> {
        serde_cbor::from_slice(&self.body)
    }
}

impl TryFrom<&mut Buffered> for Response {
    type Error = HttpClientError;

    fn try_from(buffered: &mut Buffered) -> Result<Self, Self::Error> {
        let (version, status) = match (buffered.version(), buffered.status()) {
            (Some(version), Some(status)) => (version, status),
            _ => {
                return Err(HttpClientError::BadResponse(anyhow!(
                    "HTTP version or status code missing in response"
                )));
            }
        };

        Ok(Self {
            head: Head {
                version,
                status,
                headers: buffered.take_headers(),
                request_info: buffered.request_context().info.clone(),
            },
            body: buffered.take_body(),
        })
    }
}

macro_rules! decode {
    ($decoder:tt, $body_stream:expr) => {{
        let body = $body_stream
            .map_ok(Cursor::new)
            .map_err(|e| futures::io::Error::new(futures::io::ErrorKind::Other, e));
        let reader = BufReader::new(StreamReader::new(body));
        ReaderStream::new($decoder::new(reader))
            .map_ok(|bytes| bytes.to_vec())
            .map_err(|e| {
                if e.get_ref()
                    .map(|i| i.is::<HttpClientError>())
                    .unwrap_or(false)
                {
                    *e.into_inner()
                        .unwrap()
                        .downcast::<HttpClientError>()
                        .unwrap()
                } else {
                    HttpClientError::DecompressionFailed(e)
                }
            })
            .boxed()
    }};
}

pub type ByteStream =
    Pin<Box<dyn Stream<Item = Result<Vec<u8>, HttpClientError>> + Send + 'static>>;

pub struct AsyncBody {
    // This is a `Result` so that an invalid Content-Encoding header does not prevent the caller
    // from accessing the raw body stream if desired. The error will only be propagated if the
    // caller actually wants to decode the body stream.
    encoding: Result<Encoding, HttpClientError>,
    body: ByteStream,
}

pub type CborStreamBody<T> = CborStream<T, ByteStream, Vec<u8>, HttpClientError>;
pub type BufferedStreamBody = BufferedStream<ByteStream, Vec<u8>, HttpClientError>;

impl AsyncBody {
    /// Get a stream of the response's body content.
    ///
    /// This method is the preferred way of accessing the response's body stream. The data will be
    /// automatically decoded based on the response's Content-Encoding header.
    pub fn decoded(self) -> ByteStream {
        let Self { encoding, body } = self;
        stream::once(async move {
            Ok(match encoding? {
                Encoding::Identity => body.boxed(),
                Encoding::Brotli => decode!(BrotliDecoder, body),
                Encoding::Deflate => decode!(DeflateDecoder, body),
                Encoding::Gzip => decode!(GzipDecoder, body),
                Encoding::Zstd => decode!(ZstdDecoder, body),
                other => {
                    return Err(HttpClientError::BadResponse(anyhow!(
                        "Unsupported Content-Encoding: {:?}",
                        other
                    )));
                }
            })
        })
        .try_flatten()
        .boxed()
    }

    /// Get a stream of the response's raw on-the-wire content.
    ///
    /// Note that the caller is responsible for decoding the response if it is compressed. Most
    /// callers will want to use the `decoded` method instead which does this automatically.
    pub fn raw(self) -> ByteStream {
        self.body
    }

    /// Attempt to deserialize the incoming data as a stream of CBOR values.
    pub fn cbor<T: DeserializeOwned>(self) -> CborStreamBody<T> {
        CborStream::new(self.decoded())
    }

    /// Attempt to deserialize the incoming data as a stream of CBOR values.
    pub fn cbor_with_buffer_size<T: DeserializeOwned>(self, size: usize) -> CborStreamBody<T> {
        CborStream::with_buffer_size(self.decoded(), size)
    }

    /// Create a buffered body stream that ensures that all yielded chunks
    /// (except the last) are at least as large as the given chunk size.
    pub fn buffered(self, size: usize) -> BufferedStreamBody {
        BufferedStream::new(self.decoded(), size)
    }
}

pub struct AsyncResponse {
    pub(crate) head: Head,
    pub(crate) body: AsyncBody,
}

impl AsyncResponse {
    pub(crate) async fn new(
        streams: ResponseStreams,
        request_info: RequestInfo,
    ) -> Result<Self, HttpClientError> {
        let ResponseStreams {
            headers_rx,
            body_rx,
            done_rx,
        } = streams;

        let header_lines = headers_rx
            .take_while(|h| future::ready(h != &Header::EndOfHeaders))
            .collect::<Vec<_>>()
            .await;

        let mut version = None;
        let mut status = None;
        let mut headers = HeaderMap::new();

        for line in header_lines {
            match line {
                Header::Status(v, s) => {
                    version = Some(v);
                    status = Some(s);
                }
                Header::Header(k, v) => {
                    headers.insert(k, v);
                }
                Header::EndOfHeaders => unreachable!(),
            }
        }

        let (version, status) = match (version, status) {
            (Some(v), Some(s)) => (v, s),
            _ => {
                // If we didn't get a status line, we most likely
                // failed to connect to the server at all. In this
                // case, we should expect to receive an error in
                // the `done_rx` stream. Even if not, a response
                // without a status line is invalid so we should
                // fail regardless.
                done_rx.await??;
                return Err(HttpClientError::BadResponse(anyhow!(
                    "HTTP version or status code missing in response"
                )));
            }
        };

        let body = body_rx
            .map(Ok)
            .chain(stream::once(async move {
                done_rx.await??;
                Ok(Vec::new())
            }))
            .try_filter(|chunk| future::ready(!chunk.is_empty()))
            .boxed();

        let head = Head {
            version,
            status,
            headers,
            request_info,
        };
        let encoding = head.encoding();
        let body = AsyncBody { encoding, body };

        Ok(Self { head, body })
    }

    /// Get the HTTP version of the response.
    pub fn version(&self) -> Version {
        self.head.version
    }

    /// Get the HTTP status code of the response.
    pub fn status(&self) -> StatusCode {
        self.head.status
    }

    /// Get the response's headers.
    pub fn headers(&self) -> &HeaderMap {
        &self.head.headers
    }

    /// Consume the response and obtain its body stream.
    pub fn into_body(self) -> AsyncBody {
        self.body
    }

    /// Split the response into its head and body.
    pub fn into_parts(self) -> (Head, AsyncBody) {
        (self.head, self.body)
    }

    pub fn url(&self) -> &Url {
        self.head.request_info.url()
    }
}

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

    use anyhow::Result;
    use futures::TryStreamExt;
    use mockito::mock;
    use url::Url;

    use crate::request::Encoding;
    use crate::request::Request;

    #[tokio::test]
    async fn test_decompression() -> Result<()> {
        let uncompressed = b"Hello, world!";
        let compressed = zstd::encode_all(Cursor::new(uncompressed), 0)?;

        let mock = mock("GET", "/test")
            .match_header("Accept-Encoding", "zstd, br, gzip, deflate")
            .with_status(200)
            .with_header("Content-Encoding", "zstd")
            .with_body(compressed)
            .create();

        let url = Url::parse(&mockito::server_url())?.join("test")?;
        let res = Request::get(url)
            .accept_encoding(Encoding::all())
            .send_async()
            .await?;

        mock.assert();

        let body = res.into_body().decoded().try_concat().await?;
        assert_eq!(&*body, &uncompressed[..]);

        Ok(())
    }
}


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

use std::fmt;
use std::mem;
use std::time::Duration;

#[derive(Clone, Debug, Default, Eq, PartialEq)]
pub struct Stats {
    pub downloaded: usize,
    pub uploaded: usize,
    pub requests: usize,
    pub time: Duration,
    pub latency: Duration,
}

impl Stats {
    pub fn time_in_seconds(&self) -> f64 {
        self.time.as_secs_f64()
    }

    pub fn bytes_per_second(&self) -> f64 {
        self.downloaded as f64 / self.time_in_seconds()
    }
}

impl fmt::Display for Stats {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(
            f,
            "Downloaded {amount} in {time:.time_prec$?} over {requests} \
            request{plural} ({rate}, latency: {latency:.latency_prec$?})",
            amount = byte_count(self.downloaded),
            time = self.time,
            time_prec = if self.time.as_secs() == 0 { 0 } else { 2 },
            requests = self.requests,
            plural = if self.requests == 1 { "" } else { "s" },
            rate = bit_rate(self.bytes_per_second() * 8.0),
            latency = self.latency,
            latency_prec = if self.latency.as_secs() == 0 { 0 } else { 2 },
        )
    }
}

/// Format a byte count using SI binary prefixes.
fn byte_count(value: usize) -> String {
    if value == 0 {
        return "0 B".into();
    }

    // Compute the base-1024 log of the value (i.e., log2(value) / log2(1024)).
    let log = ((8 * mem::size_of::<usize>()) - value.leading_zeros() as usize - 1) / 10;

    // Shift value down. (Use floating-point division to preserve decimals.)
    let shifted = value as f64 / (1 << (log * 10)) as f64;

    // Determine unit and precision to display.
    let unit = ["B", "kiB", "MiB", "GiB", "TiB", "PiB", "EiB"][log];
    let prec = if log > 1 { 2 } else { 0 };

    format!("{:.*} {}", prec, shifted, unit)
}

/// Format a bit rate using decimal prefixes.
fn bit_rate(rate: f64) -> String {
    // Guard against zero, NaN, infinity, etc.
    if !rate.is_normal() {
        return "0 b/s".into();
    }

    // Divide by the base-1000 log of the value to bring it under 1000.
    let log = (rate.log10() / 3.0).floor() as usize;
    let shifted = rate / 1000f64.powi(log as i32);

    // Determine unit and precision to display.
    let unit = ["b/s", "kb/s", "Mb/s", "Gb/s", "Tb/s", "Pb/s", "Eb/s"][log];
    let prec = if log > 1 { 2 } else { 0 };

    format!("{:.*} {}", prec, shifted, unit)
}

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

    #[test]
    fn test_stats_formatting() {
        let stats = Stats {
            downloaded: 10 * 1024 * 1024 + 600 * 1024, // 10.586 MiB
            uploaded: 1024,
            requests: 5,
            time: Duration::from_millis(12345),
            latency: Duration::from_micros(123456),
        };

        let expected = "Downloaded 10.59 MiB in 12.35s over 5 requests (7.19 Mb/s, latency: 123ms)";
        assert_eq!(expected, &stats.to_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.
 */

use std::cmp;
use std::marker::PhantomData;
use std::mem;
use std::pin::Pin;

use futures::prelude::*;
use futures::ready;
use futures::task::Context;
use futures::task::Poll;
use pin_project::pin_project;
use serde::de::DeserializeOwned;
use serde_cbor::Deserializer;
use thiserror::Error;

#[derive(Error, Debug)]
pub enum CborStreamError {
    #[error(transparent)]
    CborError(serde_cbor::Error),

    #[error("Unexpected trailing data found at end of CBOR stream ({0} bytes)")]
    TrailingData(usize),
}

const MEDIUM_BUFFER_SIZE: usize = 1024 * 1024;
const SMALL_BUFFER_SIZE: usize = 1024 * 8;

/// A wrapper around a `TryStream` of bytes that will attempt to deserialize
/// CBOR-encoded values from the data stream as it is received.
///
/// Data from the underlying stream will be buffered until a sufficient amount
/// has accumulated before attempting deserialization. This is important for
/// efficiency since libcurl will typically return a stream of small packets.
#[pin_project]
#[must_use = "streams do nothing unless polled"]
pub struct CborStream<T, S, B, E> {
    #[pin]
    incoming: BufferedStream<S, B, E>,
    buffer: Vec<u8>,
    threshold: usize,
    position: usize,
    terminated: bool,
    _phantom: PhantomData<(T, B, E)>,
}

impl<T, S, B, E> CborStream<T, S, B, E> {
    pub(crate) fn new(body: S) -> Self {
        Self::with_buffer_size(body, SMALL_BUFFER_SIZE)
    }

    pub(crate) fn with_buffer_size(body: S, size: usize) -> Self {
        Self {
            incoming: BufferedStream::new(body, size),
            buffer: Vec::new(),
            threshold: size,
            position: 0,
            terminated: false,
            _phantom: PhantomData,
        }
    }
}

impl<T, S, B, E> Stream for CborStream<T, S, B, E>
where
    T: DeserializeOwned,
    S: Stream<Item = Result<B, E>> + Send + 'static,
    B: AsRef<[u8]>,
    E: From<CborStreamError>,
{
    type Item = Result<T, E>;

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

        if *this.terminated {
            return Poll::Ready(None);
        }

        loop {
            // Attempt to deserialize a single item from the buffer.
            let mut de = Deserializer::from_slice(&this.buffer[*this.position..]);
            match T::deserialize(&mut de) {
                Ok(value) => {
                    *this.position += de.byte_offset();
                    // Reset the buffer threshold in case we had expanded it to fit a large value.
                    // If we indeed had to expand it, do not shrink it to smaller size than the MEDIUM_BUFFER_SIZE since
                    // it could be traffic pattern rather than a single large item.
                    if this.incoming.threshold() > (*this.threshold) {
                        this.incoming
                            .as_mut()
                            .set_threshold(cmp::max(*this.threshold, MEDIUM_BUFFER_SIZE));
                    }
                    return Poll::Ready(Some(Ok(value)));
                }
                Err(e) if !e.is_eof() => {
                    *this.terminated = true;
                    let e = CborStreamError::CborError(e);
                    return Poll::Ready(Some(Err(E::from(e))));
                }
                _ => {}
            }

            let pos = *this.position;
            if pos > 0 {
                // At this point we've deserialized most of the data in the buffer. Any remaining
                // trailing data is the prefix of a single, incomplete item. We should move it to
                // the front of the buffer to reclaim the space from the items that have already
                // been deserialized.
                let len = this.buffer.len() - pos;
                for i in 0..len {
                    this.buffer[i] = this.buffer[pos + i];
                }
                this.buffer.truncate(len);
                *this.position = 0;
            } else if this.buffer.len() >= this.incoming.threshold() {
                // If we get here, that means that we're attempting to deserialize a single value
                // that is larger than the current buffer size. If the value is significantly larger
                // than the current buffer size, this will result in accidentally quadratic behavior
                // as we repeatedly attempt to deserialize the partial value whenever a new chunk
                // comes in.
                //
                // To prevent this situation, whenever we encounter an item that exceeds the current
                // buffer size, we simply double it. This means that we'll only need to do O(log(n))
                // deserialization attempts for very large values.
                //
                // We will skip doubling small buffers, and go strait to 1MB buffer size.
                let new_threshold = cmp::max(MEDIUM_BUFFER_SIZE, 2 * this.incoming.threshold());
                this.incoming.as_mut().set_threshold(new_threshold);
            }

            // Poll the underlying stream for more incoming data.
            match ready!(this.incoming.as_mut().poll_next(cx)) {
                Some(Ok(chunk)) => this.buffer.extend_from_slice(chunk.as_ref()),
                Some(Err(e)) => {
                    return Poll::Ready(Some(Err(e)));
                }
                None => {
                    // At this point the stream is complete, so we expect to have read everything.
                    // If we haven't, then something went wrong during the transfer and the data
                    // we're handling seems corrupted: raise an error in that case.
                    if !this.buffer.is_empty() {
                        let e = CborStreamError::TrailingData(this.buffer.len());
                        return Poll::Ready(Some(Err(E::from(e))));
                    }
                    return Poll::Ready(None);
                }
            }
        }
    }
}

/// A wrapper around a byte stream that buffers the data until it exceeds
/// the given size threshold. This can improve the efficiency of processing
/// data from bytes streams that consist largely of small chunks.
#[pin_project]
#[must_use = "streams do nothing unless polled"]
pub struct BufferedStream<S, B, E> {
    #[pin]
    inner: S,
    buffer: Vec<u8>,
    threshold: usize,
    done: bool,
    _phantom: PhantomData<(B, E)>,
}

impl<S, B, E> BufferedStream<S, B, E> {
    pub(crate) fn new(inner: S, n: usize) -> Self {
        Self {
            inner,
            buffer: Vec::with_capacity(n),
            threshold: n,
            done: false,
            _phantom: PhantomData,
        }
    }

    fn threshold(&self) -> usize {
        self.threshold
    }

    /// Update the buffer size in the middle of polling the stream.
    fn set_threshold(self: Pin<&mut Self>, threshold: usize) {
        *self.project().threshold = threshold;
    }
}

impl<S, B, E> Stream for BufferedStream<S, B, E>
where
    S: Stream<Item = Result<B, E>> + Send + 'static,
    B: AsRef<[u8]>,
{
    type Item = Result<Vec<u8>, E>;

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

        if *this.done {
            return Poll::Ready(None);
        }

        while this.buffer.len() < *this.threshold {
            match ready!(this.inner.as_mut().poll_next(cx)) {
                Some(Ok(bytes)) => this.buffer.extend_from_slice(bytes.as_ref()),
                Some(Err(e)) => return Poll::Ready(Some(Err(e))),
                None => {
                    *this.done = true;
                    break;
                }
            }
        }

        let buf = mem::replace(this.buffer, Vec::with_capacity(*this.threshold));
        Poll::Ready(Some(Ok(buf)))
    }
}

#[cfg(test)]
mod tests {
    use anyhow::anyhow;
    use anyhow::Result;
    use serde::Deserialize;
    use serde::Serialize;

    use super::*;

    #[derive(Clone, Debug, Eq, PartialEq, Deserialize, Serialize)]
    struct TestItem(String);

    impl TestItem {
        fn new(s: &str) -> Self {
            TestItem(s.into())
        }
    }

    #[tokio::test]
    async fn test_single_item() -> Result<()> {
        let item = TestItem::new("hello");
        let bytes = serde_cbor::to_vec(&item)?;

        let byte_stream = stream::once(async {
            let res: Result<_> = Ok(bytes); // Need to assign for type hint.
            res
        });
        let mut cbor_stream = Box::pin(CborStream::new(byte_stream));

        let res = cbor_stream.try_next().await?;
        assert_eq!(res, Some(item));

        Ok(())
    }

    #[tokio::test]
    async fn test_single_error() -> Result<()> {
        let byte_stream = stream::once(async {
            let res: Result<Vec<u8>> = Err(anyhow!("error")); // Need to assign for type hint.
            res
        });
        let mut cbor_stream = Box::pin(CborStream::new(byte_stream));

        let res: Result<Option<TestItem>> = cbor_stream.try_next().await;
        assert!(res.is_err());

        Ok(())
    }

    #[tokio::test]
    async fn test_complete_items() -> Result<()> {
        let items = vec![
            TestItem::new("foo"),
            TestItem::new("bar"),
            TestItem::new("baz"),
        ];

        let incoming = items
            .clone()
            .into_iter()
            .map(|v| -> Result<Vec<u8>> { Ok(serde_cbor::to_vec(&v)?) });
        let cbor_stream = Box::pin(CborStream::new(stream::iter(incoming)));

        let res: Vec<TestItem> = cbor_stream.try_collect().await?;
        assert_eq!(res, items);

        Ok(())
    }

    #[tokio::test]
    async fn test_mid_stream_error() -> Result<()> {
        let items = vec![TestItem::new("foo"), TestItem::new("bar")];

        let bytes_and_errors = vec![
            Ok(serde_cbor::to_vec(&items[0])?),
            Err(anyhow!("mid-stream error")),
            Ok(serde_cbor::to_vec(&items[1])?),
        ];
        let mut cbor_stream = Box::pin(CborStream::new(stream::iter(bytes_and_errors)));

        let res1: Result<Option<TestItem>> = cbor_stream.try_next().await;
        let res2: Result<Option<TestItem>> = cbor_stream.try_next().await;
        let res3: Result<Option<TestItem>> = cbor_stream.try_next().await;

        // Due to internal buffering, the error will be returned first
        // even though it was second in the input stream.
        assert!(res1.is_err());
        assert_eq!(res2?, Some(items[0].clone()));
        assert_eq!(res3?, Some(items[1].clone()));

        Ok(())
    }

    #[tokio::test]
    async fn test_partial_item() -> Result<()> {
        let items = vec![TestItem::new("hello"), TestItem::new("world")];

        let mut incoming: Vec<Result<Vec<u8>>> = Vec::new();
        for i in &items {
            let bytes = serde_cbor::to_vec(&i)?;

            let mid = bytes.len() / 2;
            let prefix = bytes[..mid].to_vec();
            let suffix = bytes[mid..].to_vec();

            incoming.push(Ok(prefix));
            incoming.push(Ok(suffix));
        }
        let cbor_stream = Box::pin(CborStream::new(stream::iter(incoming)));

        let res: Vec<TestItem> = cbor_stream.try_collect().await?;
        assert_eq!(res, items);

        Ok(())
    }

    #[tokio::test]
    async fn test_concatenated_items() -> Result<()> {
        let items = vec![
            TestItem::new("test1"),
            TestItem::new("test2"),
            TestItem::new("test3"),
            TestItem::new("test4"),
            TestItem::new("test5"),
        ];

        let mut concat = Vec::new();
        for i in &items {
            concat.extend(serde_cbor::to_vec(&i)?);
        }

        let mid = concat.len() / 2;
        let chunk1 = concat[..mid].to_vec();
        let chunk2 = concat[mid..].to_vec();
        let chunks: Vec<Result<Vec<u8>>> = vec![Ok(chunk1), Ok(chunk2)];
        let cbor_stream = Box::pin(CborStream::new(stream::iter(chunks)));

        let res: Vec<TestItem> = cbor_stream.try_collect().await?;
        assert_eq!(res, items);

        Ok(())
    }

    #[tokio::test]
    async fn test_truncated_data() -> Result<()> {
        // Add a valid chunk, but truncate it, though in practice, the truncated data could be
        // stuff that looks like it could be CBOR but doesn't decode properly (like an error
        // string).
        let mut chunk = serde_cbor::to_vec(&TestItem::new("test1"))?;
        chunk.pop();
        let len = chunk.len();

        let chunks: Vec<Result<Vec<u8>, CborStreamError>> = vec![Ok(chunk)];
        let mut cbor_stream = Box::pin(CborStream::new(stream::iter(chunks)));

        let res1: Result<Option<TestItem>, _> = cbor_stream.try_next().await;

        match res1 {
            Err(CborStreamError::TrailingData(l)) if l == len => {}
            other => panic!("Unexpected result on trailing data: {:?}", other),
        };

        Ok(())
    }

    #[tokio::test]
    async fn test_buffered_stream() -> Result<()> {
        let words: Vec<Result<Vec<u8>>> = vec![
            Ok(b"The ".to_vec()),
            Ok(b"quick ".to_vec()),
            Ok(b"brown ".to_vec()),
            Ok(b"fox ".to_vec()),
            Ok(b"jumped ".to_vec()),
            Ok(b"over ".to_vec()),
            Ok(b"the ".to_vec()),
            Ok(b"lazy ".to_vec()),
            Ok(b"dog".to_vec()),
        ];

        let input = stream::iter(words);
        let buf_size = 20;
        let buffered = Box::pin(BufferedStream::new(input, buf_size));
        let output = buffered.try_collect::<Vec<_>>().await?;

        assert_eq!(output.len(), 3);

        // Exactly 20 bytes long.
        assert_eq!(&*output[0], &b"The quick brown fox "[..]);

        // 21 bytes long because last chunk pushed buffer over limit.
        assert_eq!(&*output[1], &b"jumped over the lazy "[..]);

        // Short last item.
        assert_eq!(&*output[2], &b"dog"[..]);

        Ok(())
    }
}


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

use std::env;
use std::io::prelude::*;
use std::io::stdin;
use std::io::stdout;
use std::io::IsTerminal;

use anyhow::Result;
use futures::prelude::*;
use http_client::AsyncResponse;
use http_client::HttpClient;
use http_client::Request;
use structopt::StructOpt;
use url::Url;

const CERT_ENV_VAR: &str = "CERT";
const KEY_ENV_VAR: &str = "KEY";
const CA_ENV_VAR: &str = "CA";

#[derive(Debug, StructOpt)]
#[structopt(name = "http_cli", about = "Send HTTP requests")]
enum Method {
    #[structopt(about = "Send a GET request")]
    Get(Args),
    #[structopt(about = "Send a HEAD request")]
    Head(Args),
    #[structopt(about = "Send a POST request")]
    Post(Args),
    #[structopt(about = "Send a PUT request")]
    Put(Args),
}

#[derive(Debug, StructOpt)]
struct Args {
    #[structopt(help = "Request URL")]
    url: String,
    #[structopt(
        long,
        short = "H",
        help = "Headers as a series of \"Header-Name: Value\" strings"
    )]
    headers: Vec<String>,
}

impl Args {
    fn url(&self) -> Result<Url> {
        Ok(Url::parse(&self.url)?)
    }
}

/// Note that we technically don't need to use async here;
/// `Request::send()` would have sufficed. However, the
/// purpose of this binary is primarily for testing and
/// debugging the library itself, so by using async we
/// can maximize the surface area exercised.
#[tokio::main]
async fn main() -> Result<()> {
    env_logger::init();
    match Method::from_args() {
        Method::Get(args) => cmd_get(args).await,
        Method::Head(args) => cmd_head(args).await,
        Method::Post(args) => cmd_post(args).await,
        Method::Put(args) => cmd_put(args).await,
    }
}

async fn cmd_get(args: Args) -> Result<()> {
    let req = HttpClient::new().get(args.url()?);
    let req = add_headers(req, &args.headers);
    let req = configure_tls(req);

    let res = req.send_async().await?;
    write_response(res).await
}

async fn cmd_head(args: Args) -> Result<()> {
    let req = HttpClient::new().head(args.url()?);
    let req = add_headers(req, &args.headers);
    let req = configure_tls(req);

    let res = req.send_async().await?;
    write_response(res).await
}

async fn cmd_post(args: Args) -> Result<()> {
    eprintln!("Reading payload from stdin");
    let body = read_input()?;

    let req = HttpClient::new().post(args.url()?).body(body);
    let req = add_headers(req, &args.headers);
    let req = configure_tls(req);

    let res = req.send_async().await?;
    write_response(res).await
}

async fn cmd_put(args: Args) -> Result<()> {
    eprintln!("Reading payload from stdin");
    let body = read_input()?;

    let req = HttpClient::new().put(args.url()?).body(body);
    let req = add_headers(req, &args.headers);
    let req = configure_tls(req);

    let res = req.send_async().await?;
    write_response(res).await
}

fn read_input() -> Result<Vec<u8>> {
    let mut buf = Vec::new();
    stdin().read_to_end(&mut buf)?;
    Ok(buf)
}

async fn write_response(res: AsyncResponse) -> Result<()> {
    eprintln!("Status: {:?} {}", res.version(), res.status());
    eprintln!("{:?}", res.headers());

    let body = res.into_body().decoded().try_concat().await?;

    if stdout().is_terminal() {
        println!("{}", String::from_utf8_lossy(&body).escape_default())
    } else {
        stdout().write_all(&body)?;
    };

    Ok(())
}

fn configure_tls(mut req: Request) -> Request {
    if let Ok(cert) = env::var(CERT_ENV_VAR) {
        req = req.cert(cert);
    }
    if let Ok(key) = env::var(KEY_ENV_VAR) {
        req = req.key(key);
    }
    if let Ok(ca) = env::var(CA_ENV_VAR) {
        req = req.cainfo(ca);
    }
    req
}

fn add_headers(mut req: Request, headers: &[String]) -> Request {
    for header in headers {
        let (name, value) = split_header(header);
        req = req.header(name, value);
    }
    req
}

fn split_header(header: &str) -> (&str, &str) {
    let parts = header.splitn(2, ':').collect::<Vec<_>>();
    if parts.len() > 1 {
        (parts[0], parts[1].trim_start())
    } else {
        (parts[0], "")
    }
}


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

use std::io::Read;
use std::io::SeekFrom;
use std::mem;

use curl::easy::Handler;
use curl::easy::ReadError;
use curl::easy::SeekResult;
use curl::easy::WriteError;
use http::header;
use http::HeaderMap;
use http::StatusCode;
use http::Version;

use super::HandlerExt;
use crate::header::Header;
use crate::progress::Progress;
use crate::RequestContext;

/// Initial buffer capacity to allocate if we don't get a Content-Length header.
/// Usually, the lack of a Content-Length header indicates a streaming response,
/// in which case the body size is expected to be relatively large.
const DEFAULT_CAPACITY: usize = 1000;

/// A simple curl Handler that buffers all received data.
pub struct Buffered {
    received: Vec<u8>,
    capacity: Option<usize>,
    version: Option<Version>,
    status: Option<StatusCode>,
    headers: HeaderMap,
    bytes_sent: usize,
    is_active: bool,
    request_context: RequestContext,
}

impl Buffered {
    pub(crate) fn new(request_context: RequestContext) -> Self {
        Self {
            received: Default::default(),
            capacity: Default::default(),
            version: Default::default(),
            status: Default::default(),
            headers: Default::default(),
            bytes_sent: Default::default(),
            is_active: false,
            request_context,
        }
    }

    pub(crate) fn version(&self) -> Option<Version> {
        self.version
    }

    pub(crate) fn status(&self) -> Option<StatusCode> {
        self.status
    }

    /// Extract the received headers.
    pub(crate) fn take_headers(&mut self) -> HeaderMap {
        mem::take(&mut self.headers)
    }

    /// Extract the received data.
    pub(crate) fn take_body(&mut self) -> Vec<u8> {
        mem::take(&mut self.received)
    }

    fn trigger_first_activity(&mut self) {
        if !self.is_active {
            self.is_active = true;
            let listeners = &self.request_context.event_listeners;
            listeners.trigger_first_activity(&self.request_context);
        }
    }
}

impl Handler for Buffered {
    fn write(&mut self, data: &[u8]) -> Result<usize, WriteError> {
        self.request_context
            .event_listeners
            .trigger_download_bytes(self.request_context(), data.len());
        // Set the buffer size based on the received Content-Length
        // header, or a default if we didn't get a Content-Length.
        self.received
            .reserve(self.capacity.unwrap_or(DEFAULT_CAPACITY));
        self.received.extend_from_slice(data);
        Ok(data.len())
    }

    fn read(&mut self, data: &mut [u8]) -> Result<usize, ReadError> {
        self.trigger_first_activity();

        Ok(if let Some(payload) = self.request_context.body.as_mut() {
            let sent = (&payload[self.bytes_sent..])
                .read(data)
                .expect("Failed to read from payload buffer");
            self.bytes_sent += sent;
            self.request_context
                .event_listeners
                .trigger_download_bytes(self.request_context(), sent);
            sent
        } else {
            0
        })
    }

    fn seek(&mut self, whence: SeekFrom) -> SeekResult {
        let size = match &self.request_context.body {
            Some(payload) => payload.len(),
            None => return SeekResult::CantSeek,
        };

        let (start, offset) = match whence {
            SeekFrom::Start(offset) => (0, offset as i64),
            SeekFrom::End(offset) => (size, offset),
            SeekFrom::Current(offset) => (self.bytes_sent, offset),
        };

        self.bytes_sent = if offset >= 0 {
            start.saturating_add(offset as usize).clamp(0, size)
        } else {
            start.saturating_sub(-offset as usize)
        };

        SeekResult::Ok
    }

    fn header(&mut self, data: &[u8]) -> bool {
        self.trigger_first_activity();

        match Header::parse(data) {
            Ok(Header::Header(name, value)) => {
                // XXX: This line triggers a lint error because `http::HeaderName`
                // is implemented using `bytes::Bytes` for custom headers, which has
                // interior mutability. There isn't anything we can really do here
                // since the problem is that the `http` crate declared this as `const`
                // instead of `static`. In this case, the lint error isn't actually
                // applicable since for standard headers like Content-Length, the
                // underlying representation is a simple enum, so there is actually
                // no interior mutability in use, so this line will not result in
                // any initialization code unintentionally running.
                #[allow(clippy::borrow_interior_mutable_const)]
                if name == header::CONTENT_LENGTH {
                    // Set the initial buffer capacity using the Content-Length header.
                    self.capacity = value.to_str().ok().and_then(|v| v.parse().ok());
                    if let Some(len) = self.capacity {
                        self.request_context
                            .event_listeners
                            .trigger_content_length(self.request_context(), len);
                    }
                }
                self.headers.insert(name, value);
            }
            Ok(Header::Status(version, code)) => {
                self.version = Some(version);
                self.status = Some(code);
            }
            Ok(Header::EndOfHeaders) => {}
            Err(e) => tracing::error!("{:?}", e),
        }
        true
    }

    fn progress(&mut self, dltotal: f64, dlnow: f64, ultotal: f64, ulnow: f64) -> bool {
        let listeners = &self.request_context.event_listeners;
        if listeners.should_trigger_progress() {
            let progress = Progress::from_curl(dltotal, dlnow, ultotal, ulnow);
            listeners.trigger_progress(self.request_context(), progress);
        }
        true
    }
}

impl HandlerExt for Buffered {
    fn request_context_mut(&mut self) -> &mut RequestContext {
        &mut self.request_context
    }

    fn request_context(&self) -> &RequestContext {
        &self.request_context
    }
}

#[cfg(test)]
mod tests {
    use assert_matches::assert_matches;
    use http::HeaderValue;

    use super::*;
    use crate::progress::ProgressReporter;

    #[test]
    fn test_read() {
        let data = [1, 2, 3, 4, 5, 6, 7, 8, 9, 0];

        let mut buf1 = [0xFF; 5];
        let mut buf2 = [0xFF; 3];
        let mut buf3 = [0xFF; 4];

        let mut handler = Buffered::new(RequestContext::dummy().body(data));

        assert_eq!(handler.read(&mut buf1[..]).unwrap(), 5);
        assert_eq!(handler.read(&mut buf2[..]).unwrap(), 3);
        assert_eq!(handler.read(&mut buf3[..]).unwrap(), 2);

        assert_eq!(&buf1[..], &[1, 2, 3, 4, 5][..]);
        assert_eq!(&buf2[..], &[6, 7, 8][..]);
        assert_eq!(&buf3[..], &[9, 0, 0xFF, 0xFF][..]);
    }

    #[test]
    fn test_write() {
        let mut handler = Buffered::new(RequestContext::dummy());

        let expected = [1, 2, 3, 4, 5, 6, 7, 8, 9, 0];

        assert_eq!(handler.write(&expected[..4]).unwrap(), 4);
        assert_eq!(handler.write(&expected[4..]).unwrap(), 6);

        let body = handler.take_body();
        assert_eq!(&expected[..], &*body);

        assert!(handler.take_body().is_empty());
    }

    #[test]
    fn test_seek() {
        let data = [1, 2, 3, 4, 5, 6, 7, 8, 9, 0];
        let mut handler = Buffered::new(RequestContext::dummy().body(data));

        assert_matches!(handler.seek(SeekFrom::Start(3)), SeekResult::Ok);
        assert_eq!(handler.bytes_sent, 3);

        assert_matches!(handler.seek(SeekFrom::End(-3)), SeekResult::Ok);
        assert_eq!(handler.bytes_sent, 7);

        assert_matches!(handler.seek(SeekFrom::End(20)), SeekResult::Ok);
        assert_eq!(handler.bytes_sent, 10);

        assert_matches!(handler.seek(SeekFrom::Current(-4)), SeekResult::Ok);
        assert_eq!(handler.bytes_sent, 6);

        assert_matches!(handler.seek(SeekFrom::Current(2)), SeekResult::Ok);
        assert_eq!(handler.bytes_sent, 8);

        assert_matches!(handler.seek(SeekFrom::Current(-20)), SeekResult::Ok);
        assert_eq!(handler.bytes_sent, 0);

        handler.request_context_mut().body = None;
        assert_matches!(handler.seek(SeekFrom::Current(0)), SeekResult::CantSeek);
    }

    #[test]
    fn test_headers() {
        let mut handler = Buffered::new(RequestContext::dummy());

        assert!(handler.header(&b"Content-Length: 1234\r\n"[..]));
        assert!(handler.header(&[0xFF, 0xFF][..])); // Invalid UTF-8 sequence.
        assert!(handler.header(&b"X-No-Value\r\n"[..]));

        let headers = handler.take_headers();

        assert_eq!(
            headers.get("Content-Length").unwrap(),
            HeaderValue::from_static("1234")
        );
        assert_eq!(
            headers.get("X-No-Value").unwrap(),
            HeaderValue::from_static("")
        );
    }

    #[test]
    fn test_capacity() {
        let mut handler = Buffered::new(RequestContext::dummy());
        let _ = handler.write(&[1, 2, 3][..]).unwrap();
        assert_eq!(handler.received.capacity(), DEFAULT_CAPACITY);

        let mut handler = Buffered::new(RequestContext::dummy());

        let _ = handler.header(&b"Content-Length: 42\r\n"[..]);
        assert_eq!(handler.capacity, Some(42));

        let _ = handler.write(&[1, 2, 3][..]).unwrap();
        assert_eq!(handler.received.capacity(), 42);
    }

    #[test]
    fn test_progress() {
        let reporter = ProgressReporter::with_callback(|_| ());

        let mut handler = Buffered::new(RequestContext::dummy());
        handler
            .request_context_mut()
            .event_listeners()
            .on_progress({
                let updater = reporter.updater();
                move |_req, p| updater.update(p)
            });
        let _ = handler.progress(1.0, 2.0, 3.0, 4.0);

        // Note that Progress struct has different argument order.
        let expected = Progress::new(2, 1, 4, 3);
        let progress = reporter.aggregate();
        assert_eq!(progress, expected);
    }
}


/*
 * 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 curl::easy::Handler;

use crate::RequestContext;

mod buffered;
mod streaming;

pub use buffered::Buffered;
pub use streaming::Streaming;

/// Extends `curl::Handler` with APIs useful for this crate.
///
/// Trait allowing a `curl::Handler` to be configured in a generic way.
/// All of the handlers used by the HTTP client need to implement this
/// trait so that they can be properly configured prior to use.
pub(crate) trait HandlerExt: Handler {
    /// Obtain the mutable `RequestContext` state.
    fn request_context_mut(&mut self) -> &mut RequestContext;

    /// Obtain the immutable `RequestContext` state.
    fn request_context(&self) -> &RequestContext;
}


/*
 * 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::Read;
use std::io::SeekFrom;

use curl::easy::Handler;
use curl::easy::ReadError;
use curl::easy::SeekResult;
use curl::easy::WriteError;

use super::HandlerExt;
use crate::header::Header;
use crate::progress::Progress;
use crate::receiver::Receiver;
use crate::RequestContext;

pub struct Streaming<R> {
    receiver: Option<R>,
    bytes_sent: usize,
    request_context: RequestContext,
    is_active: bool,
}

impl<R> Streaming<R> {
    pub(crate) fn new(receiver: R, request_context: RequestContext) -> Self {
        Self {
            receiver: Some(receiver),
            bytes_sent: 0,
            request_context,
            is_active: false,
        }
    }

    pub fn take_receiver(&mut self) -> Option<R> {
        self.receiver.take()
    }

    fn trigger_first_activity(&mut self) {
        if !self.is_active {
            self.is_active = true;
            let listeners = &self.request_context.event_listeners;
            listeners.trigger_first_activity(&self.request_context);
        }
    }
}

impl<R: Receiver> Handler for Streaming<R> {
    fn write(&mut self, data: &[u8]) -> Result<usize, WriteError> {
        self.request_context
            .event_listeners
            .trigger_download_bytes(self.request_context(), data.len());
        if let Some(ref mut receiver) = self.receiver {
            if receiver.chunk(data.into()).is_err() {
                // WriteError can only return "Pause", so instead we need to return an incorrect
                // number of bytes written, which will trigger curl to end the request.
                return Ok(0);
            }
        }
        Ok(data.len())
    }

    fn read(&mut self, data: &mut [u8]) -> Result<usize, ReadError> {
        self.trigger_first_activity();

        Ok(if let Some(payload) = self.request_context.body.as_mut() {
            let sent = (&payload[self.bytes_sent..])
                .read(data)
                .expect("Failed to read from payload buffer");
            self.bytes_sent += sent;
            self.request_context
                .event_listeners
                .trigger_download_bytes(self.request_context(), sent);
            sent
        } else {
            0
        })
    }

    fn seek(&mut self, whence: SeekFrom) -> SeekResult {
        let size = match &self.request_context.body {
            Some(payload) => payload.len(),
            None => return SeekResult::CantSeek,
        };

        let (start, offset) = match whence {
            SeekFrom::Start(offset) => (0, offset as i64),
            SeekFrom::End(offset) => (size, offset),
            SeekFrom::Current(offset) => (self.bytes_sent, offset),
        };

        self.bytes_sent = if offset >= 0 {
            start.saturating_add(offset as usize).clamp(0, size)
        } else {
            start.saturating_sub(-offset as usize)
        };

        SeekResult::Ok
    }

    fn header(&mut self, data: &[u8]) -> bool {
        self.trigger_first_activity();

        match Header::parse(data) {
            Ok(header) => {
                if let Some(ref mut receiver) = self.receiver {
                    if receiver.header(header).is_err() {
                        return false;
                    }
                }
            }
            Err(e) => {
                tracing::error!("{:?}", e);
            }
        }
        true
    }

    fn progress(&mut self, dltotal: f64, dlnow: f64, ultotal: f64, ulnow: f64) -> bool {
        let progress = Progress::from_curl(dltotal, dlnow, ultotal, ulnow);

        if let Some(ref mut receiver) = self.receiver {
            receiver.progress(progress);
        }

        let listeners = &self.request_context.event_listeners;
        if listeners.should_trigger_progress() {
            listeners.trigger_progress(self.request_context(), progress)
        }

        true
    }
}

impl<R: Receiver> HandlerExt for Streaming<R> {
    fn request_context_mut(&mut self) -> &mut RequestContext {
        &mut self.request_context
    }

    fn request_context(&self) -> &RequestContext {
        &self.request_context
    }
}

#[cfg(test)]
mod tests {
    use assert_matches::assert_matches;
    use http::header;
    use http::header::HeaderName;
    use http::header::HeaderValue;

    use super::*;
    use crate::progress::ProgressReporter;
    use crate::receiver::testutil::NullReceiver;
    use crate::receiver::testutil::TestReceiver;

    #[test]
    fn test_read() {
        let data = [1, 2, 3, 4, 5, 6, 7, 8, 9, 0];

        let mut buf1 = [0xFF; 5];
        let mut buf2 = [0xFF; 3];
        let mut buf3 = [0xFF; 4];

        let mut handler = Streaming::new(NullReceiver, RequestContext::dummy().body(data));

        assert_eq!(handler.read(&mut buf1[..]).unwrap(), 5);
        assert_eq!(handler.read(&mut buf2[..]).unwrap(), 3);
        assert_eq!(handler.read(&mut buf3[..]).unwrap(), 2);

        assert_eq!(&buf1[..], &[1, 2, 3, 4, 5][..]);
        assert_eq!(&buf2[..], &[6, 7, 8][..]);
        assert_eq!(&buf3[..], &[9, 0, 0xFF, 0xFF][..]);
    }

    #[test]
    fn test_write() {
        let receiver = TestReceiver::new();
        let mut handler = Streaming::new(receiver.clone(), RequestContext::dummy());

        let chunks = vec![vec![1, 2, 3], vec![5, 6], vec![7, 8, 9, 0]];

        assert_eq!(handler.write(&chunks[0]).unwrap(), 3);
        assert_eq!(handler.write(&chunks[1]).unwrap(), 2);
        assert_eq!(handler.write(&chunks[2]).unwrap(), 4);

        assert_eq!(receiver.chunks(), chunks);
    }

    #[test]
    fn test_seek() {
        let data = [1, 2, 3, 4, 5, 6, 7, 8, 9, 0];
        let mut handler = Streaming::new(NullReceiver, RequestContext::dummy().body(data));

        assert_matches!(handler.seek(SeekFrom::Start(3)), SeekResult::Ok);
        assert_eq!(handler.bytes_sent, 3);

        assert_matches!(handler.seek(SeekFrom::End(-3)), SeekResult::Ok);
        assert_eq!(handler.bytes_sent, 7);

        assert_matches!(handler.seek(SeekFrom::End(20)), SeekResult::Ok);
        assert_eq!(handler.bytes_sent, 10);

        assert_matches!(handler.seek(SeekFrom::Current(-4)), SeekResult::Ok);
        assert_eq!(handler.bytes_sent, 6);

        assert_matches!(handler.seek(SeekFrom::Current(2)), SeekResult::Ok);
        assert_eq!(handler.bytes_sent, 8);

        assert_matches!(handler.seek(SeekFrom::Current(-20)), SeekResult::Ok);
        assert_eq!(handler.bytes_sent, 0);

        handler.request_context_mut().body = None;
        assert_matches!(handler.seek(SeekFrom::Current(0)), SeekResult::CantSeek);
    }

    #[test]
    fn test_headers() {
        let receiver = TestReceiver::new();
        let mut handler = Streaming::new(receiver.clone(), RequestContext::dummy());

        assert!(handler.header(&b"Content-Length: 1234\r\n"[..]));
        assert!(handler.header(&[1, 2, 58, 3, 4][..])); // Valid UTF-8 but not alphanumeric.
        assert!(handler.header(&[0xFF, 0xFF][..])); // Invalid UTF-8 sequence.
        assert!(handler.header(&b"X-No-Value\r\n"[..]));

        let expected = vec![
            (header::CONTENT_LENGTH, HeaderValue::from_static("1234")),
            (
                HeaderName::from_static("x-no-value"),
                HeaderValue::from_static(""),
            ),
        ];

        assert_eq!(receiver.headers(), expected);
    }

    #[test]
    fn test_progress() {
        let receiver = TestReceiver::new();
        let mut handler = Streaming::new(receiver.clone(), RequestContext::dummy());

        let reporter = ProgressReporter::with_callback(|_| ());
        handler
            .request_context_mut()
            .event_listeners()
            .on_progress({
                let updater = reporter.updater();
                move |_req, p| updater.update(p)
            });
        let _ = handler.progress(0.0, 0.0, 0.0, 0.0);
        let _ = handler.progress(1.0, 2.0, 3.0, 4.0);

        // Note that Progress struct has different argument order.
        let expected = Progress::new(2, 1, 4, 3);

        // Check that the TestReceiver got the value.
        assert_eq!(receiver.progress().unwrap(), expected);

        // Check that ProgressReporter also got the value.
        assert_eq!(reporter.aggregate(), expected);
    }
}


/*
 * 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::Result;
use futures::channel::mpsc;
use futures::channel::oneshot;

use crate::errors::Abort;
use crate::errors::HttpClientError;
use crate::header::Header;
use crate::receiver::Receiver;

type Headers = mpsc::UnboundedReceiver<Header>;
type Chunks = mpsc::UnboundedReceiver<Vec<u8>>;
type Done = oneshot::Receiver<Result<(), HttpClientError>>;

/// The receiving end of a `ChannelReceiver`.
pub struct ResponseStreams {
    pub headers_rx: Headers,
    pub body_rx: Chunks,
    pub done_rx: Done,
}

/// A `Receiver` that forwards all received data into channels.
pub struct ChannelReceiver {
    headers_tx: mpsc::UnboundedSender<Header>,
    body_tx: mpsc::UnboundedSender<Vec<u8>>,
    done_tx: oneshot::Sender<Result<(), HttpClientError>>,
}

impl ChannelReceiver {
    pub fn new() -> (Self, ResponseStreams) {
        let (headers_tx, headers_rx) = mpsc::unbounded();
        let (body_tx, body_rx) = mpsc::unbounded();
        let (done_tx, done_rx) = oneshot::channel();

        let senders = Self {
            headers_tx,
            body_tx,
            done_tx,
        };

        let streams = ResponseStreams {
            headers_rx,
            body_rx,
            done_rx,
        };

        (senders, streams)
    }
}

impl Receiver for ChannelReceiver {
    fn chunk(&mut self, chunk: Vec<u8>) -> Result<()> {
        self.body_tx.unbounded_send(chunk).map_err(|e| e.into())
    }

    fn header(&mut self, header: Header) -> Result<()> {
        self.headers_tx.unbounded_send(header).map_err(|e| e.into())
    }

    fn done(self, res: Result<(), HttpClientError>) -> Result<(), Abort> {
        let _ = self.done_tx.send(res);
        Ok(())
    }
}


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

use anyhow::Result;

use crate::errors::Abort;
use crate::errors::HttpClientError;
use crate::header::Header;
use crate::progress::Progress;

pub mod channel;

pub use channel::ChannelReceiver;
pub use channel::ResponseStreams;

/// Interface for streaming HTTP response handlers.
pub trait Receiver: Sized {
    /// Handle received chunk of the response body.
    fn chunk(&mut self, chunk: Vec<u8>) -> Result<()>;

    /// Handle a received header.
    fn header(&mut self, header: Header) -> Result<()>;

    /// Get progress updates for this transfer.
    /// This function will be called whenever the underlying
    /// transfer makes progress.
    fn progress(&mut self, _progress: Progress) {}

    /// Called when the transfer has completed (successfully or not).
    ///
    /// If a fatal error occured while performing the transfer, the error
    /// will be passed to this method so that the `Receiver` can decide how
    /// to proceed. If the `Receiver` returns an `Abort`, all other ongoing
    /// transfers will be aborted and the operation will return early.
    fn done(self, _res: Result<(), HttpClientError>) -> Result<(), Abort> {
        Ok(())
    }
}

#[cfg(test)]
pub(crate) mod testutil {
    #![allow(dead_code)]

    use std::cell::RefCell;
    use std::rc::Rc;

    use http::header::HeaderName;
    use http::header::HeaderValue;
    use http::StatusCode;

    use super::*;

    /// Simple receiver for use in tests.
    #[derive(Clone, Debug)]
    pub struct TestReceiver {
        inner: Rc<RefCell<TestReceiverInner>>,
    }

    impl TestReceiver {
        pub fn new() -> Self {
            Self {
                inner: Rc::new(RefCell::new(Default::default())),
            }
        }

        pub fn status(&self) -> Option<StatusCode> {
            self.inner.borrow().status
        }

        pub fn headers(&self) -> Vec<(HeaderName, HeaderValue)> {
            self.inner.borrow().headers.clone()
        }

        pub fn chunks(&self) -> Vec<Vec<u8>> {
            self.inner.borrow().chunks.clone()
        }

        pub fn progress(&self) -> Option<Progress> {
            self.inner.borrow().progress
        }
    }

    #[derive(Debug, Default)]
    struct TestReceiverInner {
        status: Option<StatusCode>,
        headers: Vec<(HeaderName, HeaderValue)>,
        chunks: Vec<Vec<u8>>,
        progress: Option<Progress>,
    }

    impl Receiver for TestReceiver {
        fn chunk(&mut self, chunk: Vec<u8>) -> Result<()> {
            self.inner.borrow_mut().chunks.push(chunk);
            Ok(())
        }

        fn header(&mut self, header: Header) -> Result<()> {
            match header {
                Header::Status(_, status) => {
                    self.inner.borrow_mut().status = Some(status);
                }
                Header::Header(name, value) => {
                    self.inner.borrow_mut().headers.push((name, value));
                }
                Header::EndOfHeaders => {}
            };
            Ok(())
        }

        fn progress(&mut self, progress: Progress) {
            self.inner.borrow_mut().progress = Some(progress);
        }
    }

    /// No-op receiver for use in tests.
    #[derive(Copy, Clone, Debug)]
    pub struct NullReceiver;

    impl Receiver for NullReceiver {
        fn chunk(&mut self, _chunk: Vec<u8>) -> Result<()> {
            Ok(())
        }
        fn header(&mut self, _header: Header) -> Result<()> {
            Ok(())
        }
    }
}