From 3c19b6d1fb9e2233583a3f2e7be6b4091b943958 Mon Sep 17 00:00:00 2001
From: Dan Allan <dallan@bnl.gov>
Date: Sat, 1 Jun 2024 11:37:53 -0400
Subject: [PATCH] Fix username logging on 500 responses. (#754)

* Fix username logging on 500 responses.

A bug causes the server to send no response, such that the client
raised:

```
RemoteProtocolError: Server disconnected without sending a response.
```

instead of returning 500 response. This was a regression
introduced in #750.

Needs test.

Closes #750

* Add comment

* Refine comments

* TST: Reproduce 'Server disconnected without sending a response' in test.

* Update CHANGELOG

* Clean up prints in test

* Refactor so server can be shut down.
---
 CHANGELOG.md                         |  7 +++
 tiled/_tests/test_server.py          | 69 ++++++++++++++++++++++++++++
 tiled/server/app.py                  |  4 ++
 tiled/server/principal_log_filter.py |  9 +++-
 4 files changed, 87 insertions(+), 2 deletions(-)
 create mode 100644 tiled/_tests/test_server.py

diff --git a/CHANGELOG.md b/CHANGELOG.md
index da80e5433..f2790c131 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -3,6 +3,13 @@ Write the date in place of the "Unreleased" in the case a new version is release
 
 # Changelog
 
+## Unreleased
+
+- Fix regression introduced in the previous release (v0.1.0b1) where exceptions
+  raised in the server sent _no_ response instead of properly sending a 500
+  response. (This presents in the client as, "Server disconnected without
+  sending a response.") A test now protects against this class of regression.
+
 ## v0.1.0b2 (2024-05-28)
 
 ## Changed
diff --git a/tiled/_tests/test_server.py b/tiled/_tests/test_server.py
new file mode 100644
index 000000000..c09cdaf5a
--- /dev/null
+++ b/tiled/_tests/test_server.py
@@ -0,0 +1,69 @@
+import contextlib
+import threading
+import time
+
+import uvicorn
+from fastapi import APIRouter
+from starlette.status import HTTP_500_INTERNAL_SERVER_ERROR
+
+from ..catalog import in_memory
+from ..client import from_uri
+from ..server.app import build_app
+from ..server.logging_config import LOGGING_CONFIG
+
+router = APIRouter()
+
+
+class Server(uvicorn.Server):
+    # https://github.com/encode/uvicorn/discussions/1103#discussioncomment-941726
+
+    def install_signal_handlers(self):
+        pass
+
+    @contextlib.contextmanager
+    def run_in_thread(self):
+        thread = threading.Thread(target=self.run)
+        thread.start()
+        try:
+            # Wait for server to start up, or raise TimeoutError.
+            for _ in range(100):
+                time.sleep(0.1)
+                if self.started:
+                    break
+            else:
+                raise TimeoutError("Server did not start in 10 seconds.")
+            host, port = self.servers[0].sockets[0].getsockname()
+            yield f"http://{host}:{port}"
+        finally:
+            self.should_exit = True
+            thread.join()
+
+
+@router.get("/error")
+def error():
+    1 / 0  # error!
+
+
+def test_500_response():
+    """
+    Test that unexpected server error returns 500 response.
+
+    This test is meant to catch regressions in which server exceptions can
+    result in the server sending no response at all, leading clients to raise
+    like:
+
+    httpx.RemoteProtocolError: Server disconnected without sending a response.
+
+    This can happen when bugs are introduced in the middleware layer.
+    """
+    API_KEY = "secret"
+    catalog = in_memory()
+    app = build_app(catalog, {"single_user_api_key": API_KEY})
+    app.include_router(router)
+    config = uvicorn.Config(app, port=0, loop="asyncio", log_config=LOGGING_CONFIG)
+    server = Server(config)
+
+    with server.run_in_thread() as url:
+        client = from_uri(url, api_key=API_KEY)
+        response = client.context.http_client.get(f"{url}/error")
+    assert response.status_code == HTTP_500_INTERNAL_SERVER_ERROR
diff --git a/tiled/server/app.py b/tiled/server/app.py
index 1ba9dfcd6..e5c953b3a 100644
--- a/tiled/server/app.py
+++ b/tiled/server/app.py
@@ -330,6 +330,10 @@ async def unsupported_query_type_exception_handler(
     async def unhandled_exception_handler(
         request: Request, exc: Exception
     ) -> JSONResponse:
+        # The current_principal_logging_filter middleware will not have
+        # had a chance to finish running, so set the principal here.
+        principal = getattr(request.state, "principal", None)
+        current_principal.set(principal)
         return await http_exception_handler(
             request,
             HTTPException(
diff --git a/tiled/server/principal_log_filter.py b/tiled/server/principal_log_filter.py
index faf4f0fd9..cf5a534f6 100644
--- a/tiled/server/principal_log_filter.py
+++ b/tiled/server/principal_log_filter.py
@@ -8,8 +8,13 @@ class PrincipalFilter(Filter):
     """Logging filter to attach username or Service Principal UUID to LogRecord"""
 
     def filter(self, record: LogRecord) -> bool:
-        principal = current_principal.get()
-        if isinstance(principal, SpecialUsers):
+        principal = current_principal.get(None)
+        if principal is None:
+            # This will only occur if an uncaught exception was raised in the
+            # server before the authentication code ran. This will always be
+            # associated with a 500 Internal Server Error response.
+            short_name = "unset"
+        elif isinstance(principal, SpecialUsers):
             short_name = f"{principal.value}"
         elif principal.type == "service":
             short_name = f"service:{principal.uuid}"