coreci
7f7868f7cf
---ci--- phase: 5 milestone: v0.1 status: complete requirements: covered: [REQ-01, REQ-02, REQ-03, REQ-04, REQ-05, REQ-06, REQ-07, REQ-08, REQ-09, REQ-10, REQ-11, REQ-12, REQ-13, REQ-14, REQ-15, REQ-16, REQ-17, REQ-18, REQ-19, REQ-20] partial: [] ---/ci---
318 lines
13 KiB
Python
318 lines
13 KiB
Python
import base64
|
|
import contextlib
|
|
import io
|
|
import json
|
|
import os
|
|
import sys
|
|
import tempfile
|
|
import threading
|
|
import unittest
|
|
import urllib.error
|
|
import urllib.request
|
|
import uuid as uuidlib
|
|
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
|
|
|
|
sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), os.pardir)))
|
|
|
|
import coreci_broker
|
|
|
|
ADMIN_USER = "admin"
|
|
ADMIN_PASS = "correct-horse-battery-staple"
|
|
REGISTER_PATH = "/api/broker/agents/register"
|
|
APPS_PATH = "/api/broker/app_registrations"
|
|
USERS_PATH = "/api/broker/user_allowlist"
|
|
STATE_PATH = "/api/broker/agents/%s/state"
|
|
STUB_PREFIX = "/api/v1/broker-agent-state/"
|
|
|
|
|
|
def _basic_header(username, password):
|
|
raw = "%s:%s" % (username, password)
|
|
return "Basic " + base64.b64encode(raw.encode("utf-8")).decode("ascii")
|
|
|
|
|
|
class _StubOwuiHandler(BaseHTTPRequestHandler):
|
|
documents = None
|
|
records = None
|
|
|
|
def log_message(self, format, *args):
|
|
pass
|
|
|
|
def do_GET(self):
|
|
agent_id = self.path[len(STUB_PREFIX):]
|
|
self.records.append(("GET", self.path, self.headers.get("Authorization"), None))
|
|
if agent_id not in self.documents:
|
|
self._send(404, {"error": 404})
|
|
return
|
|
self._send(200, self.documents[agent_id])
|
|
|
|
def do_PUT(self):
|
|
agent_id = self.path[len(STUB_PREFIX):]
|
|
length = int(self.headers.get("Content-Length") or "0")
|
|
raw = self.rfile.read(length)
|
|
self.records.append(("PUT", self.path, self.headers.get("Authorization"), raw))
|
|
self.documents[agent_id] = json.loads(raw.decode("utf-8"))
|
|
self._send(200, {"ok": True})
|
|
|
|
def _send(self, status, payload):
|
|
data = json.dumps(payload).encode("utf-8")
|
|
self.send_response(status)
|
|
self.send_header("Content-Type", "application/json")
|
|
self.send_header("Content-Length", str(len(data)))
|
|
self.end_headers()
|
|
self.wfile.write(data)
|
|
|
|
|
|
class StubOwuiServer:
|
|
def __init__(self):
|
|
self.documents, self.records = {}, []
|
|
handler = type("BoundStubHandler", (_StubOwuiHandler,),
|
|
{"documents": self.documents, "records": self.records})
|
|
self.server = ThreadingHTTPServer(("127.0.0.1", 0), handler)
|
|
threading.Thread(target=self.server.serve_forever, daemon=True).start()
|
|
|
|
@property
|
|
def base_url(self):
|
|
return "http://127.0.0.1:%d" % self.server.server_address[1]
|
|
|
|
def close(self):
|
|
self.server.shutdown()
|
|
self.server.server_close()
|
|
|
|
|
|
class _StateCase(unittest.TestCase):
|
|
backend = "internal"
|
|
|
|
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"),
|
|
"BROKER_ADMIN_USER": ADMIN_USER,
|
|
"BROKER_ADMIN_PASS": ADMIN_PASS,
|
|
"BROKER_AGENT_STATE_BACKEND": self.backend,
|
|
}
|
|
self.env.update(getattr(self, "extra_env", {}))
|
|
self._start_broker()
|
|
self.admin = {"Authorization": _basic_header(ADMIN_USER, ADMIN_PASS)}
|
|
self.app_uuid = str(uuidlib.uuid4())
|
|
self.user_id = "ceo"
|
|
self._seed(self.app_uuid, self.user_id)
|
|
self.registered = self._register("agent-1", ["b", "a"])
|
|
self.app = {"Authorization": "App %s %s" % (self.app_uuid, self.registered["token"])}
|
|
|
|
def _start_broker(self):
|
|
with contextlib.redirect_stdout(io.StringIO()):
|
|
self.broker = coreci_broker.boot(dict(self.env))
|
|
self.addCleanup(self.broker.audit.close)
|
|
self.addCleanup(self.broker.store.close)
|
|
self.server = coreci_broker.serve("127.0.0.1", 0, self.broker)
|
|
self.addCleanup(self.server.shutdown)
|
|
self.addCleanup(self.server.server_close)
|
|
self.base = "http://127.0.0.1:%d" % self.server.server_address[1]
|
|
|
|
def _request(self, method, path, payload=None, headers=None):
|
|
data = None if payload is None else json.dumps(payload).encode("utf-8")
|
|
request = urllib.request.Request(self.base + path, data=data, method=method)
|
|
for key, value in (headers or {}).items():
|
|
request.add_header(key, value)
|
|
try:
|
|
with urllib.request.urlopen(request, timeout=10) as response:
|
|
raw = response.read().decode("utf-8")
|
|
return response.status, json.loads(raw) if raw else None
|
|
except urllib.error.HTTPError as error:
|
|
raw = error.read().decode("utf-8")
|
|
return error.code, json.loads(raw) if raw else None
|
|
|
|
def _state(self, method, agent_id, payload=None, headers=None):
|
|
return self._request(method, STATE_PATH % agent_id, payload,
|
|
self.app if headers is None else headers)
|
|
|
|
def _seed(self, app_uuid, user_id):
|
|
status, _ = self._request("POST", APPS_PATH, {"uuid": app_uuid}, self.admin)
|
|
self.assertEqual(201, status)
|
|
status, _ = self._request("POST", USERS_PATH, {"user_id": user_id}, self.admin)
|
|
self.assertEqual(201, status)
|
|
|
|
def _register(self, agent_id, tags):
|
|
payload = {"agent_id": agent_id, "app_registration_uuid": self.app_uuid,
|
|
"user_id": self.user_id, "tags": list(tags)}
|
|
status, body = self._request("POST", REGISTER_PATH, payload)
|
|
self.assertEqual(201, status)
|
|
return body
|
|
|
|
|
|
class InternalStateTests(_StateCase):
|
|
def test_put_then_get_roundtrip_is_document_faithful(self):
|
|
state = {"agent_id": "agent-1", "app_registration_uuid": self.app_uuid,
|
|
"user_id": self.user_id, "tags": ["b", "a"],
|
|
"pepper": self.registered["pepper"],
|
|
"current_token": self.registered["token"],
|
|
"token_expires_at": self.registered["token_expires_at"],
|
|
"first_registered_at": "t0", "last_seen_at": "t1",
|
|
"nested": {"flags": [1, 2.5, None, True], "deep": {"k": None}}, "z": 1}
|
|
status, body = self._state("PUT", "agent-1", state)
|
|
self.assertEqual(200, status)
|
|
self.assertEqual({"ok": True}, body)
|
|
status, body = self._state("GET", "agent-1")
|
|
self.assertEqual(200, status)
|
|
self.assertEqual(state, body)
|
|
self.assertEqual(list(state), list(body))
|
|
self.assertEqual(state, self.broker.store.get_agent_state("agent-1", self.app_uuid))
|
|
|
|
def test_state_survives_broker_restart(self):
|
|
state = {"counter": 7, "notes": ["x"], "pepper": "keep-me"}
|
|
status, _ = self._state("PUT", "agent-1", state)
|
|
self.assertEqual(200, status)
|
|
self.server.shutdown()
|
|
self.server.server_close()
|
|
self.broker.store.close()
|
|
with contextlib.redirect_stdout(io.StringIO()):
|
|
second = coreci_broker.boot(dict(self.env))
|
|
self.addCleanup(second.audit.close)
|
|
self.addCleanup(second.store.close)
|
|
server = coreci_broker.serve("127.0.0.1", 0, second)
|
|
self.addCleanup(server.shutdown)
|
|
self.addCleanup(server.server_close)
|
|
self.base = "http://127.0.0.1:%d" % server.server_address[1]
|
|
status, body = self._state("GET", "agent-1")
|
|
self.assertEqual(200, status)
|
|
self.assertEqual(state, body)
|
|
row = second.store.get_agent("agent-1", self.app_uuid)
|
|
self.assertTrue(coreci_broker.verify_token(self.registered["token"], row[4]))
|
|
|
|
def test_identity_mismatch_is_403_on_get_and_put(self):
|
|
second = self._register("agent-2", ["coreci"])
|
|
for method, payload in (("GET", None), ("PUT", {"who": "agent-2"})):
|
|
status, body = self._state(method, "agent-2", payload)
|
|
self.assertEqual(403, status, method)
|
|
self.assertEqual(403, body["error"])
|
|
self.assertTrue(body["message"])
|
|
status, body = self._state("GET", "agent-2", headers={
|
|
"Authorization": "App %s %s" % (self.app_uuid, second["token"])})
|
|
self.assertEqual(404, status)
|
|
|
|
def test_unregistered_agent_id_is_404_on_get(self):
|
|
status, body = self._state("GET", "ghost-agent")
|
|
self.assertEqual(404, status)
|
|
self.assertEqual({"error": 404, "message": "not found"}, body)
|
|
|
|
def test_get_without_state_is_404(self):
|
|
status, body = self._state("GET", "agent-1")
|
|
self.assertEqual(404, status)
|
|
self.assertEqual({"error": 404, "message": "not found"}, body)
|
|
|
|
def test_put_non_dict_is_400(self):
|
|
for payload in ([1, 2], "nope", 5, None):
|
|
status, body = self._state("PUT", "agent-1", payload)
|
|
self.assertEqual(400, status, repr(payload))
|
|
self.assertEqual(400, body["error"])
|
|
status, _ = self._state("GET", "agent-1")
|
|
self.assertEqual(404, status)
|
|
|
|
def test_put_above_256kib_is_400(self):
|
|
status, body = self._state("PUT", "agent-1", {"pad": "x" * (300 * 1024)})
|
|
self.assertEqual(400, status)
|
|
self.assertEqual(400, body["error"])
|
|
status, _ = self._state("GET", "agent-1")
|
|
self.assertEqual(404, status)
|
|
|
|
def test_put_at_exactly_256kib_is_stored(self):
|
|
pad = "x" * (256 * 1024 - len(json.dumps({"pad": ""})))
|
|
status, body = self._state("PUT", "agent-1", {"pad": pad})
|
|
self.assertEqual(200, status)
|
|
status, body = self._state("GET", "agent-1")
|
|
self.assertEqual(200, status)
|
|
self.assertEqual({"pad": pad}, body)
|
|
|
|
def test_missing_app_header_is_401(self):
|
|
status, body = self._state("GET", "agent-1", headers={})
|
|
self.assertEqual(401, status)
|
|
self.assertEqual({"error": 401, "message": "invalid app credentials"}, body)
|
|
|
|
|
|
class OwuiStateTests(_StateCase):
|
|
backend = "owui"
|
|
|
|
def setUp(self):
|
|
self.stub = StubOwuiServer()
|
|
self.addCleanup(self.stub.close)
|
|
self.extra_env = {
|
|
"BROKER_OWUI_BASE_URL": self.stub.base_url,
|
|
"BROKER_OWUI_API_KEY": "test-key",
|
|
}
|
|
super().setUp()
|
|
|
|
def test_put_then_get_roundtrip_flows_through_stub(self):
|
|
state = {"agent_id": "agent-1", "current_token": "tok", "n": 3}
|
|
status, body = self._state("PUT", "agent-1", state)
|
|
self.assertEqual(200, status)
|
|
self.assertEqual({"ok": True}, body)
|
|
self.assertEqual(state, self.stub.documents["agent-1"])
|
|
status, body = self._state("GET", "agent-1")
|
|
self.assertEqual(200, status)
|
|
self.assertEqual(state, body)
|
|
puts = [record for record in self.stub.records if record[0] == "PUT"]
|
|
gets = [record for record in self.stub.records if record[0] == "GET"]
|
|
self.assertEqual(1, len(puts))
|
|
self.assertEqual(1, len(gets))
|
|
for _, path, auth, _ in self.stub.records:
|
|
self.assertEqual("Bearer test-key", auth)
|
|
self.assertEqual(STUB_PREFIX + "agent-1", puts[0][1])
|
|
self.assertEqual(STUB_PREFIX + "agent-1", gets[0][1])
|
|
self.assertEqual(state, json.loads(puts[0][3].decode("utf-8")))
|
|
|
|
def test_stub_miss_maps_to_broker_404(self):
|
|
status, body = self._state("GET", "agent-1")
|
|
self.assertEqual(404, status)
|
|
self.assertEqual({"error": 404, "message": "not found"}, body)
|
|
|
|
def test_identity_enforced_before_backend_call(self):
|
|
self._register("agent-2", ["coreci"])
|
|
status, body = self._state("GET", "agent-2")
|
|
self.assertEqual(403, status)
|
|
status, body = self._state("PUT", "agent-2", {"x": 1})
|
|
self.assertEqual(403, status)
|
|
self.assertEqual([], self.stub.records)
|
|
|
|
|
|
class BackendToggleTests(unittest.TestCase):
|
|
def _boot(self, extra):
|
|
tmp = tempfile.TemporaryDirectory()
|
|
self.addCleanup(tmp.cleanup)
|
|
env = {
|
|
"BROKER_ENV_PATH": os.path.join(tmp.name, ".env"),
|
|
"BROKER_DB_PATH": os.path.join(tmp.name, "broker.db"),
|
|
"BROKER_AUDIT_LOG_PATH": os.path.join(tmp.name, "broker.audit.log"),
|
|
}
|
|
env.update(extra)
|
|
with contextlib.redirect_stdout(io.StringIO()):
|
|
broker = coreci_broker.boot(env)
|
|
self.addCleanup(broker.audit.close)
|
|
self.addCleanup(broker.store.close)
|
|
return broker
|
|
|
|
def test_default_backend_is_internal(self):
|
|
broker = self._boot({})
|
|
self.assertIsInstance(broker.agent_state_backend, coreci_broker.InternalAgentStateBackend)
|
|
|
|
def test_owui_backend_selected_by_config(self):
|
|
broker = self._boot({
|
|
"BROKER_AGENT_STATE_BACKEND": "owui",
|
|
"BROKER_OWUI_BASE_URL": "http://127.0.0.1:9999/",
|
|
"BROKER_OWUI_API_KEY": "k",
|
|
})
|
|
backend = broker.agent_state_backend
|
|
self.assertIsInstance(backend, coreci_broker.OwuiAgentStateBackend)
|
|
self.assertEqual("http://127.0.0.1:9999", backend.base_url)
|
|
self.assertEqual("k", backend.api_key)
|
|
|
|
def test_unknown_backend_value_falls_back_to_internal(self):
|
|
broker = self._boot({"BROKER_AGENT_STATE_BACKEND": "banana"})
|
|
self.assertIsInstance(broker.agent_state_backend, coreci_broker.InternalAgentStateBackend)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main() |