av-stream-info-rust 0.13.2

use crate::request_error::RequestError;

use futures_util::StreamExt;
use tokio::io::{AsyncRead, AsyncWriteExt, AsyncBufReadExt, AsyncReadExt};
use tokio::io::BufReader;
use tokio::net::TcpStream;
use tokio_rustls::rustls::{ClientConfig};
use tokio_rustls::TlsConnector;
use encoding_rs::{UTF_8, WINDOWS_1252};
use std::sync::Arc;
use std::{collections::HashMap, str::Utf8Error, time::Duration};
use url::Url;

#[derive(Debug)]
pub struct HttpHeaders {
    pub code: u32,
    pub headers: HashMap<String, String>,
}

#[derive(Debug)]
pub struct Request {
    pub info: HttpHeaders,
    content_vec: Vec<u8>,
    content_length: Option<u64>,
    has_http: bool,
    has_https: bool,
    ssl_error: bool,
}

impl Request {
    pub async fn new(url_str: &str, agent: &str, timeout: u32) -> Result<Request, RequestError> {
        let result = Request::new_default(url_str, agent, timeout, false, false).await;
        match result {
            Ok(_) => result,
            Err(err) => {
                warn!("err1: {:?}", err);
                let result2 = Request::new_default(url_str, agent, timeout, true, false).await;
                match result2 {
                    Ok(_) => result2,
                    Err(err) => {
                        warn!("err2: {:?}", err);
                        let result3 = Request::new_fallback(url_str, agent, timeout).await;
                        match result3 {
                            Ok(_) => result3,
                            Err(err) => {
                                warn!("err3: {:?}", err);
                                Request::new_default(url_str, agent, timeout, true, true).await
                            }
                        }
                    }
                }
            }
        }
    }

    async fn new_fallback(url: &str, agent: &str, timeout: u32) -> Result<Request, RequestError> {
        trace!("new_fallback({})", url);
        let (headers, body) = fetch_with_limit(url, timeout, 4096, agent).await?;
        let mut headers_map = HashMap::new();
        let mut status_code = 200; // default fallback

        for (key, value) in headers.iter() {
            if key.to_lowercase() == "status" {
                // Extract status code if present in headers
                if let Ok(code) = value.trim().parse::<u32>() {
                    status_code = code;
                }
            }
            headers_map.insert(key.to_lowercase().to_string(), value.trim().to_string());
        }

        let headers = HttpHeaders {
            code: status_code,
            headers: headers_map,
        };
        let bodylen = body.len();
        let req = Request {
            info: headers,
            content_vec: body,
            content_length: Some(bodylen as u64),
            has_http: true,
            has_https: false,
            ssl_error: false,
        };
        Ok(req)
    }

