codex-app-server/cloud/tests/test_push_toctou.py

"""推送 DNS TOCTOU 单测（残留批次 A 任务 3）。

旧版：validate_push_target 解析一次，urlopen 建连时又解析
一次；两次解析结果不一致（DNS 重绑定）时可能拨向未验证
的内网 IP。新版：建连时解析一次并判段、只拨判过的 IP，
Host 头与 TLS SNI 仍为原域名。

- 两次解析不一致（先公网、后内网）时：发送失败，且一次
  都没有拨向内网 IP；
- 混合解析结果（先内网、后公网）时：只拨公网 IP；
- HTTP 真服务器：域名经 resolver 钉到 127.0.0.1，服务器
  收到的 Host 仍为原域名；
- HTTPS 连接类：拨号地址为判过的 IP，wrap_socket 的
  server_hostname 仍为原域名（SNI 不变）。

测试用的 IP、域名、key 全是假值；除本地回环服务器外不联网。
"""

import http.client
import http.server
import os
import sys
import threading
import unittest
from pathlib import Path
from unittest import mock

try:
    from cloud.push import (
        _checked_connection_class, _post_json, send_webhook)
except ImportError:  # discover -s cloud/tests 时把项目根目录补进 sys.path
    sys.path.insert(0, str(Path(__file__).resolve().parents[2]))
    from cloud.push import (
        _checked_connection_class, _post_json, send_webhook)

PUBLIC_IP = "8.8.8.8"
PRIVATE_IP = "169.254.169.254"


class _FlipFlopResolver:
    """第一次返回公网 IP，之后返回内网 IP：模拟重绑定。

    旧版两次解析恰好踩中这个窗口；新版第二次解析判段后
    直接拒绝，绝不拨号。
    """

    def __init__(self):
        self.calls = 0

    def __call__(self, host):
        self.calls += 1
        return [PUBLIC_IP] if self.calls == 1 else [PRIVATE_IP]


class _DialRecorder:
    """记录拨号地址后直接失败：测试只断言拨向，不跑完 HTTP。"""

    def __init__(self):
        self.calls = []

    def __call__(self, address, timeout=None, source_address=None):
        self.calls.append(address)
        raise OSError("dial stopped by test")


class _DummySocket:
    def setsockopt(self, *args):
        pass


class _RecordingHandler(http.server.BaseHTTPRequestHandler):
    requests = []

    def do_POST(self):
        length = int(self.headers.get("Content-Length", 0))
        body = self.rfile.read(length) if length else b""
        _RecordingHandler.requests.append({
            "host": self.headers.get("Host"),
            "body": body,
        })
        self.send_response(200)
        self.send_header("Content-Length", "2")
        self.end_headers()
        self.wfile.write(b"{}")

    def log_message(self, *args):
        pass


class PushToctouTestCase(unittest.TestCase):
    def setUp(self):
        self._old_allow = os.environ.get("CODEX_PUSH_ALLOW_PRIVATE_TARGETS")
        os.environ.pop("CODEX_PUSH_ALLOW_PRIVATE_TARGETS", None)
        _RecordingHandler.requests = []

    def tearDown(self):
        if self._old_allow is None:
            os.environ.pop("CODEX_PUSH_ALLOW_PRIVATE_TARGETS", None)
        else:
            os.environ["CODEX_PUSH_ALLOW_PRIVATE_TARGETS"] = self._old_allow

    def _message(self, target, resolver=None):
        msg = {"event_type": "completed", "project_id": "proj-a",
               "payload": {"task_id": "t-1"}, "target": target}
        if resolver is not None:
            msg["_resolver"] = resolver
        return msg

    def test_rebinding_second_resolution_never_dialed(self):
        resolver = _FlipFlopResolver()
        dials = _DialRecorder()
        with mock.patch("socket.create_connection", dials):
            with self.assertRaises(RuntimeError):
                send_webhook(self._message(
                    "http://hook.example.invalid/hook", resolver))
        # 两次解析确实都发生了（一次预检、一次建连），但建连侧
        # 判段后拒绝：内网 IP 一次都没被拨。
        self.assertEqual(resolver.calls, 2)
        self.assertEqual(dials.calls, [])

    def test_mixed_resolution_dials_only_public(self):
        dials = _DialRecorder()
        with mock.patch("socket.create_connection", dials):
            ok, _detail = _post_json(
                "http://hook.example.invalid/hook", {"a": 1},
                resolver=lambda host: [PRIVATE_IP, PUBLIC_IP])
        self.assertFalse(ok)
        self.assertEqual(dials.calls, [(PUBLIC_IP, 80)])

    def test_http_pinned_host_header_preserved(self):
        os.environ["CODEX_PUSH_ALLOW_PRIVATE_TARGETS"] = "1"
        server = http.server.HTTPServer(("127.0.0.1", 0),
                                        _RecordingHandler)
        port = server.server_address[1]
        thread = threading.Thread(target=server.serve_forever, daemon=True)
        thread.start()
        try:
            send_webhook(self._message(
                "http://webhook.test.invalid:%d/hook" % port,
                resolver=lambda host: ["127.0.0.1"]))
        finally:
            server.shutdown()
            server.server_close()
            thread.join(timeout=5)
        self.assertEqual(len(_RecordingHandler.requests), 1)
        host_header = _RecordingHandler.requests[0]["host"]
        # 拨的是 127.0.0.1，但 Host 仍为原域名（虚拟主机依赖它）。
        self.assertEqual(host_header.split(":")[0],
                         "webhook.test.invalid")

    def test_https_sni_preserved_with_pinned_ip(self):
        wrap_calls = []

        class FakeContext:
            def wrap_socket(self, sock, server_hostname=None):
                wrap_calls.append(server_hostname)
                return sock

        cls = _checked_connection_class(http.client.HTTPSConnection,
                                        lambda host: [PUBLIC_IP])
        conn = cls("example.com", 443)
        conn._context = FakeContext()
        dials = []

        def fake_create(address, timeout=None, source_address=None):
            dials.append(address)
            return _DummySocket()

        with mock.patch("socket.create_connection", fake_create):
            conn.connect()
        self.assertEqual(dials, [(PUBLIC_IP, 443)])
        self.assertEqual(wrap_calls, ["example.com"])