firegex/tests/integration/test_udp.py

"""UDP, on whichever layer.

Both carry it and the same filters work on both, which is the whole point of the model.
They get there differently: the proxy relays datagrams with **one socket per protected
address** — `SO_ORIGINAL_DST` is TCP-only, so there is nothing to recover per datagram and
nothing to recover it from — while NFQUEUE inspects the real datagram in place.

What UDP keeps on both: per-flow filter state, released after a timeout because a datagram
has no close to observe, and the client's own address, which the proxy preserves with
`IP_TRANSPARENT` rather than giving up.
"""

import time

import pytest

from integration import filter_code
from integration.conftest import (Layer, add_python_filter, add_regex_filter,
                                  start_and_settle)

pytestmark = pytest.mark.instance

UDP_LAYERS = [
    pytest.param(Layer("proxy", ipv6=False), id="proxy-ipv4"),
    pytest.param(Layer("proxy", ipv6=True), id="proxy-ipv6", marks=pytest.mark.ipv6),
    pytest.param(Layer("nfqueue", ipv6=False), id="nfqueue-ipv4"),
    pytest.param(Layer("nfqueue", ipv6=True), id="nfqueue-ipv6", marks=pytest.mark.ipv6),
]


@pytest.fixture(params=UDP_LAYERS)
def udp_layer(request) -> Layer:
    return request.param


@pytest.fixture
def udp_service(api, service, udp_stand_in, udp_layer):
    """A running UDP service refusing anything carrying DENYME."""
    echo = udp_stand_in(udp_layer.ipv6)
    service_id = service(f"udp-{echo.port}", udp_layer.ip, echo.port,
                         udp_layer.transport, proto="udp")
    add_regex_filter(api, service_id, "DENYME")
    start_and_settle(api, service_id, wait=1.2)
    return service_id, echo


def test_a_benign_datagram_is_relayed_and_answered(udp_service):
    _, echo = udp_service
    assert echo.exchange(b"hello there") == b"hello there"


def test_a_matching_datagram_is_refused(udp_service):
    _, echo = udp_service
    assert echo.exchange(b"carrying DENYME") is None


def test_the_flow_survives_a_refused_datagram(udp_service):
    """`Verdict::Reject` drops the datagram rather than closing anything, because there
    is nothing to close."""
    _, echo = udp_service
    assert echo.exchange(b"carrying DENYME") is None
    assert echo.exchange(b"still here") == b"still here"


def test_the_block_is_counted_against_the_pattern(api, udp_service):
    service_id, echo = udp_service
    echo.exchange(b"carrying DENYME")
    time.sleep(0.6)
    filter_id = api.services_filters(service_id)[0]["filter_id"]
    rows = api.services_regexes(service_id, filter_id)
    assert any(r["blocked"] >= 1 for r in rows), str(rows)


def test_the_reply_appears_to_come_from_the_address_the_client_dialled(udp_service):
    """Replies leave through the *listener* socket so conntrack rewrites them.

    Sending from the upstream socket would reach a client not expecting that source, and
    the client would simply drop it.
    """
    _, echo = udp_service
    answered_from = echo.answered_from(b"who answered")
    assert answered_from is not None, "no answer came back at all"
    assert answered_from[1] == echo.port, \
        f"the answer came from {answered_from}, not the port the client dialled"


def test_a_second_udp_address_is_relayed_without_restarting(api, udp_service,
                                                            udp_stand_in, udp_layer):
    """New relays are bound on the fly over the control channel (`ADD_UDP`), so adding a
    UDP address drops nothing — the same promise TCP has always had here."""
    service_id, first = udp_service
    second = udp_stand_in(udp_layer.ipv6)

    why = api.services_add_address_error(service_id, udp_layer.ip, second.port)
    assert why is None, why
    time.sleep(1.2)

    assert second.exchange(b"hello second") == b"hello second"
    assert second.exchange(b"carrying DENYME") is None, \
        "the same chain is not filtering the new UDP address"
    assert first.exchange(b"still here") == b"still here", \
        "the original UDP address stopped answering"

    entries = api.services_logs(service_id)
    assert any("also protecting" in e["text"] for e in entries), str(entries[-4:])


def test_a_python_filter_works_on_datagrams(api, udp_service):
    """`cpproxy` gained `filter_action_udp` — no stream follower, no ack fixing, because
    a datagram is complete in itself."""
    service_id, echo = udp_service
    assert api.services_stop(service_id)
    add_python_filter(api, service_id, filter_code.UDP, name="datagrams")
    start_and_settle(api, service_id, wait=1.2)

    assert echo.exchange(b"carrying PYDENY") is None
    assert echo.exchange(b"perfectly fine") == b"perfectly fine"


