🐛(pymta) save the origin IP across STARTTLS restarts

src/mta-in/src/pymta/handler.py

@@ -35,6 +35,15 @@ _ENVELOPES_ATTR = "_pymta_envelopes"
 _SOFT_ERRORS_ATTR = "_pymta_soft_errors"
 _RCPT_MISSES_ATTR = "_pymta_rcpt_misses"
 
+# The PROXY-protocol source is stashed on the *server* (the per-connection SMTP
+# protocol instance), NOT on the session. aiosmtpd rebuilds ``session`` from
+# scratch when the client issues STARTTLS (connection_made -> _create_session),
+# which drops ``session.proxy_data`` and resets ``session.peer`` to the raw TCP
+# peer (the load balancer). The server instance survives that transport swap,
+# so a value stashed there is the only copy of the real client IP that outlives
+# STARTTLS. Holds a ``(addr, port)`` tuple; ``port`` may be None.
+_PROXY_SRC_ATTR = "_pymta_proxy_src"
+
 # Sentinel for the RFC 5321 null sender (MAIL FROM:<>). aiosmtpd's
 # ``smtp_RCPT`` rejects with 503 when ``envelope.mail_from`` is falsy, which
 # would block legitimate bounces. We keep the sentinel internally and rewrite
@@ -70,7 +79,14 @@ def _bump_rcpt_misses(session) -> int:
     return n
 
 
-def _peer_ip(session) -> str | None:
+def _peer_ip(session, server=None) -> str | None:
+    # Prefer the PROXY source captured at connect time and stashed on the
+    # server: it is the only copy that survives the STARTTLS session rebuild
+    # (see _PROXY_SRC_ATTR). Fall back to session.proxy_data for the pre-TLS
+    # window, then to the raw TCP peer when PROXY protocol is off.
+    stashed = getattr(server, _PROXY_SRC_ATTR, None) if server is not None else None
+    if stashed is not None and stashed[0]:
+        return str(stashed[0])
     proxy_data = getattr(session, "proxy_data", None)
     if proxy_data is not None and getattr(proxy_data, "src_addr", None):
         return str(proxy_data.src_addr)
@@ -80,7 +96,10 @@ def _peer_ip(session) -> str | None:
     return None
 
 
-def _peer_port(session) -> str | None:
+def _peer_port(session, server=None) -> str | None:
+    stashed = getattr(server, _PROXY_SRC_ATTR, None) if server is not None else None
+    if stashed is not None and stashed[1] is not None:
+        return str(stashed[1])
     proxy_data = getattr(session, "proxy_data", None)
     if proxy_data is not None and getattr(proxy_data, "src_port", None) is not None:
         return str(proxy_data.src_port)
@@ -277,8 +296,8 @@ class InboundHandler:
                     message=content,
                     sender=sender,
                     original_recipients=list(envelope.rcpt_tos),
-                    client_address=_peer_ip(session),
-                    client_port=_peer_port(session),
+                    client_address=_peer_ip(session, server),
+                    client_port=_peer_port(session, server),
                     # We do not reverse-DNS ourselves: the MDA inserts its
                     # own Received header using metadata and can decide what
                     # to do with the missing hostname.
@@ -294,7 +313,7 @@ class InboundHandler:
             logger.warning(
                 "DATA deliver deadline exceeded (%ds) for peer %s",
                 settings.PYMTA_DATA_TIMEOUT,
-                _peer_ip(session),
+                _peer_ip(session, server),
             )
             return "451 4.3.0 Delivery timed out, please retry"
 
@@ -325,6 +344,28 @@ class InboundHandler:
         real_ip = "unknown"
         if proxy_data is not None and getattr(proxy_data, "src_addr", None):
             real_ip = str(proxy_data.src_addr)
