feat: enhance authentication flow (#15)

atyrode · web-flow · commit 496b43a146fa · 2025-04-24T08:04:14.000+02:00
* feat: add static file serving for static assets in backend

* feat: add authentication completion popup HTML file

* fix: prevent workspace state polling when unauthenticated

* refactor: adjust workspace state polling and retry behavior based on authentication

* feat: enhance authentication flow with popup support and improved logout handling

- Added support for popup authentication in the login process.
- Updated the callback endpoint to handle popup state and return a specific response.
- Improved logout functionality to clear session cookies without redirecting.
- Invalidated authentication queries on logout to refresh the UI state.
diff --git a/src/backend/main.py b/src/backend/main.py
@@ -44,6 +44,7 @@ async def lifespan(_: FastAPI):
 )
 
 app.mount("/assets", StaticFiles(directory=ASSETS_DIR), name="assets")
+app.mount("/static", StaticFiles(directory=STATIC_DIR), name="static")
 
 @app.get("/")
 async def read_root(request: Request, auth: Optional[SessionData] = Depends(optional_auth)):
diff --git a/src/backend/routers/auth.py b/src/backend/routers/auth.py
@@ -2,28 +2,32 @@
 import jwt
 import httpx
 from fastapi import APIRouter, Request, HTTPException, Depends
-from fastapi.responses import RedirectResponse
+from fastapi.responses import RedirectResponse, FileResponse
+import os
 
-from config import get_auth_url, get_token_url, OIDC_CONFIG, sessions
+from config import get_auth_url, get_token_url, OIDC_CONFIG, sessions, STATIC_DIR
 from dependencies import SessionData, require_auth
 from coder import CoderAPI
 
 auth_router = APIRouter()
 coder_api = CoderAPI()
 
 @auth_router.get("/login")
-async def login(request: Request, kc_idp_hint: str = None):
+async def login(request: Request, kc_idp_hint: str = None, popup: str = None):
     session_id = secrets.token_urlsafe(32)
     auth_url = get_auth_url()
+    state = "popup" if popup == "1" else "default"
     if kc_idp_hint:
         auth_url = f"{auth_url}&kc_idp_hint={kc_idp_hint}"
+    # Add state param to OIDC URL
+    auth_url = f"{auth_url}&state={state}"
     response = RedirectResponse(auth_url)
     response.set_cookie('session_id', session_id)
 
     return response
 
 @auth_router.get("/callback")
-async def callback(request: Request, code: str):
+async def callback(request: Request, code: str, state: str = "default"):
     session_id = request.cookies.get('session_id')
     if not session_id:
         raise HTTPException(status_code=400, detail="No session")
@@ -53,19 +57,32 @@ async def callback(request: Request, code: str):
                 user_info
             )
             coder_api.ensure_workspace_exists(user_data['username'])
-            
-                
         except Exception as e:
             print(f"Error in user/workspace setup: {str(e)}")
             # Continue with login even if Coder API fails
-    
-    return RedirectResponse('/')
+
+    if state == "popup":
+        return FileResponse(os.path.join(STATIC_DIR, "auth/popup-close.html"))
+    else:
+        return RedirectResponse('/')
 
 @auth_router.get("/logout")
 async def logout(request: Request):
     session_id = request.cookies.get('session_id')
     if session_id in sessions:
         del sessions[session_id]
-    response = RedirectResponse('/')
-    response.delete_cookie('session_id')
+    
+    # Create a response that doesn't redirect but still clears the cookie
+    from fastapi.responses import JSONResponse
+    response = JSONResponse({"status": "success", "message": "Logged out successfully"})
+    
+    # Clear the session_id cookie with all necessary parameters
+    response.delete_cookie(
+        key="session_id",
+        path="/",
+        domain=None,  # Use None to match the current domain
+        secure=request.url.scheme == "https",
+        httponly=True
+    )
+    
     return response
diff --git a/src/frontend/public/auth/popup-close.html b/src/frontend/public/auth/popup-close.html
@@ -0,0 +1,34 @@
+<!DOCTYPE html>
+<html lang="en">
+<head>
+  <meta charset="UTF-8" />
+  <title>Authentication Complete</title>
+  <style>
+    body {
+      background: #18181b;
+      color: #fff;
+      font-family: sans-serif;
+      display: flex;
+      flex-direction: column;
+      align-items: center;
+      justify-content: center;
+      height: 100vh;
+      margin: 0;
+    }
+    .message {
+      margin-bottom: 2rem;
+      font-size: 1.2rem;
+    }
+  </style>
+</head>
+<body>
+  <div class="message">Authentication complete! You may close this window.</div>
+  <script>
+    localStorage.setItem('auth_completed', Date.now().toString());
+    // Close the window after a short delay to ensure message is sent
+    setTimeout(() => {
+      window.close();
+    }, 100);
+  </script>
+</body>
+</html>
diff --git a/src/frontend/src/api/hooks.ts b/src/frontend/src/api/hooks.ts
@@ -132,10 +132,16 @@ export function useUserProfile(options?: UseQueryOptions<UserProfile>) {
 }
 
 export function useWorkspaceState(options?: UseQueryOptions<WorkspaceState>) {
+  // Get the current auth state from the query cache
+  const authState = queryClient.getQueryData<boolean>(['auth']);
+  
   return useQuery({
     queryKey: ['workspaceState'],
     queryFn: api.getWorkspaceState,
-    refetchInterval: 5000, // Poll every 5 seconds
+    // Only poll if authenticated
+    refetchInterval: authState === true ? 5000 : false, // Poll every 5 seconds if authenticated, otherwise don't poll
+    // Don't retry on error if not authenticated
+    retry: authState === true ? 1 : false,
     ...options,
   });
 }
