ZaStoGram_desktop/Telegram/SourceFiles/tests/test_proxy_capability_cache.py

from pathlib import Path


SOURCE_DIR = Path(__file__).resolve().parents[1]
ROOT = SOURCE_DIR.parents[1]
CMAKE = ROOT / "Telegram" / "CMakeLists.txt"
PROXY_DIR = SOURCE_DIR / "mtproto" / "proxy"
CAPABILITIES_H = PROXY_DIR / "capabilities.h"
CAPABILITIES_CPP = PROXY_DIR / "capabilities.cpp"
PROXY_SERVICES_CPP = PROXY_DIR / "proxy_services.cpp"
PROXY_SERVICES_H = PROXY_DIR / "proxy_services.h"
RUNTIME_CPP = SOURCE_DIR / "mtproto" / "runtime" / "runtime_environment.cpp"
TRANSPORT_POLICY_CPP = PROXY_DIR / "transport_policy.cpp"
ENDPOINT_HEALTH_LIFECYCLE_CPP = (
    PROXY_DIR / "mtproxy" / "endpoint_health_lifecycle.cpp")
ENDPOINT_HEALTH_CAPABILITIES_CPP = (
    PROXY_DIR / "mtproxy" / "endpoint_health_capabilities.cpp")
TLS_SOCKET_CPP = PROXY_DIR / "mtproxy" / "tls_socket.cpp"
TLS_SOCKET_RECORDS_CPP = PROXY_DIR / "mtproxy" / "tls_socket_records.cpp"


def read(path):
    assert path.exists(), f"missing expected source file: {path}"
    return path.read_text(encoding="utf-8")


def read_endpoint_health_sources():
    return "\n".join(read(path) for path in (
        ENDPOINT_HEALTH_CPP,
        ENDPOINT_HEALTH_LIFECYCLE_CPP,
    ))


def function_body(source, signature):
    start = source.index(signature)
    brace = source.index("{", start)
    depth = 0
    for i in range(brace, len(source)):
        if source[i] == "{":
            depth += 1
        elif source[i] == "}":
            depth -= 1
            if depth == 0:
                return source[brace:i + 1]
    raise AssertionError(f"function body not found: {signature}")


def test_capability_cache_module_is_file_backed_and_registered():
    header = read(CAPABILITIES_H)
    source = read(CAPABILITIES_CPP)
    runtime = read(RUNTIME_CPP)
    services = read(PROXY_SERVICES_CPP)
    cmake = read(CMAKE)

    assert "mtproto/proxy/capabilities.cpp" in cmake
    assert "mtproto/proxy/capabilities.h" in cmake
    assert "struct ProxyCapabilityCard" in header
    assert "class ProxyCapabilityCache final" in header
    assert "ProxyCapabilityCache &Instance()" not in header
    assert "ProxyCapabilityCache _capabilities;" in read(PROXY_SERVICES_H)
    assert "QJsonDocument" in source
    assert "QSaveFile" in source
    assert 'u"proxy-capabilities.json"_q' in runtime
    assert "cWorkingDir() + u\"tdata/\"_q" in runtime
    assert "QDir().mkpath(" in runtime
    assert "_capabilities(runtime->proxyCapabilities().path)" in services
    assert "SetProxyCapabilityPathProvider" not in runtime
    assert "CapabilitiesPath()" not in source
    assert "QString ProxyCapabilityCache::path() const" in source
    assert "load()" in source
    assert "save()" in source


def test_capability_card_contains_transport_profile_flags_and_routes():
    header = read(CAPABILITIES_H)
    source = read(CAPABILITIES_CPP)

    for field in (
        "QString proxyKey;",
        "ProxyCapabilityTransport lastGoodTransport",
        "QString lastGoodRoute;",
        "ProxyTlsProfile lastGoodProfile",
        "int lastGoodRecipeLevel",
        "bool relayProven",
        "bool autoRotateAllowed",
        "bool wssAllowed",
        "bool syntheticPskAllowed",
        "bool fragmentationAllowed",
        "crl::time lastSuccessAt",
        "crl::time relayProvenAt",
        "QString lastFailureClass;",
        "std::vector<QString> badRoutes;",
        "std::vector<QString> goodRoutes;",
        "crl::time wssBlockedUntil",
    ):
        assert field in header
    for json_key in (
        '"proxyKey"',
        '"lastGoodTransport"',
        '"lastGoodRoute"',
        '"lastGoodProfile"',
        '"lastGoodRecipeLevel"',
        '"relayProven"',
        '"autoRotateAllowed"',
        '"wssAllowed"',
        '"syntheticPskAllowed"',
        '"fragmentationAllowed"',
        '"lastSuccessAt"',
        '"relayProvenAt"',
        '"lastFailureClass"',
        '"badRoutes"',
        '"goodRoutes"',
        '"wssBlockedUntil"',
    ):
        assert json_key in source
    assert 'u"MtproxyFakeTlsTcp"_q' in source
    assert 'u"Wss"_q' in source


