ninlight-w600/tools/station/tests/test_server.py

import hashlib
import http.client
import json
import os
import select
import socket
import stat
import subprocess
import sys
import threading
import time
from pathlib import Path

import pytest

SERVER = Path(__file__).resolve().parents[1] / "server.py"
TRAILER_SIZE = 512
RECENT_MAX = 8


def make_image(path: Path, device: str, version: str, payload: bytes) -> bytes:
    """Builds a file with the same trailer layout the firmware verifies."""
    trailer = bytearray(b"\xff" * TRAILER_SIZE)
    trailer[0:4] = b"NINL"
    trailer[4] = 1
    trailer[5:21] = device.encode().ljust(16, b"\x00")
    trailer[21:37] = version.encode().ljust(16, b"\x00")
    trailer[37:69] = hashlib.sha256(payload).digest()
    data = payload + bytes(trailer)
    path.write_bytes(data)
    return data


def free_port() -> int:
    s = socket.socket()
    s.bind(("127.0.0.1", 0))
    port = s.getsockname()[1]
    s.close()
    return port


class Station:
    def __init__(self, tmp_path, *extra):
        self.image_dir = tmp_path / "images"
        self.image_dir.mkdir()
        self.state = tmp_path / "status.json"
        self.port = free_port()
        self.proc = subprocess.Popen(
            [sys.executable, str(SERVER), "--image-dir", str(self.image_dir),
             "--bind", "127.0.0.1", "--port", str(self.port),
             "--state", str(self.state), *extra],
            stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True,
        )
        deadline = time.time() + 10
        while time.time() < deadline:
            try:
                socket.create_connection(("127.0.0.1", self.port), 0.2).close()
                return
            except OSError:
                if self.proc.poll() is not None:
                    raise RuntimeError("server died: " + (self.proc.stdout.read() or ""))
                time.sleep(0.05)
        raise RuntimeError("server did not come up")

    def conn(self):
        return http.client.HTTPConnection("127.0.0.1", self.port, timeout=5)

    def read_state(self, predicate=None):
        """Reads status.json, optionally waiting for it to satisfy `predicate`.

        The server only writes the entry for a request after that request's response has
        actually finished (or failed) on the wire - by design, so it never claims success
        before the last byte is confirmed sent (see server.py). A client that has just
        finished reading its own response can therefore still see a status.json that is a
        few milliseconds stale. Rather than have the server promise something before it is
        true, we wait here for the state to catch up.
        """
        last = None
        for _ in range(100):
            if self.state.exists():
                try:
                    last = json.loads(self.state.read_text())
                    if predicate is None or predicate(last):
                        return last
                except json.JSONDecodeError:
                    pass
            time.sleep(0.05)
        if last is None:
            raise AssertionError("no state file written")
        raise AssertionError(f"state never satisfied predicate, last seen: {last!r}")

    def stop(self):
        self.proc.terminate()
        self.proc.wait(timeout=5)

    def wait_for_log_line(self, contains, timeout=3.0):
        """Reads the server's merged stdout/stderr (see __init__) until a line containing
        `contains` shows up, skipping unrelated lines (e.g. the startup banner) - used to
        check that a request actually produced a log line (F3: every request must,
        including rejected ones and errors the base handler catches before routing).

        Reads the raw fd directly with os.read() instead of self.proc.stdout.readline():
        readline() goes through TextIOWrapper's internal BufferedReader, whose underlying
        read() can slurp everything currently sitting in the pipe (banner *and* the log
        line we're waiting for) into its own buffer on the very first call, while only
        handing back the first line. select() then reports the fd as not-readable on the
        next iteration - the data isn't gone, it is just buffered somewhere select() can't
        see - and a fixed-size test timeout would flake on exactly the line it's checking
        for.
        """
        fd = self.proc.stdout.fileno()
        buf = getattr(self, "_log_buf", "")
        deadline = time.time() + timeout
        while True:
            while "\n" in buf:
                line, buf = buf.split("\n", 1)
                self._log_buf = buf
                if contains in line:
                    return line + "\n"
            remaining = deadline - time.time()
            if remaining <= 0:
                self._log_buf = buf
                return None
            r, _, _ = select.select([fd], [], [], remaining)
            if not r:
                self._log_buf = buf
                return None
            chunk = os.read(fd, 65536)
            if not chunk:
                self._log_buf = buf
                return None
            buf += chunk.decode("utf-8", "replace")