+            # Stash on the server so the real client IP outlives the STARTTLS
+            # session rebuild that would otherwise drop session.proxy_data.
+            setattr(
+                server,
+                _PROXY_SRC_ATTR,
+                (str(proxy_data.src_addr), getattr(proxy_data, "src_port", None)),
+            )
+        if proxy_data is not None:
+            # Permanent forensic record: ties the SMTP session to the real
+            # origin IP carried in the PROXY header. Every other mail.log line
+            # is keyed on session.peer (the load balancer), so this is the only
+            # place the true client IP is recorded. Logging peer alongside src
+            # also surfaces misconfigurations at a glance: src == peer means the
+            # header is not carrying a real origin.
+            logger.info(
+                "PROXY header: src=%s:%s peer=%r version=%r protocol=%r",
+                getattr(proxy_data, "src_addr", None),
+                getattr(proxy_data, "src_port", None),
+                getattr(session, "peer", None),
+                getattr(proxy_data, "version", None),
+                getattr(proxy_data, "protocol", None),
+            )
         return await server.acquire_gate_post_proxy(real_ip)
 
     # ------------------------------------------------------------------ misc

src/mta-in/tests/test_handler.py

@@ -8,6 +8,7 @@ stand-ins — no Docker stack, no real SMTP traffic.
 from __future__ import annotations
 
 import types
+from ipaddress import ip_address
 
 import pytest
 
@@ -28,10 +29,17 @@ class _FakeMDA:
         self.check_result = check_result or MDAResult(
             ok=True, temp_fail=False, payload={}, status_code=200
         )
+        self.deliver_kwargs: dict | None = None
 
     async def check_recipient(self, address: str) -> MDAResult:
         return self.check_result
 
+    async def deliver(self, **kwargs) -> MDAResult:
+        self.deliver_kwargs = kwargs
+        return MDAResult(
+            ok=True, temp_fail=False, payload={"status": "ok"}, status_code=200
+        )
+
 
 def _session():
     return types.SimpleNamespace(host_name=None, peer=("203.0.113.5", 12345))
@@ -180,3 +188,59 @@ async def test_null_sender_round_trip_via_sentinel():
     reply = await _handler().handle_MAIL(None, session, envelope, "<>", [])
     assert reply.startswith("250")
     assert envelope.mail_from == NULL_SENDER_SENTINEL
+
+
+# ---------------------------------------------------------------------------
+# PROXY-protocol source survives the STARTTLS session rebuild.
+#
+# aiosmtpd rebuilds ``session`` from scratch when the client issues STARTTLS,
+# dropping ``session.proxy_data`` and resetting ``session.peer`` to the raw TCP
+# peer (the load balancer). The real client IP must still reach the MDA on the
+# post-TLS DATA command. Regression guard for the "client_ip == LB internal IP"
+# bug seen in production behind HAProxy.
+# ---------------------------------------------------------------------------
+
+
+class _FakeServer:
+    """Per-connection SMTP protocol stand-in (survives the STARTTLS swap)."""
+
+    async def acquire_gate_post_proxy(self, ip: str) -> bool:
+        return True
+
+
+@pytest.mark.asyncio
+async def test_proxy_source_survives_starttls_and_reaches_mda():
+    real_client = ip_address("203.0.113.9")
+    lb_peer = ("10.89.0.2", 43154)  # HAProxy/podman gateway — NOT the client
+
+    server = _FakeServer()
+    mda = _FakeMDA()
+    handler = _handler(mda)
+
+    # 1. PROXY header parsed on the plaintext connection, before STARTTLS.
+    proxy_data = types.SimpleNamespace(
+        src_addr=real_client, src_port=52000, version=2, protocol=1
+    )
+    session_pre_tls = types.SimpleNamespace(
+        host_name=None, peer=lb_peer, proxy_data=proxy_data
+    )
+    gate = await handler.handle_PROXY(server, session_pre_tls, _envelope(), proxy_data)
+    assert gate is True
+
+    # 2. STARTTLS rebuilds the session: proxy_data gone, peer is the LB again.
+    #    Same server instance carries over.
+    session_post_tls = types.SimpleNamespace(
+        host_name=None, peer=lb_peer, proxy_data=None
+    )
+
+    # 3. DATA delivers using the post-TLS session.
+    envelope = _envelope()
+    envelope.mail_from = "sender@example.com"
+    envelope.rcpt_tos = ["rcpt@example.com"]
+    envelope.content = b"Subject: hi\r\n\r\nbody\r\n"
+    reply = await handler.handle_DATA(server, session_post_tls, envelope)
+
+    assert reply.startswith("250"), reply
+    assert mda.deliver_kwargs is not None
+    assert mda.deliver_kwargs["client_address"] == "203.0.113.9"
+    assert mda.deliver_kwargs["client_port"] == "52000"