rws-storage 18.28.0

use super::{S3Config, S3Storage};
use crate::storage::Storage;
use std::io::{Read, Write};
use std::net::TcpListener;
use std::sync::{Arc, Mutex, OnceLock};
use std::thread;

// ── mock S3 server ───────────────────────────────────────────────────────────

struct CapturedRequest {
    method: String,
    path: String,
    headers: Vec<(String, String)>,
    body: Vec<u8>,
}

fn find_subslice(haystack: &[u8], needle: &[u8]) -> Option<usize> {
    if haystack.len() < needle.len() {
        return None;
    }
    haystack.windows(needle.len()).position(|w| w == needle)
}

/// Spawns a one-shot mock S3 endpoint: accepts one connection, captures the
/// request line/headers/body, replies with `status_line` + `response_body`,
/// then closes.
fn spawn_mock_s3(
    status_line: &'static str,
    response_body: &'static [u8],
) -> (u16, Arc<Mutex<Option<CapturedRequest>>>) {
    let listener = TcpListener::bind("127.0.0.1:0").expect("bind mock S3 server");
    let port = listener.local_addr().unwrap().port();
    let captured = Arc::new(Mutex::new(None));
    let captured2 = Arc::clone(&captured);

    thread::spawn(move || {
        if let Ok((mut stream, _)) = listener.accept() {
            let mut data = Vec::new();
            let mut buf = [0u8; 8192];
            let header_end = loop {
                let n = stream.read(&mut buf).unwrap_or(0);
                if n == 0 {
                    return;
                }
                data.extend_from_slice(&buf[..n]);
                if let Some(pos) = find_subslice(&data, b"\r\n\r\n") {
                    break pos + 4;
                }
            };

            let header_str = String::from_utf8_lossy(&data[..header_end]).to_string();
            let mut lines = header_str.lines();
            let request_line = lines.next().unwrap_or("").to_string();
            let mut rl_parts = request_line.split_whitespace();
            let method = rl_parts.next().unwrap_or("").to_string();
            let path = rl_parts.next().unwrap_or("").to_string();

            let mut headers = Vec::new();
            let mut content_length = 0usize;
            for line in lines {
                if let Some((k, v)) = line.split_once(':') {
                    let k = k.trim().to_string();
                    let v = v.trim().to_string();
                    if k.eq_ignore_ascii_case("content-length") {
                        content_length = v.parse().unwrap_or(0);
                    }
                    headers.push((k, v));
                }
            }

            while data.len() < header_end + content_length {
                let n = stream.read(&mut buf).unwrap_or(0);
                if n == 0 {
                    break;
                }
                data.extend_from_slice(&buf[..n]);
            }
            let available = data.len().saturating_sub(header_end);
            let body = data[header_end..header_end + content_length.min(available)].to_vec();

            *captured2.lock().unwrap() = Some(CapturedRequest { method, path, headers, body });

            let resp = format!(
                "{status_line}\r\nContent-Type: text/plain\r\nContent-Length: {}\r\nConnection: close\r\n\r\n",
                response_body.len()
            );
            let _ = stream.write_all(resp.as_bytes());
            let _ = stream.write_all(response_body);
        }
    });

    (port, captured)
}

fn test_config(port: u16) -> S3Config {
    S3Config {
        bucket: "test-bucket".to_string(),
        region: "us-east-1".to_string(),
        access_key: "AKIAIOSFODNN7EXAMPLE".to_string(),
        secret_key: "wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY".to_string(),
        endpoint: format!("http://127.0.0.1:{port}"),
    }
}

fn header<'a>(req: &'a CapturedRequest, name: &str) -> Option<&'a str> {
    req.headers.iter().find(|(k, _)| k.eq_ignore_ascii_case(name)).map(|(_, v)| v.as_str())
}

// ── tests ────────────────────────────────────────────────────────────────────

