update authentication backend to respect cookies and bearer authorization

Author: felixevers

backend/auth.py

@@ -3,7 +3,7 @@
 
 import requests
 from config import CLIENT_ID, ISSUER_URI, LOGGER, PUBLIC_ISSUER_URI
-from fastapi import HTTPException, Request, status
+from fastapi import Request
 from fastapi.responses import RedirectResponse
 from jose import jwk, jwt
 from starlette.requests import HTTPConnection
@@ -25,37 +25,17 @@ def delete_auth_cookie(response):
     )
 
 
-async def authenticate_connection(connection, token: str | None):
-    user_payload = None
+def get_user_payload(connection: HTTPConnection) -> dict | None:
+    token = get_token_source(connection)
 
-    if token:
-        try:
-            user_payload = verify_token(token)
-        except Exception as e:
-            logger.warning(f"Token validation failed: {e}")
+    if not token:
+        return None
 
-    if user_payload:
-        return user_payload
-
-    if connection.scope["type"] == "http":
-        response = RedirectResponse("/login", status_code=302)
-        delete_auth_cookie(response)
-
-        accept_header = connection.headers.get("accept", "")
-
-        if "text/html" in accept_header:
-            raise UnauthenticatedRedirect(response=response)
-
-        raise HTTPException(
-            status_code=401,
-            detail="Not authenticated",
-            headers={"WWW-Authenticate": "Bearer"},
-        )
-
-    if connection.scope["type"] == "websocket":
-        raise HTTPException(status_code=403, detail="Not authenticated")
-
-    raise RuntimeError("Unsupported connection type")
+    try:
+        return verify_token(token)
+    except Exception as e:
+        logger.warning(f"Token validation failed: {e}")
+        return None
 
 
 def get_public_key(token: str) -> Any:
@@ -71,7 +51,6 @@ def get_public_key(token: str) -> Any:
         if kid in jwks_cache:
             return jwks_cache[kid]
 
-        # TODO use openid config endpoint to obtain well-known endpoints in the future
         jwks_uri = f"{ISSUER_URI}/protocol/openid-connect/certs"
 
         response = requests.get(jwks_uri, timeout=5)
@@ -87,12 +66,8 @@ def get_public_key(token: str) -> Any:
         raise Exception("Public key not found in JWKS")
 
     except Exception as e:
-        print(f"Auth Error: {e}")
-        raise HTTPException(
-            status_code=status.HTTP_401_UNAUTHORIZED,
-            detail="Could not validate credentials",
-            headers={"WWW-Authenticate": "Bearer"},
-        )
+        logger.error(f"Auth Error: {e}")
+        raise e
 
 
 def verify_token(token: str) -> dict:
@@ -124,16 +99,14 @@ def verify_token(token: str) -> dict:
             f"Invalid audience/azp. Expected {CLIENT_ID}, got azp={azp}, aud={aud}",
         )
     except jwt.ExpiredSignatureError:
-        raise HTTPException(status_code=401, detail="Token has expired")
+        raise Exception("Token has expired")
     except jwt.JWTError as e:
-        raise HTTPException(status_code=401, detail=f"Invalid token: {e!s}")
-    except Exception:
-        raise HTTPException(status_code=401, detail="Authentication failed")
+        raise Exception(f"Invalid token: {e!s}")
+    except Exception as e:
+        raise Exception(f"Authentication failed: {e!s}")
 
 
-async def get_token_source(
-    connection: HTTPConnection,
-) -> str | None:
+def get_token_source(connection: HTTPConnection) -> str | None:
     auth_header = connection.headers.get("authorization")
     if auth_header and auth_header.startswith("Bearer "):
         return auth_header.split(" ")[1]

backend/main.py

