ZaStoGram_desktop/Telegram/SourceFiles/tests/test_mtproxy_target_flow.py

from pathlib import Path
from session_private_sources import read_session_private_sources


SOURCE_DIR = Path(__file__).resolve().parents[1]
PROXY_DIR = SOURCE_DIR / "mtproto" / "proxy"
MTPROXY_DIR = PROXY_DIR / "mtproxy"

SESSION_CPP = SOURCE_DIR / "mtproto" / "session" / "session.cpp"
SESSION_PRIVATE_CPP = SOURCE_DIR / "mtproto" / "session" / "private" / "session_private.cpp"
ABSTRACT_CONNECTION_CPP = SOURCE_DIR / "mtproto" / "transport" / "connection_abstract.cpp"
TRANSPORT_POLICY_CPP = PROXY_DIR / "transport_policy.cpp"
RESOLVING_CPP = PROXY_DIR / "resolving_connection.cpp"
CHECK_H = PROXY_DIR / "check.h"
CHECK_CPP = PROXY_DIR / "check.cpp"
TLS_SOCKET_CPP = MTPROXY_DIR / "tls_socket.cpp"
TLS_SOCKET_HANDSHAKE_CPP = MTPROXY_DIR / "tls_socket_handshake.cpp"
TLS_SOCKET_RECORDS_CPP = MTPROXY_DIR / "tls_socket_records.cpp"
STATUS_CPP = PROXY_DIR / "status.cpp"
WINDOW_CONNECTING_CPP = SOURCE_DIR / "window" / "window_connecting_widget.cpp"


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


def compact(text):
    return " ".join(text.split())


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


def test_user_proxy_selection_uses_fixed_strict_mtproxy_plan():
    session = read(SESSION_CPP)
    policy = read(TRANSPORT_POLICY_CPP)
    refresh = function_body(session, "void Session::refreshOptions(")
    effective = function_body(policy, "ProxyStealthOptions EffectiveProxyStealthOptions(")
    mtproxy_branch = effective.split(
        "proxy.type == ProxyData::Type::Mtproto) {", 1)[1].split(
        "if (settings == ProxyData::Settings::Enabled", 1)[0]

    # The ClientHello profile is whatever the user configured, fixed: a
    # fingerprint filter is deterministic, so steering the profile from the
    # outcome of earlier dials only hands the other side more profiles.
    assert "EffectiveProxyStealthOptions(" in refresh
    assert mtproxy_branch.index("result.transport = ProxyTransport::Tcp;") < (
        mtproxy_branch.index("BoringMtproxyStealthOptions(std::move(result))"))
    boring = function_body(policy, "ProxyStealthOptions BoringMtproxyStealthOptions(")
    assert boring.index("const auto profile = result.tlsProfile;") < (
        boring.index("CompatStrictProxyStealthOptions(std::move(result));"))
    assert boring.index("CompatStrictProxyStealthOptions(std::move(result));") < (
        boring.index("result.tlsProfile = profile;"))
    assert "ProxyTlsProfile::" not in boring
    assert "capabilities()" not in mtproxy_branch
    assert "capability." not in mtproxy_branch
    assert "ProxyTransport::Wss" not in mtproxy_branch
    assert "ApplyProxyStealthLevel(" not in mtproxy_branch
    assert "syntheticPskAllowed" not in mtproxy_branch
    assert "fragmentationAllowed" not in mtproxy_branch


def test_dns_singleflight_precedes_route_open_and_route_racing_is_bounded():
    abstract = read(ABSTRACT_CONNECTION_CPP)
    resolving = read(RESOLVING_CPP)
    connect = function_body(
        resolving,
        "void ResolvingConnection::connectToServer(")
    start_resolving = function_body(
        resolving,
        "void ResolvingConnection::startResolving()")
    route_order = function_body(
        resolving,
        "std::vector<int> ResolvingConnection::routeOrder(")
    add_route = function_body(
        resolving,
        "void ResolvingConnection::addRouteAttempt(")

    assert "proxy.tryCustomResolve()" in abstract
    assert "ConnectionPointer::New<ResolvingConnection>" in abstract
    assert connect.index("_mtproxyAttempt = context.mtproxyAttempt;") < (
        connect.index("startResolving();"))
    assert "_runtime->proxyServices().dnsResolver().request(" in start_resolving
    assert "domainResolved(" in start_resolving
    mtproxy_order = compact(route_order).split(
        "if (_proxy.type != ProxyData::Type::Mtproto) {", 1)[1].split(
        "return result; }", 1)[1]
    # Addresses that answered last time go first, untried ones next and the
    # ones that already failed last.
    good = mtproxy_order.index("listed(capability.goodRoutes, keyOf(index))")
    untried = mtproxy_order.index("!listed(capability.badRoutes, key)")
    assert mtproxy_order.index(
        "_runtime->proxyServices().capabilities().lookup( _proxy)") < good
    assert good < untried < mtproxy_order.rindex(
        "for (auto index = 0; index != int(_proxy.resolvedIPs.size()); ++index) { append(index); }")
    assert "MtProxy::RouteKey(MtProxy::RouteEndpointFromAddress(" in mtproxy_order
    assert "attempt.child = _child->clone(ToDirectIpProxy(_proxy, ipIndex));" in add_route
    assert add_route.index("ToDirectIpProxy(_proxy, ipIndex)") < (
        add_route.index("stored.child->connectToServer("))
    assert "kRouteRaceDelay = crl::time(300)" in resolving
    assert "kMaxParallelRouteAttempts = 2" in resolving
    assert "activeRouteAttempts() >= kMaxParallelRouteAttempts" in resolving