def test_the_udp_block_is_counted_against_the_function_that_made_it(api, udp_service):
    service_id, echo = udp_service
    assert api.services_stop(service_id)
    filter_id = add_python_filter(api, service_id, filter_code.UDP, name="datagrams")
    start_and_settle(api, service_id, wait=1.2)

    echo.exchange(b"carrying PYDENY")
    time.sleep(0.8)
    functions = api.services_functions(service_id, filter_id)
    assert any(f["blocked"] >= 1 for f in functions), str(functions)


def test_a_filter_needing_a_tcp_stream_is_refused_on_a_udp_service(api, udp_service):
    """`RawPacket` is the only model that survives a datagram.

    Every other one declines to be built when the traffic is not TCP, so a filter asking
    for one would sit in the chain and never be called. The refusal names both the
    service's protocol and the model that cannot be built on it.
    """
    service_id, _ = udp_service
    assert api.services_stop(service_id)
    assert api.services_add_filter(service_id, "pyfilter", "wants-http")
    filter_id = [f["filter_id"] for f in api.services_filters(service_id)
                 if f["name"] == "wants-http"][0]

    why = api.services_set_code_error(service_id, filter_id, filter_code.HTTP)
    assert why is not None, "a filter needing a TCP stream was saved on a UDP service"
    assert "UDP" in why and "HttpRequest" in why, why


def test_a_tcp_service_and_a_udp_one_may_share_an_address(api, service, stand_in,
                                                          udp_stand_in):
    """Exactly as the kernel allows, which is what makes `(ip, port, proto)` the key.

    The protocol is denormalised onto each address for this reason; `tls` is stored as
    the TCP it is on the wire, or a TCP service and a TLS one could claim one `ip:port`
    between them.
    """
    tcp_server = stand_in()
    udp_echo = udp_stand_in(port=tcp_server.port)

    tcp_id = service(f"shared-tcp-{tcp_server.port}", "127.0.0.1", tcp_server.port,
                     "proxy")
    udp_id = service(f"shared-udp-{udp_echo.port}", "127.0.0.1", udp_echo.port,
                     "proxy", proto="udp")
    assert tcp_id and udp_id
    start_and_settle(api, tcp_id)
    start_and_settle(api, udp_id, wait=1.2)
    assert udp_echo.exchange(b"both at once") == b"both at once"


def test_udp_on_an_interface_needs_an_address_to_bind_a_relay_to(api, service,
                                                                 udp_stand_in):
    """A relay binds somewhere concrete, and an interface name is not somewhere.

    The interface's own address stands in — IPv4 first and link-local last, because an
    interface almost always carries an `fe80::` address it was never configured with, and
    binding there would put the service where nobody dials.
    """
    echo = udp_stand_in()
    service_id = service(f"udpif-{echo.port}", "lo", echo.port, "proxy", proto="udp")
    started = api.services_start(service_id)
    if not started:
        pytest.skip("this host has no `lo` to resolve, or would not start the service")
    time.sleep(1.2)
    assert echo.exchange(b"through the interface") == b"through the interface"


# --- how many flows keep their state, on NFQUEUE -------------------------------------

REMEMBERS = """
from firegex.pyfilters import pyfilter, ACCEPT, REJECT
from firegex.pyfilters.models import RawPacket

# Module level, so it is per flow: each one gets globals of its own.
marked = False

@pyfilter
def remembers(packet: RawPacket):
    global marked
    if b"MARK" in packet.data:
        marked = True
        return ACCEPT
    return REJECT if marked else ACCEPT
"""


def _from(port: int, echo, payload: bytes) -> bytes | None:
    """One datagram from a fixed source port, so several exchanges are one flow."""
    import socket
    sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
    sock.settimeout(1.5)
    try:
        sock.bind(("127.0.0.1", port))
        sock.sendto(payload, (echo.host, echo.port))
        return sock.recvfrom(65535)[0]
    except (TimeoutError, OSError):
        return None
    finally:
        sock.close()


