serve: reverse-proxy Host allowlist + quiet prefill keepalive (#597)

c/openai_server.py

@@ -1304,7 +1304,7 @@ class APIServer(ThreadingHTTPServer):
 
     def __init__(self, address, engine, model_id, api_key=None, max_tokens=1024,
                  cors_origins=DEFAULT_CORS_ORIGINS, max_queue=8, queue_timeout=300,
-                 kv_slots=1):
+                 kv_slots=1, allowed_hosts=()):
         super().__init__(address, APIHandler)
         self.engine = engine
         self.model_id = model_id
@@ -1313,6 +1313,11 @@ def __init__(self, address, engine, model_id, api_key=None, max_tokens=1024,
         self.scheduler = GenerationScheduler(max_queue, queue_timeout, kv_slots)
         self.kv_slots = kv_slots
         self.cors_origins = tuple(cors_origins)
+        # Extra Host header values trusted past the DNS-rebinding guard, for a
+        # reverse proxy / MagicDNS in front of the loopback bind (#597). Explicit
+        # opt-in only: no wildcard, default stays loopback + bind address.
+        self.allowed_hosts = tuple(
+            h.strip().lower() for h in allowed_hosts if h and h.strip())
         self.created = int(time.time())
 
 
@@ -1383,6 +1388,7 @@ def _check_host(self):
             name = host                                            # bare host / bracketless ipv6
         name = name.strip().lower()
         allowed = set(self.LOOPBACK_HOSTS)
+        allowed.update(self.server.allowed_hosts)          # #597: operator-trusted reverse-proxy names
         try:
             allowed.add(str(self.server.server_address[0]).strip("[]").lower())
         except Exception:
@@ -1657,7 +1663,14 @@ def event(choices, usage_marker=False):
                         connected = False
 
             def _keepalive():
-                ping = [{"index": 0, "delta": ({"reasoning_content": "."} if chat else {"content": ""}),
+                # #597: an empty delta already resets the client's idle timer without
+                # painting hundreds of dots in the reasoning panel during a minutes-long
+                # cold prefill. COLI_VISIBLE_KEEPALIVE=1 restores the old visible "." for
+                # diagnosing whether keepalives are being delivered at all.
+                visible = os.environ.get("COLI_VISIBLE_KEEPALIVE") == "1"
+                ping = [{"index": 0,
+                         "delta": ({"reasoning_content": "." if visible else ""} if chat
+                                   else {"content": ""}),
                          "logprobs": None, "finish_reason": None}]
                 while not ka_stop.wait(1.0):
                     if not connected:
@@ -2028,7 +2041,7 @@ def completion(self, body, request_id):
 
 def serve(model, host="127.0.0.1", port=8000, model_id="glm-5.2-colibri", api_key=None,
           cap=8, max_tokens=1024, engine=None, env=None, cors_origins=None,
-          max_queue=8, queue_timeout=300, kv_slots=1):
+          max_queue=8, queue_timeout=300, kv_slots=1, allowed_hosts=()):
     if engine is None:
         engine = default_engine()
     if not 1 <= max_tokens:
@@ -2055,7 +2068,7 @@ def serve(model, host="127.0.0.1", port=8000, model_id="glm-5.2-colibri", api_ke
     # Bind before starting the 744B engine. A stale/occupied port must fail in
     # milliseconds rather than loading hundreds of GB and leaking a child.
     server = APIServer((host, port), None, model_id, api_key, max_tokens, origins,
-                       max_queue, queue_timeout, kv_slots)
+                       max_queue, queue_timeout, kv_slots, allowed_hosts=allowed_hosts)
     runtime = None
     previous_sigterm = signal.getsignal(signal.SIGTERM)
     try:
@@ -2090,6 +2103,11 @@ def main():
     parser.add_argument("--queue-timeout", type=float,
                         default=float(os.environ.get("COLI_QUEUE_TIMEOUT", "300")))
     parser.add_argument("--kv-slots", type=int, default=int(os.environ.get("COLI_KV_SLOTS", "1")))
+    parser.add_argument("--allowed-host", action="append",
+        default=[h.strip() for h in os.environ.get("COLI_ALLOWED_HOSTS", "").split(",") if h.strip()],
+        help="additional Host header value accepted by the DNS-rebinding guard "
+             "(reverse proxy / MagicDNS in front of the loopback bind); repeat as needed, "
+             "or set COLI_ALLOWED_HOSTS as a comma-separated list")
     args = parser.parse_args()
     global ARCH
     ARCH = args.arch
@@ -2103,7 +2121,8 @@ def main():
         args.model_id = "inkling-colibri" if ARCH == "inkling" else "glm-5.2-colibri"
     serve(args.model, args.host, args.port, args.model_id, args.api_key,
           args.cap,args.max_tokens,args.engine,cors_origins=args.cors_origin,
-          max_queue=args.max_queue,queue_timeout=args.queue_timeout,kv_slots=args.kv_slots)
+          max_queue=args.max_queue,queue_timeout=args.queue_timeout,kv_slots=args.kv_slots,
+          allowed_hosts=args.allowed_host)
 
 
 if __name__ == "__main__":

c/tests/test_openai_server.py

@@ -1,3 +1,4 @@
+import http.client
 import io
 import json
 import math
@@ -957,5 +958,52 @@ def test_rejects_tool_choice_without_tools(self):
             generation_options({"messages": [], "tool_choice": "required"}, 128)
 
 
+class AllowedHostsTest(unittest.TestCase):
+    """#597: the DNS-rebinding guard must accept operator-trusted reverse-proxy
+    Host values, while still rejecting everything else by default."""
+
+    def _make_server(self, allowed_hosts=()):
+        server = APIServer(("127.0.0.1", 0), FakeEngine(), "test-model",
+                           allowed_hosts=allowed_hosts)
+        thread = threading.Thread(target=server.serve_forever, daemon=True)
+        thread.start()
+        self.addCleanup(thread.join, 2)
+        self.addCleanup(server.server_close)
+        self.addCleanup(server.shutdown)
+        self.addCleanup(server.scheduler.close)
+        return server
+
+    def _get_models(self, port, host_header):
+        # http.client lets us set an arbitrary Host header (urlopen forces its own).
+        conn = http.client.HTTPConnection("127.0.0.1", port, timeout=2)
+        try:
+            conn.putrequest("GET", "/v1/models", skip_host=True)
+            conn.putheader("Host", host_header)
+            conn.endheaders()
+            return conn.getresponse().status
+        finally:
+            conn.close()
+
+    def test_allowlist_wiring_normalises_and_filters(self):
+        server = self._make_server(allowed_hosts=("  Proxy.Example.TS.net ", "", "   "))
+        self.assertEqual(server.allowed_hosts, ("proxy.example.ts.net",))
+
+    def test_trusted_reverse_proxy_host_is_accepted(self):
+        server = self._make_server(allowed_hosts=("proxy.example.ts.net",))
+        port = server.server_port
+        # trusted name (with a port suffix, case-insensitive) passes the guard
+        self.assertEqual(self._get_models(port, "Proxy.Example.TS.net:8000"), 200)
+        # loopback still works, unaffected by the allowlist
+        self.assertEqual(self._get_models(port, "localhost"), 200)
+
+    def test_untrusted_host_is_rejected_by_default(self):
+        server = self._make_server()               # no allowlist: loopback/bind only
+        self.assertEqual(self._get_models(server.server_port, "evil.example.com"), 403)
+
+    def test_untrusted_host_still_rejected_with_allowlist(self):
+        server = self._make_server(allowed_hosts=("proxy.example.ts.net",))
+        self.assertEqual(self._get_models(server.server_port, "evil.example.com"), 403)
+
+
 if __name__ == "__main__":
     unittest.main()