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

import hashlib
import http.client
import json
import os
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)


@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_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_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