ProxyWarden/src-tauri/tests/privileged_package_trust_tests.rs

use proxywarden_lib::component_catalog::ComponentId;
use proxywarden_lib::component_packages::{
    ComponentPackageService, ComponentPackagesError, GithubReleaseDigestProof,
    NativePrivilegedBundleVerifier, PackageSignatureVerifier, PackageSource,
    PrivilegedBundleVerificationError, PrivilegedBundleVerifier, PrivilegedCachedUpdatePlan,
    SignaturePublisher, SignatureVerifierError, UpdateRequestKind, UpdateTransport,
    UpdateTransportError, UpdateTransportRequest, UpdateTransportResponse,
};
use proxywarden_lib::safe_fs::protect_path_for_owner_admin_system;
use proxywarden_lib::storage::StoragePaths;
use serde_json::json;
use sha2::{Digest, Sha256};
use std::fs;
use std::io::Cursor;
use std::path::{Path, PathBuf};
use std::sync::atomic::{AtomicUsize, Ordering};
use uuid::Uuid;

#[test]
fn privileged_fresh_install_ignores_forged_newer_cache_and_has_no_transport() {
    let workspace = TestWorkspace::new();
    let paths = workspace.storage_paths();
    let forged = paths
        .packages_dir
        .join(ComponentId::SingBox.as_str())
        .join("999.0.0");
    fs::create_dir_all(&forged).expect("create forged cache");
    fs::write(forged.join("forged.zip"), b"not an official package").expect("write forged cache");
    let service =
        ComponentPackageService::open(bundled_root(), &paths).expect("open component packages");
    let verifier = ExplicitTestBundleVerifier::default();

    let lease = service
        .lease_bundled_for_privileged_install(ComponentId::SingBox, &verifier)
        .expect("lease immutable bundled sing-box");

    assert_eq!(lease.proof().source, PackageSource::Bundled);
    assert_eq!(lease.proof().version, "1.13.19");
    assert_eq!(lease.proof().independent_proof, None);
    assert_eq!(verifier.calls.load(Ordering::SeqCst), 1);
    assert!(lease.asset_path().starts_with(bundled_root()));
    // There is intentionally no transport argument on the fresh-install API.
}

#[test]
fn native_bundle_verifier_rejects_user_owned_or_noncanonical_bundle() {
    let workspace = TestWorkspace::new();
    let service = ComponentPackageService::open(bundled_root(), &workspace.storage_paths())
        .expect("open component packages");

    assert!(matches!(
        service.lease_bundled_for_privileged_install(
            ComponentId::SingBox,
            &NativePrivilegedBundleVerifier
        ),
        Err(ComponentPackagesError::UntrustedBundleRoot)
    ));
}

#[test]
fn privileged_update_requires_the_exact_live_latest_proof_before_staging() {
    let workspace = TestWorkspace::new();
    let service = ComponentPackageService::open(bundled_root(), &workspace.storage_paths())
        .expect("open component packages");
    let plan = sing_box_plan(b"cached sing-box update");
    let mut wrong_latest = plan.clone();
    wrong_latest.independent_proof.release_id += 1;
    let transport = MetadataTransport::success(metadata_for(&wrong_latest));

    assert!(matches!(
        service.lease_cached_update_for_privileged_install(
            &plan,
            workspace.path(),
            &transport,
            &NoopVerifier,
        ),
        Err(ComponentPackagesError::InvalidPrivilegedUpdatePlan)
    ));
    assert_eq!(transport.calls.load(Ordering::SeqCst), 1);
    assert_no_privileged_staging(workspace.path());

    let timeout = MetadataTransport::failure();
    assert!(matches!(
        service.lease_cached_update_for_privileged_install(
            &plan,
            workspace.path(),
            &timeout,
            &NoopVerifier,
        ),
        Err(ComponentPackagesError::Transport)
    ));
    assert_eq!(timeout.calls.load(Ordering::SeqCst), 1);
    assert_no_privileged_staging(workspace.path());
}