@pytest.mark.parametrize("limit,kept", [(0, True), (2, False)], ids=["no-limit", "limit-2"])
def test_a_queued_udp_service_keeps_flows_up_to_its_own_limit(api, service, udp_stand_in,
                                                              limit, kept):
    """`cpproxy` held at most 4096 flows' state, whatever the operator wanted, and let go
    of an arbitrary one past that — a TCP stream's included. The number is the service's
    own now, none by default, and past it the flow quiet the longest starts over."""
    from helpers.net import free_port

    echo = udp_stand_in()
    service_id = service(f"flows-{limit}-{echo.port}", "127.0.0.1", echo.port, "nfqueue",
                         proto="udp", max_connections=limit)
    add_python_filter(api, service_id, REMEMBERS, name="remembers")
    start_and_settle(api, service_id, wait=1.2)

    first = free_port(udp=True)
    assert _from(first, echo, b"MARK") == b"MARK"
    assert _from(first, echo, b"again") is None, "the flow did not keep its state at all"
    # Enough of them that some land on the first flow's thread: flows are shared out
    # between the queue threads by a hash, and so is the limit. Each from a port handed
    # out once, rather than one the kernel picks: that can be the first flow's own, now
    # that nothing holds it, and a datagram sent from it is the first flow — the one the
    # filter was told to refuse — which failed this one run in a few hundred.
    for _ in range(200):
        assert _from(free_port(udp=True), echo, b"another flow") == b"another flow"

    answered = _from(first, echo, b"once more")
    if kept:
        assert answered is None, "with no limit, a flow lost its state to two others"
    else:
        assert answered == b"once more", "past the limit, the oldest flow kept its state"


# --- flows conntrack is still carrying -------------------------------------------------
#
# Conntrack translates a flow at its first packet, and a UDP flow lasts for as long as
# datagrams keep arriving. So on the proxy layer, which is a translation, what a flow
# already had decided where each of its datagrams went — until firegex started forgetting
# the flows whose answer the rules had just changed.


def _proxy_only(layer: Layer):
    if layer.transport != "proxy":
        pytest.skip("NFQUEUE queues every packet as it goes by; conntrack decides nothing")


def test_a_flow_open_before_the_service_started_meets_its_filters(api, service,
                                                                 udp_stand_in, udp_layer):
    """Nothing redirected it, and nothing would: it went on straight to the service,
    unfiltered, for as long as it kept talking — a way past the protection for anybody
    who opened one before it went up."""
    echo = udp_stand_in(udp_layer.ipv6)
    service_id = service(f"udp-early-{echo.port}", udp_layer.ip, echo.port,
                         udp_layer.transport, proto="udp")
    add_regex_filter(api, service_id, "DENYME")
    with echo.flow() as flow:
        assert flow.exchange(b"before the service") == b"before the service"
        start_and_settle(api, service_id, wait=1.2)
        assert flow.exchange(b"carrying DENYME") is None, \
            "a flow older than the service went on unfiltered"
        assert flow.exchange(b"still talking") == b"still talking", \
            "the flow stopped being answered once the service started"


def test_a_flow_carried_across_a_restart_meets_the_new_filters(api, udp_service, udp_layer):
    """A restart hands the connections over to the engine replacing the old one; a flow
    stayed with the old one instead, under the filters it started with, for as long as
    it kept talking."""
    _proxy_only(udp_layer)
    service_id, echo = udp_service
    with echo.flow() as flow:
        assert flow.exchange(b"before the restart") == b"before the restart"
        # Anything but its name rebuilds the datapath.
        assert api.services_edit(service_id, max_connections=500)
        time.sleep(1.2)
        add_regex_filter(api, service_id, "LATECOMER", name="later")
        time.sleep(0.5)
        assert flow.exchange(b"carrying LATECOMER") is None, \
            "the flow was still carried by the engine from before the restart"
        assert flow.exchange(b"after the restart") == b"after the restart"


def _engines_relaying_to(port: int) -> list[int]:
    """The proxy engines holding a UDP relay towards `port`, by their environment."""
    import os
    found = []
    for pid in filter(str.isdigit, os.listdir("/proc")):
        try:
            with open(f"/proc/{pid}/environ", "rb") as f:
                environ = f.read().split(b"\0")
        except OSError:
            continue
        for entry in environ:
            if entry.startswith(b"FGEX_PROXY_UDP=") and f":{port}|".encode() in entry:
                found.append(int(pid))
    return found


@pytest.mark.root
def test_a_flow_is_answered_again_after_its_engine_is_killed(api, udp_service, udp_layer):
    """Conntrack went on sending the flow to the relay port of an engine that no longer
    existed, every datagram refreshing the entry that did so: a client that kept talking
    was cut off from the service until it fell silent for two minutes."""
    import os
    import signal

    _proxy_only(udp_layer)
    _, echo = udp_service
    with echo.flow() as flow:
        assert flow.exchange(b"before the crash") == b"before the crash"
        engines = _engines_relaying_to(echo.port)
        if not engines:
            pytest.skip("cannot find the engine to kill; is the instance elsewhere?")
        for pid in engines:
            os.kill(pid, signal.SIGKILL)
        answered, deadline, n = None, time.monotonic() + 10, 0
        while answered is None and time.monotonic() < deadline:
            n += 1
            answered = flow.exchange(f"after the crash {n}".encode())
        assert answered is not None, "the flow was never answered again after the engine died"
        assert flow.exchange(b"carrying DENYME") is None, \
            "the engine that came back is not filtering the flow"