diff --git a/api/apps/connector_app.py b/api/apps/connector_app.py index ba78d6d6f..23965e617 100644 --- a/api/apps/connector_app.py +++ b/api/apps/connector_app.py @@ -13,16 +13,26 @@ # See the License for the specific language governing permissions and # limitations under the License. # +import json +import logging import time +import uuid +from html import escape +from typing import Any -from flask import request -from flask_login import login_required, current_user +from flask import make_response, request +from flask_login import current_user, login_required +from google_auth_oauthlib.flow import Flow from api.db import InputType from api.db.services.connector_service import ConnectorService, SyncLogsService -from api.utils.api_utils import get_json_result, validate_request, get_data_error_result -from common.misc_utils import get_uuid +from api.utils.api_utils import get_data_error_result, get_json_result, validate_request from common.constants import RetCode, TaskStatus +from common.data_source.config import GOOGLE_DRIVE_WEB_OAUTH_REDIRECT_URI, DocumentSource +from common.data_source.google_util.constant import GOOGLE_DRIVE_WEB_OAUTH_POPUP_TEMPLATE, GOOGLE_SCOPES +from common.misc_utils import get_uuid +from rag.utils.redis_conn import REDIS_CONN + @manager.route("/set", methods=["POST"]) # noqa: F821 @login_required @@ -42,8 +52,8 @@ def set_connector(): "config": req["config"], "refresh_freq": int(req.get("refresh_freq", 30)), "prune_freq": int(req.get("prune_freq", 720)), - "timeout_secs": int(req.get("timeout_secs", 60*29)), - "status": TaskStatus.SCHEDULE + "timeout_secs": int(req.get("timeout_secs", 60 * 29)), + "status": TaskStatus.SCHEDULE, } conn["status"] = TaskStatus.SCHEDULE ConnectorService.save(**conn) @@ -105,3 +115,181 @@ def rm_connector(connector_id): ConnectorService.resume(connector_id, TaskStatus.CANCEL) ConnectorService.delete_by_id(connector_id) return get_json_result(data=True) + + +GOOGLE_WEB_FLOW_STATE_PREFIX = "google_drive_web_flow_state" +GOOGLE_WEB_FLOW_RESULT_PREFIX = "google_drive_web_flow_result" +WEB_FLOW_TTL_SECS = 15 * 60 + + +def _web_state_cache_key(flow_id: str) -> str: + return f"{GOOGLE_WEB_FLOW_STATE_PREFIX}:{flow_id}" + + +def _web_result_cache_key(flow_id: str) -> str: + return f"{GOOGLE_WEB_FLOW_RESULT_PREFIX}:{flow_id}" + + +def _load_credentials(payload: str | dict[str, Any]) -> dict[str, Any]: + if isinstance(payload, dict): + return payload + try: + return json.loads(payload) + except json.JSONDecodeError as exc: # pragma: no cover - defensive + raise ValueError("Invalid Google credentials JSON.") from exc + + +def _get_web_client_config(credentials: dict[str, Any]) -> dict[str, Any]: + web_section = credentials.get("web") + if not isinstance(web_section, dict): + raise ValueError("Google OAuth JSON must include a 'web' client configuration to use browser-based authorization.") + return {"web": web_section} + + +def _render_web_oauth_popup(flow_id: str, success: bool, message: str): + status = "success" if success else "error" + auto_close = "window.close();" if success else "" + escaped_message = escape(message) + payload_json = json.dumps( + { + "type": "ragflow-google-drive-oauth", + "status": status, + "flowId": flow_id or "", + "message": message, + } + ) + html = GOOGLE_DRIVE_WEB_OAUTH_POPUP_TEMPLATE.format( + heading="Authorization complete" if success else "Authorization failed", + message=escaped_message, + payload_json=payload_json, + auto_close=auto_close, + ) + response = make_response(html, 200) + response.headers["Content-Type"] = "text/html; charset=utf-8" + return response + + +@manager.route("/google-drive/oauth/web/start", methods=["POST"]) # noqa: F821 +@login_required +@validate_request("credentials") +def start_google_drive_web_oauth(): + if not GOOGLE_DRIVE_WEB_OAUTH_REDIRECT_URI: + return get_json_result( + code=RetCode.SERVER_ERROR, + message="Google Drive OAuth redirect URI is not configured on the server.", + ) + + req = request.json or {} + raw_credentials = req.get("credentials", "") + try: + credentials = _load_credentials(raw_credentials) + except ValueError as exc: + return get_json_result(code=RetCode.ARGUMENT_ERROR, message=str(exc)) + + if credentials.get("refresh_token"): + return get_json_result( + code=RetCode.ARGUMENT_ERROR, + message="Uploaded credentials already include a refresh token.", + ) + + try: + client_config = _get_web_client_config(credentials) + except ValueError as exc: + return get_json_result(code=RetCode.ARGUMENT_ERROR, message=str(exc)) + + flow_id = str(uuid.uuid4()) + try: + flow = Flow.from_client_config(client_config, scopes=GOOGLE_SCOPES[DocumentSource.GOOGLE_DRIVE]) + flow.redirect_uri = GOOGLE_DRIVE_WEB_OAUTH_REDIRECT_URI + authorization_url, _ = flow.authorization_url( + access_type="offline", + include_granted_scopes="true", + prompt="consent", + state=flow_id, + ) + except Exception as exc: # pragma: no cover - defensive + logging.exception("Failed to create Google OAuth flow: %s", exc) + return get_json_result( + code=RetCode.SERVER_ERROR, + message="Failed to initialize Google OAuth flow. Please verify the uploaded client configuration.", + ) + + cache_payload = { + "user_id": current_user.id, + "client_config": client_config, + "created_at": int(time.time()), + } + REDIS_CONN.set_obj(_web_state_cache_key(flow_id), cache_payload, WEB_FLOW_TTL_SECS) + + return get_json_result( + data={ + "flow_id": flow_id, + "authorization_url": authorization_url, + "expires_in": WEB_FLOW_TTL_SECS, + } + ) + + +@manager.route("/google-drive/oauth/web/callback", methods=["GET"]) # noqa: F821 +def google_drive_web_oauth_callback(): + state_id = request.args.get("state") + error = request.args.get("error") + error_description = request.args.get("error_description") or error + + if not state_id: + return _render_web_oauth_popup("", False, "Missing OAuth state parameter.") + + state_cache = REDIS_CONN.get(_web_state_cache_key(state_id)) + if not state_cache: + return _render_web_oauth_popup(state_id, False, "Authorization session expired. Please restart from the main window.") + + state_obj = json.loads(state_cache) + client_config = state_obj.get("client_config") + if not client_config: + REDIS_CONN.delete(_web_state_cache_key(state_id)) + return _render_web_oauth_popup(state_id, False, "Authorization session was invalid. Please retry.") + + if error: + REDIS_CONN.delete(_web_state_cache_key(state_id)) + return _render_web_oauth_popup(state_id, False, error_description or "Authorization was cancelled.") + + code = request.args.get("code") + if not code: + return _render_web_oauth_popup(state_id, False, "Missing authorization code from Google.") + + try: + flow = Flow.from_client_config(client_config, scopes=GOOGLE_SCOPES[DocumentSource.GOOGLE_DRIVE]) + flow.redirect_uri = GOOGLE_DRIVE_WEB_OAUTH_REDIRECT_URI + flow.fetch_token(code=code) + except Exception as exc: # pragma: no cover - defensive + logging.exception("Failed to exchange Google OAuth code: %s", exc) + REDIS_CONN.delete(_web_state_cache_key(state_id)) + return _render_web_oauth_popup(state_id, False, "Failed to exchange tokens with Google. Please retry.") + + creds_json = flow.credentials.to_json() + result_payload = { + "user_id": state_obj.get("user_id"), + "credentials": creds_json, + } + REDIS_CONN.set_obj(_web_result_cache_key(state_id), result_payload, WEB_FLOW_TTL_SECS) + REDIS_CONN.delete(_web_state_cache_key(state_id)) + + return _render_web_oauth_popup(state_id, True, "Authorization completed successfully.") + + +@manager.route("/google-drive/oauth/web/result", methods=["POST"]) # noqa: F821 +@login_required +@validate_request("flow_id") +def poll_google_drive_web_result(): + req = request.json or {} + flow_id = req.get("flow_id") + cache_raw = REDIS_CONN.get(_web_result_cache_key(flow_id)) + if not cache_raw: + return get_json_result(code=RetCode.RUNNING, message="Authorization is still pending.") + + result = json.loads(cache_raw) + if result.get("user_id") != current_user.id: + return get_json_result(code=RetCode.PERMISSION_ERROR, message="You are not allowed to access this authorization result.") + + REDIS_CONN.delete(_web_result_cache_key(flow_id)) + return get_json_result(data={"credentials": result.get("credentials")}) diff --git a/common/data_source/config.py b/common/data_source/config.py index 5eea4f6e4..196d9ed3e 100644 --- a/common/data_source/config.py +++ b/common/data_source/config.py @@ -190,6 +190,7 @@ OAUTH_GOOGLE_DRIVE_CLIENT_ID = os.environ.get("OAUTH_GOOGLE_DRIVE_CLIENT_ID", "" OAUTH_GOOGLE_DRIVE_CLIENT_SECRET = os.environ.get( "OAUTH_GOOGLE_DRIVE_CLIENT_SECRET", "" ) +GOOGLE_DRIVE_WEB_OAUTH_REDIRECT_URI = os.environ.get("GOOGLE_DRIVE_WEB_OAUTH_REDIRECT_URI", "http://localhost:9380/v1/connector/google-drive/oauth/web/callback") CONFLUENCE_OAUTH_TOKEN_URL = "https://auth.atlassian.com/oauth/token" RATE_LIMIT_MESSAGE_LOWERCASE = "Rate limit exceeded".lower() diff --git a/common/data_source/google_util/constant.py b/common/data_source/google_util/constant.py index c0b7f0711..8ab75fa14 100644 --- a/common/data_source/google_util/constant.py +++ b/common/data_source/google_util/constant.py @@ -47,3 +47,57 @@ USER_FIELDS = "nextPageToken, users(primaryEmail)" # Error message substrings MISSING_SCOPES_ERROR_STR = "client not authorized for any of the scopes requested" SCOPE_INSTRUCTIONS = "" + + +GOOGLE_DRIVE_WEB_OAUTH_POPUP_TEMPLATE = """ + + + + Google Drive Authorization + + + +
+