#[cfg(windows)]
#[test]
fn tampered_cache_fails_before_any_admin_staging_is_created() {
    let workspace = TestWorkspace::new();
    let paths = workspace.storage_paths();
    let service =
        ComponentPackageService::open(bundled_root(), &paths).expect("open component packages");
    let plan = sing_box_plan(b"official update bytes");
    write_cache(&paths, &plan, b"tampered cache bytes");
    let transport = MetadataTransport::success(metadata_for(&plan));

    assert!(matches!(
        service.lease_cached_update_for_privileged_install(
            &plan,
            workspace.path(),
            &transport,
            &NoopVerifier,
        ),
        Err(ComponentPackagesError::DigestMismatch)
    ));
    assert_eq!(transport.calls.load(Ordering::SeqCst), 1);
    assert_no_privileged_staging(workspace.path());
}

#[test]
fn winsw_and_vc_runtime_make_zero_privileged_update_transport_calls() {
    let workspace = TestWorkspace::new();
    let service = ComponentPackageService::open(bundled_root(), &workspace.storage_paths())
        .expect("open component packages");

    for component_id in [ComponentId::Winsw, ComponentId::VcRuntime] {
        let transport = MetadataTransport::failure();
        let mut plan = sing_box_plan(b"unused");
        plan.component_id = component_id;
        assert!(matches!(
            service.lease_cached_update_for_privileged_install(
                &plan,
                workspace.path(),
                &transport,
                &NoopVerifier,
            ),
            Err(ComponentPackagesError::NoTrustedUpdate)
        ));
        assert_eq!(transport.calls.load(Ordering::SeqCst), 0);
    }
}

#[derive(Default)]
struct ExplicitTestBundleVerifier {
    calls: AtomicUsize,
}

struct NoopVerifier;

impl PackageSignatureVerifier for NoopVerifier {
    fn verify(&self, _path: &Path) -> Result<SignaturePublisher, SignatureVerifierError> {
        Err(SignatureVerifierError)
    }
}

impl PrivilegedBundleVerifier for ExplicitTestBundleVerifier {
    fn verify(
        &self,
        bundled_root: &Path,
        catalog_path: &Path,
        asset_path: &Path,
    ) -> Result<(), PrivilegedBundleVerificationError> {
        self.calls.fetch_add(1, Ordering::SeqCst);
        if catalog_path != bundled_root.join("catalog.json")
            || !asset_path.starts_with(bundled_root)
        {
            return Err(PrivilegedBundleVerificationError);
        }
        Ok(())
    }
}

struct MetadataTransport {
    response: Result<Vec<u8>, UpdateTransportError>,
    calls: AtomicUsize,
}

impl MetadataTransport {
    fn success(response: Vec<u8>) -> Self {
        Self {
            response: Ok(response),
            calls: AtomicUsize::new(0),
        }
    }

    fn failure() -> Self {
        Self {
            response: Err(UpdateTransportError::RequestFailed),
            calls: AtomicUsize::new(0),
        }
    }
}

impl UpdateTransport for MetadataTransport {
    fn get(
        &self,
        request: &UpdateTransportRequest,
    ) -> Result<UpdateTransportResponse, UpdateTransportError> {
        self.calls.fetch_add(1, Ordering::SeqCst);
        assert_eq!(request.kind, UpdateRequestKind::GithubReleaseMetadata);
        assert_eq!(
            request.url,
            "https://api.github.com/repos/SagerNet/sing-box/releases/latest"
        );
        let body = self.response.clone()?;
        Ok(UpdateTransportResponse {
            status: 200,
            location: None,
            content_length: Some(body.len() as u64),
            body: Box::new(Cursor::new(body)),
        })
    }
}

fn sing_box_plan(bytes: &[u8]) -> PrivilegedCachedUpdatePlan {
    let version = "1.14.0";
    let asset_name = format!("sing-box-{version}-windows-amd64.zip");
    let sha256 = format!("{:x}", Sha256::digest(bytes));
    PrivilegedCachedUpdatePlan {
        component_id: ComponentId::SingBox,
        version: version.to_string(),
        independent_proof: GithubReleaseDigestProof {
            repository: "SagerNet/sing-box".to_string(),
            release_id: 700,
            asset_id: 701,
            stable_tag: format!("v{version}"),
            asset_name,
            size: bytes.len() as u64,
            sha256_from_api: sha256,
            verified_signatures: Vec::new(),
        },
    }
}

