Merge pull request #131 from dolfinus/master

Fix handling default server routes in TraefikTomlProxy

.gitignore

@@ -102,3 +102,5 @@ venv.bak/
 
 # mypy
 .mypy_cache/
+
+/bin

jupyterhub_traefik_proxy/consul.py

@@ -165,7 +165,9 @@ async def _kv_atomic_delete_route_parts(self, jupyterhub_routespec, route_keys):
 
         index, v = await self.kv_client.kv.get(escaped_jupyterhub_routespec)
         if v is None:
-            self.log.warning("Route %s doesn't exist. Nothing to delete", jupyterhub_routespec)
+            self.log.warning(
+                "Route %s doesn't exist. Nothing to delete", jupyterhub_routespec
+            )
             return True, None
         target = v["Value"]
         escaped_target = escapism.escape(target, safe=self.key_safe_chars)

jupyterhub_traefik_proxy/etcd.py

@@ -147,7 +147,9 @@ async def _kv_atomic_add_route_parts(
     async def _kv_atomic_delete_route_parts(self, jupyterhub_routespec, route_keys):
         value = await maybe_future(self._etcd_get(jupyterhub_routespec))
         if value is None:
-            self.log.warning("Route %s doesn't exist. Nothing to delete", jupyterhub_routespec)
+            self.log.warning(
+                "Route %s doesn't exist. Nothing to delete", jupyterhub_routespec
+            )
             return True, None
 
         target = value.decode()

jupyterhub_traefik_proxy/proxy.py

@@ -248,20 +248,17 @@ async def _setup_traefik_static_config(self):
         self.static_config["api"] = {"dashboard": True, "entrypoint": "auth_api"}
         self.static_config["wss"] = {"protocol": "http"}
 
-
     def _routespec_to_traefik_path(self, routespec):
         path = self.validate_routespec(routespec)
-        if path != '/' and path.endswith('/'):
-            path = path.rstrip('/')
+        if path != "/" and path.endswith("/"):
+            path = path.rstrip("/")
         return path
 
-
     def _routespec_from_traefik_path(self, routespec):
-        if not routespec.endswith('/'):
-            routespec = routespec + '/'
+        if not routespec.endswith("/"):
+            routespec = routespec + "/"
         return routespec
 
-
     async def start(self):
         """Start the proxy.
 

jupyterhub_traefik_proxy/toml.py

@@ -51,9 +51,9 @@ def __init__(self, **kwargs):
             self.routes_cache = traefik_utils.load_routes(self.toml_dynamic_config_file)
         except FileNotFoundError:
             self.routes_cache = {}
-        finally:
-            if not self.routes_cache:
-                self.routes_cache = {"backends": {}, "frontends": {}}
+
+        if not self.routes_cache:
+            self.routes_cache = {"backends": {}, "frontends": {}}
 
     async def _setup_traefik_static_config(self):
         await super()._setup_traefik_static_config()
@@ -100,8 +100,8 @@ def _clean_resources(self):
             raise
 
     def _get_route_unsafe(self, traefik_routespec):
-        safe = string.ascii_letters + string.digits + "-"
-        escaped_routespec = escapism.escape(traefik_routespec, safe=safe)
+        backend_alias = traefik_utils.generate_alias(traefik_routespec, "backend")
+        frontend_alias = traefik_utils.generate_alias(traefik_routespec, "frontend")
         routespec = self._routespec_from_traefik_path(traefik_routespec)
         result = {"data": "", "target": "", "routespec": routespec}
 
@@ -118,12 +118,12 @@ def get_target_data(d, to_find):
                 if isinstance(v, dict):
                     get_target_data(v, to_find)
 
-        for key, value in self.routes_cache["backends"].items():
-            if escaped_routespec in key:
-                get_target_data(value, "url")
-        for key, value in self.routes_cache["frontends"].items():
-            if escaped_routespec in key:
-                get_target_data(value, "data")
+        if backend_alias in self.routes_cache["backends"]:
+            get_target_data(self.routes_cache["backends"][backend_alias], "url")
+
+        if frontend_alias in self.routes_cache["frontends"]:
+            get_target_data(self.routes_cache["frontends"][frontend_alias], "data")
+
         if not result["data"] and not result["target"]:
             self.log.info("No route for {} found!".format(routespec))
             result = None
@@ -214,18 +214,13 @@ async def delete_route(self, routespec):
         **Subclasses must define this method**
         """
         routespec = self._routespec_to_traefik_path(routespec)
-        safe = string.ascii_letters + string.digits + "-"
-        escaped_routespec = escapism.escape(routespec, safe=safe)
+        backend_alias = traefik_utils.generate_alias(routespec, "backend")
+        frontend_alias = traefik_utils.generate_alias(routespec, "frontend")
 
         async with self.mutex:
-            for key, value in self.routes_cache["frontends"].items():
-                if escaped_routespec in key:
-                    del self.routes_cache["frontends"][key]
-                    break
-            for key, value in self.routes_cache["backends"].items():
-                if escaped_routespec in key:
-                    del self.routes_cache["backends"][key]
-                    break
+            self.routes_cache["frontends"].pop(frontend_alias, None)
+            self.routes_cache["backends"].pop(backend_alias, None)
+
         traefik_utils.persist_routes(self.toml_dynamic_config_file, self.routes_cache)
 
     async def get_all_routes(self):

performance/check_perf.py

@@ -42,9 +42,9 @@ async def add_route_perf(proxy, route_idx, stdout_print):
 async def delete_route_perf(proxy, route_idx, stdout_print):
     """
     Computes time taken (ns) to delete the "route_idx"
-    route from the proxy's routing table (which 
+    route from the proxy's routing table (which
     contains route_idx + 1 routes when ran sequentially).
-    
+
     It assumes the route to be deleted exists.
 
     Returns a tuple:

tests/proxytest.py

@@ -101,95 +101,186 @@ async def wait_for_services(urls):
 
 
 @pytest.mark.parametrize(
-    "routespec",
+    "routespec, existing_routes",
     [
-        "/",  # default route
-        "/has%20space/foo/",
-        "/missing-trailing/slash",
-        "/has/@/",
-        "/has/" + quote("üñîçø∂é"),
-        "host.name/path/",
-        "other.host/path/no/slash",
+        # default route
+        (
+            "/",
+            [
+                "/abc",
+                "/has%20space/",
+                "/has%20space/foo/",
+                "/missing-trailing/",
+                "/missing-trailing/slash",
+                "/has/",
+                "/has/@/",
+                "host.name/",
+                "host.name/path/",
+                "other.host/",
+                "other.host/path/",
+                "other.host/path/no/",
+                "other.host/path/no/slash",
+            ],
+        ),
+        ("/has%20space/foo/", ["/", "/has%20space/", "/has%20space/foo/abc/"]),
+        (
+            "/missing-trailing/slash",
+            ["/", "/missing-trailing/", "/missing-trailing/slash/abc"],
+        ),
+        ("/has/@/", ["/", "/has/", "/has/@/abc/"]),
+        (
+            "/has/" + quote("üñîçø∂é"),
+            ["/", "/has/", "/has/" + quote("üñîçø∂é") + "/abc/"],
+        ),
+        ("host.name/path/", ["/", "host.name/", "host.name/path/abc/"]),
+        (
+            "other.host/path/no/slash",
+            [
+                "/",
+                "other.host/",
+                "other.host/path/",
+                "other.host/path/no/",
+                "other.host/path/no/slash/abc/",
+            ],
+        ),
     ],
 )
-async def test_add_get_delete(proxy, launch_backend, routespec):
-    target = "http://127.0.0.1:9000"
+async def test_add_get_delete(
+    request, proxy, launch_backend, routespec, existing_routes, event_loop
+):
+    default_target = "http://127.0.0.1:9000"
     data = {"test": "test1", "user": "username"}
-    expected_output = {
-        "routespec": routespec if routespec.endswith("/") else routespec + "/",
-        "target": target,
-        "data": data,
-    }
 
-    backend_port = urlparse(target).port
-    launch_backend(backend_port)
-    await wait_for_services([proxy.public_url, target])
+    default_backend = urlparse(default_target)
+    extra_backends = []
 
-    # host-routes when not host-routing raises an error
-    # and vice versa
-    subdomain_host = False
-    try:
-        subdomain_host = bool(proxy.app.subdomain_host)
-    except AttributeError:
-        pass
-    finally:
-        expect_value_error = subdomain_host ^ (not routespec.startswith("/"))
+    proxy_url = proxy.public_url.rstrip("/")
+
+    def normalize_spec(spec):
+        return proxy.validate_routespec(spec)
+
+    def expected_output(spec, url):
+        return {
+            "routespec": normalize_spec(spec),
+            "target": url,
+            "data": data,
+        }
+
+    # just use existing Jupyterhub check instead of making own one
+    def expect_value_error(spec):
+        try:
+            normalize_spec(spec)
+        except ValueError:
+            return True
+
+        return False
 
     @contextmanager
-    def context():
-        if expect_value_error:
+    def context(spec):
+        if expect_value_error(spec):
             with pytest.raises(ValueError):
                 yield
         else:
             yield
 
-    """ Test add and get """
-    with context():
-        await proxy.add_route(routespec, target, copy.copy(data))
-        # Make sure the route was added
-        route = await proxy.get_route(routespec)
-    if not expect_value_error:
+    async def test_route_exist(spec, backend):
+        with context(spec):
+            route = await proxy.get_route(spec)
+
+        if not expect_value_error(spec):
+            try:
+                del route["data"]["last_activity"]  # CHP
+            except KeyError:
+                pass
+
+            assert route == expected_output(spec, backend.geturl())
+
+            # Test the actual routing
+            responding_backend1 = await utils.get_responding_backend_port(
+                proxy_url, normalize_spec(spec)
+            )
+            responding_backend2 = await utils.get_responding_backend_port(
+                proxy_url, normalize_spec(spec) + "something"
+            )
+            assert (
+                responding_backend1 == backend.port
+                and responding_backend2 == backend.port
+            )
+
+    for i, spec in enumerate(existing_routes, start=1):
+        backend = default_backend._replace(
+            netloc=f"{default_backend.hostname}:{default_backend.port+i}"
+        )
+        launch_backend(backend.port, backend.scheme)
+        extra_backends.append(backend)
+
+    launch_backend(default_backend.port, default_backend.scheme)
+    await wait_for_services(
+        [proxy.public_url, default_backend.geturl()]
+        + [backend.geturl() for backend in extra_backends]
+    )
+
+    # Create existing routes
+    for i, spec in enumerate(existing_routes):
         try:
-            del route["data"]["last_activity"]  # CHP
-        except KeyError:
+            await proxy.add_route(spec, extra_backends[i].geturl(), copy.copy(data))
+        except Exception:
             pass
-        finally:
-            assert route == expected_output
-        if proxy.public_url.endswith("/"):
-            req_url = proxy.public_url[:-1]
-        else:
-            req_url = proxy.public_url
-        # Test the actual routing
-        responding_backend1 = await utils.get_responding_backend_port(
-            req_url, routespec
-        )
-        responding_backend2 = await utils.get_responding_backend_port(
-            req_url, routespec + "/something"
-        )
-        assert (
-            responding_backend1 == backend_port and responding_backend2 == backend_port
-        )
 
-    """ Test delete + get """
-    with context():
+    def finalizer():
+        async def cleanup():
+            """ Cleanup """
+            for spec in existing_routes:
+                try:
+                    await proxy.delete_route(spec)
+                except Exception:
+                    pass
+
+        event_loop.run_until_complete(cleanup())
+
+    request.addfinalizer(finalizer)
+
+    # Test add
+    with context(routespec):
+        await proxy.add_route(routespec, default_backend.geturl(), copy.copy(data))
+
+    # Test get
+    await test_route_exist(routespec, default_backend)
+    for i, spec in enumerate(existing_routes):
+        await test_route_exist(spec, extra_backends[i])
+
+    # Test delete
+    with context(routespec):
         await proxy.delete_route(routespec)
         route = await proxy.get_route(routespec)
-    if not expect_value_error:
+
+    # Test that deleted route does not exist anymore
+    if not expect_value_error(routespec):
         assert route == None
 
         async def _wait_for_deletion():
-            deleted = False
-            try:
-                await utils.get_responding_backend_port(req_url, routespec)
-            except HTTPClientError:
-                deleted = True
-            finally:
-                return deleted
-
-        """ If this raises a TimeoutError, the route wasn't properly deleted,
-        thus the proxy still has a route for the given routespec"""
+            deleted = 0
+            for spec in [
+                normalize_spec(routespec),
+                normalize_spec(routespec) + "something",
+            ]:
+                try:
+                    result = await utils.get_responding_backend_port(proxy_url, spec)
+                    if result != default_backend.port:
+                        deleted += 1
+                except HTTPClientError:
+                    deleted += 1
+
+            return deleted == 2
+
+        # If this raises a TimeoutError, the route wasn't properly deleted,
+        # thus the proxy still has a route for the given routespec
         await exponential_backoff(_wait_for_deletion, "Route still exists")
 
+    # Test that other routes are still exist
+    for i, spec in enumerate(existing_routes):
+        await test_route_exist(spec, extra_backends[i])
+
 
 async def test_get_all_routes(proxy, launch_backend):
     routespecs = ["/proxy/path1", "/proxy/path2/", "/proxy/path3/"]
@@ -233,8 +324,8 @@ async def test_get_all_routes(proxy, launch_backend):
             del routes[route_key]["data"]["last_activity"]  # CHP
     except KeyError:
         pass
-    finally:
-        assert routes == expected_output
+
+    assert routes == expected_output
 
 
 async def test_host_origin_headers(proxy, launch_backend):