#[test]
fn put_sends_signed_request_with_correct_path_and_body() {
    // `RWS_S3_CREDENTIAL_SOURCE` is a process-wide override that takes
    // priority over static keys (see `aws_credentials::CredentialsProvider::
    // detect`), so any test constructing an `S3Storage` — even with static
    // keys — races with tests that set it unless they share this lock.
    let _g = crate::storage::aws_credentials::credential_env_lock().lock().unwrap();
    let (port, captured) = spawn_mock_s3("HTTP/1.1 200 OK", b"");
    let store = S3Storage::new(test_config(port));

    let key = store.put("uploads/photo.png", b"binary-data", "image/png").unwrap();
    assert_eq!("uploads/photo.png", key);

    let req = captured.lock().unwrap().take().unwrap();
    assert_eq!("PUT", req.method);
    assert_eq!("/test-bucket/uploads/photo.png", req.path);
    assert_eq!(b"binary-data".to_vec(), req.body);

    assert!(header(&req, "Authorization").unwrap().starts_with(
        "AWS4-HMAC-SHA256 Credential=AKIAIOSFODNN7EXAMPLE/"
    ));
    assert!(header(&req, "x-amz-date").is_some());
    assert!(header(&req, "x-amz-content-sha256").is_some());
    assert_eq!(Some("image/png"), header(&req, "Content-Type"));

    // Exactly one Host header — no duplicate from the signed header list.
    let host_count = req.headers.iter().filter(|(k, _)| k.eq_ignore_ascii_case("host")).count();
    assert_eq!(1, host_count);

    // Regression guard: static credentials never carry a session token, so
    // no x-amz-security-token header should appear on the wire.
    assert!(header(&req, "x-amz-security-token").is_none());
}

#[test]
fn get_returns_body_on_success() {
    let _g = crate::storage::aws_credentials::credential_env_lock().lock().unwrap();
    let (port, _captured) = spawn_mock_s3("HTTP/1.1 200 OK", b"file contents");
    let store = S3Storage::new(test_config(port));
    let bytes = store.get("uploads/photo.png").unwrap();
    assert_eq!(b"file contents".to_vec(), bytes);
}

#[test]
fn get_returns_error_on_404() {
    let _g = crate::storage::aws_credentials::credential_env_lock().lock().unwrap();
    let (port, _captured) = spawn_mock_s3("HTTP/1.1 404 Not Found", b"NoSuchKey");
    let store = S3Storage::new(test_config(port));
    let err = store.get("missing.png").unwrap_err();
    assert!(err.to_string().contains("404"));
}

#[test]
fn delete_sends_delete_method() {
    let _g = crate::storage::aws_credentials::credential_env_lock().lock().unwrap();
    let (port, captured) = spawn_mock_s3("HTTP/1.1 204 No Content", b"");
    let store = S3Storage::new(test_config(port));
    store.delete("uploads/photo.png").unwrap();
    let req = captured.lock().unwrap().take().unwrap();
    assert_eq!("DELETE", req.method);
    assert_eq!("/test-bucket/uploads/photo.png", req.path);
}

#[test]
fn url_uses_path_style_addressing() {
    let store = S3Storage::new(test_config(9999));
    assert_eq!("http://127.0.0.1:9999/test-bucket/uploads/photo.png", store.url("uploads/photo.png"));
}

#[test]
fn key_with_special_characters_is_percent_encoded_in_path() {
    let _g = crate::storage::aws_credentials::credential_env_lock().lock().unwrap();
    let (port, captured) = spawn_mock_s3("HTTP/1.1 200 OK", b"");
    let store = S3Storage::new(test_config(port));
    store.put("a file.txt", b"x", "text/plain").unwrap();
    let req = captured.lock().unwrap().take().unwrap();
    assert_eq!("/test-bucket/a%20file.txt", req.path);
}

// ── Dynamic (workload-identity) credentials ─────────────────────────────────

