From c9deae940c9b4f0caaaf9edd2ad8c20e6c733618 Mon Sep 17 00:00:00 2001 From: Anish Sarkar <104695310+AnishSarkar22@users.noreply.github.com> Date: Thu, 19 Mar 2026 17:51:59 +0530 Subject: [PATCH] feat: implement re-authentication flows for Google Drive, Gmail, and Calendar connectors - Added re-authentication endpoints for Google Drive, Gmail, and Calendar connectors to handle expired authentication. - Enhanced the UI to prompt users for re-authentication when their credentials are expired. - Updated backend logic to mark connectors as 'auth_expired' and manage re-authentication requests effectively. - Improved error handling for authentication failures across Google connectors. --- .../google_calendar_add_connector_route.py | 98 +++++++++++++++ .../google_drive_add_connector_route.py | 15 +++ .../google_gmail_add_connector_route.py | 98 +++++++++++++++ .../routes/search_source_connectors_routes.py | 51 ++++++-- .../assistant-ui/connector-popup.tsx | 60 +++++++--- .../components/google-drive-config.tsx | 88 ++++++++++++-- .../views/connector-edit-view.tsx | 5 +- .../views/connector-accounts-list-view.tsx | 3 + surfsense_web/hooks/use-google-picker.ts | 112 +++++++++--------- 9 files changed, 432 insertions(+), 98 deletions(-) diff --git a/surfsense_backend/app/routes/google_calendar_add_connector_route.py b/surfsense_backend/app/routes/google_calendar_add_connector_route.py index 6b74bc05d..3f3201043 100644 --- a/surfsense_backend/app/routes/google_calendar_add_connector_route.py +++ b/surfsense_backend/app/routes/google_calendar_add_connector_route.py @@ -10,8 +10,10 @@ from fastapi import APIRouter, Depends, HTTPException, Request from fastapi.responses import RedirectResponse from google_auth_oauthlib.flow import Flow from pydantic import ValidationError +from sqlalchemy import select from sqlalchemy.exc import IntegrityError from sqlalchemy.ext.asyncio import AsyncSession +from sqlalchemy.orm.attributes import flag_modified from app.config import config from app.connectors.google_gmail_connector import fetch_google_user_email @@ -111,6 +113,66 @@ async def connect_calendar(space_id: int, user: User = Depends(current_active_us ) from e +@router.get("/auth/google/calendar/connector/reauth") +async def reauth_calendar( + space_id: int, + connector_id: int, + return_url: str | None = None, + user: User = Depends(current_active_user), + session: AsyncSession = Depends(get_async_session), +): + """Initiate Google Calendar re-authentication for an existing connector.""" + try: + result = await session.execute( + select(SearchSourceConnector).filter( + SearchSourceConnector.id == connector_id, + SearchSourceConnector.user_id == user.id, + SearchSourceConnector.search_space_id == space_id, + SearchSourceConnector.connector_type + == SearchSourceConnectorType.GOOGLE_CALENDAR_CONNECTOR, + ) + ) + connector = result.scalars().first() + if not connector: + raise HTTPException( + status_code=404, + detail="Google Calendar connector not found or access denied", + ) + + if not config.SECRET_KEY: + raise HTTPException( + status_code=500, detail="SECRET_KEY not configured for OAuth security." + ) + + flow = get_google_flow() + + state_manager = get_state_manager() + extra: dict = {"connector_id": connector_id} + if return_url and return_url.startswith("/"): + extra["return_url"] = return_url + state_encoded = state_manager.generate_secure_state(space_id, user.id, **extra) + + auth_url, _ = flow.authorization_url( + access_type="offline", + prompt="consent", + include_granted_scopes="true", + state=state_encoded, + ) + + logger.info( + f"Initiating Google Calendar re-auth for user {user.id}, connector {connector_id}" + ) + return {"auth_url": auth_url} + + except HTTPException: + raise + except Exception as e: + logger.error(f"Failed to initiate Calendar re-auth: {e!s}", exc_info=True) + raise HTTPException( + status_code=500, detail=f"Failed to initiate Calendar re-auth: {e!s}" + ) from e + + @router.get("/auth/google/calendar/connector/callback") async def calendar_callback( request: Request, @@ -197,6 +259,42 @@ async def calendar_callback( # Mark that credentials are encrypted for backward compatibility creds_dict["_token_encrypted"] = True + reauth_connector_id = data.get("connector_id") + reauth_return_url = data.get("return_url") + + if reauth_connector_id: + result = await session.execute( + select(SearchSourceConnector).filter( + SearchSourceConnector.id == reauth_connector_id, + SearchSourceConnector.user_id == user_id, + SearchSourceConnector.search_space_id == space_id, + SearchSourceConnector.connector_type + == SearchSourceConnectorType.GOOGLE_CALENDAR_CONNECTOR, + ) + ) + db_connector = result.scalars().first() + if not db_connector: + raise HTTPException( + status_code=404, + detail="Connector not found or access denied during re-auth", + ) + + db_connector.config = {**creds_dict} + flag_modified(db_connector, "config") + await session.commit() + await session.refresh(db_connector) + + logger.info( + f"Re-authenticated Calendar connector {db_connector.id} for user {user_id}" + ) + if reauth_return_url and reauth_return_url.startswith("/"): + return RedirectResponse( + url=f"{config.NEXT_FRONTEND_URL}{reauth_return_url}" + ) + return RedirectResponse( + url=f"{config.NEXT_FRONTEND_URL}/dashboard/{space_id}/connectors/callback?success=true&connector=google-calendar-connector&connectorId={db_connector.id}" + ) + # Check for duplicate connector (same account already connected) is_duplicate = await check_duplicate_connector( session, diff --git a/surfsense_backend/app/routes/google_drive_add_connector_route.py b/surfsense_backend/app/routes/google_drive_add_connector_route.py index 23fba8770..e21b95dfd 100644 --- a/surfsense_backend/app/routes/google_drive_add_connector_route.py +++ b/surfsense_backend/app/routes/google_drive_add_connector_route.py @@ -525,6 +525,21 @@ async def list_google_drive_folders( raise except Exception as e: logger.error(f"Error listing Drive contents: {e!s}", exc_info=True) + error_lower = str(e).lower() + if "invalid_grant" in error_lower or "token has been expired or revoked" in error_lower or "authentication failed" in error_lower: + from sqlalchemy.orm.attributes import flag_modified + + try: + if connector and not connector.config.get("auth_expired"): + connector.config = {**connector.config, "auth_expired": True} + flag_modified(connector, "config") + await session.commit() + logger.info(f"Marked connector {connector_id} as auth_expired") + except Exception: + logger.warning(f"Failed to persist auth_expired for connector {connector_id}", exc_info=True) + raise HTTPException( + status_code=400, detail="Google Drive authentication expired. Please re-authenticate." + ) from e raise HTTPException( status_code=500, detail=f"Failed to list Drive contents: {e!s}" ) from e diff --git a/surfsense_backend/app/routes/google_gmail_add_connector_route.py b/surfsense_backend/app/routes/google_gmail_add_connector_route.py index aafe4d271..dd8e6bfac 100644 --- a/surfsense_backend/app/routes/google_gmail_add_connector_route.py +++ b/surfsense_backend/app/routes/google_gmail_add_connector_route.py @@ -10,8 +10,10 @@ from fastapi import APIRouter, Depends, HTTPException, Request from fastapi.responses import RedirectResponse from google_auth_oauthlib.flow import Flow from pydantic import ValidationError +from sqlalchemy import select from sqlalchemy.exc import IntegrityError from sqlalchemy.ext.asyncio import AsyncSession +from sqlalchemy.orm.attributes import flag_modified from app.config import config from app.connectors.google_gmail_connector import fetch_google_user_email @@ -129,6 +131,66 @@ async def connect_gmail(space_id: int, user: User = Depends(current_active_user) ) from e +@router.get("/auth/google/gmail/connector/reauth") +async def reauth_gmail( + space_id: int, + connector_id: int, + return_url: str | None = None, + user: User = Depends(current_active_user), + session: AsyncSession = Depends(get_async_session), +): + """Initiate Gmail re-authentication for an existing connector.""" + try: + result = await session.execute( + select(SearchSourceConnector).filter( + SearchSourceConnector.id == connector_id, + SearchSourceConnector.user_id == user.id, + SearchSourceConnector.search_space_id == space_id, + SearchSourceConnector.connector_type + == SearchSourceConnectorType.GOOGLE_GMAIL_CONNECTOR, + ) + ) + connector = result.scalars().first() + if not connector: + raise HTTPException( + status_code=404, + detail="Gmail connector not found or access denied", + ) + + if not config.SECRET_KEY: + raise HTTPException( + status_code=500, detail="SECRET_KEY not configured for OAuth security." + ) + + flow = get_google_flow() + + state_manager = get_state_manager() + extra: dict = {"connector_id": connector_id} + if return_url and return_url.startswith("/"): + extra["return_url"] = return_url + state_encoded = state_manager.generate_secure_state(space_id, user.id, **extra) + + auth_url, _ = flow.authorization_url( + access_type="offline", + prompt="consent", + include_granted_scopes="true", + state=state_encoded, + ) + + logger.info( + f"Initiating Gmail re-auth for user {user.id}, connector {connector_id}" + ) + return {"auth_url": auth_url} + + except HTTPException: + raise + except Exception as e: + logger.error(f"Failed to initiate Gmail re-auth: {e!s}", exc_info=True) + raise HTTPException( + status_code=500, detail=f"Failed to initiate Gmail re-auth: {e!s}" + ) from e + + @router.get("/auth/google/gmail/connector/callback") async def gmail_callback( request: Request, @@ -228,6 +290,42 @@ async def gmail_callback( # Mark that credentials are encrypted for backward compatibility creds_dict["_token_encrypted"] = True + reauth_connector_id = data.get("connector_id") + reauth_return_url = data.get("return_url") + + if reauth_connector_id: + result = await session.execute( + select(SearchSourceConnector).filter( + SearchSourceConnector.id == reauth_connector_id, + SearchSourceConnector.user_id == user_id, + SearchSourceConnector.search_space_id == space_id, + SearchSourceConnector.connector_type + == SearchSourceConnectorType.GOOGLE_GMAIL_CONNECTOR, + ) + ) + db_connector = result.scalars().first() + if not db_connector: + raise HTTPException( + status_code=404, + detail="Connector not found or access denied during re-auth", + ) + + db_connector.config = {**creds_dict} + flag_modified(db_connector, "config") + await session.commit() + await session.refresh(db_connector) + + logger.info( + f"Re-authenticated Gmail connector {db_connector.id} for user {user_id}" + ) + if reauth_return_url and reauth_return_url.startswith("/"): + return RedirectResponse( + url=f"{config.NEXT_FRONTEND_URL}{reauth_return_url}" + ) + return RedirectResponse( + url=f"{config.NEXT_FRONTEND_URL}/dashboard/{space_id}/connectors/callback?success=true&connector=google-gmail-connector&connectorId={db_connector.id}" + ) + # Check for duplicate connector (same account already connected) is_duplicate = await check_duplicate_connector( session, diff --git a/surfsense_backend/app/routes/search_source_connectors_routes.py b/surfsense_backend/app/routes/search_source_connectors_routes.py index 28bb95e5e..8b2209cc8 100644 --- a/surfsense_backend/app/routes/search_source_connectors_routes.py +++ b/surfsense_backend/app/routes/search_source_connectors_routes.py @@ -956,23 +956,46 @@ async def index_connector_content( index_google_drive_files_task, ) - if not drive_items or not drive_items.has_items(): - raise HTTPException( - status_code=400, - detail="Google Drive indexing requires drive_items body parameter with folders or files", + if drive_items and drive_items.has_items(): + logger.info( + f"Triggering Google Drive indexing for connector {connector_id} into search space {search_space_id}, " + f"folders: {len(drive_items.folders)}, files: {len(drive_items.files)}" + ) + items_dict = drive_items.model_dump() + else: + # Quick Index / periodic sync: fall back to stored config + config = connector.config or {} + selected_folders = config.get("selected_folders", []) + selected_files = config.get("selected_files", []) + if not selected_folders and not selected_files: + raise HTTPException( + status_code=400, + detail="Google Drive indexing requires folders or files to be configured. " + "Please select folders/files to index.", + ) + indexing_options = config.get( + "indexing_options", + { + "max_files_per_folder": 100, + "incremental_sync": True, + "include_subfolders": True, + }, + ) + items_dict = { + "folders": selected_folders, + "files": selected_files, + "indexing_options": indexing_options, + } + logger.info( + f"Triggering Google Drive indexing for connector {connector_id} into search space {search_space_id} " + f"using existing config" ) - logger.info( - f"Triggering Google Drive indexing for connector {connector_id} into search space {search_space_id}, " - f"folders: {len(drive_items.folders)}, files: {len(drive_items.files)}" - ) - - # Pass structured data to Celery task index_google_drive_files_task.delay( connector_id, search_space_id, str(user.id), - drive_items.model_dump(), # Convert to dict for JSON serialization + items_dict, ) response_message = "Google Drive indexing started in the background." @@ -3233,6 +3256,12 @@ async def get_drive_picker_token( raise except Exception as e: logger.error(f"Failed to get Drive picker token: {e!s}", exc_info=True) + if _is_auth_error(str(e)): + await _persist_auth_expired(session, connector_id) + raise HTTPException( + status_code=400, + detail="Google Drive authentication expired. Please re-authenticate.", + ) from e raise HTTPException( status_code=500, detail="Failed to retrieve access token. Check server logs for details.", diff --git a/surfsense_web/components/assistant-ui/connector-popup.tsx b/surfsense_web/components/assistant-ui/connector-popup.tsx index e065ce72d..cadf976d8 100644 --- a/surfsense_web/components/assistant-ui/connector-popup.tsx +++ b/surfsense_web/components/assistant-ui/connector-popup.tsx @@ -3,6 +3,7 @@ import { useAtomValue, useSetAtom } from "jotai"; import { AlertTriangle, Cable, Settings } from "lucide-react"; import { forwardRef, useEffect, useImperativeHandle, useMemo, useState } from "react"; +import { createPortal } from "react-dom"; import { documentTypeCountsAtom } from "@/atoms/documents/document-query.atoms"; import { statusInboxItemsAtom } from "@/atoms/inbox/status-inbox.atom"; import { @@ -215,11 +216,8 @@ export const ConnectorIndicator = forwardRef { - if (!open && pickerOpen) return; - handleOpenChange(open); - }} - modal={!pickerOpen} + modal={false} + onOpenChange={handleOpenChange} > {showTrigger && ( )} + {isOpen && + createPortal( + )} - + - {pickerError &&

{pickerError}

} + {(pickerError || isAuthExpired) && ( +
+ {pickerError && !isAuthExpired && ( +

{pickerError}

+ )} + {isAuthExpired && ( +
+

+ Your Google Drive authentication has expired. +

+ +
+ )} +
+ )} {/* Indexing Options */} diff --git a/surfsense_web/components/assistant-ui/connector-popup/connector-configs/views/connector-edit-view.tsx b/surfsense_web/components/assistant-ui/connector-popup/connector-configs/views/connector-edit-view.tsx index 536044df4..f4ac27504 100644 --- a/surfsense_web/components/assistant-ui/connector-popup/connector-configs/views/connector-edit-view.tsx +++ b/surfsense_web/components/assistant-ui/connector-popup/connector-configs/views/connector-edit-view.tsx @@ -169,10 +169,9 @@ export const ConnectorEditView: FC = ({

- {/* Quick Index Button - only show for indexable connectors, but not for Google Drive (requires folder selection) */} + {/* Quick Index Button - shown for indexable connectors that have been configured */} {connector.is_indexable && - onQuickIndex && - connector.connector_type !== "GOOGLE_DRIVE_CONNECTOR" && ( + onQuickIndex && (