def test_logs_and_left_proxy_shield_expose_target_flow_state():
    resolving = read(RESOLVING_CPP)
    tls_handshake = read(TLS_SOCKET_HANDSHAKE_CPP)
    window = read(WINDOW_CONNECTING_CPP)
    route_event = function_body(resolving, "void ReportRouteEvent(")

    assert "MtProxy::ToLegacyDiagnostic(reason)" in route_event
    assert 'u"tcp_not_connected"_q' in route_event
    assert "ProxyDiagnosticsPhase::RouteSelected" in resolving
    assert "ProxyDiagnosticsPhase::RouteFailed" in resolving
    assert "ProxyDiagnosticsPhase::ClientHelloSent" in tls_handshake
    assert "reportTransportEvent(" in tls_handshake
    assert "_proxyIcon->moveToLeft(xShift, yShift);" in window
    assert "const auto progressVisible = visible && !_currentLayout.proxyEnabled;" in window
    assert "_proxyIcon->setVisible(_currentLayout.proxyEnabled);" in window
    assert "_cacheErrorDns" in window
    assert "_cacheErrorTimeout" in window
    assert "_cacheErrorHandshake" in window
    assert "_cacheErrorData" in window


def test_first_app_data_reported_once_per_tls_socket():
    tls = read(TLS_SOCKET_RECORDS_CPP)
    body = function_body(tls, "bool TlsSocket::checkNextPacket(")

    # FirstDataReceived progress and the data-path success must fire once
    # per socket, not on every incoming TLS record: repeated reports knock
    # the proxy status from Connected back to CheckingTelegram.
    guarded = function_body(body, "if (!_firstAppDataReceived) {")
    assert "_firstAppDataReceived = true;" in guarded
    assert "connectionProgress(_phase);" in guarded
    assert "u\"mtproxy first tls appdata received\"_q" in guarded
    assert "syntheticPsks().noteDataPathSuccess(" in guarded
    assert body.count("connectionProgress(") == 1
    assert body.count("noteDataPathSuccess(") == 1


def test_resolving_connection_forwards_timeout_to_route_attempts():
    resolving = read(RESOLVING_CPP)
    timed_out = function_body(
        resolving,
        "void ResolvingConnection::timedOut(")

    # The owner times out on the ResolvingConnection wrapper; without
    # forwarding, the TlsSocket doing the actual connect never reports
    # its failure to EndpointHealth and domain proxies never degrade.
    assert "attempt.child->timedOut();" in timed_out
    assert "_child->timedOut();" in timed_out


def test_stealth_option_changes_restart_proxy_connections():
    app = read(SOURCE_DIR / "core" / "application.cpp")
    box = read(SOURCE_DIR / "boxes" / "connection_box.cpp")
    apply_body = function_body(
        app,
        "void Application::applyProxyStealthOptions(")
    restart_body = function_body(
        app,
        "void Application::restartProxyConnections(")

    # Sessions read stealth options only when (re)connecting: writing
    # the setting without restarting MTP leaves e.g. a WSS transport
    # switch inert until something else (like toggling IPv6) restarts
    # connections.
    assert "setProxyStealthOptions(options);" in apply_body
    assert "_proxyRestartTimer" in apply_body
    assert "_proxyChanges.fire" in restart_body
    assert "Core::App().settings().setProxyStealthOptions(" not in box
    assert "Core::App().applyProxyStealthOptions(o);" in box
    assert "Core::App().restartProxyConnections();" in box


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