def test_proxy_capability_key_uses_canonical_identity_not_route_ip():
    source = read(CAPABILITIES_CPP)
    body = function_body(source, "QString ProxyCapabilityKey(")

    assert "ProxyCapabilityHost(proxy)" in body
    assert "proxy.originalHost.isEmpty()" in source
    assert "? proxy.host" in source
    assert ": proxy.originalHost" in source
    assert "proxy.host +" not in body
    assert "QString::number(int(proxy.type))" in body
    assert "QString::number(proxy.port)" in body
    assert "ProxyCapabilitySecretHash(proxy)" in body
    assert "QCryptographicHash::Sha256" in source


def test_wss_remote_closed_is_persisted_with_ttl_per_proxy():
    source = read(TRANSPORT_POLICY_CPP)
    capabilities = read(CAPABILITIES_CPP)
    note_body = function_body(source, "void NoteProxyWssRemoteClosed(")
    allowed_body = function_body(source, "bool ProxyWssAllowed(")

    assert "kWssRemoteClosedTtl = crl::time(" in source
    assert "runtime->proxyServices().capabilities().noteWssRemoteClosed(" in (
        note_body)
    assert "kWssRemoteClosedTtl" in note_body
    assert "runtime->proxyServices().capabilities().wssAllowed(proxy)" in (
        allowed_body)
    assert "std::set<QString> WssForbiddenProxyKeys" not in source
    assert "QMutex WssForbiddenProxyKeysMutex" not in source
    assert "card.wssAllowed = false;" in capabilities
    assert "card.wssBlockedUntil = crl::now() + ttl;" in capabilities


def test_mtproxy_relay_success_persists_boring_last_good_path():
    header = read(CAPABILITIES_H)
    source = read(CAPABILITIES_CPP)
    success = function_body(source, "void ProxyCapabilityCache::noteMtproxySuccess(")
    read_card = function_body(source, "ProxyCapabilityCard ReadCard(")
    write_card = function_body(source, "QJsonObject WriteCard(")

    assert "QString lastGoodRoute;" in header
    assert "int lastGoodRecipeLevel = 0;" in header
    assert "bool relayProven = false;" in header
    assert "crl::time relayProvenAt = 0;" in header
    assert "bool autoRotateAllowed = true;" in header

    for json_key in (
        '"lastGoodRoute"',
        '"lastGoodRecipeLevel"',
        '"relayProven"',
        '"relayProvenAt"',
        '"autoRotateAllowed"',
    ):
        assert json_key in read_card
        assert json_key in write_card

    assert "const QString &lastGoodRoute" in source
    assert "int recipeLevel" in source
    assert "bool relayProven" in source
    assert "const auto now = crl::now();" in success
    assert "card.lastGoodRoute = lastGoodRoute;" in success
    assert "card.lastGoodRecipeLevel = recipeLevel;" in success
    assert "card.relayProven = relayProven;" in success
    assert "card.relayProvenAt = relayProven ? now : 0;" in success
    assert "card.autoRotateAllowed = false;" in success
    assert "card.syntheticPskAllowed = stealth.syntheticPsk;" in success
    assert ("card.fragmentationAllowed = (stealth.clientHelloFragmentation\n"
        "\t\t!= ProxyClientHelloFragmentation::Off);") in success
    assert "if (stealth.syntheticPsk) {" not in success
    assert "card.syntheticPskAllowed = true;" not in success
    assert "card.fragmentationAllowed = true;" not in success


def test_legacy_relay_cache_uses_last_success_as_proof_time():
    source = read(CAPABILITIES_CPP)
    read_card = function_body(source, "ProxyCapabilityCard ReadCard(")

    assert "const auto relayProvenAt = crl::time(" in read_card
    assert "object.value(\"relayProvenAt\").toDouble()" in read_card
    assert "result.relayProven" in read_card
    assert "result.lastSuccessAt" in read_card


def test_cached_capability_does_not_steer_mtproxy_stealth_options():
    source = read(TRANSPORT_POLICY_CPP)
    body = function_body(source, "ProxyStealthOptions EffectiveProxyStealthOptions(")
    mtproxy_branch = body.split(
        "proxy.type == ProxyData::Type::Mtproto) {", 1)[1].split(
        "if (settings == ProxyData::Settings::Enabled", 1)[0]

    # The cache still records what worked (relay proof, routes), but the
    # MTProxy ClientHello plan no longer follows it: the profile stays the
    # one the user configured.
    assert "FreshMtproxyRelayProof(card)" in read(CAPABILITIES_CPP)
    assert "BoringMtproxyStealthOptions(std::move(result))" in mtproxy_branch
    assert "capabilities().lookup(" not in mtproxy_branch
    assert "capability.lastGoodProfile" not in body
    assert "CompatStrictProxyStealthOptions(std::move(result))" in body
    assert "capability.syntheticPskAllowed" not in mtproxy_branch
    assert "capability.fragmentationAllowed" not in mtproxy_branch
    assert "result.level == ProxyStealthLevel::Experimental" not in mtproxy_branch


if __name__ == "__main__":
    for name, value in list(globals().items()):
        if name.startswith("test_") and callable(value):
            value()
    print("ok")