@pytest.fixture
def station(tmp_path):
    s = Station(tmp_path)
    yield s
    s.stop()


def test_serves_a_known_image(station):
    data = make_image(station.image_dir / "ninlight_evk_ota.img", "evk", "1.2.3", b"payload" * 500)
    c = station.conn()
    c.request("GET", "/ninlight_evk_ota.img")
    r = c.getresponse()
    body = r.read()
    assert r.status == 200
    assert int(r.getheader("Content-Length")) == len(data)
    assert body == data


def test_missing_image_is_404(station):
    c = station.conn()
    c.request("GET", "/ninlight_g95_cw_ota.img")
    assert c.getresponse().status == 404


def test_unknown_target_is_404(station):
    make_image(station.image_dir / "ninlight_evk_ota.img", "evk", "1.0.0", b"x" * 100)
    for path in ("/ninlight_nosuch_ota.img", "/ninlight_evk_ota.img.bak", "/", "/images/"):
        c = station.conn()
        c.request("GET", path)
        assert c.getresponse().status == 404, path


def test_path_traversal_is_refused(station):
    secret = station.image_dir.parent / "secret.txt"
    secret.write_text("nicht ausliefern")
    c = station.conn()
    c.request("GET", "/../secret.txt")
    assert c.getresponse().status == 404


def test_other_methods_are_405(station):
    c = station.conn()
    c.request("POST", "/ninlight_evk_ota.img", body=b"")
    assert c.getresponse().status == 405


def test_rejected_method_produces_a_log_line(station):
    """F3: do_POST (and PUT/DELETE/HEAD) used to call reject() without logging anything -
    the station is the only diagnostic tool during a rescue, and the spec requires a line
    per request, rejections included."""
    c = station.conn()
    c.request("POST", "/ninlight_evk_ota.img", body=b"")
    assert c.getresponse().status == 405
    line = station.wait_for_log_line(" 405 ")
    assert line is not None, "kein Log-Eintrag fuer die abgelehnte POST-Anfrage"
    assert "/ninlight_evk_ota.img" in line


def test_log_error_is_not_silenced_by_log_message(capsys):
    """F3: BaseHTTPRequestHandler.log_error() calls log_message() by default (used e.g. for
    a malformed request line or a client timeout, cases that never reach do_GET/do_POST).
    server.py silences log_message() in favour of its own per-request log() calls - which
    used to silence log_error() right along with it, since it wasn't overridden separately.
    This exercises Handler.log_error() directly rather than over a raw socket: this
    sandbox's network layer intercepts/rewrites malformed non-HTTP payloads before they
    reach the server (observed via a spoofed 'Server: aiohttp' response), which makes a
    genuine malformed-request-line test unreliable here."""
    import importlib.util

    spec = importlib.util.spec_from_file_location("ninlight_station_server", SERVER)
    module = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(module)

    class _Stub:
        client_address = ("203.0.113.5", 54321)

    module.Handler.log_error(_Stub(), "code %d, message %s", 400, "Bad request version")
    out = capsys.readouterr().out
    assert "ERROR" in out
    assert "203.0.113.5" in out
    assert "code 400, message Bad request version" in out


def test_range_request_returns_206_and_the_right_slice(station):
    data = make_image(station.image_dir / "ninlight_e27_rgb_ota.img", "e27_rgb", "2.0.0", b"abc" * 2000)
    c = station.conn()
    c.request("GET", "/ninlight_e27_rgb_ota.img", headers={"Range": "bytes=1000-"})
    r = c.getresponse()
    body = r.read()
    assert r.status == 206
    assert r.getheader("Content-Range") == f"bytes 1000-{len(data)-1}/{len(data)}"
    assert body == data[1000:]


def test_state_file_records_a_delivery(station):
    make_image(station.image_dir / "ninlight_evk_ota.img", "evk", "0.3.4", b"z" * 1000)
    c = station.conn()
    c.request("GET", "/ninlight_evk_ota.img")
    c.getresponse().read()
    state = station.read_state(lambda s: s["served"] >= 1)
    assert state["served"] == 1
    entry = state["recent"][-1]
    assert entry["target"] == "evk"
    assert entry["version"] == "0.3.4"
    assert entry["status"] == "ok"
    assert entry["ip"] == "127.0.0.1"
    assert entry["bytes"] == 1000 + TRAILER_SIZE


