User reverse proxy, feed user UUID and roles into header

eceo-epfl · May 23, 2024 · 9a55f6a · 9a55f6a
1 parent 2a27c21
commit 9a55f6a
Show file tree

Hide file tree

Showing 4 changed files with 31 additions and 91 deletions.
diff --git a/app/auth.py b/app/auth.py
@@ -53,7 +53,7 @@ async def get_payload(token: str = Security(oauth2_scheme)) -> dict:
 # Get user infos from the payload
 async def get_user_info(payload: dict = Depends(get_payload)) -> User:
     try:
-        return User(
+        user = User(
             id=payload.get("sub"),
             username=payload.get("preferred_username"),
             email=payload.get("email"),
@@ -62,6 +62,7 @@ async def get_user_info(payload: dict = Depends(get_payload)) -> User:
             realm_roles=payload.get("realm_access", {}).get("roles", []),
             client_roles=payload.get("realm_access", {}).get("roles", []),
         )
+        return user
     except Exception as e:
         raise HTTPException(
             status_code=status.HTTP_400_BAD_REQUEST,

diff --git a/app/objects.py b/app/objects.py
@@ -2,7 +2,7 @@
 from fastapi import Depends, APIRouter, Query, Response, Body
 from fastapi.responses import StreamingResponse
 from app.config import config
-from app.utils import get_async_client
+from app.utils import get_async_client, _reverse_proxy
 import httpx
 from uuid import UUID
 from app.models.user import User
@@ -133,40 +133,21 @@ async def check_uploaded_chunks(
 
 @router.get("/{object_id}")
 async def get_object(
-    client: httpx.AsyncClient = Depends(get_async_client),
-    *,
     object_id: UUID,
-    user: User = Depends(get_user_info),
+    reverse_proxy: Any = Depends(_reverse_proxy),
 ) -> Any:
     """Get a object by id"""
 
-    res = await client.get(
-        f"{config.DEEPREEFMAP_API_URL}/v1/objects/{object_id}",
-    )
-
-    return res.json()
+    return reverse_proxy
 
 
 @router.get("")
 async def get_objects(
-    response: Response,
-    *,
-    filter: str = Query(None),
-    sort: str = Query(None),
-    range: str = Query(None),
-    client: httpx.AsyncClient = Depends(get_async_client),
-    user: User = Depends(get_user_info),
+    reverse_proxy: Any = Depends(_reverse_proxy),
 ) -> Any:
     """Get all objects"""
 
-    res = await client.get(
-        f"{config.DEEPREEFMAP_API_URL}/v1/objects",
-        params={"sort": sort, "range": range, "filter": filter},
-    )
-    response.headers["Access-Control-Expose-Headers"] = "Content-Range"
-    response.headers["Content-Range"] = res.headers["Content-Range"]
-
-    return res.json()
+    return reverse_proxy
 
 
 @router.post("/{object_id}")
@@ -187,30 +168,18 @@ async def regenerate_statistics(
 @router.put("/{object_id}")
 async def update_object(
     object_id: UUID,
-    object: Any = Body(...),
-    client: httpx.AsyncClient = Depends(get_async_client),
-    admin_user: User = Depends(require_admin),
+    reverse_proxy: Any = Depends(_reverse_proxy),
 ) -> Any:
     """ "Updates an object by id"""
 
-    res = await client.put(
-        f"{config.DEEPREEFMAP_API_URL}/v1/objects/{object_id}",
-        json=object,
-    )
-
-    return res.json()
+    return reverse_proxy
 
 
 @router.delete("/{object_id}")
 async def delete_object(
     object_id: UUID,
-    client: httpx.AsyncClient = Depends(get_async_client),
-    admin_user: User = Depends(require_admin),
+    reverse_proxy: Any = Depends(_reverse_proxy),
 ) -> None:
     """Delete an object by id"""
 
-    res = await client.delete(
-        f"{config.DEEPREEFMAP_API_URL}/v1/objects/{object_id}"
-    )
-
-    return res.json()
+    return reverse_proxy
diff --git a/app/submissions.py b/app/submissions.py
@@ -1,7 +1,7 @@
 from typing import Any
 from fastapi import Depends, APIRouter, Query, Response, Body, HTTPException
 from app.config import config
-from app.utils import get_async_client
+from app.utils import get_async_client, _reverse_proxy
 import httpx
 from uuid import UUID
 from app.models.user import User
@@ -94,18 +94,12 @@ async def get_submission_output_file(
 
 @router.get("/{submission_id}")
 async def get_submission(
-    client: httpx.AsyncClient = Depends(get_async_client),
-    *,
     submission_id: UUID,
-    user: User = Depends(get_user_info),
+    reverse_proxy: Any = Depends(_reverse_proxy),
 ) -> Any:
     """Get a submission by id"""
 
-    res = await client.get(
-        f"{config.DEEPREEFMAP_API_URL}/v1/submissions/{submission_id}",
-    )
-
-    return res.json()
+    return reverse_proxy
 
 
 @router.get("/{submission_id}/{filename}", response_model=DownloadToken)
@@ -155,71 +149,40 @@ async def execute_submission(
 
 @router.get("")
 async def get_submissions(
-    response: Response,
-    *,
-    filter: str = Query(None),
-    sort: str = Query(None),
-    range: str = Query(None),
-    client: httpx.AsyncClient = Depends(get_async_client),
-    user: User = Depends(get_user_info),
+    reverse_proxy: Any = Depends(_reverse_proxy),
 ) -> Any:
     """Get all submissions"""
 
-    res = await client.get(
-        f"{config.DEEPREEFMAP_API_URL}/v1/submissions",
-        params={"sort": sort, "range": range, "filter": filter},
-    )
-    response.headers["Access-Control-Expose-Headers"] = "Content-Range"
-    response.headers["Content-Range"] = res.headers["Content-Range"]
-
-    return res.json()
+    return reverse_proxy
 
 
 @router.post("")
 async def create_submission(
-    submission: Any = Body(...),
-    client: httpx.AsyncClient = Depends(get_async_client),
-    admin_user: User = Depends(require_admin),
+    reverse_proxy: Any = Depends(_reverse_proxy),
 ) -> Any:
     """Creates an submission
 
     Creates a new submission with the given files
     """
 
-    res = await client.post(
-        f"{config.DEEPREEFMAP_API_URL}/v1/submissions",
-        json=submission,
-    )
-    return res.json()
+    return reverse_proxy
 
 
 @router.put("/{submission_id}")
 async def update_submission(
     submission_id: UUID,
-    submission: Any = Body(...),
-    client: httpx.AsyncClient = Depends(get_async_client),
-    admin_user: User = Depends(require_admin),
+    reverse_proxy: Any = Depends(_reverse_proxy),
 ) -> Any:
     """ "Updates an submission by id"""
 
-    res = await client.put(
-        f"{config.DEEPREEFMAP_API_URL}/v1/submissions/{submission_id}",
-        json=submission,
-    )
-
-    return res.json()
+    return reverse_proxy
 
 
 @router.delete("/{submission_id}")
 async def delete_submission(
     submission_id: UUID,
-    client: httpx.AsyncClient = Depends(get_async_client),
-    admin_user: User = Depends(require_admin),
+    reverse_proxy: Any = Depends(_reverse_proxy),
 ) -> None:
     """Delete an submission by id"""
 
-    res = await client.delete(
-        f"{config.DEEPREEFMAP_API_URL}/v1/submissions/{submission_id}"
-    )
-
-    return res.json()
+    return reverse_proxy
diff --git a/app/utils.py b/app/utils.py
@@ -31,18 +31,25 @@ async def lifespan(app: FastAPI):
 
 async def _reverse_proxy(
     request: Request,
+    user: User = Depends(get_user_info),
 ):
     client = request.state.client
-
     path = request.url.path.replace("/api", "/v1")
     url = httpx.URL(
         path=path,
         query=request.url.query.encode("utf-8"),
     )
+    headers = {
+        key.decode(): value.decode() for key, value in request.headers.raw
+    }
+    headers.update(  # Add user ID and roles to the headers
+        {"User-ID": user.id, "User-Roles": ",".join(user.realm_roles)}
+    )
+
     req = client.build_request(
         request.method,
         url,
-        headers=request.headers.raw,
+        headers=headers,
         content=request.stream(),
     )
     r = await client.send(req, stream=True)