const DYNAMIC_CREDENTIAL_ENV_VARS: &[&str] = &[
    "RWS_S3_CREDENTIAL_SOURCE",
    "AWS_ROLE_ARN",
    "AWS_WEB_IDENTITY_TOKEN_FILE",
    "AWS_CONTAINER_CREDENTIALS_FULL_URI",
    "AWS_CONTAINER_CREDENTIALS_RELATIVE_URI",
    "AWS_CONTAINER_AUTHORIZATION_TOKEN",
];

fn clear_dynamic_credential_env_vars() {
    for v in DYNAMIC_CREDENTIAL_ENV_VARS {
        std::env::remove_var(v);
    }
}

#[test]
fn put_sends_x_amz_security_token_header_when_using_dynamic_credentials() {
    // Shared with `aws_credentials::tests` — both files mutate the same
    // process-wide AWS_*/RWS_S3_CREDENTIAL_SOURCE env vars.
    let _g = crate::storage::aws_credentials::credential_env_lock().lock().unwrap();
    clear_dynamic_credential_env_vars();

    // Stand-in ECS task-role metadata endpoint returning temporary
    // credentials with a session token.
    let creds_json = br#"{"Code":"Success","AccessKeyId":"AKIADYNAMIC","SecretAccessKey":"secretdynamic","Token":"dynamictoken","Expiration":"2099-01-01T00:00:00Z"}"#;
    let (ecs_port, _ecs_captured) = spawn_mock_s3("HTTP/1.1 200 OK", creds_json);
    std::env::set_var("AWS_CONTAINER_CREDENTIALS_FULL_URI", format!("http://127.0.0.1:{ecs_port}/creds"));

    let (s3_port, s3_captured) = spawn_mock_s3("HTTP/1.1 200 OK", b"");
    let mut config = test_config(s3_port);
    config.access_key = String::new();
    config.secret_key = String::new();
    let store = S3Storage::new(config);

    store.put("uploads/photo.png", b"data", "image/png").unwrap();

    let req = s3_captured.lock().unwrap().take().unwrap();
    assert_eq!(Some("dynamictoken"), header(&req, "x-amz-security-token"));
    assert!(header(&req, "Authorization").unwrap().starts_with("AWS4-HMAC-SHA256 Credential=AKIADYNAMIC/"));

    clear_dynamic_credential_env_vars();
}

#[test]
fn put_returns_storage_error_when_credential_source_is_unreachable() {
    let _g = crate::storage::aws_credentials::credential_env_lock().lock().unwrap();
    clear_dynamic_credential_env_vars();
    // Force IMDS with nothing listening — the request must fail cleanly
    // rather than hang or panic.
    std::env::set_var("RWS_S3_CREDENTIAL_SOURCE", "imds");

    let (s3_port, _s3_captured) = spawn_mock_s3("HTTP/1.1 200 OK", b"");
    let mut config = test_config(s3_port);
    config.access_key = String::new();
    config.secret_key = String::new();
    let store = S3Storage::new(config);

    let err = store.put("uploads/photo.png", b"data", "image/png").unwrap_err();
    assert!(err.to_string().contains("IMDSv2"));

    clear_dynamic_credential_env_vars();
}

// ── S3Config::from_env ───────────────────────────────────────────────────────

fn env_lock() -> &'static Mutex<()> {
    static LOCK: OnceLock<Mutex<()>> = OnceLock::new();
    LOCK.get_or_init(|| Mutex::new(()))
}

#[test]
fn from_env_requires_bucket() {
    let _g = env_lock().lock().unwrap();
    std::env::remove_var("RWS_S3_BUCKET");
    std::env::remove_var("RWS_S3_ACCESS_KEY");
    std::env::remove_var("RWS_S3_SECRET_KEY");
    assert!(S3Config::from_env().is_err());
}