def test_state_file_records_a_missing_image(station):
    c = station.conn()
    c.request("GET", "/ninlight_g95_cw_ota.img")
    c.getresponse().read()
    state = station.read_state(lambda s: len(s["recent"]) >= 1)
    assert state["served"] == 0
    assert state["recent"][-1]["status"] == "notfound"
    assert state["recent"][-1]["target"] == "g95_cw"


def test_state_file_is_world_readable_even_under_a_restrictive_umask(tmp_path):
    """status.json is written by a root service but read by a normal-user CLI.

    tempfile.mkstemp() creates its file 0600 and os.replace() preserves whatever mode
    the temp file had, so the server must explicitly widen the mode before the rename -
    and must do so with chmod() rather than relying on the mode passed to mkstemp/open,
    because that one is still narrowed by the process umask. A restrictive umask (as
    root services on the Pi may well run under) must not bring the 0600 back.
    """
    old_umask = os.umask(0o077)
    try:
        s = Station(tmp_path)
    finally:
        os.umask(old_umask)
    try:
        make_image(s.image_dir / "ninlight_evk_ota.img", "evk", "1.0.0", b"w" * 100)

        def world_readable() -> bool:
            mode = stat.S_IMODE(s.state.stat().st_mode)
            return bool(mode & stat.S_IROTH)

        # Written once already at startup (Station.__init__ -> write_state()).
        s.read_state(lambda _: world_readable())
        assert world_readable(), "status.json is not world-readable after the initial write"

        # A second write (a real delivery) must not lose the mode either - e.g. if the
        # widening were only applied to the file's very first creation instead of to
        # every temp file, mode 0600 would silently come back on every later write.
        c = s.conn()
        c.request("GET", "/ninlight_evk_ota.img")
        c.getresponse().read()
        s.read_state(lambda st: st["served"] >= 1)
        assert world_readable(), "status.json lost its mode on a later overwrite"
    finally:
        s.stop()


def test_recent_is_capped_at_eight(station):
    make_image(station.image_dir / "ninlight_evk_ota.img", "evk", "1.0.0", b"q" * 10)
    for _ in range(12):
        c = station.conn()
        c.request("GET", "/ninlight_evk_ota.img")
        c.getresponse().read()
    state = station.read_state(lambda s: s["served"] >= 12)
    assert state["served"] == 12
    assert len(state["recent"]) == RECENT_MAX


def test_aborted_transfer_is_recorded(tmp_path):
    s = Station(tmp_path, "--stall-after", "256")
    try:
        make_image(s.image_dir / "ninlight_evk_ota.img", "evk", "1.0.0", b"y" * 100000)
        sock = socket.create_connection(("127.0.0.1", s.port), timeout=5)
        sock.sendall(b"GET /ninlight_evk_ota.img HTTP/1.1\r\nHost: x\r\n\r\n")
        sock.recv(4096)
        time.sleep(0.3)
        sock.close()
        deadline = time.time() + 10
        while time.time() < deadline:
            try:
                entry = json.loads(s.state.read_text())["recent"][-1]
                if entry["status"] == "aborted":
                    assert entry["bytes"] < 100000
                    return
            except (OSError, json.JSONDecodeError, IndexError, KeyError):
                pass
            time.sleep(0.1)
        raise AssertionError("abort was not recorded")
    finally:
        s.stop()


def test_stall_after_stops_sending_but_keeps_the_connection(tmp_path):
    s = Station(tmp_path, "--stall-after", "512")
    try:
        make_image(s.image_dir / "ninlight_evk_ota.img", "evk", "1.0.0", b"y" * 200000)
        sock = socket.create_connection(("127.0.0.1", s.port), timeout=5)
        sock.sendall(b"GET /ninlight_evk_ota.img HTTP/1.1\r\nHost: x\r\n\r\n")
        got = b""
        sock.settimeout(2)
        try:
            while len(got) < 200000:
                chunk = sock.recv(65536)
                if not chunk:
                    break
                got += chunk
        except socket.timeout:
            pass
        # far less than the full image
        assert len(got) < 20000
        # and the connection is genuinely still open: a server that (wrongly) closed it
        # after the stall boundary would make this recv() return b"" (EOF) instead of
        # timing out. Just counting received bytes above would not catch that - a server
        # closing the connection right after --stall-after also sends far less than
        # 20000 bytes.
        sock.settimeout(1)
        with pytest.raises(socket.timeout):
            sock.recv(1)
        sock.close()
    finally:
        s.stop()