diff --git a/src/frontend/src/auth/AuthModal.tsx b/src/frontend/src/auth/AuthModal.tsx
@@ -2,6 +2,7 @@ import React, { useState, useEffect } from "react";
 import ReactDOM from "react-dom";
 import { capture } from "../utils/posthog";
 import { Mail } from "lucide-react";
+import { queryClient } from "../api/queryClient";
 import "../styles/AuthModal.scss";
 
 interface AuthModalProps {
@@ -16,11 +17,22 @@ const AuthModal: React.FC<AuthModalProps> = ({
   useEffect(() => {
     setIsMounted(true);
     capture("auth_modal_shown");
-    // Prevent scrolling when modal is open
-    document.body.style.overflow = "hidden";
+  }, []);
 
+  useEffect(() => {
+    const checkLocalStorage = () => {
+      const authCompleted = localStorage.getItem('auth_completed');
+      if (authCompleted) {
+        localStorage.removeItem('auth_completed');
+        queryClient.invalidateQueries({ queryKey: ['auth'] });
+        clearInterval(intervalId);
+      }
+    };
+    
+    const intervalId = setInterval(checkLocalStorage, 500);
+    
     return () => {
-      document.body.style.overflow = "auto";
+      clearInterval(intervalId);
     };
   }, []);
 
@@ -62,7 +74,11 @@ const AuthModal: React.FC<AuthModalProps> = ({
               <button
                 className="auth-modal-button auth-modal-button-primary"
                 onClick={() => {
-                  window.location.href = "/auth/login?kc_idp_hint=google";
+                  window.open(
+                    "/auth/login?kc_idp_hint=google&popup=1",
+                    "authPopup",
+                    "width=500,height=700,noopener,noreferrer"
+                  );
                 }}
               >
                 <svg
@@ -96,7 +112,11 @@ const AuthModal: React.FC<AuthModalProps> = ({
               <button
                 className="auth-modal-button auth-modal-button-outline"
                 onClick={() => {
-                  window.location.href = "/auth/login?kc_idp_hint=github";
+                  window.open(
+                    "/auth/login?kc_idp_hint=github&popup=1",
+                    "authPopup",
+                    "width=500,height=700,noopener,noreferrer"
+                  );
                 }}
               >
                 <svg
diff --git a/src/frontend/src/pad/controls/StateIndicator.tsx b/src/frontend/src/pad/controls/StateIndicator.tsx
@@ -4,9 +4,13 @@ import '../styles/index.scss';
 
 export const StateIndicator: React.FC = () => {
   const { data: isAuthenticated, isLoading: isAuthLoading } = useAuthCheck();
+  
+  // Only fetch workspace state if authenticated
   const { data: workspaceState, isLoading: isWorkspaceLoading } = useWorkspaceState({
     queryKey: ['workspaceState'],
     enabled: isAuthenticated === true && !isAuthLoading,
+    // Explicitly set refetchInterval to false when not authenticated
+    refetchInterval: isAuthenticated === true ? undefined : false,
   });
 
   const getStateClassName = () => {
diff --git a/src/frontend/src/ui/MainMenu.tsx b/src/frontend/src/ui/MainMenu.tsx
@@ -7,6 +7,7 @@ import { LogOut, SquarePlus, LayoutDashboard, SquareCode, Eye, Coffee, Grid2x2,
 import { capture } from '../utils/posthog';
 import { ExcalidrawElementFactory, PlacementMode } from '../lib/ExcalidrawElementFactory';
 import { useUserProfile } from "../api/hooks";
+import { queryClient } from "../api/queryClient";
 
 interface MainMenuConfigProps {
   MainMenu: typeof MainMenuType;
@@ -171,9 +172,27 @@ export const MainMenuConfig: React.FC<MainMenuConfigProps> = ({
 
       <MainMenu.Item
           icon={<LogOut />}
-          onClick={() => {
+          onClick={async () => {
             capture('logout_clicked');
-            window.location.href = "/auth/logout";
+            
+            try {
+              // Call the logout endpoint but don't follow the redirect
+              await fetch('/auth/logout', { 
+                method: 'GET',
+                credentials: 'include' 
+              });
+              
+              // Clear the session_id cookie client-side
+              document.cookie = "session_id=; expires=Thu, 01 Jan 1970 00:00:00 UTC; path=/;";
+              
+              // Invalidate auth query to show the AuthModal
+              queryClient.invalidateQueries({ queryKey: ['auth'] });
+              queryClient.invalidateQueries({ queryKey: ['userProfile'] });
+              
+              console.log("Logged out successfully");
+            } catch (error) {
+              console.error("Logout failed:", error);
+            }
           }}
         >
           Logout

Original file line number	Diff line number	Diff line change
`@@ -44,6 +44,7 @@ async def lifespan(_: FastAPI):`
`44`	`44`	`)`
`45`	`45`
`46`	`46`	`app.mount("/assets", StaticFiles(directory=ASSETS_DIR), name="assets")`
	`47`	`+app.mount("/static", StaticFiles(directory=STATIC_DIR), name="static")`
`47`	`48`
`48`	`49`	`@app.get("/")`
`49`	`50`	`async def read_root(request: Request, auth: Optional[SessionData] = Depends(optional_auth)):`