broker/tests/test_hardening.py

import contextlib
import io
import json
import os
import sqlite3
import sys
import tempfile
import threading
import unittest
import urllib.error
import urllib.request
from unittest import mock

sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), os.pardir)))

import coreci_broker

THREAD_COUNT = 8
WRITES_PER_THREAD = 16
MAX_BODY_BYTES = 1024 * 1024


class HardeningTests(unittest.TestCase):
    def setUp(self):
        tmp = tempfile.TemporaryDirectory()
        self.addCleanup(tmp.cleanup)
        self.tmp = tmp.name
        self.env = {
            "BROKER_ENV_PATH": os.path.join(self.tmp, ".env"),
            "BROKER_DB_PATH": os.path.join(self.tmp, "broker.db"),
            "BROKER_AUDIT_LOG_PATH": os.path.join(self.tmp, "broker.audit.log"),
        }

    def _boot_and_serve(self):
        buffer = io.StringIO()
        with contextlib.redirect_stdout(buffer):
            broker = coreci_broker.boot(dict(self.env))
        server = coreci_broker.serve("127.0.0.1", 0, broker)
        self.addCleanup(broker.audit.close)
        self.addCleanup(broker.store.close)
        self.addCleanup(server.server_close)
        self.addCleanup(server.shutdown)
        return broker, server

    def test_oversized_declared_content_length_post_rejected(self):
        broker, server = self._boot_and_serve()
        port = server.server_address[1]
        request = urllib.request.Request(
            "http://127.0.0.1:%d/healthz" % port,
            data=b"",
            headers={"Content-Length": str(MAX_BODY_BYTES + 1)},
            method="POST",
        )
        with self.assertRaises(urllib.error.HTTPError) as raised:
            urllib.request.urlopen(request, timeout=5)
        error = raised.exception
        self.assertEqual(400, error.code)
        payload = json.loads(error.read().decode("utf-8"))
        self.assertEqual({"error": 400, "message": "body too large"}, payload)

    def test_handler_timeout_attribute_is_30(self):
        self.assertEqual(30, coreci_broker.BrokerRequestHandler.timeout)
        broker, server = self._boot_and_serve()
        self.assertEqual(30, server.RequestHandlerClass.timeout)

    def test_concurrent_store_writes_from_eight_threads(self):
        db_path = os.path.join(self.tmp, "concurrent.db")
        store = coreci_broker.Store(db_path)
        self.addCleanup(store.close)
        errors = []
        barrier = threading.Barrier(THREAD_COUNT)

        def worker(index):
            try:
                barrier.wait(timeout=5)
                for slot in range(WRITES_PER_THREAD):
                    store.set_config(
                        "key-%d-%d" % (index, slot), "value-%d-%d" % (index, slot)
                    )
                store.get_config("key-%d-0" % index)
            except BaseException as error:
                errors.append(error)

        threads = [
            threading.Thread(target=worker, args=(index,))
            for index in range(THREAD_COUNT)
        ]
        for thread in threads:
            thread.start()
        for thread in threads:
            thread.join(timeout=30)
        self.assertEqual([], errors)
        for index in range(THREAD_COUNT):
            for slot in range(WRITES_PER_THREAD):
                key = "key-%d-%d" % (index, slot)
                expected = "value-%d-%d" % (index, slot)
                self.assertEqual(expected, store.get_config(key))
        conn = sqlite3.connect(db_path)
        self.addCleanup(conn.close)
        count = conn.execute("SELECT COUNT(*) FROM broker_config").fetchone()[0]
        self.assertEqual(THREAD_COUNT * WRITES_PER_THREAD, count)

    def test_first_boot_env_write_path_still_works(self):
        buffer = io.StringIO()
        with contextlib.redirect_stdout(buffer):
            config = coreci_broker.load_config(dict(self.env))
        env_path = self.env["BROKER_ENV_PATH"]
        self.assertTrue(os.path.isfile(env_path))
        self.assertEqual(0o600, os.stat(env_path).st_mode & 0o777)
        stored = {}
        with open(env_path, "r", encoding="utf-8") as handle:
            for raw in handle:
                line = raw.strip()
                if not line or line.startswith("#") or "=" not in line:
                    continue
                key, _, value = line.partition("=")
                stored[key.strip()] = value.strip()
        self.assertEqual(config["BROKER_ADMIN_PASS"], stored["BROKER_ADMIN_PASS"])
        self.assertEqual(
            config["BROKER_INSTANCE_PEPPER"], stored["BROKER_INSTANCE_PEPPER"]
        )
        self.assertEqual(1, buffer.getvalue().count(config["BROKER_ADMIN_PASS"]))
        second = io.StringIO()
        with contextlib.redirect_stdout(second):
            coreci_broker.load_config(dict(self.env))
        self.assertEqual("", second.getvalue())

    def test_append_env_entries_write_error_propagates_with_single_close(self):
        env_path = os.path.join(self.tmp, ".env")
        with open(env_path, "w", encoding="utf-8") as handle:
            handle.write("seed=1\n")
        closes = []

        class FailingHandle:
            def __init__(self, fd):
                self.fd = fd

            def write(self, data):
                raise OSError("simulated write failure")

            def __enter__(self):
                return self

            def __exit__(self, exc_type, exc_value, traceback):
                os.close(self.fd)
                closes.append(self.fd)
                return False

        def fake_fdopen(fd, *args, **kwargs):
            return FailingHandle(fd)

        with mock.patch.object(os, "fdopen", fake_fdopen):
            with self.assertRaises(OSError) as raised:
                coreci_broker._append_env_entries(env_path, [("NEW", "value")])
        self.assertEqual(1, len(closes))
        self.assertIn("simulated write failure", str(raised.exception))
        with open(env_path, "r", encoding="utf-8") as handle:
            self.assertEqual("seed=1\n", handle.read())


if __name__ == "__main__":
    unittest.main()