    async fn new_default(
        url_str: &str,
        agent: &str,
        timeout: u32,
        ignore_tls_errors: bool,
        ignore_header_errors: bool,
    ) -> Result<Request, RequestError> {
        trace!(
            "new_default({}, ignore_tls_errors={})",
            url_str, ignore_tls_errors
        );
        let client = reqwest::Client::builder()
            .user_agent(agent)
            .tls_danger_accept_invalid_certs(ignore_tls_errors)
            .tls_danger_accept_invalid_hostnames(ignore_tls_errors)
            .http1_ignore_invalid_headers_in_responses(ignore_header_errors)
            .http1_allow_obsolete_multiline_headers_in_responses(ignore_header_errors)
            .http1_allow_spaces_after_header_name_in_responses(ignore_header_errors)
            .timeout(std::time::Duration::from_secs(timeout as u64))
            .build()?;

        let response = client.get(url_str).send().await?;

        let status = response.status();
        let content_length = response.content_length();
        let code = status.as_u16() as u32;

        let mut headers_map = HashMap::new();
        for (key, value) in response.headers() {
            headers_map.insert(key.to_string(), value.to_str()?.to_string());
        }

        let mut body_stream = response.bytes_stream();
        let mut content_vec = Vec::new();
        let mut total_read = 0;
        let max_bytes: usize = content_length.unwrap_or(4096) as usize;

        while let Some(chunk_result) = body_stream.next().await {
            let chunk = chunk_result?;
            let remain = max_bytes.saturating_sub(total_read);

            if remain == 0 {
                break;
            }
            let to_take = chunk.len().min(remain);
            content_vec.extend_from_slice(&chunk[..to_take]);
            total_read += to_take;
        }

        let http_headers = HttpHeaders {
            code,
            headers: headers_map,
        };

        let mut has_http = url_str.to_lowercase().starts_with("http://");
        let mut has_https = url_str.to_lowercase().starts_with("https://");

        // check http/https separately
        {
            let other_url = swap_scheme(url_str);
            trace!("Check http/https {}", other_url);

            let response = client.get(other_url).send().await;
            if let Ok(response) = response {
                let other_status = response.status();
                let other_content_length = response.content_length();
                let other_code = other_status.as_u16() as u32;

                // if the second request is also ok then mark has_xxxx as ok
                trace!("{}, {:?}", other_code, other_content_length);
                if code == other_code && content_length == other_content_length {
                    if !has_http {
                        has_http = true;
                    } else {
                        has_https = true;
                    }
                }
            }
        };

        let req = Request {
            info: http_headers,
            content_vec,
            content_length,
            has_http,
            has_https,
            ssl_error: ignore_tls_errors,
        };

        Ok(req)
    }

    pub fn text(&self) -> Result<&str, Utf8Error> {
        std::str::from_utf8(&self.content_vec)
    }

    pub fn raw(&self) -> &Vec<u8> {
        &self.content_vec
    }