#[test]
fn from_env_no_longer_requires_access_and_secret_key() {
    // Static keys are optional: when absent, credentials are auto-detected
    // from workload identity (IRSA/ECS/IMDS) instead — see
    // `crate::storage::aws_credentials`. Construction itself does no I/O
    // and must not fail just because the keys are unset.
    let _g = env_lock().lock().unwrap();
    std::env::set_var("RWS_S3_BUCKET", "my-bucket");
    std::env::remove_var("RWS_S3_ACCESS_KEY");
    std::env::remove_var("RWS_S3_SECRET_KEY");

    let cfg = S3Config::from_env().unwrap();
    assert_eq!("", cfg.access_key);
    assert_eq!("", cfg.secret_key);

    std::env::remove_var("RWS_S3_BUCKET");
}

#[test]
fn from_env_applies_defaults() {
    let _g = env_lock().lock().unwrap();
    std::env::set_var("RWS_S3_BUCKET", "my-bucket");
    std::env::set_var("RWS_S3_ACCESS_KEY", "AK");
    std::env::set_var("RWS_S3_SECRET_KEY", "SK");
    std::env::remove_var("RWS_S3_REGION");
    std::env::remove_var("RWS_S3_ENDPOINT");

    let cfg = S3Config::from_env().unwrap();
    assert_eq!("my-bucket", cfg.bucket);
    assert_eq!("us-east-1", cfg.region);
    assert_eq!("https://s3.us-east-1.amazonaws.com", cfg.endpoint);

    std::env::remove_var("RWS_S3_BUCKET");
    std::env::remove_var("RWS_S3_ACCESS_KEY");
    std::env::remove_var("RWS_S3_SECRET_KEY");
}

#[test]
fn from_env_respects_custom_endpoint_and_region() {
    let _g = env_lock().lock().unwrap();
    std::env::set_var("RWS_S3_BUCKET", "my-bucket");
    std::env::set_var("RWS_S3_ACCESS_KEY", "AK");
    std::env::set_var("RWS_S3_SECRET_KEY", "SK");
    std::env::set_var("RWS_S3_REGION", "eu-west-1");
    std::env::set_var("RWS_S3_ENDPOINT", "https://accountid.r2.cloudflarestorage.com");

    let cfg = S3Config::from_env().unwrap();
    assert_eq!("eu-west-1", cfg.region);
    assert_eq!("https://accountid.r2.cloudflarestorage.com", cfg.endpoint);

    std::env::remove_var("RWS_S3_BUCKET");
    std::env::remove_var("RWS_S3_ACCESS_KEY");
    std::env::remove_var("RWS_S3_SECRET_KEY");
    std::env::remove_var("RWS_S3_REGION");
    std::env::remove_var("RWS_S3_ENDPOINT");
}

// ── #31: streamed uploads (multipart) ────────────────────────────────────────

/// A reader of `len` deterministic bytes.
struct Pattern {
    len: usize,
    pos: usize,
}

impl std::io::Read for Pattern {
    fn read(&mut self, buf: &mut [u8]) -> std::io::Result<usize> {
        let n = buf.len().min(self.len - self.pos);
        for (i, b) in buf[..n].iter_mut().enumerate() {
            *b = ((self.pos + i) % 251) as u8;
        }
        self.pos += n;
        Ok(n)
    }
}

fn peak_rss_kib() -> u64 {
    std::fs::read_to_string("/proc/self/status").ok()
        .and_then(|s| s.lines().find(|l| l.starts_with("VmHWM:")).and_then(|l| l.split_whitespace().nth(1)?.parse().ok()))
        .unwrap_or(0)
}