{heading}

+

{message}

+

You can close this window.

+
+ + + +""" diff --git a/common/data_source/google_util/oauth_flow.py b/common/data_source/google_util/oauth_flow.py index edf9a24fd..7e39e5283 100644 --- a/common/data_source/google_util/oauth_flow.py +++ b/common/data_source/google_util/oauth_flow.py @@ -3,9 +3,15 @@ import os import threading from typing import Any, Callable +import requests + from common.data_source.config import DocumentSource from common.data_source.google_util.constant import GOOGLE_SCOPES +GOOGLE_DEVICE_CODE_URL = "https://oauth2.googleapis.com/device/code" +GOOGLE_DEVICE_TOKEN_URL = "https://oauth2.googleapis.com/token" +DEFAULT_DEVICE_INTERVAL = 5 + def _get_requested_scopes(source: DocumentSource) -> list[str]: """Return the scopes to request, honoring an optional override env var.""" @@ -49,6 +55,62 @@ def _run_with_timeout(func: Callable[[], Any], timeout_secs: int, timeout_messag return result.get("value") +def _extract_client_info(credentials: dict[str, Any]) -> tuple[str, str | None]: + if "client_id" in credentials: + return credentials["client_id"], credentials.get("client_secret") + for key in ("installed", "web"): + if key in credentials and isinstance(credentials[key], dict): + nested = credentials[key] + if "client_id" not in nested: + break + return nested["client_id"], nested.get("client_secret") + raise ValueError("Provided Google OAuth credentials are missing client_id.") + + +def start_device_authorization_flow( + credentials: dict[str, Any], + source: DocumentSource, +) -> tuple[dict[str, Any], dict[str, Any]]: + client_id, client_secret = _extract_client_info(credentials) + data = { + "client_id": client_id, + "scope": " ".join(_get_requested_scopes(source)), + } + if client_secret: + data["client_secret"] = client_secret + resp = requests.post(GOOGLE_DEVICE_CODE_URL, data=data, timeout=15) + resp.raise_for_status() + payload = resp.json() + state = { + "client_id": client_id, + "client_secret": client_secret, + "device_code": payload.get("device_code"), + "interval": payload.get("interval", DEFAULT_DEVICE_INTERVAL), + } + response_data = { + "user_code": payload.get("user_code"), + "verification_url": payload.get("verification_url") or payload.get("verification_uri"), + "verification_url_complete": payload.get("verification_url_complete") + or payload.get("verification_uri_complete"), + "expires_in": payload.get("expires_in"), + "interval": state["interval"], + } + return state, response_data + + +def poll_device_authorization_flow(state: dict[str, Any]) -> dict[str, Any]: + data = { + "client_id": state["client_id"], + "device_code": state["device_code"], + "grant_type": "urn:ietf:params:oauth:grant-type:device_code", + } + if state.get("client_secret"): + data["client_secret"] = state["client_secret"] + resp = requests.post(GOOGLE_DEVICE_TOKEN_URL, data=data, timeout=20) + resp.raise_for_status() + return resp.json() + + def _run_local_server_flow(client_config: dict[str, Any], source: DocumentSource) -> dict[str, Any]: """Launch the standard Google OAuth local-server flow to mint user tokens.""" from google_auth_oauthlib.flow import InstalledAppFlow # type: ignore diff --git a/web/src/pages/user-setting/data-source/component/google-drive-token-field.tsx b/web/src/pages/user-setting/data-source/component/google-drive-token-field.tsx index 6d212e971..1f0d91617 100644 --- a/web/src/pages/user-setting/data-source/component/google-drive-token-field.tsx +++ b/web/src/pages/user-setting/data-source/component/google-drive-token-field.tsx @@ -1,64 +1,383 @@ -import { useMemo, useState } from 'react'; +import { useCallback, useEffect, useMemo, useRef, useState } from 'react'; import { FileUploader } from '@/components/file-uploader'; +import { Button } from '@/components/ui/button'; +import { + Dialog, + DialogContent, + DialogDescription, + DialogFooter, + DialogHeader, + DialogTitle, +} from '@/components/ui/dialog'; import message from '@/components/ui/message'; -import { Textarea } from '@/components/ui/textarea'; import { FileMimeType } from '@/constants/common'; +import { + pollGoogleDriveWebAuthResult, + startGoogleDriveWebAuth, +} from '@/services/data-source-service'; +import { Loader2 } from 'lucide-react'; type GoogleDriveTokenFieldProps = { value?: string; onChange: (value: any) => void; - placeholder?: string; +}; + +const credentialHasRefreshToken = (content: string) => { + try { + const parsed = JSON.parse(content); + return Boolean(parsed?.refresh_token); + } catch { + return false; + } +}; + +const describeCredentials = (content?: string) => { + if (!content) return ''; + try { + const parsed = JSON.parse(content); + if (parsed?.refresh_token) { + return 'Uploaded OAuth tokens with a refresh token.'; + } + if (parsed?.installed || parsed?.web) { + return 'Client credentials detected. Complete verification to mint long-lived tokens.'; + } + return 'Stored Google credential JSON.'; + } catch { + return ''; + } }; const GoogleDriveTokenField = ({ value, onChange, - placeholder, }: GoogleDriveTokenFieldProps) => { const [files, setFiles] = useState([]); + const [pendingCredentials, setPendingCredentials] = useState(''); + const [dialogOpen, setDialogOpen] = useState(false); + const [webAuthLoading, setWebAuthLoading] = useState(false); + const [webFlowId, setWebFlowId] = useState(null); + const [webStatus, setWebStatus] = useState< + 'idle' | 'waiting' | 'success' | 'error' + >('idle'); + const [webStatusMessage, setWebStatusMessage] = useState(''); + const webFlowIdRef = useRef(null); + const webPollTimerRef = useRef | null>(null); - const handleValueChange = useMemo( - () => (nextFiles: File[]) => { + const clearWebState = useCallback(() => { + if (webPollTimerRef.current) { + clearTimeout(webPollTimerRef.current); + webPollTimerRef.current = null; + } + webFlowIdRef.current = null; + setWebFlowId(null); + setWebStatus('idle'); + setWebStatusMessage(''); + }, []); + + useEffect(() => { + return () => { + if (webPollTimerRef.current) { + clearTimeout(webPollTimerRef.current); + } + }; + }, []); + + useEffect(() => { + webFlowIdRef.current = webFlowId; + }, [webFlowId]); + + const credentialSummary = useMemo(() => describeCredentials(value), [value]); + const hasVerifiedTokens = useMemo( + () => Boolean(value && credentialHasRefreshToken(value)), + [value], + ); + const hasUploadedButUnverified = useMemo( + () => Boolean(value && !hasVerifiedTokens), + [hasVerifiedTokens, value], + ); + + const resetDialog = useCallback( + (shouldResetState: boolean) => { + setDialogOpen(false); + clearWebState(); + if (shouldResetState) { + setPendingCredentials(''); + setFiles([]); + } + }, + [clearWebState], + ); + + const fetchWebResult = useCallback( + async (flowId: string) => { + try { + const { data } = await pollGoogleDriveWebAuthResult({ + flow_id: flowId, + }); + if (data.code === 0 && data.data?.credentials) { + onChange(data.data.credentials); + setPendingCredentials(''); + message.success('Google Drive credentials verified.'); + resetDialog(false); + return; + } + if (data.code === 106) { + setWebStatus('waiting'); + setWebStatusMessage('Authorization confirmed. Finalizing tokens...'); + if (webPollTimerRef.current) { + clearTimeout(webPollTimerRef.current); + } + webPollTimerRef.current = setTimeout( + () => fetchWebResult(flowId), + 1500, + ); + return; + } + message.error(data.message || 'Authorization failed.'); + clearWebState(); + } catch (err) { + message.error('Unable to retrieve authorization result.'); + clearWebState(); + } + }, + [clearWebState, onChange, resetDialog], + ); + + useEffect(() => { + const handler = (event: MessageEvent) => { + const payload = event.data; + if (!payload || payload.type !== 'ragflow-google-drive-oauth') { + return; + } + if (!payload.flowId) { + return; + } + if (webFlowIdRef.current && webFlowIdRef.current !== payload.flowId) { + return; + } + + if (payload.status === 'success') { + setWebStatus('waiting'); + setWebStatusMessage('Authorization confirmed. Finalizing tokens...'); + fetchWebResult(payload.flowId); + } else { + message.error( + payload.message || 'Authorization window reported an error.', + ); + clearWebState(); + } + }; + + window.addEventListener('message', handler); + return () => window.removeEventListener('message', handler); + }, [clearWebState, fetchWebResult]); + + const handleValueChange = useCallback( + (nextFiles: File[]) => { if (!nextFiles.length) { setFiles([]); + onChange(''); + setPendingCredentials(''); + clearWebState(); return; } const file = nextFiles[nextFiles.length - 1]; file .text() .then((text) => { - JSON.parse(text); - onChange(text); + try { + JSON.parse(text); + } catch { + message.error('Invalid JSON file.'); + setFiles([]); + clearWebState(); + return; + } setFiles([file]); - message.success('JSON uploaded'); + clearWebState(); + if (credentialHasRefreshToken(text)) { + onChange(text); + setPendingCredentials(''); + message.success('OAuth credentials uploaded.'); + return; + } + setPendingCredentials(text); + setDialogOpen(true); + message.info( + 'Client configuration uploaded. Verification is required to finish setup.', + ); }) .catch(() => { - message.error('Invalid JSON file.'); + message.error('Unable to read the uploaded file.'); setFiles([]); }); }, - [onChange], + [clearWebState, onChange], ); - return ( -
-