diff --git a/README.md b/README.md index a5d7b51..c3e6335 100644 --- a/README.md +++ b/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). diff --git a/crab/router.py b/crab/router.py index f1d8eb7..32c9f34 100644 --- a/crab/router.py +++ b/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 ) diff --git a/tests/test_router.py b/tests/test_router.py index 2e4f149..7b87dc6 100644 --- a/tests/test_router.py +++ b/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()