megacode-agent-sdk 0.31.0

//! The release download (spec slice-ar R12.2): every hop https (http only to
//! a loopback host), no user info in any URL, at most 5 redirects, the whole
//! download within 120 s, the archive at most 64 MiB, `Accept-Encoding:
//! identity`; the digest is of the exact bytes received.

use std::net::IpAddr;
use std::time::Duration;

use reqwest::Url;
use reqwest::redirect::{Attempt, Policy};
use sha2::{Digest, Sha256};

use super::hex;

/// The largest archive accepted.
const MAX_ARCHIVE: u64 = 64 * 1024 * 1024;
const MAX_REDIRECTS: usize = 5;
const WHOLE_DOWNLOAD: Duration = Duration::from_secs(120);

/// Whether `url` may be fetched: https, or http to a loopback host; no user
/// info.
fn allowed(url: &Url) -> Result<(), String> {
    if !url.username().is_empty() || url.password().is_some() {
        return Err(format!(
            "{} carries user info",
            url.host_str().unwrap_or("a URL")
        ));
    }
    match url.scheme() {
        "https" => Ok(()),
        "http" if loopback(url) => Ok(()),
        _ => Err(format!("{url} is not https")),
    }
}

fn loopback(url: &Url) -> bool {
    let Some(host) = url.host_str() else {
        return false;
    };
    let bare = host.trim_start_matches('[').trim_end_matches(']');
    match bare.parse::<IpAddr>() {
        Ok(ip) => ip.is_loopback(),
        Err(_) => bare.eq_ignore_ascii_case("localhost"),
    }
}

fn redirect(attempt: Attempt<'_>) -> reqwest::redirect::Action {
    // `previous` holds every URL already requested, the first included.
    if attempt.previous().len() > MAX_REDIRECTS {
        return attempt.error(format!("more than {MAX_REDIRECTS} redirects"));
    }
    match allowed(attempt.url()) {
        Ok(()) => attempt.follow(),
        Err(why) => {
            let why = format!("redirected: {why}");
            attempt.error(why)
        }
    }
}

/// `e` and its sources, joined.
fn chain(e: &dyn std::error::Error) -> String {
    let mut out = e.to_string();
    let mut source = e.source();
    while let Some(s) = source {
        let text = s.to_string();
        if !out.contains(&text) {
            out.push_str(": ");
            out.push_str(&text);
        }
        source = s.source();
    }
    out
}

/// The archive at `url`, its SHA-256 checked against `archive_sha256`.
pub(crate) async fn archive(url: &str, archive_sha256: &str) -> Result<Vec<u8>, String> {
    let fail = |why: String| format!("cannot download {url}: {why}; nothing was cached");
    let parsed = Url::parse(url).map_err(|e| fail(e.to_string()))?;
    allowed(&parsed).map_err(fail)?;
    let mut builder = reqwest::Client::builder();
    if loopback(&parsed) {
        // A loopback release is never reached through a proxy.
        builder = builder.no_proxy();
    }
    let client = builder
        .no_gzip()
        .no_brotli()
        .no_deflate()
        .no_zstd()
        .redirect(Policy::custom(redirect))
        .timeout(WHOLE_DOWNLOAD)
        .build()
        .map_err(|e| fail(chain(&e)))?;
    let mut response = client
        .get(parsed)
        .header(reqwest::header::ACCEPT_ENCODING, "identity")
        .send()
        .await
        .map_err(|e| fail(chain(&e)))?;
    if !response.status().is_success() {
        return Err(fail(response.status().to_string()));
    }
    let over = || fail("the archive is over 64 MiB".into());
    if response.content_length().is_some_and(|n| n > MAX_ARCHIVE) {
        return Err(over());
    }
    let mut body = Vec::new();
    while let Some(chunk) = response.chunk().await.map_err(|e| fail(chain(&e)))? {
        if (body.len() + chunk.len()) as u64 > MAX_ARCHIVE {
            return Err(over());
        }
        body.extend_from_slice(&chunk);
    }
    let digest = hex(&Sha256::digest(&body));
    if digest != archive_sha256 {
        return Err(fail(format!(
            "the archive does not match its pinned SHA-256 (got {digest}, pinned {archive_sha256})"
        )));
    }
    Ok(body)
}

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

    #[test]
    fn only_https_or_loopback_http_without_user_info() {
        let ok = |u: &str| allowed(&Url::parse(u).unwrap()).is_ok();
        assert!(ok("https://megacode-harness.vercel.app/dl"));
        assert!(ok("http://127.0.0.1:8080/dl"));
        assert!(ok("http://127.8.0.1/dl"));
        assert!(ok("http://[::1]:80/dl"));
        assert!(ok("http://localhost/dl"));
        assert!(!ok("http://example.com/dl"));
        assert!(!ok("http://10.0.0.1/dl"));
        assert!(!ok("ftp://127.0.0.1/dl"));
        assert!(!ok("https://user@example.com/dl"));
        assert!(!ok("https://:pw@example.com/dl"));
    }
}