Skip to content

Conversation

@christopherholland-workday
Copy link
Contributor

Create basic guardrails for GraphCypherQAChain, this change does the following:

  1. Validates the users prompt against basic prompt injection patterns
  2. Sanitizes the users prompt against basic prompt injection patterns
  3. Validates the cypher query against known modification patterns
  4. Adds a maximum return size to prevent data exfiltration

@gemini-code-assist
Copy link
Contributor

Summary of Changes

Hello @christopherholland-workday, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request significantly enhances the security and robustness of the GraphCypherQAChain by introducing multiple layers of defense. It focuses on preventing prompt injection, sanitizing user input, blocking malicious Cypher write operations, and limiting data exfiltration, thereby making the chain more resilient to various attack vectors and ensuring safer interaction with graph databases.

Highlights

  • Prompt Injection Validation: Implemented robust validation against basic and sophisticated prompt injection patterns to protect the LLM from manipulation.
  • User Input Sanitization: Added a sanitization step for user prompts, normalizing Unicode, removing control characters, comments, semicolons, and collapsing whitespace to harden input.
  • Cypher Query Validation: Introduced validation for generated Cypher queries to explicitly block any write operations (e.g., CREATE, DELETE, SET) for enhanced security, ensuring read-only access.
  • Data Exfiltration Prevention: Enforced a maximum return size (100 results) for Cypher queries to prevent large-scale data exfiltration.
  • Input Length Limiting: Added a hardcoded maximum input length (2000 characters) to prevent abuse and resource exhaustion.
Changelog
  • .eslintrc.js
    • Disabled the no-control-regex ESLint rule to allow matching control regex in user input.
  • packages/components/jest.config.js
    • Updated Jest configuration to include the test directory in the root paths for test discovery.
  • packages/components/nodes/chains/GraphCypherQAChain/GraphCypherQAChain.ts
    • Introduced validateCypherQuery to prevent write operations in LLM-generated Cypher queries.
    • Implemented sanitizeUserInput for normalizing and hardening user input against common abuse patterns.
    • Added detectPromptInjection with multi-layered heuristics to identify prompt manipulation and Cypher injection attempts.
    • Integrated Cypher query validation and result limiting (max 100) into the graph.query execution within constructGraphCypherQAChain.
    • Applied input length validation (max 2000 characters), prompt injection detection, and user input sanitization in the run method.
    • Exported security utility functions (sanitizeUserInput, detectPromptInjection, validateCypherQuery) for external testing.
  • packages/components/test/nodes/chains/GraphCypherQAChain/GraphCypherQAChain.test.ts
    • Added comprehensive unit tests for the new security functions: sanitizeUserInput, detectPromptInjection, and validateCypherQuery, covering various attack patterns and legitimate use cases.
Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature Command Description
Code Review /gemini review Performs a code review for the current pull request in its current state.
Pull Request Summary /gemini summary Provides a summary of the current pull request in its current state.
Comment @gemini-code-assist Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help /gemini help Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Footnotes

  1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution.

Copy link
Contributor

@gemini-code-assist gemini-code-assist bot left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Code Review

This pull request introduces important security guardrails for the GraphCypherQAChain, including input sanitization, prompt injection detection, validation of generated Cypher queries to prevent write operations, and limiting query results to prevent data exfiltration. The changes are well-implemented and include a comprehensive suite of tests. My feedback focuses on improving the robustness of the prompt injection heuristics to avoid false positives and making hardcoded limits configurable for better flexibility.

Comment on lines +157 to +168
// Check for excessive special characters (potential obfuscation)
const specialCharCount = (input.match(/[{}()[\];|&$`\\]/g) || []).length
if (specialCharCount > 5) {
return true
}

// Check for suspicious Cypher keywords in close proximity
const cypherKeywords = ['MATCH', 'CREATE', 'MERGE', 'DELETE', 'DETACH', 'SET', 'REMOVE', 'RETURN', 'WHERE', 'WITH']
const foundKeywords = cypherKeywords.filter((keyword) => lowerInput.includes(keyword.toLowerCase()))
if (foundKeywords.length >= 3) {
return true
}
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

The heuristic checks for special characters and keyword clustering are a good idea for defense-in-depth, but the current thresholds might be too sensitive and lead to false positives for legitimate user questions, especially if they are asking about code or Cypher queries.

  • Special Character Count: The limit of 5 special characters (/[{}()[\];|&$`]/g`) is quite low. A user asking a question about a JSON object or a code snippet could easily exceed this.
  • Keyword Clustering: The check for 3 or more Cypher keywords can be triggered by natural language questions about how to use Cypher.

I recommend increasing these thresholds (e.g., 10 for special characters, 4 for keywords) or making the logic more sophisticated to reduce false positives. For example, you could check for keywords in combination with query-like structures rather than just their presence.

const cypherModel = nodeData.inputs?.cypherModel
const qaModel = nodeData.inputs?.qaModel
const graph = nodeData.inputs?.graph
const maxResults = 100 // Hardcoded limit to prevent data exfiltration
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

The maxResults limit is a great security measure to prevent data exfiltration. However, hardcoding it to 100 might be too restrictive for some use cases. I suggest making this a configurable parameter on the node (by adding it to the inputs array in the constructor), with 100 as the default value. This would allow administrators to adjust it based on their specific needs while still providing a safe default.

Suggested change
const maxResults = 100 // Hardcoded limit to prevent data exfiltration
const maxResults = (nodeData.inputs?.maxResults as number) ?? 100 // Hardcoded limit to prevent data exfiltration

const chain = nodeData.instance as GraphCypherQAChain
const moderations = nodeData.inputs?.inputModeration as Moderation[]
const returnDirect = nodeData.inputs?.returnDirect as boolean
const maxInputLength = 2000 // Hardcoded limit to prevent abuse
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

Similar to maxResults, the maxInputLength is hardcoded. While this is a good guardrail against abuse, making it configurable would provide more flexibility. I recommend exposing this as a node parameter (by adding it to the inputs array in the constructor) with a default of 2000.

Suggested change
const maxInputLength = 2000 // Hardcoded limit to prevent abuse
const maxInputLength = (nodeData.inputs?.maxInputLength as number) ?? 2000 // Hardcoded limit to prevent abuse

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants