Crude working collections filter

Author: alukach

src/stac_auth_proxy/app.py

@@ -6,14 +6,16 @@
 """
 
 import logging
-from typing import Optional
+from typing import Optional, Annotated
 
-from fastapi import Depends, FastAPI
+from fastapi import FastAPI, Security, Request, Depends
+from cql2 import Expr
 
 from .auth import OpenIdConnectAuth
 from .config import Settings
 from .handlers import OpenApiSpecHandler, ReverseProxyHandler
 from .middleware import AddProcessTimeHeaderMiddleware
+from .utils import apply_filter
 
 logger = logging.getLogger(__name__)
 
@@ -28,8 +30,8 @@ def create_app(settings: Optional[Settings] = None) -> FastAPI:
     app.add_middleware(AddProcessTimeHeaderMiddleware)
 
     auth_scheme = OpenIdConnectAuth(
-        openid_configuration_url=str(settings.oidc_discovery_url)
-    ).valid_token_dependency
+        openid_configuration_url=settings.oidc_discovery_url
+    )
 
     if settings.debug:
         app.add_api_route(
@@ -38,12 +40,40 @@ def create_app(settings: Optional[Settings] = None) -> FastAPI:
             methods=["GET"],
         )
 
-    proxy_handler = ReverseProxyHandler(upstream=str(settings.upstream_url))
+    collections_filter = (
+        settings.collections_filter(auth_scheme.maybe_validated_user)
+        if settings.collections_filter
+        else None
+    )
+    items_filter = (
+        settings.items_filter(auth_scheme.maybe_validated_user)
+        if settings.items_filter
+        else None
+    )
+    proxy_handler = ReverseProxyHandler(
+        upstream=str(settings.upstream_url),
+        collections_filter=collections_filter,
+        items_filter=items_filter,
+    )
     openapi_handler = OpenApiSpecHandler(
         proxy=proxy_handler,
         oidc_config_url=str(settings.oidc_discovery_url),
     )
 
+    # @app.get("/collections")
+    # async def collections(
+    #     request: Request,
+    #     filter: Annotated[Optional[Expr], Depends(collections_filter.dependency)],
+    # ):
+    #     # if filter:
+    #     #     print(f"{request.receive=}")
+    #     #     request = await apply_filter(
+    #     #         request,
+    #     #         filter,
+    #     #     )
+    #     #     print(f"{request.receive=}")
+    #     return await proxy_handler.stream(request=request)
+
     # Endpoints that are explicitely marked private
     for path, methods in settings.private_endpoints.items():
         app.add_api_route(
@@ -54,7 +84,7 @@ def create_app(settings: Optional[Settings] = None) -> FastAPI:
                 else openapi_handler.dispatch
             ),
             methods=methods,
-            dependencies=[Depends(auth_scheme)],
+            dependencies=[Security(auth_scheme.validated_user)],
         )
 
     # Endpoints that are explicitely marked as public
@@ -67,14 +97,23 @@ def create_app(settings: Optional[Settings] = None) -> FastAPI:
                 else openapi_handler.dispatch
             ),
             methods=methods,
+            dependencies=[Security(auth_scheme.maybe_validated_user)],
         )
 
     # Catchall for remainder of the endpoints
     app.add_api_route(
         "/{path:path}",
         proxy_handler.stream,
         methods=["GET", "POST", "PUT", "PATCH", "DELETE"],
-        dependencies=([] if settings.default_public else [Depends(auth_scheme)]),
+        dependencies=(
+            [
+                Security(
+                    auth_scheme.maybe_validated_user
+                    if settings.default_public
+                    else auth_scheme.validated_user
+                )
+            ]
+        ),
     )
 
     return app

src/stac_auth_proxy/auth.py

@@ -31,7 +31,7 @@ class OpenIdConnectAuth:
     def __post_init__(self):
         """Initialize the OIDC authentication class."""
         logger.debug("Requesting OIDC config")
-        origin_url = (
+        origin_url = str(
             self.openid_configuration_internal_url or self.openid_configuration_url
         )
         with urllib.request.urlopen(origin_url) as response:

src/stac_auth_proxy/filters/template.py

@@ -19,31 +19,37 @@ class Template:
 
     # Generated attributes
     env: Environment = field(init=False)
+    dependency: Callable[[Request, Security], Expr] = field(init=False)
 
     def __post_init__(self):
         """Initialize the Jinja2 environment."""
         self.env = Environment(loader=BaseLoader).from_string(self.template_str)
-        self.render.__annotations__["auth_token"] = Security(self.token_dependency)
-
-    async def cql2(self, request: Request, auth_token=Security(...)) -> Expr:
-        """Render a CQL2 filter expression with the request and auth token."""
-        # TODO: How to handle the case where auth_token is null?
-        context = {
-            "req": {
-                "path": request.url.path,
-                "method": request.method,
-                "query_params": dict(request.query_params),
-                "path_params": extract_variables(request.url.path),
-                "headers": dict(request.headers),
-                "body": (
-                    await request.json()
-                    if request.headers.get("content-type") == "application/json"
-                    else (await request.body()).decode()
-                ),
-            },
-            "token": auth_token,
-        }
-        cql2_str = self.env.render(**context)
-        cql2_expr = Expr(cql2_str)
-        cql2_expr.validate()
-        return cql2_expr
+        self.dependency = self.build()
+
+    def build(self):
+        async def dependency(
+            request: Request, auth_token=Security(self.token_dependency)
+        ) -> Expr:
+            """Render a CQL2 filter expression with the request and auth token."""
+            # TODO: How to handle the case where auth_token is null?
+            context = {
+                "req": {
+                    "path": request.url.path,
+                    "method": request.method,
+                    "query_params": dict(request.query_params),
+                    "path_params": extract_variables(request.url.path),
+                    "headers": dict(request.headers),
+                    "body": (
+                        await request.json()
+                        if request.headers.get("content-type") == "application/json"
+                        else (await request.body()).decode()
+                    ),
+                },
+                "token": auth_token,
+            }
+            cql2_str = self.env.render(**context)
+            cql2_expr = Expr(cql2_str)
+            cql2_expr.validate()
+            return cql2_expr
+
+        return dependency

src/stac_auth_proxy/handlers/reverse_proxy.py

@@ -3,13 +3,17 @@
 import logging
 import time
 from dataclasses import dataclass
+from typing import Optional, Annotated
 
+from cql2 import Expr
 import httpx
-from fastapi import Request
+from fastapi import Request, Depends
 from starlette.background import BackgroundTask
 from starlette.datastructures import MutableHeaders
 from starlette.responses import StreamingResponse
 
+from ..utils import update_qs
+
 logger = logging.getLogger(__name__)
 
 
@@ -19,6 +23,8 @@ class ReverseProxyHandler:
 
     upstream: str
     client: httpx.AsyncClient = None
+    collections_filter: Optional[callable] = None
+    items_filter: Optional[callable] = None
 
     def __post_init__(self):
         """Initialize the HTTP client."""
@@ -27,17 +33,41 @@ def __post_init__(self):
             timeout=httpx.Timeout(timeout=15.0),
         )
 
-    async def proxy_request(self, request: Request, *, stream=False) -> httpx.Response:
+        self.proxy_request.__annotations__["collections_filter"] = Annotated[
+            Optional[Expr], Depends(self.collections_filter.dependency)
+        ]
+        self.stream.__annotations__["collections_filter"] = Annotated[
+            Optional[Expr], Depends(self.collections_filter.dependency)
+        ]
+
+    async def proxy_request(
+        self,
+        request: Request,
+        *,
+        collections_filter: Annotated[Optional[Expr], Depends(...)],
+        stream=False,
+    ) -> httpx.Response:
         """Proxy a request to the upstream STAC API."""
         headers = MutableHeaders(request.headers)
         headers.setdefault("X-Forwarded-For", request.client.host)
         headers.setdefault("X-Forwarded-Host", request.url.hostname)
 
+        path = request.url.path
+        query = request.url.query.encode("utf-8")
+
         # https://github.com/fastapi/fastapi/discussions/7382#discussioncomment-5136466
+        # TODO: Examine filters
+        if collections_filter:
+            if request.method == "GET" and path == "/collections":
+                query += b"&" + update_qs(
+                    request.query_params, filter=collections_filter.to_text()
+                )
+
         url = httpx.URL(
-            path=request.url.path,
-            query=request.url.query.encode("utf-8"),
+            path=path,
+            query=query,
         )
+
         rp_req = self.client.build_request(
             request.method,
             url=url,
@@ -56,9 +86,17 @@ async def proxy_request(self, request: Request, *, stream=False) -> httpx.Respon
         rp_resp.headers["X-Upstream-Time"] = f"{proxy_time:.3f}"
         return rp_resp
 
-    async def stream(self, request: Request) -> StreamingResponse:
+    async def stream(
+        self,
+        request: Request,
+        collections_filter: Annotated[Optional[Expr], Depends(...)],
+    ) -> StreamingResponse:
         """Transparently proxy a request to the upstream STAC API."""
-        rp_resp = await self.proxy_request(request, stream=True)
+        rp_resp = await self.proxy_request(
+            request,
+            collections_filter=collections_filter,
+            stream=True,
+        )
         return StreamingResponse(
             rp_resp.aiter_raw(),
             status_code=rp_resp.status_code,

src/stac_auth_proxy/utils.py

@@ -3,7 +3,10 @@
 import re
 from urllib.parse import urlparse
 
+from cql2 import Expr
+from fastapi import Request
 from fastapi.dependencies.models import Dependant
+from starlette.datastructures import QueryParams
 from httpx import Headers
 
 
@@ -42,3 +45,47 @@ def has_any_security_requirements(dependency: Dependant) -> bool:
     return any(
         has_any_security_requirements(sub_dep) for sub_dep in dependency.dependencies
     )
+
+
+async def apply_filter(request: Request, filter: Expr) -> Request:
+    """Apply a CQL2 filter to a request."""
+    req_filter = request.query_params.get("filter") or (
+        (await request.json()).get("filter")
+        if request.headers.get("content-length")
+        else None
+    )
+
+    new_filter = Expr(" AND ".join(e.to_text() for e in [req_filter, filter] if e))
+    new_filter.validate()
+
+    if request.method == "GET":
+        updated_scope = request.scope.copy()
+        updated_scope["query_string"] = update_qs(
+            request.query_params,
+            filter=new_filter.to_text(),
+        )
+        return Request(
+            scope=updated_scope,
+            receive=request.receive,
+            # send=request._send,
+        )
+
+    # TODO: Support POST/PUT/PATCH
+    # elif request.method == "POST":
+    #     request_body = await request.body()
+    #     query = request.url.query
+    #     query += "&" if query else "?"
+    #     query += f"filter={filter}"
+    #     request.url.query = query
+
+    return request
+
+
+def update_qs(query_params: QueryParams, **updates) -> bytes:
+    query_dict = {
+        **query_params,
+        **updates,
+    }
+    return "&".join(f"{key}={value}" for key, value in query_dict.items()).encode(
+        "utf-8"
+    )