Return as soon as we have a port for the current hostname

README.md

@@ -98,6 +98,12 @@ You can set this environment variable in the `.env` file for your project e.g.
 VIRTUAL_HOST=mywebsite.localhost
 ```
 
+You can also specify multiple hostnames:
+
+```
+VIRTUAL_HOST=mywebsite.localhost,myotherwebsite.localhost
+```
+
 Then you can start (or restart) your project, visit `http://mywebsite.localhost:8080` in your browser, and the traffic will magically be routed to the right place.
 
 (Note that at least Chrome automatically routes everything with the TLD `.localhost` to 127.0.0.1. Other browsers may or may not follow this standard).

crab/router.py

@@ -7,16 +7,13 @@
 import sys
 
 
-def get_routes():
-    routes = {}
+def get_route_for_hostname(hostname):
     for process in psutil.process_iter(attrs=["environ"]):
-        try:
-            host = process.info["environ"]["VIRTUAL_HOST"]
-            port = process.info["environ"]["PORT"]
-            routes[host] = port
-        except:
-            pass
-    return routes
+        process_env = process.info["environ"]
+        if not process_env or "PORT" not in process_env:
+            continue
+        if hostname in process_env.get("VIRTUAL_HOST", "").split(","):
+            return process_env["PORT"]
 
 
 app = Flask(__name__, static_folder=None)
@@ -26,14 +23,14 @@ def get_routes():
 
 @app.endpoint("proxy")
 def proxy(path):
-    routes = get_routes()
     hostname = urlparse(request.base_url).hostname
-    if hostname not in routes:
+    upstream_port = get_route_for_hostname(hostname)
+    if upstream_port is None:
         app.logger.warn("No backend for %s", hostname)
         abort(502)
 
     path = request.full_path if request.args else request.path
-    target_url = "http://localhost:" + routes[hostname] + path
+    target_url = "http://localhost:" + upstream_port + path
     app.logger.info(
         "Routing request to backend - %s %s%s", request.method, hostname, path
     )

tests/test_router.py

@@ -1,4 +1,4 @@
-from crab.router import get_routes
+from crab.router import get_route_for_hostname
 from unittest import TestCase
 import subprocess
 import os
@@ -10,6 +10,20 @@ def test_get_routes(self):
             ["sleep", "5"],
             env={**os.environ, "VIRTUAL_HOST": "test.localhost", "PORT": "1234"},
         ) as subproc:
-            routes = get_routes()
-            self.assertEqual(routes, {"test.localhost": "1234"})
+            self.assertEqual(get_route_for_hostname("test.localhost"), "1234")
+            self.assertIsNone(get_route_for_hostname("not-test.localhost"))
+            subproc.terminate()
+
+    def test_multiple_routes(self):
+        with subprocess.Popen(
+            ["sleep", "5"],
+            env={
+                **os.environ,
+                "VIRTUAL_HOST": "test.localhost,other-test.localhost",
+                "PORT": "1234",
+            },
+        ) as subproc:
+            self.assertEqual(get_route_for_hostname("test.localhost"), "1234")
+            self.assertEqual(get_route_for_hostname("other-test.localhost"), "1234")
+            self.assertIsNone(get_route_for_hostname("not-test.localhost"))
             subproc.terminate()