diff --git a/ansible/ansible.cfg b/ansible/ansible.cfg index d5db45cc6..6d99d1494 100644 --- a/ansible/ansible.cfg +++ b/ansible/ansible.cfg @@ -1,10 +1,10 @@ [defaults] inventory = inventory.ini -host_key_checking = False +host_key_checking = True retry_files_enabled = False stdout_callback = default gathering = smart fact_caching = memory [ssh_connection] -ssh_args = -o ControlMaster=auto -o ControlPersist=60s -o UserKnownHostsFile=/dev/null -o StrictHostKeyChecking=no +ssh_args = -o ControlMaster=auto -o ControlPersist=60s -o StrictHostKeyChecking=accept-new diff --git a/ansible/playbooks/openrag.yml b/ansible/playbooks/openrag.yml index 71bb73c0a..db2b125a8 100644 --- a/ansible/playbooks/openrag.yml +++ b/ansible/playbooks/openrag.yml @@ -104,7 +104,7 @@ dest: "{{ project_path }}/.env" owner: "{{ project_user }}" group: "{{ project_user }}" - mode: "0644" + mode: "0600" when: - not env_file.stat.exists - local_env_file.stat.exists @@ -124,8 +124,11 @@ - name: Setup Python environment block: - name: Install uv (Python package manager) - shell: curl -LsSf https://astral.sh/uv/install.sh | sh + shell: | + set -euo pipefail + curl -LsSf https://astral.sh/uv/0.5.11/install.sh | sh args: + executable: /bin/bash creates: "/home/{{ project_user }}/.cargo/bin/uv" - name: Add uv to PATH in .bashrc diff --git a/docs/assets/compose_linux_gpu.yaml b/docs/assets/compose_linux_gpu.yaml index f13d68cf7..b3db5cfde 100644 --- a/docs/assets/compose_linux_gpu.yaml +++ b/docs/assets/compose_linux_gpu.yaml @@ -16,7 +16,7 @@ x-openrag: &openrag_template - ./ray_mount/logs:/app/logs ports: - ${APP_PORT:-8080}:${APP_iPORT:-8080} - - ${RAY_DASHBOARD_PORT:-8265}:8265 # Disable when in cluster mode + - 127.0.0.1:${RAY_DASHBOARD_PORT:-8265}:8265 # Localhost only: Ray dashboard/Jobs API is unauthenticated. Disable when in cluster mode networks: default: aliases: diff --git a/openrag/components/auth/oidc_client.py b/openrag/components/auth/oidc_client.py index 06b95303a..cdede77c3 100644 --- a/openrag/components/auth/oidc_client.py +++ b/openrag/components/auth/oidc_client.py @@ -290,6 +290,8 @@ async def _verify_id_token(self, token: str, *, expected_nonce: str | None) -> d if isinstance(aud, list): if self.client_id not in aud: raise ValueError(f"ID token aud {aud!r} does not contain client_id {self.client_id!r}") + if len(aud) > 1 and decoded.get("azp") != self.client_id: + raise ValueError(f"ID token azp {decoded.get('azp')!r} != client_id {self.client_id!r}") elif aud != self.client_id: raise ValueError(f"ID token aud {aud!r} != client_id {self.client_id!r}") @@ -349,6 +351,8 @@ async def verify_logout_token(self, token: str) -> LogoutTokenClaims: if isinstance(aud, list): if self.client_id not in aud: raise ValueError(f"logout_token aud {aud!r} does not contain client_id {self.client_id!r}") + if len(aud) > 1 and decoded.get("azp") != self.client_id: + raise ValueError(f"logout_token azp {decoded.get('azp')!r} != client_id {self.client_id!r}") elif aud != self.client_id: raise ValueError(f"logout_token aud {aud!r} != client_id {self.client_id!r}") diff --git a/openrag/components/indexer/loaders/CustomDocLoader.py b/openrag/components/indexer/loaders/CustomDocLoader.py index b91206e81..3238d2c26 100644 --- a/openrag/components/indexer/loaders/CustomDocLoader.py +++ b/openrag/components/indexer/loaders/CustomDocLoader.py @@ -34,6 +34,6 @@ async def aload_document(self, file_path, metadata: dict = None): s = "" for page_num, p in enumerate(pages, start=1): - s = p.page_content.strip() + f"\n[PAGE_{page_num}]\n" + s += p.page_content.strip() + f"\n[PAGE_{page_num}]\n" return Document(page_content=s, metadata=metadata) diff --git a/openrag/components/indexer/utils/text_sanitizer.py b/openrag/components/indexer/utils/text_sanitizer.py index 4a308ce23..c0479d6bf 100644 --- a/openrag/components/indexer/utils/text_sanitizer.py +++ b/openrag/components/indexer/utils/text_sanitizer.py @@ -98,7 +98,11 @@ def sanitize_text( # Control tokens a document could embed to forge citations or source boundaries: # "[Source N]" (block marker), "[Sources: 1, 3]" / "Sources: 1, 3" (answer tag), # and "----------" (separator). Neutralize them in untrusted text. -_INJECT_SOURCE_BLOCK_RE = re.compile(r"\[\s*(sources?)\b", re.IGNORECASE) +# Capture an optional trailing colon: dropping it also breaks the answer-tag form +# "[Sources: 1, 2]". The output parser's bracket is optional (\[?), so neutralizing +# only the "[" would leave "(Sources: 1, 2]" — still a parser match. Removing the +# colon defangs the keyword the parser keys on. +_INJECT_SOURCE_BLOCK_RE = re.compile(r"\[\s*(sources?)\b(\s*:)?", re.IGNORECASE) _INJECT_SOURCES_TAG_RE = re.compile(r"(?im)^([ \t]*)(sources?)(\s*:\s*)(\[?[\d,\s]+\]?)[ \t]*$") _INJECT_SEPARATOR_RE = re.compile(r"-{4,}") @@ -107,8 +111,9 @@ def neutralize_prompt_control_tokens(text: str) -> str: """Defang RAG control tokens in untrusted text so they can't fake markers.""" if not text: return text - # "[Source...]" / "[Sources...]" -> open paren so it can't start a marker. - text = _INJECT_SOURCE_BLOCK_RE.sub(r"(\1", text) + # "[Source...]" / "[Sources...]" -> open paren so it can't start a marker, and + # drop any "[Sources:" colon so the answer-tag parser can't match the remainder. + text = _INJECT_SOURCE_BLOCK_RE.sub(lambda m: "(" + m.group(1) + (" " if m.group(2) else ""), text) # Break the unbracketed "Sources: 1, 2" line form the parser also matches. text = _INJECT_SOURCES_TAG_RE.sub(r"\1\2 \4", text) # Cap long hyphen runs so they can't reproduce the separator. diff --git a/openrag/components/indexer/vectordb/utils.py b/openrag/components/indexer/vectordb/utils.py index 7bc0983c9..24fe211e2 100644 --- a/openrag/components/indexer/vectordb/utils.py +++ b/openrag/components/indexer/vectordb/utils.py @@ -54,6 +54,7 @@ def __init__(self, database_url: str, logger=logger): ) def _ensure_admin_user(self, admin_token: str): + token_provided = bool(admin_token) if not admin_token: admin_token = f"or-{secrets.token_hex(16)}" hashed_token = self.hash_token(admin_token) @@ -70,9 +71,13 @@ def _ensure_admin_user(self, admin_token: str): self.logger.info("Created admin user") else: admin.is_admin = True - admin.token = hashed_token + # Only (re)set the token when one was explicitly supplied via + # AUTH_TOKEN; otherwise keep the stored token so it doesn't + # rotate on every startup and invalidate existing clients. + if token_provided: + admin.token = hashed_token s.commit() - self.logger.info("Upgraded existing user to admin") + self.logger.info("Ensured admin user (id=1)") def list_partition_files(self, partition: str, limit: int | None = None): """List files in a partition with optional limit - Optimized by querying File table directly""" @@ -435,6 +440,9 @@ def regenerate_user_token(self, user_id: int) -> dict | None: user.token = hashed_token s.commit() s.refresh(user) + # Rotating the API token also invalidates the user's active OIDC + # browser sessions so they can't outlive the rotation. + self.revoke_oidc_sessions_by_user_id(user_id) return { "id": user.id, @@ -1113,6 +1121,22 @@ def revoke_oidc_session_by_id(self, session_id: int) -> None: s.commit() self.logger.bind(session_id=session_id).info("Revoked OIDC session") + def revoke_oidc_sessions_by_user_id(self, user_id: int) -> int: + """Revoke all active OIDC sessions for a user (e.g. on API-token rotation).""" + now = datetime.now() + with self.Session() as s: + stmt = ( + update(OIDCSession) + .where(OIDCSession.user_id == user_id) + .where(OIDCSession.revoked_at.is_(None)) + .values(revoked_at=now) + ) + result = s.execute(stmt) + s.commit() + count = result.rowcount or 0 + self.logger.bind(user_id=user_id, count=count).info("Revoked OIDC sessions by user_id") + return count + def cleanup_expired_oidc_sessions(self) -> int: """Delete rows whose ``session_expires_at`` is older than 7 days. diff --git a/openrag/models/user.py b/openrag/models/user.py index 56a49f8d5..87fc8ac31 100644 --- a/openrag/models/user.py +++ b/openrag/models/user.py @@ -1,4 +1,4 @@ -from pydantic import BaseModel, ConfigDict, Field +from pydantic import BaseModel, ConfigDict, Field, field_validator class UserBase(BaseModel): @@ -8,6 +8,15 @@ class UserBase(BaseModel): is_admin: bool = False file_quota: int | None = Field(default=10) + @field_validator("external_user_id", mode="before") + @classmethod + def _empty_external_id_to_none(cls, v): + # Coerce "" / whitespace to NULL so it can't collide on the unique index + # (Postgres allows many NULLs but only one empty string). + if isinstance(v, str) and not v.strip(): + return None + return v + class UserCreate(UserBase): # Reject unknown fields so callers cannot smuggle extra column names. diff --git a/openrag/routers/auth.py b/openrag/routers/auth.py index d5a09ca15..91b3eadf8 100644 --- a/openrag/routers/auth.py +++ b/openrag/routers/auth.py @@ -195,6 +195,10 @@ def _sanitize_next_url(next_url: str | None) -> str: """ if not next_url: return "/" + # Reject backslashes (browsers treat "\" as "/", so "/\evil.com" resolves + # protocol-relative) and any control chars (CR/LF header injection). + if "\\" in next_url or any(ord(c) < 0x20 or ord(c) == 0x7F for c in next_url): + return "/" if next_url.startswith("/") and not next_url.startswith("//"): return next_url # Absolute URL: only allow whitelisted origins. @@ -443,6 +447,16 @@ async def callback(request: Request, code: str | None = None, state: str | None "Failed to fetch userinfo from IdP.", delete_state_cookie=True, ) + # Bind userinfo to the verified ID token: per the OIDC spec the + # userinfo `sub` MUST equal the ID token `sub`, else the response + # could describe a different principal (token substitution). + if claims_for_mapping.get("sub") != sub: + logger.warning(f"OIDC userinfo sub mismatch for user_id={user['id']}") + return _json_error( + status.HTTP_400_BAD_REQUEST, + "userinfo sub does not match ID token.", + delete_state_cookie=True, + ) else: claims_for_mapping = bundle.claims diff --git a/openrag/routers/indexer.py b/openrag/routers/indexer.py index 5204c3c38..482ad9c2c 100644 --- a/openrag/routers/indexer.py +++ b/openrag/routers/indexer.py @@ -26,6 +26,7 @@ current_user_partitions, ensure_partition_role, human_readable_size, + is_file_id_valid, require_partition_editor, require_task_owner, validate_file_format, @@ -41,7 +42,6 @@ DATA_DIR = config.paths.data_dir VECTORDB_TIMEOUT = config.ray.indexer.vectordb_timeout -FORBIDDEN_CHARS_IN_FILE_ID = set("/") # set('"<>#%{}|\\^`[]') LOG_FILE = Path(config.paths.log_dir or "logs") / "app.json" # supported file formats or mimetypes @@ -407,6 +407,11 @@ async def copy_file_between_partitions( user_partitions=Depends(current_user_partitions), _quota_check=Depends(check_user_file_quota), ): + if not is_file_id_valid(source_file_id): + raise HTTPException( + status_code=status.HTTP_400_BAD_REQUEST, + detail="Source file ID may only contain letters, digits, '.', '_', ':' and '-'.", + ) # Make sure user has access to destination partition await ensure_partition_role( partition=source_partition, diff --git a/openrag/routers/openai.py b/openrag/routers/openai.py index 2f7d13cb9..a1f241411 100644 --- a/openrag/routers/openai.py +++ b/openrag/routers/openai.py @@ -354,7 +354,13 @@ async def stream_response(): return StreamingResponse(stream_response(), media_type="text/event-stream") else: - chunk = await llm_output.__anext__() + try: + chunk = await llm_output.__anext__() + except StopAsyncIteration: + raise HTTPException( + status_code=status.HTTP_502_BAD_GATEWAY, + detail="Empty response from upstream LLM", + ) chunk["model"] = model_name content = chunk.get("choices", [{}])[0].get("message", {}).get("content", "") or "" @@ -430,7 +436,13 @@ async def openai_completion( sources = __prepare_sources(request2, docs) - complete_response = await llm_output.__anext__() + try: + complete_response = await llm_output.__anext__() + except StopAsyncIteration: + raise HTTPException( + status_code=status.HTTP_502_BAD_GATEWAY, + detail="Empty response from upstream LLM", + ) text = complete_response.get("choices", [{}])[0].get("text", "") or "" clean_text, citations = extract_and_strip_sources_block(text) diff --git a/openrag/routers/search.py b/openrag/routers/search.py index 545d27acd..27a683f97 100644 --- a/openrag/routers/search.py +++ b/openrag/routers/search.py @@ -126,7 +126,7 @@ async def search_multiple_partitions( log = logger.bind( partitions=partitions, - query=search_params.text, + query_len=len(search_params.text), top_k=search_params.top_k, workspace=workspace, include_related=related_params.include_related, @@ -236,7 +236,7 @@ async def search_one_partition( ): log = logger.bind( partition=partition, - query=search_params.text, + query_len=len(search_params.text), top_k=search_params.top_k, workspace=workspace, include_related=related_params.include_related, @@ -336,7 +336,9 @@ async def search_file( vectordb=Depends(get_vectordb), partition_viewer=Depends(require_partition_viewer), ): - log = logger.bind(partition=partition, file_id=file_id, query=search_params.text, top_k=search_params.top_k) + log = logger.bind( + partition=partition, file_id=file_id, query_len=len(search_params.text), top_k=search_params.top_k + ) filter = "file_id == {_file_id}" + (f" AND {search_params.filter}" if search_params.filter else "") params = {"_file_id": file_id} diff --git a/tests/api_tests/test_indexer.py b/tests/api_tests/test_indexer.py index 1cc86f855..2d44b4959 100644 --- a/tests/api_tests/test_indexer.py +++ b/tests/api_tests/test_indexer.py @@ -784,11 +784,18 @@ def test_cancel_increments_total_cancelled(self, api_client, created_partition, cancel_response = api_client.delete(f"/indexer/task/{task_id}") assert cancel_response.status_code == 200 - info_after = api_client.get("/queue/info") - assert info_after.status_code == 200 - cancelled_after = info_after.json()["tasks"]["total_cancelled"] + # The counter is updated asynchronously, so poll until it increments. + cancelled_after = cancelled_before + start = time.time() + while time.time() - start < TASK_TIMEOUT: + info_after = api_client.get("/queue/info") + assert info_after.status_code == 200 + cancelled_after = info_after.json()["tasks"]["total_cancelled"] + if cancelled_after >= cancelled_before + 1: + break + time.sleep(0.5) - assert cancelled_after == cancelled_before + 1 + assert cancelled_after >= cancelled_before + 1 def test_cancel_nonexistent_task_returns_404(self, api_client): """Cancelling a task that does not exist must return 404."""