def test_stall_after_budget_survives_a_resumed_request(tmp_path):
    """--stall-after must stop a client for good, not just for one request.

    On real hardware, a lamp whose download stalls hits its own receive timeout, gives up
    on the connection, and opens a new one with a Range header picking up right where it
    left off - exactly what test_stall_after_stops_sending_but_keeps_the_connection above
    exercises for a single request. A per-request byte counter would hand that resumed
    request a brand-new --stall-after allowance, so the "stall" would only ever delay
    delivery instead of actually stopping it - which is what was observed on the EVK: the
    server stalled at 65536 bytes each time, but the lamp reconnected five times and walked
    away with the complete image. This test reproduces that reconnect and checks that the
    second, resumed request gets no further bytes at all.
    """
    s = Station(tmp_path, "--stall-after", "1000")
    try:
        make_image(s.image_dir / "ninlight_evk_ota.img", "evk", "1.0.0", b"y" * 200000)

        # First request: burns through the whole --stall-after budget, then the server
        # holds the connection open and sends nothing further.
        sock = socket.create_connection(("127.0.0.1", s.port), timeout=5)
        sock.sendall(b"GET /ninlight_evk_ota.img HTTP/1.1\r\nHost: x\r\n\r\n")
        got = b""
        sock.settimeout(2)
        try:
            while len(got) < 200000:
                chunk = sock.recv(65536)
                if not chunk:
                    break
                got += chunk
        except socket.timeout:
            pass
        assert len(got) < 20000  # far less than the full image, same bound as above
        sock.close()  # the lamp's watchdog gives up and reconnects, like on real hardware

        # Second request, resuming with Range exactly like a real lamp would. This
        # client's cumulative budget is already spent, so it must get zero more bytes of
        # body - not a fresh --stall-after allowance just because it's a new connection.
        sock2 = socket.create_connection(("127.0.0.1", s.port), timeout=5)
        sock2.sendall(b"GET /ninlight_evk_ota.img HTTP/1.1\r\nHost: x\r\nRange: bytes=1000-\r\n\r\n")
        sock2.settimeout(2)
        header = b""
        while b"\r\n\r\n" not in header:
            header += sock2.recv(4096)
        assert b" 206 " in header.split(b"\r\n", 1)[0]
        body_after_header = header.split(b"\r\n\r\n", 1)[1]
        assert body_after_header == b"", "resumed request must not receive any body bytes"
        try:
            more = sock2.recv(1)
        except socket.timeout:
            more = None
        assert more in (None, b""), "server sent data after the client's stall budget was spent"
        sock2.close()
    finally:
        s.stop()


def test_concurrent_downloads_are_all_recorded(station):
    """Several lamps hitting the station at once must not lose or corrupt status.json.

    Station.served is a shared counter and Station.recent a shared list, mutated from one
    handler thread per connection under ThreadingHTTPServer. Without a lock around the
    read-modify-write of `served` and the recent-list mutation + serialize-to-disk step,
    concurrent requests can lose increments (two threads read the same old value) or have
    one thread's status.json write clobber another's.
    """
    make_image(station.image_dir / "ninlight_evk_ota.img", "evk", "1.0.0", b"c" * 5000)
    expected_len = 5000 + TRAILER_SIZE
    n_workers = 20
    errors = []
    results_lock = threading.Lock()

    def worker():
        try:
            c = station.conn()
            c.request("GET", "/ninlight_evk_ota.img")
            r = c.getresponse()
            body = r.read()
            if r.status != 200 or len(body) != expected_len:
                with results_lock:
                    errors.append((r.status, len(body)))
        except Exception as e:  # pragma: no cover - surfaced via the errors list below
            with results_lock:
                errors.append(e)

    threads = [threading.Thread(target=worker) for _ in range(n_workers)]
    for t in threads:
        t.start()
    for t in threads:
        t.join(timeout=10)

    assert not errors, f"some downloads failed or returned the wrong body: {errors}"

    state = station.read_state(lambda s: s["served"] >= n_workers)
    assert state["served"] == n_workers
    assert len(state["recent"]) == min(n_workers, RECENT_MAX)
    for entry in state["recent"]:
        assert entry["status"] == "ok"
        assert entry["target"] == "evk"
        assert entry["bytes"] == expected_len