@@ -5,8 +5,7 @@
 from api.resolvers import Mutation, Query, Subscription
 from auth import (
     UnauthenticatedRedirect,
-    authenticate_connection,
-    get_token_source,
+    get_user_payload,
     unauthenticated_redirect_handler,
 )
 from config import (
@@ -19,55 +18,87 @@
 from database.models.user import User
 from database.session import get_db_session
 from fastapi import Depends, FastAPI, HTTPException, Request
-from fastapi.responses import RedirectResponse
+from fastapi.responses import HTMLResponse, RedirectResponse
+from graphql import FieldNode
 from sqlalchemy import select
 from starlette.requests import HTTPConnection
-from starlette.status import HTTP_400_BAD_REQUEST
 from strawberry import Schema
+from strawberry.extensions import SchemaExtension
 from strawberry.fastapi import GraphQLRouter
 
 logger = logging.getLogger(LOGGER)
 
 
+class GlobalAuthExtension(SchemaExtension):
+    def on_execute(self):
+        execution_context = self.execution_context
+        user = execution_context.context.user
+
+        if user:
+            yield
+            return
+
+        document = execution_context.graphql_document
+        if document:
+            for definition in document.definitions:
+                if definition.kind == "operation_definition":
+                    for selection in definition.selection_set.selections:
+                        if not isinstance(selection, FieldNode):
+                            raise HTTPException(
+                                status_code=401,
+                                detail="Not authenticated",
+                            )
+                        if not selection.name.value.startswith("__"):
+                            raise HTTPException(
+                                status_code=401,
+                                detail="Not authenticated",
+                            )
+
+        yield
+
+
 async def get_context(
     connection: HTTPConnection,
-    token: str | None = Depends(get_token_source),
     session=Depends(get_db_session),
 ) -> Context:
-    user_payload = await authenticate_connection(connection, token)
+    user_payload = get_user_payload(connection)
 
-    user_id = user_payload.get("sub")
-    username = user_payload.get("preferred_username") or user_payload.get(
-        "name",
-    )
-    firstname = user_payload.get("given_name")
-    lastname = user_payload.get("family_name")
+    db_user = None
 
-    if not (user_id and username and firstname and lastname):
-        raise HTTPException(
-            status_code=HTTP_400_BAD_REQUEST,
-            detail="Missing required user details.",
+    if user_payload:
+        user_id = user_payload.get("sub")
+        username = user_payload.get("preferred_username") or user_payload.get(
+            "name",
         )
+        firstname = user_payload.get("given_name")
+        lastname = user_payload.get("family_name")
 
-    result = await session.execute(select(User).where(User.id == user_id))
-    db_user = result.scalars().first()
+        if user_id:
+            result = await session.execute(
+                select(User).where(User.id == user_id),
+            )
+            db_user = result.scalars().first()
 
-    if not db_user:
-        db_user = User(
-            id=user_id,
-            name=username,
-            firstname=firstname,
-            lastname=lastname,
-            title="User",
-        )
-        session.add(db_user)
-    elif db_user.name != username or db_user.firstname != firstname:
-        db_user.name = username
-        db_user.firstname = firstname
-        db_user.lastname = lastname
+            if not db_user:
+                db_user = User(
+                    id=user_id,
+                    name=username,
+                    firstname=firstname,
+                    lastname=lastname,
+                    title="User",
+                )
+                session.add(db_user)
+            elif (
+                db_user.name != username
+                or db_user.firstname != firstname
+                or db_user.lastname != lastname
+            ):
+                db_user.name = username
+                db_user.firstname = firstname
+                db_user.lastname = lastname
 
-    await session.commit()
-    await session.refresh(db_user)
+            await session.commit()
+            await session.refresh(db_user)
 
     return Context(db=session, user=db_user)
 
@@ -76,8 +107,107 @@ async def get_context(
     query=Query,
     mutation=Mutation,
     subscription=Subscription,
+    extensions=[GlobalAuthExtension],
 )
-graphql_app = GraphQLRouter(
+
+
+class AuthedGraphQLRouter(GraphQLRouter):
+    async def render_graphql_ide(self, request: Request) -> HTMLResponse:
+        response = await super().render_graphql_ide(request)
+
+        redirect_uri = f"{request.base_url}callback"
+        login_url = (
+            f"{ISSUER_URI}/protocol/openid-connect/auth"
+            f"?client_id={CLIENT_ID}"
+            f"&response_type=code"
+            f"&scope=openid profile email"
+            f"&redirect_uri={redirect_uri}"
+        )
+        logout_url = f"{request.base_url}logout"
+
+        is_authenticated = (
+            "true" if request.cookies.get("access_token") else "false"
+        )
+
+        injection_script = f"""
+        <script>
+            (function() {{
+                var loginUrl = "{login_url}";
+                var logoutUrl = "{logout_url}";
+                var isAuthenticated = {is_authenticated};
+                
+                function injectAuthButton() {{
+                    var sidebars = document.querySelectorAll('.graphiql-sidebar-section');
+                    var sidebar = sidebars[0]; 
+                    
+                    if (sidebar && !document.getElementById('custom-auth-button')) {{
+                        var button = document.createElement('button');
+                        button.id = 'custom-auth-button';
+                        button.className = 'graphiql-un-styled';
+                        button.type = 'button';
+                        
+                        if (isAuthenticated) {{
+                            button.setAttribute('aria-label', 'Logout');
+                            button.title = 'Logout';
+                            button.innerHTML = `
+                                <svg height="1em" viewBox="0 0 24 24" fill="none" 
+                                     stroke="currentColor" stroke-width="1.5" 
+                                     stroke-linecap="round" stroke-linejoin="round" 
+                                     xmlns="http://www.w3.org/2000/svg">
+                                    <path d="M9 21H5a2 2 0 0 1-2-2V5a2 2 0 0 1 2-2h4"></path>
+                                    <polyline points="16 17 21 12 16 7"></polyline>
+                                    <line x1="21" y1="12" x2="9" y2="12"></line>
+                                </svg>
+                            `;
+                            button.onclick = function() {{
+                                window.location.href = logoutUrl;
+                            }};
+                        }} else {{
+                            button.setAttribute('aria-label', 'Login with OIDC');
+                            button.title = 'Login with OIDC';
+                            button.innerHTML = `
+                                <svg height="1em" viewBox="0 0 24 24" fill="none" 
+                                     stroke="currentColor" stroke-width="1.5" 
+                                     stroke-linecap="round" stroke-linejoin="round" 
+                                     xmlns="http://www.w3.org/2000/svg">
+                                    <path d="M15 3h4a2 2 0 0 1 2 2v14a2 2 0 0 1-2 2h-4"></path>
+                                    <polyline points="10 17 15 12 10 7"></polyline>
+                                    <line x1="15" y1="12" x2="3" y2="12"></line>
+                                </svg>
+                            `;
+                            button.onclick = function() {{
+                                window.location.href = loginUrl;
+                            }};
+                        }}
+
+                        sidebar.appendChild(button);
+                    }}
+                }}
+
+                var observer = new MutationObserver(function(mutations) {{
+                    injectAuthButton();
+                }});
+
+                observer.observe(document.body, {{
+                    childList: true,
+                    subtree: true
+                }});
+                
+                injectAuthButton();
+            }})();
+        </script>
+        """
+
+        html_content = response.body.decode("utf-8")
+        new_html = html_content.replace(
+            "</body>",
+            f"{injection_script}</body>",
+        )
+
+        return HTMLResponse(new_html)
+
+
+graphql_app = AuthedGraphQLRouter(
     schema,
     context_getter=get_context,
     graphql_ide=IS_DEV,
@@ -97,6 +227,19 @@ async def get_context(
 )
 
 
+@app.get("/logout")
+def logout(_: Request):
+    response = RedirectResponse(url="/graphql")
+    response.delete_cookie(
+        "access_token",
+        path="/",
+        secure=True,
+        httponly=True,
+        samesite="lax",
+    )
+    return response
+
+
 @app.get("/callback")
 def oauth_callback(code: str, request: Request):
     token_endpoint = f"{ISSUER_URI}/protocol/openid-connect/token"
@@ -123,6 +266,8 @@ def oauth_callback(code: str, request: Request):
             access_token,
             httponly=True,
             max_age=3600,
+            samesite="lax",
+            secure=True,
         )
         return resp