/// Against a real S3-compatible server, configured by `RWS_S3_E2E_*`
/// (endpoint, bucket, access/secret key): a 100 MiB streamed upload goes up
/// in parts with peak memory well under the object's size, and reads back
/// byte for byte. `cargo test --features storage-s3 --lib s3_multipart -- --ignored`.
#[test]
#[ignore]
fn s3_multipart_upload_streams_in_constant_memory() {
    let env = |name: &str| std::env::var(format!("RWS_S3_E2E_{name}")).expect("set RWS_S3_E2E_{ENDPOINT,BUCKET,ACCESS_KEY,SECRET_KEY}");
    let storage = S3Storage::new(S3Config {
        bucket: env("BUCKET"),
        region: "us-east-1".to_string(),
        access_key: env("ACCESS_KEY"),
        secret_key: env("SECRET_KEY"),
        endpoint: env("ENDPOINT"),
    });
    const LEN: usize = 100 * 1024 * 1024 + 12345; // 13 parts, the last short
    let before = peak_rss_kib();
    storage.put_stream("rws-e2e/big.bin", &mut Pattern { len: LEN, pos: 0 }, "application/octet-stream").unwrap();
    let grown_mib = peak_rss_kib().saturating_sub(before) / 1024;
    assert!(grown_mib < 48, "a 100 MiB streamed upload grew the peak by {grown_mib} MiB");
    let back = storage.get("rws-e2e/big.bin").unwrap();
    assert_eq!(LEN, back.len());
    assert!(back.iter().enumerate().all(|(i, b)| *b == (i % 251) as u8), "content differs");
    // A small object takes the single-PUT path.
    storage.put_stream("rws-e2e/small.txt", &mut &b"small"[..], "text/plain").unwrap();
    assert_eq!(b"small".to_vec(), storage.get("rws-e2e/small.txt").unwrap());
    storage.delete("rws-e2e/big.bin").unwrap();
    storage.delete("rws-e2e/small.txt").unwrap();
}

#[test]
fn part_reading_and_xml_helpers() {
    let mut reader = Pattern { len: super::S3_PART_SIZE + 10, pos: 0 };
    assert_eq!(super::S3_PART_SIZE, super::read_part(&mut reader, "k").unwrap().len());
    assert_eq!(10, super::read_part(&mut reader, "k").unwrap().len());
    assert!(super::read_part(&mut reader, "k").unwrap().is_empty());
    let xml = "<InitiateMultipartUploadResult><Bucket>b</Bucket><UploadId>abc-123</UploadId></InitiateMultipartUploadResult>";
    assert_eq!(Some("abc-123".to_string()), super::xml_element(xml, "UploadId"));
    assert_eq!(None, super::xml_element(xml, "Missing"));
    assert_eq!("&quot;e&amp;t&quot;", super::xml_escape("\"e&t\""));
}

#[test]
fn the_signed_host_is_the_host_header_sent() {
    let _g = crate::storage::aws_credentials::credential_env_lock().lock().unwrap();
    let (port, captured) = spawn_mock_s3("HTTP/1.1 200 OK", b"");
    let store = S3Storage::new(test_config(port));
    store.put("a.txt", b"x", "text/plain").unwrap();
    let req = captured.lock().unwrap().take().unwrap();
    assert_eq!(Some(store.host().as_str()), header(&req, "host"));
    assert_eq!(format!("127.0.0.1:{port}"), store.host());
    for (endpoint, host) in [
        ("https://s3.eu-west-1.amazonaws.com", "s3.eu-west-1.amazonaws.com"),
        ("https://s3.example:443/", "s3.example"),
        ("http://rook-ceph-rgw.rook-ceph.svc:80", "rook-ceph-rgw.rook-ceph.svc"),
        ("http://minio:9000", "minio:9000"),
        ("https://minio:9000", "minio:9000"),
    ] {
        let store = S3Storage::new(S3Config { endpoint: endpoint.to_string(), ..test_config(1) });
        assert_eq!(host, store.host(), "{endpoint}");
    }
}

#[test]
fn a_ca_file_is_read_or_refused() {
    let _g = crate::storage::aws_credentials::credential_env_lock().lock().unwrap();
    let store = S3Storage::new(test_config(1));
    assert!(store.with_ca_file("/nonexistent/ca.pem").is_err());
    let empty = crate::test_env::TempPath::new("s3-empty-ca.pem");
    std::fs::write(&*empty, "not a certificate").unwrap();
    assert!(S3Storage::new(test_config(1)).with_ca_file(&empty.to_string_lossy()).is_err());
}