Add RAG sensitive data exposure#19
Conversation
📝 WalkthroughWalkthroughA new RAG Sensitive Data Exposure lab is introduced with a FAISS+SQLite vector store, three vulnerability levels simulating data leakage, FastAPI endpoints, and a complete frontend UI for interactive testing and secret validation. ChangesRAG Sensitive Data Exposure Lab
Sequence Diagram(s)sequenceDiagram
participant Browser as Browser / User
participant Controller as RAG Controller
participant Lab as Lab Pipeline
participant Store as Vector Store
participant Embedder as Embedding Service
participant LLM as LLM Chat
Browser->>Controller: POST /level{N} with user_input
Controller->>Lab: evaluate_level(level, user_input, model)
Lab->>Lab: validate input length & denylist
Lab->>Store: ensure_level_indexed(level)
Store->>Embedder: embed all doc chunks
Embedder->>Store: embeddings + metadata
Store->>Store: index via FAISS
Lab->>Embedder: embed user_input
Lab->>Store: search(embedding, top_k)
Store->>Lab: ranked matches + metadata
Lab->>Lab: format context from matches
Lab->>LLM: call with system_prompt + context
LLM->>Lab: model output
Lab->>Lab: detect secret in output
Lab->>Controller: response with retrieved_docs, leak reason
Controller->>Browser: JSON: {retrieved_docs, input_accepted, secret_exposed}
Browser->>Browser: render docs list, show leak feedback
Estimated code review effort🎯 4 (Complex) | ⏱️ ~60 minutes The PR introduces a substantial new lab with vector store implementation (FAISS+SQLite), multi-level configuration with per-level secrets and denylist behavior, embedding/retrieval logic with metadata filtering, and a complete frontend UI. While individual files are well-scoped, the changes span framework extensions, service layer, API controller, test data, and frontend, requiring review of dense logic in the lab module, controller error handling, and frontend state management across multiple files. Possibly related PRs
Poem
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
Lexical denylist (password, secret, admin) bypassed by semantic paraphrase retrieval. - Add L2 corpus and denylist input handler - Add level2 endpoint and locale entries - Make the facade level-aware
…tive-data-exposure
d3b5626 to
5ef3d01
Compare
There was a problem hiding this comment.
Actionable comments posted: 6
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@src/controllers/rag_data_exposure_controller.py`:
- Around line 47-53: The controller currently treats any non-"validate" action
as "generate"; add an explicit allowlist check for action values ("validate" or
"generate") and reject others with an error response. In the handler around the
action variable, validate that action is one of {"validate","generate"} and
return a 4xx error (or raise the appropriate HTTPException) for invalid values;
keep existing calls to validate_rag_data_exposure_secret when action ==
"validate" and to evaluate_rag_data_exposure_level when action == "generate"
(use the same user_input/model extraction). Ensure you reference the action
variable, validate_rag_data_exposure_secret, and
evaluate_rag_data_exposure_level when implementing this guard so bad clients
fail fast instead of following the wrong path.
- Around line 42-44: In the _handle_level method move the JSON parsing (await
request.json()) inside the existing try/except so malformed JSON can't bubble
out; specifically, call await request.json() within the try block at the start
of _handle_level, validate/normalize action = str(data.get("action",
"generate")).strip().lower() there, and in the except handle
JSONDecodeError/ValueError by returning a controlled error dict (and/or
appropriate status payload) instead of letting an unhandled 500 occur.
In `@src/framework/registry.py`:
- Around line 99-100: The branch comparing controller_name uses a hyphenated
string ("rag-sensitive-data-exposure") but other dispatches expect underscored
names (e.g., "rag_sensitive_data_exposure"), causing mismatches; update the
dispatch logic in the function that handles controller_name (the branch that
calls validate_rag_data_exposure_secret) to normalize controller_name (e.g.,
controller_name = controller_name.replace('-', '_') or compare against both
forms) before the if checks so the compare will match and
validate_rag_data_exposure_secret(level_number, candidate_secret) is reached for
registrations named rag_sensitive_data_exposure.
In `@src/service/vulnerabilities/rag_data_exposure_lab.py`:
- Around line 195-231: The current flow appends vectors to FAISS before
inserting rows into SQLite which can cause FAISS/DB divergence on
unique-constraint failures; change the order so you build the DB rows and
perform the INSERT (using self._connect() and handling
sqlite3.IntegrityError/unique-constraint races by skipping or deduping duplicate
rows) before mutating FAISS, then compute start_vector_id = int(index.ntotal),
call faiss.normalize_L2(matrix) and index.add(matrix); apply the same
swap-and-tolerate-duplicate-insert logic to the other identical block referenced
around lines 461-471 (use symbols: index.add, faiss.normalize_L2,
start_vector_id, self._connect, rows, documents).
In `@src/static/facade/rag_data_exposure_template.html`:
- Around line 11-21: The textarea (id="ragExposurePrompt") and the secret input
lack programmatic labels; add accessible labels by either inserting <label
for="ragExposurePrompt">Prompt</label> tied to the textarea and a corresponding
<label for="..."> for the secret input, or by adding clear aria-label attributes
to those elements, and if needed use a visually-hidden utility class to keep
visual layout unchanged; ensure the labels reference the exact element ids used
in this template so screen readers can announce the fields.
In `@src/static/facade/rag_data_exposure_template.js`:
- Around line 72-78: The UI currently accepts any integer level >=1 (from the
level parsing code) but endpointForLevel only implements level1..level3, causing
404s for higher levels; clamp the parsed level to the supported range (1..3) or
clamp inside endpointForLevel so that any incoming level is reduced to
Math.min(Math.max(level, 1), 3) before building the URL. Update the level
parsing function or endpointForLevel (referencing endpointForLevel and the
level-parsing code that sets const level = Number(match[1])) to enforce this
clamp so requests always target /level1, /level2, or /level3.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro Plus
Run ID: 14462f52-df10-4037-8223-dc7e7b71afe4
📒 Files selected for processing (13)
locale/messages_us.propertiessrc/app.pysrc/controllers/rag_data_exposure_controller.pysrc/framework/decorators.pysrc/framework/registry.pysrc/service/vulnerabilities/__init__.pysrc/service/vulnerabilities/docs/RAG_DATA_EXPOSURE/LEVEL1/documents.jsonsrc/service/vulnerabilities/docs/RAG_DATA_EXPOSURE/LEVEL2/documents.jsonsrc/service/vulnerabilities/docs/RAG_DATA_EXPOSURE/LEVEL3/documents.jsonsrc/service/vulnerabilities/rag_data_exposure_lab.pysrc/static/facade/rag_data_exposure_template.csssrc/static/facade/rag_data_exposure_template.htmlsrc/static/facade/rag_data_exposure_template.js
| async def _handle_level(self, level: int, request: Request) -> dict: | ||
| data = await request.json() | ||
| action = str(data.get("action", "generate")).strip().lower() |
There was a problem hiding this comment.
Handle malformed JSON inside guarded error flow.
On Line 43, await request.json() runs before the try block, so invalid JSON can escape your handler and return an unhandled 500. Move parsing into try and return a controlled error payload.
Proposed fix
async def _handle_level(self, level: int, request: Request) -> dict:
- data = await request.json()
- action = str(data.get("action", "generate")).strip().lower()
-
try:
+ data = await request.json()
+ action = str(data.get("action", "generate")).strip().lower()
if action == "validate":
candidate_secret = str(data.get("candidate_secret", ""))
return validate_rag_data_exposure_secret(level, candidate_secret)
@@
except ValueError as exc:
return {"error": str(exc)}
+ except Exception:
+ return {"error": "Invalid JSON payload"}
except httpx.HTTPError:
return {"error": "Model or embedding service unavailable"}🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@src/controllers/rag_data_exposure_controller.py` around lines 42 - 44, In the
_handle_level method move the JSON parsing (await request.json()) inside the
existing try/except so malformed JSON can't bubble out; specifically, call await
request.json() within the try block at the start of _handle_level,
validate/normalize action = str(data.get("action", "generate")).strip().lower()
there, and in the except handle JSONDecodeError/ValueError by returning a
controlled error dict (and/or appropriate status payload) instead of letting an
unhandled 500 occur.
| if action == "validate": | ||
| candidate_secret = str(data.get("candidate_secret", "")) | ||
| return validate_rag_data_exposure_secret(level, candidate_secret) | ||
|
|
||
| user_input = str(data.get("user_input", "")) | ||
| model = data.get("model") | ||
| return await evaluate_rag_data_exposure_level(level, user_input, model=model) |
There was a problem hiding this comment.
Reject unsupported action values explicitly.
Right now, any non-validate action is treated as generate. On Line 47, add an allowlist check (generate / validate) so bad clients fail fast instead of silently executing the wrong path.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@src/controllers/rag_data_exposure_controller.py` around lines 47 - 53, The
controller currently treats any non-"validate" action as "generate"; add an
explicit allowlist check for action values ("validate" or "generate") and reject
others with an error response. In the handler around the action variable,
validate that action is one of {"validate","generate"} and return a 4xx error
(or raise the appropriate HTTPException) for invalid values; keep existing calls
to validate_rag_data_exposure_secret when action == "validate" and to
evaluate_rag_data_exposure_level when action == "generate" (use the same
user_input/model extraction). Ensure you reference the action variable,
validate_rag_data_exposure_secret, and evaluate_rag_data_exposure_level when
implementing this guard so bad clients fail fast instead of following the wrong
path.
| if controller_name == "rag-sensitive-data-exposure": | ||
| return validate_rag_data_exposure_secret(level_number, candidate_secret) |
There was a problem hiding this comment.
Normalize controller-name format in secret dispatch.
Line 99 uses a hyphenated controller id, while dispatch in this function otherwise relies on underscored names. If registration.name is rag_sensitive_data_exposure, this branch won’t match and verify-secret returns 404.
Suggested fix
def _verify_secret(controller_name: str, level_number: int, secret_token: str | None, candidate_secret: str) -> dict:
- if controller_name == "prompt_injection":
+ normalized_controller_name = controller_name.replace("-", "_")
+ if normalized_controller_name == "prompt_injection":
return verify_level_secret(level_number, secret_token, candidate_secret)
- if controller_name == "indirect_prompt_injection":
+ if normalized_controller_name == "indirect_prompt_injection":
return verify_indirect_level_secret(level_number, secret_token, candidate_secret)
- if controller_name == "rag-sensitive-data-exposure":
+ if normalized_controller_name == "rag_sensitive_data_exposure":
return validate_rag_data_exposure_secret(level_number, candidate_secret)
raise HTTPException(status_code=404, detail="Unsupported controller")🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@src/framework/registry.py` around lines 99 - 100, The branch comparing
controller_name uses a hyphenated string ("rag-sensitive-data-exposure") but
other dispatches expect underscored names (e.g., "rag_sensitive_data_exposure"),
causing mismatches; update the dispatch logic in the function that handles
controller_name (the branch that calls validate_rag_data_exposure_secret) to
normalize controller_name (e.g., controller_name = controller_name.replace('-',
'_') or compare against both forms) before the if checks so the compare will
match and validate_rag_data_exposure_secret(level_number, candidate_secret) is
reached for registrations named rag_sensitive_data_exposure.
| start_vector_id = int(index.ntotal) | ||
| faiss.normalize_L2(matrix) | ||
| index.add(matrix) | ||
|
|
||
| self._ensure_schema() | ||
| rows = [] | ||
| for offset, document in enumerate(documents): | ||
| metadata = dict(document.get("metadata") or {}) | ||
| rows.append( | ||
| { | ||
| "namespace": namespace, | ||
| "vector_id": start_vector_id + offset, | ||
| "doc_id": document["doc_id"], | ||
| "chunk_id": document["chunk_id"], | ||
| "title": document["title"], | ||
| "source": document["source"], | ||
| "content": document["text"], | ||
| "sensitivity": document["sensitivity"], | ||
| "metadata_json": json.dumps(metadata, ensure_ascii=True, sort_keys=True), | ||
| } | ||
| ) | ||
|
|
||
| with self._connect() as connection: | ||
| connection.executemany( | ||
| """ | ||
| INSERT INTO rag_data_exposure_chunks ( | ||
| namespace, vector_id, doc_id, chunk_id, title, source, | ||
| content, sensitivity, metadata_json | ||
| ) | ||
| VALUES ( | ||
| :namespace, :vector_id, :doc_id, :chunk_id, :title, :source, | ||
| :content, :sensitivity, :metadata_json | ||
| ) | ||
| """, | ||
| rows, | ||
| ) | ||
|
|
There was a problem hiding this comment.
Prevent FAISS/SQLite divergence during concurrent first-time indexing.
Two concurrent calls can both pass ensure_level_indexed pre-check. In add, vectors are appended to FAISS before SQLite insert; if SQLite then hits the unique constraint, FAISS and SQLite drift out of sync.
💡 Suggested fix (insert first, then mutate FAISS; tolerate duplicate insert race)
@@
- start_vector_id = int(index.ntotal)
- faiss.normalize_L2(matrix)
- index.add(matrix)
-
self._ensure_schema()
+ start_vector_id = int(index.ntotal)
rows = []
@@
- with self._connect() as connection:
- connection.executemany(
- """
- INSERT INTO rag_data_exposure_chunks (
- namespace, vector_id, doc_id, chunk_id, title, source,
- content, sensitivity, metadata_json
- )
- VALUES (
- :namespace, :vector_id, :doc_id, :chunk_id, :title, :source,
- :content, :sensitivity, :metadata_json
- )
- """,
- rows,
- )
+ try:
+ with self._connect() as connection:
+ connection.executemany(
+ """
+ INSERT INTO rag_data_exposure_chunks (
+ namespace, vector_id, doc_id, chunk_id, title, source,
+ content, sensitivity, metadata_json
+ )
+ VALUES (
+ :namespace, :vector_id, :doc_id, :chunk_id, :title, :source,
+ :content, :sensitivity, :metadata_json
+ )
+ """,
+ rows,
+ )
+ except sqlite3.IntegrityError:
+ return {
+ "added": 0,
+ "total_documents": self.namespace_count(namespace),
+ "dimension": self._dimensions[namespace],
+ "namespace": namespace,
+ }
+
+ faiss.normalize_L2(matrix)
+ index.add(matrix)Also applies to: 461-471
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@src/service/vulnerabilities/rag_data_exposure_lab.py` around lines 195 - 231,
The current flow appends vectors to FAISS before inserting rows into SQLite
which can cause FAISS/DB divergence on unique-constraint failures; change the
order so you build the DB rows and perform the INSERT (using self._connect() and
handling sqlite3.IntegrityError/unique-constraint races by skipping or deduping
duplicate rows) before mutating FAISS, then compute start_vector_id =
int(index.ntotal), call faiss.normalize_L2(matrix) and index.add(matrix); apply
the same swap-and-tolerate-duplicate-insert logic to the other identical block
referenced around lines 461-471 (use symbols: index.add, faiss.normalize_L2,
start_vector_id, self._connect, rows, documents).
| <h3>Prompt</h3> | ||
| <span id="queryCounter" class="counter">0 / 240</span> | ||
| </div> | ||
|
|
||
| <textarea | ||
| id="ragExposurePrompt" | ||
| class="rag-input" | ||
| maxlength="240" | ||
| rows="4" | ||
| spellcheck="false" | ||
| ></textarea> |
There was a problem hiding this comment.
Add explicit labels for form controls (a11y).
The prompt textarea and secret input rely on surrounding text/placeholder, but neither has a programmatic <label> (or aria-label). Screen-reader users may not get reliable field names.
Proposed fix
- <h3>Prompt</h3>
+ <h3><label for="ragExposurePrompt">Prompt</label></h3>
@@
- <input
+ <label for="ragExposureSecret" class="sr-only">Secret to verify</label>
+ <input
id="ragExposureSecret"Also applies to: 40-48
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@src/static/facade/rag_data_exposure_template.html` around lines 11 - 21, The
textarea (id="ragExposurePrompt") and the secret input lack programmatic labels;
add accessible labels by either inserting <label
for="ragExposurePrompt">Prompt</label> tied to the textarea and a corresponding
<label for="..."> for the secret input, or by adding clear aria-label attributes
to those elements, and if needed use a visually-hidden utility class to keep
visual layout unchanged; ensure the labels reference the exact element ids used
in this template so screen readers can announce the fields.
| const level = Number(match[1]); | ||
| return Number.isInteger(level) && level >= 1 ? level : 1; | ||
| } | ||
|
|
||
| function endpointForLevel(level) { | ||
| return apiPrefix + "/api/v1/vulnerabilities/rag-sensitive-data-exposure/level" + level; | ||
| } |
There was a problem hiding this comment.
Clamp level to implemented endpoints to avoid guaranteed 404s.
On Line 73, any LEVEL_n (n >= 1) is accepted, but this UI only implements level1..level3. If LEVEL_4/LEVEL_5 is selected upstream, requests go to missing endpoints.
Proposed fix
+ const SUPPORTED_LEVELS = [1, 2, 3];
@@
function levelFromGlobalState() {
@@
const level = Number(match[1]);
- return Number.isInteger(level) && level >= 1 ? level : 1;
+ return Number.isInteger(level) && SUPPORTED_LEVELS.includes(level) ? level : 1;
}🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@src/static/facade/rag_data_exposure_template.js` around lines 72 - 78, The UI
currently accepts any integer level >=1 (from the level parsing code) but
endpointForLevel only implements level1..level3, causing 404s for higher levels;
clamp the parsed level to the supported range (1..3) or clamp inside
endpointForLevel so that any incoming level is reduced to
Math.min(Math.max(level, 1), 3) before building the URL. Update the level
parsing function or endpointForLevel (referencing endpointForLevel and the
level-parsing code that sets const level = Number(match[1])) to enforce this
clamp so requests always target /level1, /level2, or /level3.
Overview
This PR implements a RAG Sensitive Data Exposure lab.
The implementation uses a
RagDataExposureVectorStore, which combines FAISS with SQLite to store vectors by namespace (one namespace per level) and persist them for later retrieval.In the interface, it adds a facade that displays the retrieved segments and validates the leaked secret against the model's responses.
How It Works
vector_id) for each namespace.Levels
Each level builds on the same pipeline, changing only the defense that sits in front of retrieval.
Level 1 — Direct Retrieval
A direct query embeds close to the sensitive chunk, FAISS retrieves it, and the model leaks the secret straight back.
Level 2 — Semantic Paraphrase Bypass
The denylist is purely lexical, but retrieval is semantic — so a paraphrase (e.g. "internal recovery value", "privileged access") still lands on the sensitive chunk and leaks it.
Level 3 — Misclassified Low-Sensitivity Document
A document tagged as
lowstill contains one chunk with a secret, so document-level tags are too coarse and the sensitive chunk passes the filter.Closes #8
Summary by CodeRabbit
Release Notes
New Features
Documentation