    pub fn had_ssl_error(&self) -> bool {
        self.ssl_error
    }

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

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

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

/// Try HTTP/1.1 first; if that fails (or times out) retry once with HTTP/0.9.
///
/// Returns a list of decoded headers (empty for HTTP 0.9) plus at most
/// `maxbytes` of body data.
///
/// Only plain-HTTP is supported; add async-tls/async-native-tls if you need HTTPS.
async fn fetch_with_limit(
    url: &str,
    timeout_secs: u32,
    maxbytes: u32,
    agent: &str,
) -> Result<(Vec<(String, String)>, Vec<u8>), RequestError> {
    trace!("fetch_with_limit({url})", url = url);
    let parsed = Url::parse(url)?;
    let scheme = parsed.scheme();
    if scheme != "http" && scheme != "https" {
        return Err(RequestError::new(
            "Only http:// or https:// URLs are supported",
        ));
    }
    let secure = scheme == "https";

    let host = parsed
        .host_str()
        .ok_or(RequestError::new("URL is missing a host"))?
        .to_owned();
    let port = parsed.port_or_known_default().unwrap_or(80);
    let mut path = parsed.path().to_owned();
    if path.is_empty() {
        path.push('/');
    }
    if let Some(q) = parsed.query() {
        path.push('?');
        path.push_str(q);
    }

    // ------------ attempt 1: HTTP/1.1 ------------
    match fetch_once(
        &host,
        port,
        &format!(
            "GET {} HTTP/1.1\r\nHost: {}\r\nUser-Agent: {}\r\nAccept: */*\r\nConnection: close\r\n\r\n",
            path, host, agent
        ),
        true,
        secure,
        timeout_secs,
        maxbytes,
    )
    .await
    {
        Ok(r) => return Ok(r),
        Err(e) => {
            // fall through to HTTP/0.9 attempt
            info!("HTTP/1.x attempt failed ({e}); retrying with HTTP/0.9 …");
        }
    }

    // ------------ attempt 2: HTTP/0.9 ------------
    fetch_once(
        &host,
        port,
        &format!("GET {}\r\n", path), // minimal request line
        false,
        secure,
        timeout_secs,
        maxbytes,
    )
    .await
}

// Custom verifier that blindly accepts any certificate
/*
struct NoVerifier;
impl rustls::client::ServerCertVerifier for NoVerifier {
    fn verify_server_cert(
        &self,
        _end_entity: &rustls::Certificate,
        _intermediates: &[rustls::Certificate],
        _server_name: &rustls::client::ServerName,
        _scts: &mut dyn Iterator<Item=&[u8]>,
        _ocsp: &[u8],
        _now: std::time::SystemTime,
    ) -> Result<rustls::client::ServerCertVerified, rustls::Error> {
        Ok(rustls::client::ServerCertVerified::assertion())
    }
}
*/

use rustls::client::danger::{
    HandshakeSignatureValid, ServerCertVerified, ServerCertVerifier,
};
use rustls::pki_types::{CertificateDer, ServerName, UnixTime};
use rustls::{DigitallySignedStruct, Error, SignatureScheme};

#[derive(Debug)]
struct NoCertificateVerification;

impl ServerCertVerifier for NoCertificateVerification {
    fn verify_server_cert(
        &self,
        _end_entity: &CertificateDer<'_>,
        _intermediates: &[CertificateDer<'_>],
        _server_name: &ServerName<'_>,
        _ocsp_response: &[u8],
        _now: UnixTime,
    ) -> Result<ServerCertVerified, Error> {
        Ok(ServerCertVerified::assertion())
    }

    fn verify_tls12_signature(
        &self,
        _message: &[u8],
        _cert: &CertificateDer<'_>,
        _dss: &DigitallySignedStruct,
    ) -> Result<HandshakeSignatureValid, Error> {
        Ok(HandshakeSignatureValid::assertion())
    }

    fn verify_tls13_signature(
        &self,
        _message: &[u8],
        _cert: &CertificateDer<'_>,
        _dss: &DigitallySignedStruct,
    ) -> Result<HandshakeSignatureValid, Error> {
        Ok(HandshakeSignatureValid::assertion())
    }

    fn supported_verify_schemes(&self) -> Vec<SignatureScheme> {
        // Return a broad set so handshakes don’t fail due to “no supported schemes”.
        vec![
            SignatureScheme::RSA_PKCS1_SHA256,
            SignatureScheme::RSA_PKCS1_SHA384,
            SignatureScheme::RSA_PKCS1_SHA512,
            SignatureScheme::RSA_PSS_SHA256,
            SignatureScheme::RSA_PSS_SHA384,
            SignatureScheme::RSA_PSS_SHA512,
            SignatureScheme::ECDSA_NISTP256_SHA256,
            SignatureScheme::ECDSA_NISTP384_SHA384,
            SignatureScheme::ECDSA_NISTP521_SHA512,
            SignatureScheme::ED25519,
            SignatureScheme::ED448,
        ]
    }
}


async fn fetch_once(
    host: &str,
    port: u16,
    request: &str,
    expect_headers: bool,
    secure: bool,
    timeout_secs: u32,
    maxbytes: u32,
) -> Result<(Vec<(String, String)>, Vec<u8>), RequestError> {
    trace!(
        "fetch_once({host}:{port})\n{request}",
        host = host,
        port = port,
        request = request
    );
    let to = Duration::from_secs(timeout_secs as u64);

    // 1. connect
    let stream = tokio::time::timeout(to, TcpStream::connect((host, port))).await??;
    if secure {
        let config = ClientConfig::builder()
            .dangerous()
            .with_custom_certificate_verifier(Arc::new(NoCertificateVerification))
            .with_no_client_auth();
        let connector: TlsConnector = Arc::new(config).into();

        let host2 = ServerName::try_from(host.to_owned())?;
        let tls = tokio::time::timeout(to, connector.connect(host2, stream)).await??;
        fetch_once_i(tls, request, to, expect_headers, maxbytes).await
    } else {
        fetch_once_i(stream, request, to, expect_headers, maxbytes).await
    }
}

async fn fetch_once_i<D>(
    mut stream: D,
    request: &str,
    to: Duration,
    expect_headers: bool,
    maxbytes: u32,
) -> Result<(Vec<(String, String)>, Vec<u8>), RequestError>
where
    D: AsyncRead + AsyncWriteExt + Unpin,
{
    // 2. send
    tokio::time::timeout(to, stream.write_all(request.as_bytes())).await??;

    // 3. read
    let mut reader = BufReader::new(stream);
    let mut headers = Vec::<(String, String)>::new();
    let mut body = Vec::with_capacity(maxbytes as usize);

    if expect_headers {
        // ---- try to read status line ----
        let mut line = String::new();
        let n = tokio::time::timeout(to, reader.read_line(&mut line)).await??;
        if n == 0 {
            return Err(RequestError::new("connection closed before status line"));
        }
        if !line.starts_with("HTTP/") && !line.starts_with("ICY") {
            // not a valid HTTP/1.x response – bail so caller can retry 0.9
            return Err(RequestError::new("response did not start with HTTP/"));
        }
        // Parse status code from status line (e.g., "HTTP/1.1 200 OK")
        if let Some(status_part) = line.split_whitespace().nth(1) {
            if let Ok(status_code) = status_part.parse::<u32>() {
                headers.push(("status".to_owned(), status_code.to_string()));
            }
        }

        line.clear();

        // ---- read header block ----
        loop {
            line.clear();
            let n = read_line_with_encoding(&mut reader, &mut line, to).await?;
            if n == 0 || line == "\r\n" {
                break;
            }
            if let Some((name, value)) = line.split_once(':') {
                trace!("{name}: {value}", name = name, value = value.trim());
                headers.push((name.trim().to_owned(), value.trim().to_owned()));
            }
        }
    }

    // ---- read up to maxbytes of body ----
    let mut buf = vec![0_u8; 4096];
    let mut remaining = maxbytes as usize;
    while remaining > 0 {
        let n = tokio::time::timeout(to, reader.read(&mut buf)).await??;
        if n == 0 {
            break; // EOF
        }
        let take = n.min(remaining);
        body.extend_from_slice(&buf[..take]);
        remaining -= take;
    }

    Ok((headers, body))
}

use regex::Regex;

async fn read_line_with_encoding<R>(
    reader: &mut BufReader<R>,
    line: &mut String,
    timeout: Duration,
) -> Result<usize, RequestError>
where
    R: AsyncRead + Unpin,
{
    let mut bytes = Vec::new();
    let mut byte = [0u8; 1];

    loop {
        let n = tokio::time::timeout(timeout, reader.read(&mut byte)).await??;
        if n == 0 {
            break;
        }

        bytes.push(byte[0]);

        if bytes.len() >= 2 && bytes[bytes.len() - 2..] == [b'\r', b'\n'] {
            break;
        }
        if bytes.len() >= 1 && bytes[bytes.len() - 1] == b'\n' {
            break;
        }
    }

    let total_bytes = bytes.len();

    if let Ok(utf8_str) = std::str::from_utf8(&bytes) {
        line.push_str(utf8_str);
        return Ok(total_bytes);
    }

    let encodings = [WINDOWS_1252, UTF_8];

    for encoding in &encodings {
        let (decoded, _, had_errors) = encoding.decode(&bytes);
        if !had_errors {
            line.push_str(&decoded);
            return Ok(total_bytes);
        }
    }

    let (decoded, _, _) = WINDOWS_1252.decode(&bytes);
    line.push_str(&decoded);
    Ok(total_bytes)
}

/// Swap `http` ↔ `https`, ignoring the original casing.
/// Leaves the rest of the URL untouched.
fn swap_scheme(url: &str) -> String {
    // (?i)  → case-insensitive
    // ^     → anchor at the very start of the string
    // (https?) captures "http" or "https"
    let re = Regex::new(r"(?i)^(https?)://").unwrap();

    re.replace(url, |caps: &regex::Captures| {
        match &caps[1].to_ascii_lowercase()[..] {
            "http" => "https://".to_string(),
            "https" => "http://".to_string(),
            _ => unreachable!(),
        }
    })
    .into_owned()
}