fn metadata_for(plan: &PrivilegedCachedUpdatePlan) -> Vec<u8> {
    let proof = &plan.independent_proof;
    serde_json::to_vec(&json!({
        "id": proof.release_id,
        "tag_name": proof.stable_tag,
        "draft": false,
        "prerelease": false,
        "assets": [{
            "id": proof.asset_id,
            "name": proof.asset_name,
            "size": proof.size,
            "digest": format!("sha256:{}", proof.sha256_from_api),
            "browser_download_url": format!(
                "https://github.com/{}/releases/download/{}/{}",
                proof.repository, proof.stable_tag, proof.asset_name
            )
        }]
    }))
    .expect("serialize GitHub metadata")
}

#[cfg(windows)]
fn write_cache(paths: &StoragePaths, plan: &PrivilegedCachedUpdatePlan, bytes: &[u8]) {
    use proxywarden_lib::component_packages::{
        PackageCacheManifest, PACKAGE_CACHE_MANIFEST_FILENAME,
        PACKAGE_CACHE_MANIFEST_SCHEMA_VERSION,
    };

    let component_root = paths.packages_dir.join(plan.component_id.as_str());
    let version_root = component_root.join(&plan.version);
    fs::create_dir_all(&version_root).expect("create cache directory");
    let asset_path = version_root.join(&plan.independent_proof.asset_name);
    fs::write(&asset_path, bytes).expect("write cache bytes");
    let manifest_path = version_root.join(PACKAGE_CACHE_MANIFEST_FILENAME);
    let manifest = PackageCacheManifest {
        schema_version: PACKAGE_CACHE_MANIFEST_SCHEMA_VERSION,
        component_id: plan.component_id,
        version: plan.version.clone(),
        asset_name: plan.independent_proof.asset_name.clone(),
        sha256: plan.independent_proof.sha256_from_api.clone(),
        size: plan.independent_proof.size,
        independent_proof: plan.independent_proof.clone(),
    };
    fs::write(
        &manifest_path,
        serde_json::to_vec_pretty(&manifest).expect("serialize cache manifest"),
    )
    .expect("write cache manifest");
    for path in [
        paths.packages_dir.as_path(),
        component_root.as_path(),
        version_root.as_path(),
        manifest_path.as_path(),
        asset_path.as_path(),
    ] {
        protect_path_for_owner_admin_system(path).expect("protect cache path");
    }
}

fn assert_no_privileged_staging(parent: &Path) {
    let staging = fs::read_dir(parent)
        .expect("read staging parent")
        .filter_map(Result::ok)
        .filter_map(|entry| entry.file_name().into_string().ok())
        .filter(|name| name.starts_with(".package-"))
        .collect::<Vec<_>>();
    assert!(
        staging.is_empty(),
        "unexpected staging entries: {staging:?}"
    );
}

fn bundled_root() -> PathBuf {
    Path::new(env!("CARGO_MANIFEST_DIR"))
        .join("bundled")
        .join("components")
}

struct TestWorkspace {
    path: PathBuf,
}

impl TestWorkspace {
    fn new() -> Self {
        let target = Path::new(env!("CARGO_MANIFEST_DIR")).join("target");
        fs::create_dir_all(&target).expect("create Cargo target directory");
        let path = target.join(format!("privileged-package-trust-{}", Uuid::new_v4()));
        fs::create_dir(&path).expect("create test workspace");
        Self { path }
    }

    fn path(&self) -> &Path {
        &self.path
    }

    fn storage_paths(&self) -> StoragePaths {
        StoragePaths::new(&self.path)
    }
}

impl Drop for TestWorkspace {
    fn drop(&mut self) {
        if self
            .path
            .file_name()
            .and_then(|name| name.to_str())
            .is_some_and(|name| name.starts_with("privileged-package-trust-"))
        {
            let _ = fs::remove_dir_all(&self.path);
        }
    }
}