package utils import ( "fmt" "regexp" "strings" pg_query "github.com/pganalyze/pg_query_go/v6" ) // This file provides comprehensive SQL validation and security features /* Example Usage: 1. Basic SQL parsing: result := ParseSQL("SELECT * FROM users WHERE age > 18") fmt.Printf("Tables: %v\n", result.TableNames) fmt.Printf("WHERE fields: %v\n", result.WhereFields) 2. Simple validation with table whitelist: parseResult, validation := ValidateSQL( "SELECT * FROM users WHERE age > 18", WithAllowedTables("users", "orders"), ) if !validation.Valid { for _, err := range validation.Errors { fmt.Printf("Error: %s - %s\n", err.Type, err.Message) } } 3. Check for SQL injection risks: parseResult, validation := ValidateSQL( "SELECT * FROM users WHERE id = 1 OR 1=1", WithInjectionRiskCheck(), ) if !validation.Valid { fmt.Println("SQL injection risk detected!") } 4. Comprehensive security validation: parseResult, validation := ValidateSQL( "SELECT * FROM users WHERE age > 18", WithInputValidation(6, 4096), WithSelectOnly(), WithSingleStatement(), WithAllowedTables("users", "orders"), WithDefaultSafeFunctions(), WithNoSubqueries(), WithNoCTEs(), WithNoSystemColumns(), ) 5. Use security defaults (recommended for production): parseResult, validation := ValidateSQL( "SELECT * FROM knowledge_bases WHERE name LIKE '%test%'", WithSecurityDefaults(tenantID), ) 6. Validate and secure SQL with tenant isolation: securedSQL, validation, err := ValidateAndSecureSQL( "SELECT * FROM knowledge_bases", WithSecurityDefaults(tenantID), ) // securedSQL will have tenant_id automatically injected: // "SELECT * FROM knowledge_bases WHERE knowledge_bases.tenant_id = 123" 7. Custom validation options: parseResult, validation := ValidateSQL( "SELECT COUNT(*), AVG(score) FROM sessions", WithAllowedTables("sessions", "messages"), WithAllowedFunctions("count", "avg", "sum"), WithTenantIsolation(tenantID, "sessions"), ) */ // SQLParseResult represents the parsed components of a SELECT SQL statement type SQLParseResult struct { IsSelect bool `json:"is_select"` // Whether the SQL is a SELECT statement TableNames []string `json:"table_names"` // List of table names in FROM clause SelectFields []string `json:"select_fields"` // List of fields in SELECT clause WhereFields []string `json:"where_fields"` // List of fields in WHERE clause WhereClause string `json:"where_clause"` // Complete WHERE clause text OriginalSQL string `json:"original_sql"` // Original SQL statement ParseError string `json:"parse_error,omitempty"` // Error message if parsing failed } // SQLValidationError represents a validation error type SQLValidationError struct { Type string `json:"type"` // Error type: "table_not_allowed", "sql_injection_risk", etc. Message string `json:"message"` // Error message Details string `json:"details"` // Additional details } // SQLValidationResult represents the result of SQL validation type SQLValidationResult struct { Valid bool `json:"valid"` // Whether the SQL passed validation Errors []SQLValidationError `json:"errors"` // List of validation errors } // SQLValidationOption is a function that configures SQL validation type SQLValidationOption func(*sqlValidator) // sqlValidator holds validation configuration type sqlValidator struct { // Basic validation checkInputValidation bool minLength int maxLength int // Statement type validation checkSelectOnly bool checkSingleStatement bool // Table validation allowedTables map[string]bool checkTableNames bool // Function validation allowedFunctions map[string]bool checkFunctionNames bool // Security checks checkInjectionRisk bool checkSubqueries bool checkCTEs bool checkSystemColumns bool checkSchemaAccess bool checkDangerousFuncs bool // Tenant isolation enableTenantInjection bool tenantID uint64 tablesWithTenantID map[string]bool // Soft delete filtering enableSoftDeleteInjection bool tablesWithDeletedAt map[string]bool // Hidden knowledge base filtering (is_temporary = false) enableHiddenKBFilter bool // Enabled chunk filtering (is_enabled = true) enableChunkEnabledFilter bool // Search scope filtering (restrict to specific KBs and knowledges) enableSearchScopeFilter bool searchScopeKBIDs []string searchScopeKnowledgeIDs []string searchScopes []SearchScope } // SearchScope describes one allowed knowledge scope for SQL query injection. // Empty KnowledgeIDs and TagIDs means the whole KB is in scope. type SearchScope struct { KnowledgeBaseID string KnowledgeIDs []string TagIDs []string } // ParseSQL parses a SQL statement using pg_query_go and extracts table names, select fields, and where fields // This uses the PostgreSQL parser for accurate SQL parsing func ParseSQL(sql string) *SQLParseResult { result := &SQLParseResult{ OriginalSQL: sql, TableNames: make([]string, 0), SelectFields: make([]string, 0), WhereFields: make([]string, 0), } // Parse the SQL using pg_query_go parseResult, err := pg_query.Parse(sql) if err != nil { result.IsSelect = false result.ParseError = fmt.Sprintf("Failed to parse SQL: %v", err) return result } // Check if it's a SELECT statement if len(parseResult.Stmts) == 0 { result.IsSelect = false result.ParseError = "No statements found in SQL" return result } // Get the first statement stmt := parseResult.Stmts[0] if stmt.Stmt == nil { result.IsSelect = false result.ParseError = "Invalid statement" return result } // Check if it's a SELECT statement selectStmt := stmt.Stmt.GetSelectStmt() if selectStmt == nil { result.IsSelect = false result.ParseError = "Not a SELECT statement" return result } result.IsSelect = true // Extract SELECT fields result.SelectFields = extractSelectFieldsFromPgQuery(selectStmt) // Extract table names from FROM clause result.TableNames = extractTableNamesFromPgQuery(selectStmt) // Extract WHERE clause fields and text whereFields, whereClause := extractWhereFromPgQuery(selectStmt, sql) result.WhereFields = whereFields result.WhereClause = whereClause return result } // extractSelectFieldsFromPgQuery extracts field names from SELECT clause using pg_query parse tree func extractSelectFieldsFromPgQuery(selectStmt *pg_query.SelectStmt) []string { fields := make([]string, 0) fieldMap := make(map[string]bool) // Avoid duplicates if selectStmt.TargetList == nil { return fields } for _, target := range selectStmt.TargetList { resTarget := target.GetResTarget() if resTarget == nil { continue } // Extract column names from the target colNames := extractColumnNamesFromNode(resTarget.Val) for _, colName := range colNames { if colName != "" && !fieldMap[colName] { fieldMap[colName] = true fields = append(fields, colName) } } } return fields } // extractTableNamesFromPgQuery extracts table names from FROM clause using pg_query parse tree func extractTableNamesFromPgQuery(selectStmt *pg_query.SelectStmt) []string { tables := make([]string, 0) tableMap := make(map[string]bool) // Avoid duplicates if selectStmt.FromClause == nil { return tables } for _, fromItem := range selectStmt.FromClause { tableNames := extractTableNamesFromNode(fromItem) for _, tableName := range tableNames { if tableName != "" && !tableMap[tableName] { tableMap[tableName] = true tables = append(tables, tableName) } } } return tables } // extractWhereFromPgQuery extracts WHERE clause fields and text using pg_query parse tree func extractWhereFromPgQuery(selectStmt *pg_query.SelectStmt, originalSQL string) ([]string, string) { fields := make([]string, 0) fieldMap := make(map[string]bool) // Avoid duplicates whereClause := "" if selectStmt.WhereClause == nil { return fields, whereClause } // Extract WHERE clause text from original SQL whereClause = extractWhereClauseText(originalSQL) // Extract column names from WHERE clause colNames := extractColumnNamesFromNode(selectStmt.WhereClause) for _, colName := range colNames { if colName != "" && !fieldMap[colName] { fieldMap[colName] = true fields = append(fields, colName) } } return fields, whereClause } // extractColumnNamesFromNode recursively extracts column names from a parse tree node func extractColumnNamesFromNode(node *pg_query.Node) []string { if node == nil { return nil } colNames := make([]string, 0) // Handle ColumnRef (column reference) if colRef := node.GetColumnRef(); colRef != nil { if len(colRef.Fields) > 0 { // Extract only the actual column name (the last part of table.column or schema.table.column) lastField := colRef.Fields[len(colRef.Fields)-1] if strNode := lastField.GetString_(); strNode != nil { if strNode.Sval != "*" { // Skip wildcard colNames = append(colNames, strNode.Sval) } } } return colNames } // Handle A_Expr (expression with operators) if aExpr := node.GetAExpr(); aExpr != nil { colNames = append(colNames, extractColumnNamesFromNode(aExpr.Lexpr)...) colNames = append(colNames, extractColumnNamesFromNode(aExpr.Rexpr)...) return colNames } // Handle BoolExpr (AND, OR, NOT) if boolExpr := node.GetBoolExpr(); boolExpr != nil { if boolExpr.Args != nil { for _, arg := range boolExpr.Args { colNames = append(colNames, extractColumnNamesFromNode(arg)...) } } return colNames } // Handle FuncCall (function calls) if funcCall := node.GetFuncCall(); funcCall != nil { if funcCall.Args != nil { for _, arg := range funcCall.Args { colNames = append(colNames, extractColumnNamesFromNode(arg)...) } } return colNames } // Handle ResTarget (result target in SELECT) if resTarget := node.GetResTarget(); resTarget != nil { colNames = append(colNames, extractColumnNamesFromNode(resTarget.Val)...) return colNames } // Handle SubLink (subquery) if subLink := node.GetSubLink(); subLink != nil { colNames = append(colNames, extractColumnNamesFromNode(subLink.Testexpr)...) return colNames } // Handle NullTest (IS NULL, IS NOT NULL) if nullTest := node.GetNullTest(); nullTest != nil { colNames = append(colNames, extractColumnNamesFromNode(nullTest.Arg)...) return colNames } // Handle CaseExpr (CASE WHEN) if caseExpr := node.GetCaseExpr(); caseExpr != nil { colNames = append(colNames, extractColumnNamesFromNode(caseExpr.Arg)...) if caseExpr.Args != nil { for _, arg := range caseExpr.Args { colNames = append(colNames, extractColumnNamesFromNode(arg)...) } } colNames = append(colNames, extractColumnNamesFromNode(caseExpr.Defresult)...) return colNames } // Handle CaseWhen (WHEN clause in CASE) if caseWhen := node.GetCaseWhen(); caseWhen != nil { colNames = append(colNames, extractColumnNamesFromNode(caseWhen.Expr)...) colNames = append(colNames, extractColumnNamesFromNode(caseWhen.Result)...) return colNames } return colNames } // extractTableNamesFromNode recursively extracts table names from a parse tree node func extractTableNamesFromNode(node *pg_query.Node) []string { if node == nil { return nil } tableNames := make([]string, 0) // Handle RangeVar (table reference) if rangeVar := node.GetRangeVar(); rangeVar != nil { if rangeVar.Relname != "" { tableNames = append(tableNames, rangeVar.Relname) } return tableNames } // Handle JoinExpr (JOIN) if joinExpr := node.GetJoinExpr(); joinExpr != nil { tableNames = append(tableNames, extractTableNamesFromNode(joinExpr.Larg)...) tableNames = append(tableNames, extractTableNamesFromNode(joinExpr.Rarg)...) return tableNames } // Handle RangeSubselect (subquery in FROM) if rangeSubselect := node.GetRangeSubselect(); rangeSubselect != nil { // We could recursively parse the subquery here if needed return tableNames } return tableNames } // extractWhereClauseText extracts the WHERE clause text from the original SQL func extractWhereClauseText(sql string) string { lowerSQL := strings.ToLower(sql) wherePos := strings.Index(lowerSQL, "where") if wherePos == -1 { return "" } // Find the end of WHERE clause whereClauseEnd := len(sql) for _, keyword := range []string{"group by", "order by", "limit", "having", "union", "intersect", "except"} { if pos := strings.Index(lowerSQL[wherePos:], keyword); pos != -1 { actualPos := wherePos + pos if actualPos < whereClauseEnd { whereClauseEnd = actualPos } } } // Extract WHERE clause (skip "WHERE" keyword) whereClause := strings.TrimSpace(sql[wherePos+5 : whereClauseEnd]) return whereClause } // WithAllowedTables creates a validation option that checks if table names are in the allowed list func WithAllowedTables(tables ...string) SQLValidationOption { return func(v *sqlValidator) { v.checkTableNames = true v.allowedTables = make(map[string]bool) for _, table := range tables { v.allowedTables[strings.ToLower(table)] = true } } } // WithInjectionRiskCheck creates a validation option that checks for SQL injection risks func WithInjectionRiskCheck() SQLValidationOption { return func(v *sqlValidator) { v.checkInjectionRisk = true } } // WithInputValidation enables basic input validation (length, null bytes, etc.) func WithInputValidation(minLen, maxLen int) SQLValidationOption { return func(v *sqlValidator) { v.checkInputValidation = true v.minLength = minLen v.maxLength = maxLen } } // WithSelectOnly ensures only SELECT statements are allowed func WithSelectOnly() SQLValidationOption { return func(v *sqlValidator) { v.checkSelectOnly = true } } // WithSingleStatement ensures only single statement is allowed (no multiple statements) func WithSingleStatement() SQLValidationOption { return func(v *sqlValidator) { v.checkSingleStatement = true } } // WithAllowedFunctions creates a validation option that checks if functions are in the allowed list func WithAllowedFunctions(functions ...string) SQLValidationOption { return func(v *sqlValidator) { v.checkFunctionNames = true v.allowedFunctions = make(map[string]bool) for _, fn := range functions { v.allowedFunctions[strings.ToLower(fn)] = true } } } // WithDefaultSafeFunctions enables a default set of safe SQL functions func WithDefaultSafeFunctions() SQLValidationOption { return func(v *sqlValidator) { v.checkFunctionNames = true v.allowedFunctions = map[string]bool{ // Aggregate functions "count": true, "sum": true, "avg": true, "min": true, "max": true, "array_agg": true, "string_agg": true, "bool_and": true, "bool_or": true, "json_agg": true, "jsonb_agg": true, "json_object_agg": true, "jsonb_object_agg": true, // Safe scalar functions "coalesce": true, "nullif": true, "greatest": true, "least": true, "abs": true, "ceil": true, "floor": true, "round": true, "length": true, "lower": true, "upper": true, "trim": true, "ltrim": true, "rtrim": true, "substring": true, "concat": true, "concat_ws": true, "replace": true, "left": true, "right": true, "now": true, "current_date": true, "current_timestamp": true, "date_trunc": true, "extract": true, "to_char": true, "to_date": true, "to_timestamp": true, "date_part": true, "age": true, } } } // WithNoSubqueries blocks all subqueries func WithNoSubqueries() SQLValidationOption { return func(v *sqlValidator) { v.checkSubqueries = true } } // WithNoCTEs blocks Common Table Expressions (WITH clause) func WithNoCTEs() SQLValidationOption { return func(v *sqlValidator) { v.checkCTEs = true } } // WithNoSystemColumns blocks access to PostgreSQL system columns func WithNoSystemColumns() SQLValidationOption { return func(v *sqlValidator) { v.checkSystemColumns = true } } // WithNoSchemaAccess blocks schema-qualified access (except public schema) func WithNoSchemaAccess() SQLValidationOption { return func(v *sqlValidator) { v.checkSchemaAccess = true } } // WithNoDangerousFunctions blocks dangerous PostgreSQL functions func WithNoDangerousFunctions() SQLValidationOption { return func(v *sqlValidator) { v.checkDangerousFuncs = true } } // WithTenantIsolation enables automatic tenant_id injection for multi-tenant security func WithTenantIsolation(tenantID uint64, tables ...string) SQLValidationOption { return func(v *sqlValidator) { v.enableTenantInjection = true v.tenantID = tenantID v.tablesWithTenantID = make(map[string]bool) if len(tables) == 0 { // Default tables with tenant_id // SECURITY: All tables with tenant_id column must be listed here // to ensure proper tenant isolation and prevent cross-tenant data access v.tablesWithTenantID = map[string]bool{ "knowledge_bases": true, "knowledges": true, "chunks": true, } } else { for _, table := range tables { v.tablesWithTenantID[strings.ToLower(table)] = true } } } } // WithSoftDeleteFilter enables automatic deleted_at IS NULL injection. func WithSoftDeleteFilter(tables ...string) SQLValidationOption { return func(v *sqlValidator) { v.enableSoftDeleteInjection = true v.tablesWithDeletedAt = make(map[string]bool) if len(tables) == 0 { // Default tables with soft-delete support. v.tablesWithDeletedAt = map[string]bool{ "knowledge_bases": true, "knowledges": true, "chunks": true, } } else { for _, table := range tables { v.tablesWithDeletedAt[strings.ToLower(table)] = true } } } } // WithHiddenKBFilter excludes internal/temporary knowledge bases (is_temporary = true) // from query results. These are system-managed KBs like __chat_history__ that should // not be visible to end users. func WithHiddenKBFilter() SQLValidationOption { return func(v *sqlValidator) { v.enableHiddenKBFilter = true } } // WithChunkEnabledFilter excludes disabled chunks from model-visible SQL // queries. It is intentionally opt-in because administrative queries may need // to inspect disabled rows. func WithChunkEnabledFilter() SQLValidationOption { return func(v *sqlValidator) { v.enableChunkEnabledFilter = true } } // WithSearchScopeFilter restricts queries to the specified knowledge bases and // (optionally) specific knowledge documents. For the knowledge_bases table it // filters by id; for knowledges it filters by knowledge_base_id (and id when // knowledgeIDs is non-empty); for chunks it filters by knowledge_base_id (and // knowledge_id when knowledgeIDs is non-empty). func WithSearchScopeFilter(kbIDs []string, knowledgeIDs []string) SQLValidationOption { return func(v *sqlValidator) { if len(kbIDs) > 0 { v.enableSearchScopeFilter = true v.searchScopeKBIDs = kbIDs v.searchScopeKnowledgeIDs = knowledgeIDs } } } // WithSearchScopes restricts queries using structured OR scopes. Each scope can // represent a full KB, specific documents, or a tag-constrained KB. func WithSearchScopes(scopes []SearchScope) SQLValidationOption { return func(v *sqlValidator) { if len(scopes) == 0 { return } v.enableSearchScopeFilter = true v.searchScopes = append([]SearchScope(nil), scopes...) } } // WithSecurityDefaults applies a comprehensive set of security validations func WithSecurityDefaults(tenantID uint64) SQLValidationOption { return func(v *sqlValidator) { // Apply all security checks WithInputValidation(6, 4096)(v) WithSelectOnly()(v) WithSingleStatement()(v) WithNoSubqueries()(v) WithNoCTEs()(v) WithNoSystemColumns()(v) WithNoSchemaAccess()(v) WithNoDangerousFunctions()(v) WithDefaultSafeFunctions()(v) WithTenantIsolation(tenantID)(v) // Default allowed tables // SECURITY: Only tables with tenant_id column should be listed here // Tables without tenant_id (messages, embeddings) are excluded to prevent // cross-tenant data access vulnerabilities (CVE: Broken Access Control) WithAllowedTables( "knowledge_bases", "knowledges", "chunks", )(v) } } // ValidateSQL validates a SQL statement with the given options func ValidateSQL(sql string, opts ...SQLValidationOption) (*SQLParseResult, *SQLValidationResult) { // Initialize validator with defaults validator := &sqlValidator{ allowedTables: make(map[string]bool), allowedFunctions: make(map[string]bool), tablesWithTenantID: make(map[string]bool), tablesWithDeletedAt: make(map[string]bool), minLength: 6, maxLength: 4096, } // Apply options for _, opt := range opts { opt(validator) } // Initialize validation result validationResult := &SQLValidationResult{ Valid: true, Errors: make([]SQLValidationError, 0), } // Phase 1: Basic input validation if validator.checkInputValidation { if err := validator.validateInput(sql); err != nil { validationResult.Valid = false validationResult.Errors = append(validationResult.Errors, SQLValidationError{ Type: "input_validation_error", Message: "Input validation failed", Details: err.Error(), }) return nil, validationResult } } // Phase 2: Parse SQL using PostgreSQL's official parser parseResult, err := pg_query.Parse(sql) if err != nil { validationResult.Valid = false validationResult.Errors = append(validationResult.Errors, SQLValidationError{ Type: "parse_error", Message: "Failed to parse SQL", Details: fmt.Sprintf("SQL parse error: %v", err), }) return &SQLParseResult{ OriginalSQL: sql, ParseError: err.Error(), }, validationResult } // Phase 3: Validate statement count if len(parseResult.Stmts) == 0 { validationResult.Valid = false validationResult.Errors = append(validationResult.Errors, SQLValidationError{ Type: "empty_query", Message: "Empty query", Details: "No statements found in SQL", }) return &SQLParseResult{ OriginalSQL: sql, ParseError: "empty query", }, validationResult } if validator.checkSingleStatement && len(parseResult.Stmts) > 1 { validationResult.Valid = false validationResult.Errors = append(validationResult.Errors, SQLValidationError{ Type: "multiple_statements", Message: "Multiple statements are not allowed", Details: fmt.Sprintf("Found %d statements, only 1 is allowed", len(parseResult.Stmts)), }) return &SQLParseResult{ OriginalSQL: sql, ParseError: "multiple statements", }, validationResult } stmt := parseResult.Stmts[0].Stmt // Phase 4: Ensure it's a SELECT statement selectStmt := stmt.GetSelectStmt() if validator.checkSelectOnly && selectStmt == nil { validationResult.Valid = false validationResult.Errors = append(validationResult.Errors, SQLValidationError{ Type: "not_select_statement", Message: "Only SELECT queries are allowed", Details: "Statement is not a SELECT query", }) return &SQLParseResult{ OriginalSQL: sql, IsSelect: false, ParseError: "not a SELECT statement", }, validationResult } // Build parse result result := &SQLParseResult{ OriginalSQL: sql, IsSelect: selectStmt != nil, TableNames: make([]string, 0), SelectFields: make([]string, 0), WhereFields: make([]string, 0), } if selectStmt != nil { // Extract SELECT fields result.SelectFields = extractSelectFieldsFromPgQuery(selectStmt) // Extract table names from FROM clause result.TableNames = extractTableNamesFromPgQuery(selectStmt) // Extract WHERE clause fields and text whereFields, whereClause := extractWhereFromPgQuery(selectStmt, sql) result.WhereFields = whereFields result.WhereClause = whereClause // Phase 5: Validate the SELECT statement with deep inspection if err := validator.validateSelectStmt(selectStmt, validationResult); err != nil { validationResult.Valid = false validationResult.Errors = append(validationResult.Errors, SQLValidationError{ Type: "statement_validation_error", Message: "Statement validation failed", Details: err.Error(), }) } // Phase 6: Validate table names if validator.checkTableNames { for _, table := range result.TableNames { if !validator.allowedTables[strings.ToLower(table)] { validationResult.Valid = false validationResult.Errors = append(validationResult.Errors, SQLValidationError{ Type: "table_not_allowed", Message: fmt.Sprintf("Table '%s' is not in the allowed list", table), Details: fmt.Sprintf("Allowed tables: %v", getMapKeys(validator.allowedTables)), }) } } } // Phase 7: Check for SQL injection risks (legacy check) if validator.checkInjectionRisk { injectionErrors := checkSQLInjectionRisks(result.WhereClause) if len(injectionErrors) > 0 { validationResult.Valid = false validationResult.Errors = append(validationResult.Errors, injectionErrors...) } } } return result, validationResult } // ValidateAndSecureSQL validates SQL and returns a secured version with tenant isolation // This is a convenience function that combines validation and SQL rewriting func ValidateAndSecureSQL(sql string, opts ...SQLValidationOption) (string, *SQLValidationResult, error) { // Parse and validate _, validationResult := ValidateSQL(sql, opts...) // If validation failed, return error if !validationResult.Valid { errMsg := "SQL validation failed" if len(validationResult.Errors) > 0 { errMsg = validationResult.Errors[0].Message } return "", validationResult, fmt.Errorf("%s", errMsg) } // Find validator config to check if tenant injection is enabled validator := &sqlValidator{ tablesWithTenantID: make(map[string]bool), tablesWithDeletedAt: make(map[string]bool), } for _, opt := range opts { opt(validator) } // If no SQL rewriting is enabled, return original SQL if !validator.enableTenantInjection && !validator.enableSoftDeleteInjection && !validator.enableHiddenKBFilter && !validator.enableChunkEnabledFilter && !validator.enableSearchScopeFilter { return sql, validationResult, nil } // Parse again to get normalized SQL result, err := pg_query.Parse(sql) if err != nil { return "", validationResult, fmt.Errorf("failed to parse SQL: %v", err) } // Normalize SQL normalizedSQL, err := pg_query.Deparse(result) if err != nil { return "", validationResult, fmt.Errorf("failed to normalize SQL: %v", err) } // Build table→alias map from parse tree (respects SQL aliases like "kb", "k") tablesInQuery := extractTableAliasMap(result) // Inject tenant conditions securedSQL := validator.injectTenantConditions(normalizedSQL, tablesInQuery) // Inject deleted_at IS NULL conditions securedSQL = validator.injectSoftDeleteConditions(securedSQL, tablesInQuery) // Inject hidden KB filter (exclude is_temporary = true knowledge bases) securedSQL = validator.injectHiddenKBFilter(securedSQL, tablesInQuery) // Exclude disabled chunks from model-visible query results. securedSQL = validator.injectChunkEnabledFilter(securedSQL, tablesInQuery) // Inject search scope filter (restrict to allowed KBs and knowledges) securedSQL = validator.injectSearchScopeConditions(securedSQL, tablesInQuery) return securedSQL, validationResult, nil } // extractTableAliasMap walks the parse tree to build a table_name→alias map. // When a table has an alias (e.g., "knowledge_bases kb"), the map entry is // {"knowledge_bases": "kb"}. Without an alias, both key and value are the table name. func extractTableAliasMap(parseResult *pg_query.ParseResult) map[string]string { m := make(map[string]string) if len(parseResult.Stmts) == 0 || parseResult.Stmts[0].Stmt == nil { return m } selectStmt := parseResult.Stmts[0].Stmt.GetSelectStmt() if selectStmt == nil { return m } for _, fromItem := range selectStmt.FromClause { collectTableAliases(fromItem, m) } return m } // collectTableAliases recursively collects table→alias mappings from FROM clause nodes. func collectTableAliases(node *pg_query.Node, m map[string]string) { if node == nil { return } if rv := node.GetRangeVar(); rv != nil { tableName := strings.ToLower(rv.Relname) alias := tableName if rv.Alias != nil && rv.Alias.Aliasname != "" { alias = strings.ToLower(rv.Alias.Aliasname) } m[tableName] = alias return } if je := node.GetJoinExpr(); je != nil { collectTableAliases(je.Larg, m) collectTableAliases(je.Rarg, m) return } } // InjectAndConditions injects filter conditions into a SQL statement using AND semantics. // If WHERE exists, the original WHERE predicates will be wrapped in parentheses. // Compiled once: the WHERE keyword and the set of clauses that may trail a // WHERE expression. reSQLTailClause is shared by InjectAndConditions for both // "where does the WHERE expression end" and "where to insert a new WHERE". var ( reSQLWhereKeyword = regexp.MustCompile(`(?i)\bWHERE\b`) reSQLTailClause = regexp.MustCompile(`(?i)\b(GROUP BY|ORDER BY|LIMIT|OFFSET|HAVING|FETCH)\b`) ) func InjectAndConditions(sql, filter string) string { filter = strings.TrimSpace(filter) if filter == "" { return sql } // Check if WHERE clause exists if loc := reSQLWhereKeyword.FindStringIndex(sql); loc != nil { // Add filter and wrap existing conditions in parentheses to prevent OR precedence issues. // The wrapping must only apply to the original WHERE expression, not trailing clauses like // ORDER BY / GROUP BY / LIMIT, otherwise it can generate invalid SQL. whereExprStart := loc[1] tailLoc := reSQLTailClause.FindStringIndex(sql[whereExprStart:]) if tailLoc == nil { originalWhereExpr := strings.TrimSpace(sql[whereExprStart:]) return fmt.Sprintf("%sWHERE %s AND (%s)", sql[:loc[0]], filter, originalWhereExpr) } whereExprEnd := whereExprStart + tailLoc[0] originalWhereExpr := strings.TrimSpace(sql[whereExprStart:whereExprEnd]) tailClause := strings.TrimLeft(sql[whereExprEnd:], " \t\r\n") return fmt.Sprintf("%sWHERE %s AND (%s) %s", sql[:loc[0]], filter, originalWhereExpr, tailClause) } // Add new WHERE clause before ORDER BY, GROUP BY, LIMIT, etc. if loc := reSQLTailClause.FindStringIndex(sql); loc != nil { prefix := strings.TrimRight(sql[:loc[0]], " \t\r\n") suffix := strings.TrimLeft(sql[loc[0]:], " \t\r\n") return fmt.Sprintf("%s WHERE %s %s", prefix, filter, suffix) } // Add WHERE clause at the end return fmt.Sprintf("%s WHERE %s", sql, filter) } // injectTenantConditions adds tenant_id filtering to the query func (v *sqlValidator) injectTenantConditions(sql string, tablesInQuery map[string]string) string { if !v.enableTenantInjection { return sql } // Build tenant conditions var conditions []string for tableName, alias := range tablesInQuery { if v.tablesWithTenantID[tableName] { if tableName == "tenants" { conditions = append(conditions, fmt.Sprintf("%s.id = %d", alias, v.tenantID)) } else { conditions = append(conditions, fmt.Sprintf("%s.tenant_id = %d", alias, v.tenantID)) } } } if len(conditions) == 0 { return sql } tenantFilter := strings.Join(conditions, " AND ") return InjectAndConditions(sql, tenantFilter) } // injectSoftDeleteConditions adds deleted_at IS NULL filtering to the query. func (v *sqlValidator) injectSoftDeleteConditions(sql string, tablesInQuery map[string]string) string { if !v.enableSoftDeleteInjection { return sql } var conditions []string for tableName, alias := range tablesInQuery { if v.tablesWithDeletedAt[tableName] { conditions = append(conditions, fmt.Sprintf("%s.deleted_at IS NULL", alias)) } } if len(conditions) == 0 { return sql } return InjectAndConditions(sql, strings.Join(conditions, " AND ")) } // injectHiddenKBFilter adds is_temporary = false filtering for the knowledge_bases table, // hiding internal/system-managed KBs (e.g., __chat_history__) from query results. func (v *sqlValidator) injectHiddenKBFilter(sql string, tablesInQuery map[string]string) string { if !v.enableHiddenKBFilter { return sql } alias, ok := tablesInQuery["knowledge_bases"] if !ok { return sql } return InjectAndConditions(sql, fmt.Sprintf("%s.is_temporary = false", alias)) } func (v *sqlValidator) injectChunkEnabledFilter(sql string, tablesInQuery map[string]string) string { if !v.enableChunkEnabledFilter { return sql } alias, ok := tablesInQuery["chunks"] if !ok { return sql } return InjectAndConditions(sql, fmt.Sprintf("%s.is_enabled = true", alias)) } // injectSearchScopeConditions restricts queries to the allowed knowledge bases // and (optionally) specific knowledge documents. func (v *sqlValidator) injectSearchScopeConditions(sql string, tablesInQuery map[string]string) string { if !v.enableSearchScopeFilter || len(v.searchScopeKBIDs) == 0 { if len(v.searchScopes) == 0 { return sql } } if len(v.searchScopes) > 0 { return v.injectStructuredSearchScopeConditions(sql, tablesInQuery) } quotedKBIDs := quoteStringSlice(v.searchScopeKBIDs) kbList := strings.Join(quotedKBIDs, ", ") var conditions []string if alias, ok := tablesInQuery["knowledge_bases"]; ok { conditions = append(conditions, fmt.Sprintf("%s.id IN (%s)", alias, kbList)) } if alias, ok := tablesInQuery["knowledges"]; ok { conditions = append(conditions, fmt.Sprintf("%s.knowledge_base_id IN (%s)", alias, kbList)) if len(v.searchScopeKnowledgeIDs) > 0 { quotedKIDs := quoteStringSlice(v.searchScopeKnowledgeIDs) conditions = append(conditions, fmt.Sprintf("%s.id IN (%s)", alias, strings.Join(quotedKIDs, ", "))) } } if alias, ok := tablesInQuery["chunks"]; ok { conditions = append(conditions, fmt.Sprintf("%s.knowledge_base_id IN (%s)", alias, kbList)) if len(v.searchScopeKnowledgeIDs) > 0 { quotedKIDs := quoteStringSlice(v.searchScopeKnowledgeIDs) conditions = append(conditions, fmt.Sprintf("%s.knowledge_id IN (%s)", alias, strings.Join(quotedKIDs, ", "))) } } if len(conditions) == 0 { return sql } return InjectAndConditions(sql, strings.Join(conditions, " AND ")) } func (v *sqlValidator) injectStructuredSearchScopeConditions(sql string, tablesInQuery map[string]string) string { var conditions []string if alias, ok := tablesInQuery["knowledge_bases"]; ok { if cond := buildKnowledgeBaseScopeCondition(alias, v.searchScopes); cond != "" { conditions = append(conditions, cond) } } if alias, ok := tablesInQuery["knowledges"]; ok { if cond := buildKnowledgeScopeCondition(alias, v.searchScopes); cond != "" { conditions = append(conditions, cond) } } if alias, ok := tablesInQuery["chunks"]; ok { if cond := buildChunkScopeCondition(alias, v.searchScopes); cond != "" { conditions = append(conditions, cond) } } if len(conditions) == 0 { return sql } return InjectAndConditions(sql, strings.Join(conditions, " AND ")) } func buildKnowledgeBaseScopeCondition(alias string, scopes []SearchScope) string { kbIDs := uniqueScopeKBIDs(scopes) if len(kbIDs) == 0 { return "" } return fmt.Sprintf("%s.id IN (%s)", alias, strings.Join(quoteStringSlice(kbIDs), ", ")) } // buildScopeClause renders one scope. A scope's document whitelist and tag // filter are ANDed: each is an independent narrowing of the same KB, so // applying only one of them would admit rows the other excludes. Scopes are // ORed against each other by joinOrClauses because they are alternatives. func buildScopeClause(alias, knowledgeIDColumn string, scope SearchScope) string { if scope.KnowledgeBaseID == "" { return "" } conditions := []string{ fmt.Sprintf("%s.knowledge_base_id = %s", alias, quoteString(scope.KnowledgeBaseID)), } if len(scope.KnowledgeIDs) > 0 { conditions = append(conditions, fmt.Sprintf( "%s.%s IN (%s)", alias, knowledgeIDColumn, strings.Join(quoteStringSlice(scope.KnowledgeIDs), ", "), )) } if len(scope.TagIDs) > 0 { conditions = append(conditions, fmt.Sprintf( "EXISTS (SELECT 1 FROM knowledge_tag_relations ktr WHERE ktr.knowledge_id = %s.%s AND ktr.tag_id IN (%s))", alias, knowledgeIDColumn, strings.Join(quoteStringSlice(scope.TagIDs), ", "), )) } if len(conditions) == 1 { return conditions[0] } return "(" + strings.Join(conditions, " AND ") + ")" } func buildKnowledgeScopeCondition(alias string, scopes []SearchScope) string { clauses := make([]string, 0, len(scopes)) for _, scope := range scopes { if clause := buildScopeClause(alias, "id", scope); clause != "" { clauses = append(clauses, clause) } } return joinOrClauses(clauses) } func buildChunkScopeCondition(alias string, scopes []SearchScope) string { clauses := make([]string, 0, len(scopes)) for _, scope := range scopes { if clause := buildScopeClause(alias, "knowledge_id", scope); clause != "" { clauses = append(clauses, clause) } } return joinOrClauses(clauses) } func uniqueScopeKBIDs(scopes []SearchScope) []string { seen := make(map[string]bool, len(scopes)) out := make([]string, 0, len(scopes)) for _, scope := range scopes { if scope.KnowledgeBaseID == "" || seen[scope.KnowledgeBaseID] { continue } seen[scope.KnowledgeBaseID] = true out = append(out, scope.KnowledgeBaseID) } return out } func joinOrClauses(clauses []string) string { if len(clauses) == 0 { return "" } if len(clauses) == 1 { return clauses[0] } return "(" + strings.Join(clauses, " OR ") + ")" } func quoteString(s string) string { return quoteStringSlice([]string{s})[0] } func quoteStringSlice(ss []string) []string { quoted := make([]string, len(ss)) for i, s := range ss { escaped := strings.ReplaceAll(s, "'", "''") quoted[i] = fmt.Sprintf("'%s'", escaped) } return quoted } // Compiled once. checkSQLInjectionRisks runs on every validated WHERE clause, // so these patterns are hoisted to package scope instead of being recompiled // per call. var ( reSQLWhitespace = regexp.MustCompile(`\s+`) sqlAlwaysTruePatterns = []struct { pattern *regexp.Regexp description string }{ { pattern: regexp.MustCompile(`(^|\s|\()(1\s*=\s*1|'1'\s*=\s*'1'|"1"\s*=\s*"1")(\s|\)|$|and|or)`), description: "Always-true condition '1=1' or similar", }, { pattern: regexp.MustCompile(`(^|\s|\()(0\s*=\s*0|'0'\s*=\s*'0'|"0"\s*=\s*"0")(\s|\)|$|and|or)`), description: "Always-true condition '0=0' or similar", }, { pattern: regexp.MustCompile(`(^|\s|\()(true)(\s|\)|$|and|or)`), description: "Always-true condition 'true'", }, { pattern: regexp.MustCompile(`(^|\s|\()('\s*'\s*=\s*'\s*'|"\s*"\s*=\s*"\s*")(\s|\)|$|and|or)`), description: "Always-true condition with empty strings", }, } sqlAlwaysFalsePatterns = []struct { pattern *regexp.Regexp description string }{ { pattern: regexp.MustCompile(`(^|\s|\()(1\s*=\s*0|0\s*=\s*1|'1'\s*=\s*'0'|"1"\s*=\s*"0")(\s|\)|$|and|or)`), description: "Always-false condition '1=0' or similar", }, { pattern: regexp.MustCompile(`(^|\s|\()(false)(\s|\)|$|and|or)`), description: "Always-false condition 'false'", }, } reSQLOrAlwaysTrue = regexp.MustCompile(`or\s+(1\s*=\s*1|'1'\s*=\s*'1'|true)`) ) // checkSQLInjectionRisks checks for common SQL injection patterns in WHERE clause func checkSQLInjectionRisks(whereClause string) []SQLValidationError { errors := make([]SQLValidationError, 0) if whereClause == "" { return errors } // Normalize the WHERE clause for checking normalizedWhere := strings.ToLower(strings.TrimSpace(whereClause)) normalizedWhere = reSQLWhitespace.ReplaceAllString(normalizedWhere, " ") // Pattern 1: Always true conditions like "1=1", "'1'='1'", "true", etc. for _, pt := range sqlAlwaysTruePatterns { if pt.pattern.MatchString(normalizedWhere) { errors = append(errors, SQLValidationError{ Type: "sql_injection_risk", Message: "Potential SQL injection risk detected", Details: fmt.Sprintf("%s found in WHERE clause: %s", pt.description, whereClause), }) } } // Pattern 2: Always false conditions that might be used for testing for _, pt := range sqlAlwaysFalsePatterns { if pt.pattern.MatchString(normalizedWhere) { errors = append(errors, SQLValidationError{ Type: "sql_injection_risk", Message: "Suspicious SQL pattern detected", Details: fmt.Sprintf("%s found in WHERE clause: %s", pt.description, whereClause), }) } } // Pattern 3: OR with always-true condition (common injection pattern) if reSQLOrAlwaysTrue.MatchString(normalizedWhere) { errors = append(errors, SQLValidationError{ Type: "sql_injection_risk", Message: "High-risk SQL injection pattern detected", Details: fmt.Sprintf("OR with always-true condition found in WHERE clause: %s", whereClause), }) } return errors } // getMapKeys returns the keys of a map as a slice func getMapKeys(m map[string]bool) []string { keys := make([]string, 0, len(m)) for k := range m { keys = append(keys, k) } return keys } // validateInput performs basic input validation func (v *sqlValidator) validateInput(sql string) error { // Check for null bytes if strings.Contains(sql, "\x00") { return fmt.Errorf("invalid character in SQL query") } // Check length limits if len(sql) < v.minLength { return fmt.Errorf("SQL query too short (min %d characters)", v.minLength) } if len(sql) > v.maxLength { return fmt.Errorf("SQL query too long (max %d characters)", v.maxLength) } return nil } // validateSelectStmt validates a SELECT statement with configured options func (v *sqlValidator) validateSelectStmt(stmt *pg_query.SelectStmt, result *SQLValidationResult) error { tablesInQuery := make(map[string]string) // table name -> alias // Check for UNION/INTERSECT/EXCEPT (compound queries) if stmt.Op != pg_query.SetOperation_SETOP_NONE { return fmt.Errorf("compound queries (UNION/INTERSECT/EXCEPT) are not allowed") } // Check for WITH clause (CTEs) if v.checkCTEs && stmt.WithClause != nil { return fmt.Errorf("WITH clause (CTEs) is not allowed") } // Check for INTO clause (SELECT INTO) if stmt.IntoClause != nil { return fmt.Errorf("SELECT INTO is not allowed") } // Check for LOCKING clause (FOR UPDATE, etc.) if len(stmt.LockingClause) > 0 { return fmt.Errorf("locking clauses (FOR UPDATE, etc.) are not allowed") } // Validate FROM clause for _, fromItem := range stmt.FromClause { if err := v.validateFromItem(fromItem, tablesInQuery, result); err != nil { return err } } // Validate target list (SELECT columns) for _, target := range stmt.TargetList { if err := v.validateNode(target, result); err != nil { return err } } // Validate WHERE clause if stmt.WhereClause != nil { if err := v.validateNode(stmt.WhereClause, result); err != nil { return err } } // Validate GROUP BY clause for _, groupBy := range stmt.GroupClause { if err := v.validateNode(groupBy, result); err != nil { return err } } // Validate HAVING clause if stmt.HavingClause != nil { if err := v.validateNode(stmt.HavingClause, result); err != nil { return err } } // Validate ORDER BY clause for _, sortBy := range stmt.SortClause { if err := v.validateNode(sortBy, result); err != nil { return err } } // Ensure at least one valid table is referenced if len(tablesInQuery) == 0 { return fmt.Errorf("no valid table found in query") } return nil } // validateFromItem validates a FROM clause item func (v *sqlValidator) validateFromItem(node *pg_query.Node, tables map[string]string, result *SQLValidationResult) error { if node == nil { return nil } // Handle RangeVar (simple table reference) if rv := node.GetRangeVar(); rv != nil { tableName := strings.ToLower(rv.Relname) // Check for schema qualification if v.checkSchemaAccess && rv.Schemaname != "" { schemaName := strings.ToLower(rv.Schemaname) if schemaName != "public" { return fmt.Errorf("access to schema '%s' is not allowed", rv.Schemaname) } } // Get alias alias := tableName if rv.Alias != nil && rv.Alias.Aliasname != "" { alias = strings.ToLower(rv.Alias.Aliasname) } tables[tableName] = alias return nil } // Handle JoinExpr (JOIN) if je := node.GetJoinExpr(); je != nil { if err := v.validateFromItem(je.Larg, tables, result); err != nil { return err } if err := v.validateFromItem(je.Rarg, tables, result); err != nil { return err } if je.Quals != nil { if err := v.validateNode(je.Quals, result); err != nil { return err } } return nil } // Handle RangeSubselect (subquery in FROM) if rss := node.GetRangeSubselect(); rss != nil { if v.checkSubqueries { return fmt.Errorf("subqueries in FROM clause are not allowed") } // SECURITY: Even when subqueries are permitted, recurse into the // subquery so dangerous constructs hidden inside it are still // validated. Without this, a FROM subquery like // (SELECT * FROM read_text('/etc/passwd')) // smuggles a RangeFunction past the check below. return v.validateSubquery(rss.Subquery, tables, result) } // Handle RangeFunction (function in FROM) if node.GetRangeFunction() != nil { return fmt.Errorf("functions in FROM clause are not allowed") } return nil } // validateSubquery validates a SELECT statement nested in a FROM subquery. // It reuses the FROM-item and expression validators so RangeFunction and // dangerous function checks apply recursively to arbitrarily nested subqueries. func (v *sqlValidator) validateSubquery(node *pg_query.Node, tables map[string]string, result *SQLValidationResult) error { if node == nil { return nil } sub := node.GetSelectStmt() if sub == nil { return nil } for _, fromItem := range sub.FromClause { if err := v.validateFromItem(fromItem, tables, result); err != nil { return err } } for _, target := range sub.TargetList { if err := v.validateNode(target, result); err != nil { return err } } if sub.WhereClause != nil { if err := v.validateNode(sub.WhereClause, result); err != nil { return err } } for _, groupBy := range sub.GroupClause { if err := v.validateNode(groupBy, result); err != nil { return err } } if sub.HavingClause != nil { if err := v.validateNode(sub.HavingClause, result); err != nil { return err } } return nil } // validateNode recursively validates AST nodes // SECURITY: This function uses a COMPREHENSIVE approach to validate ALL node types. // Any node type that contains child expressions MUST be handled to prevent bypass attacks. // The principle is: if we don't know how to validate a node type, we REJECT it. func (v *sqlValidator) validateNode(node *pg_query.Node, result *SQLValidationResult) error { if node == nil { return nil } // Check for subqueries (SubLink) if v.checkSubqueries { if sl := node.GetSubLink(); sl != nil { return fmt.Errorf("subqueries are not allowed") } } // Check for function calls if fc := node.GetFuncCall(); fc != nil { if err := v.validateFuncCall(fc, result); err != nil { return err } } // Check for column references if cr := node.GetColumnRef(); cr != nil { if err := v.validateColumnRef(cr); err != nil { return err } } // Check for type casts if tc := node.GetTypeCast(); tc != nil { if err := v.validateNode(tc.Arg, result); err != nil { return err } if tc.TypeName != nil { typeName := v.getTypeName(tc.TypeName) if strings.HasPrefix(strings.ToLower(typeName), "pg_") { return fmt.Errorf("casting to system type '%s' is not allowed", typeName) } } } // Recursively check A_Expr (expressions) if ae := node.GetAExpr(); ae != nil { if err := v.validateNode(ae.Lexpr, result); err != nil { return err } if err := v.validateNode(ae.Rexpr, result); err != nil { return err } } // Check BoolExpr (AND, OR, NOT) if be := node.GetBoolExpr(); be != nil { for _, arg := range be.Args { if err := v.validateNode(arg, result); err != nil { return err } } } // Check NullTest if nt := node.GetNullTest(); nt != nil { if err := v.validateNode(nt.Arg, result); err != nil { return err } } // Check CoalesceExpr if ce := node.GetCoalesceExpr(); ce != nil { for _, arg := range ce.Args { if err := v.validateNode(arg, result); err != nil { return err } } } // Check CaseExpr if caseExpr := node.GetCaseExpr(); caseExpr != nil { if err := v.validateNode(caseExpr.Arg, result); err != nil { return err } for _, when := range caseExpr.Args { if err := v.validateNode(when, result); err != nil { return err } } if err := v.validateNode(caseExpr.Defresult, result); err != nil { return err } } // Check CaseWhen if cw := node.GetCaseWhen(); cw != nil { if err := v.validateNode(cw.Expr, result); err != nil { return err } if err := v.validateNode(cw.Result, result); err != nil { return err } } // Check ResTarget (SELECT list items) if rt := node.GetResTarget(); rt != nil { if err := v.validateNode(rt.Val, result); err != nil { return err } } // Check SortBy (ORDER BY items) if sb := node.GetSortBy(); sb != nil { if err := v.validateNode(sb.Node, result); err != nil { return err } } // Check List if list := node.GetList(); list != nil { for _, item := range list.Items { if err := v.validateNode(item, result); err != nil { return err } } } // ============================================================ // SECURITY FIX: Comprehensive handling of ALL expression types // that can contain child nodes (potential bypass vectors) // ============================================================ // ArrayExpr (ARRAY[...] expressions) // Attack: SELECT ARRAY[pg_read_file('/etc/passwd')] FROM table if ae := node.GetAArrayExpr(); ae != nil { for _, elem := range ae.Elements { if err := v.validateNode(elem, result); err != nil { return err } } } // RowExpr (ROW(...) expressions) // Attack: SELECT ROW(pg_read_file('/etc/passwd')) FROM table if re := node.GetRowExpr(); re != nil { for _, arg := range re.Args { if err := v.validateNode(arg, result); err != nil { return err } } } // MinMaxExpr (GREATEST/LEAST expressions) if mm := node.GetMinMaxExpr(); mm != nil { for _, arg := range mm.Args { if err := v.validateNode(arg, result); err != nil { return err } } } // NullIfExpr (NULLIF expressions) if ni := node.GetNullIfExpr(); ni != nil { for _, arg := range ni.Args { if err := v.validateNode(arg, result); err != nil { return err } } } // ScalarArrayOpExpr (IN, ANY, ALL with arrays) if sao := node.GetScalarArrayOpExpr(); sao != nil { for _, arg := range sao.Args { if err := v.validateNode(arg, result); err != nil { return err } } } // ArrayCoerceExpr if ace := node.GetArrayCoerceExpr(); ace != nil { if err := v.validateNode(ace.Arg, result); err != nil { return err } } // CoerceViaIO (type coercion via I/O) if cvi := node.GetCoerceViaIo(); cvi != nil { if err := v.validateNode(cvi.Arg, result); err != nil { return err } } // CollateExpr (COLLATE expressions) if ce := node.GetCollateExpr(); ce != nil { if err := v.validateNode(ce.Arg, result); err != nil { return err } } // SubLink (subqueries) - validate child expressions even if subqueries are allowed if sl := node.GetSubLink(); sl != nil { if err := v.validateNode(sl.Testexpr, result); err != nil { return err } } // OpExpr (operator expressions) if oe := node.GetOpExpr(); oe != nil { for _, arg := range oe.Args { if err := v.validateNode(arg, result); err != nil { return err } } } // DistinctExpr (IS DISTINCT FROM) if de := node.GetDistinctExpr(); de != nil { for _, arg := range de.Args { if err := v.validateNode(arg, result); err != nil { return err } } } // XmlExpr (XML expressions) if xe := node.GetXmlExpr(); xe != nil { for _, arg := range xe.Args { if err := v.validateNode(arg, result); err != nil { return err } } for _, arg := range xe.NamedArgs { if err := v.validateNode(arg, result); err != nil { return err } } } // JsonConstructorExpr if jce := node.GetJsonConstructorExpr(); jce != nil { for _, arg := range jce.Args { if err := v.validateNode(arg, result); err != nil { return err } } } // ============================================================ // Additional expression types that need recursive validation // ============================================================ // FuncExpr (different from FuncCall - internal function representation) if fe := node.GetFuncExpr(); fe != nil { for _, arg := range fe.Args { if err := v.validateNode(arg, result); err != nil { return err } } } // Aggref (aggregate function reference) if ag := node.GetAggref(); ag != nil { for _, arg := range ag.Args { if err := v.validateNode(arg, result); err != nil { return err } } for _, arg := range ag.Aggdirectargs { if err := v.validateNode(arg, result); err != nil { return err } } if ag.Aggfilter != nil { if err := v.validateNode(ag.Aggfilter, result); err != nil { return err } } } // WindowFunc if wf := node.GetWindowFunc(); wf != nil { for _, arg := range wf.Args { if err := v.validateNode(arg, result); err != nil { return err } } if wf.Aggfilter != nil { if err := v.validateNode(wf.Aggfilter, result); err != nil { return err } } } // SubscriptingRef (array subscripting like arr[1]) if sr := node.GetSubscriptingRef(); sr != nil { for _, idx := range sr.Refupperindexpr { if err := v.validateNode(idx, result); err != nil { return err } } for _, idx := range sr.Reflowerindexpr { if err := v.validateNode(idx, result); err != nil { return err } } if err := v.validateNode(sr.Refexpr, result); err != nil { return err } if err := v.validateNode(sr.Refassgnexpr, result); err != nil { return err } } // NamedArgExpr (named arguments in function calls) if nae := node.GetNamedArgExpr(); nae != nil { if err := v.validateNode(nae.Arg, result); err != nil { return err } } // FieldSelect (field selection from composite type) if fs := node.GetFieldSelect(); fs != nil { if err := v.validateNode(fs.Arg, result); err != nil { return err } } // FieldStore if fs := node.GetFieldStore(); fs != nil { if err := v.validateNode(fs.Arg, result); err != nil { return err } for _, newval := range fs.Newvals { if err := v.validateNode(newval, result); err != nil { return err } } } // RelabelType (type relabeling) if rt := node.GetRelabelType(); rt != nil { if err := v.validateNode(rt.Arg, result); err != nil { return err } } // ConvertRowtypeExpr if cre := node.GetConvertRowtypeExpr(); cre != nil { if err := v.validateNode(cre.Arg, result); err != nil { return err } } // RowCompareExpr if rce := node.GetRowCompareExpr(); rce != nil { for _, arg := range rce.Largs { if err := v.validateNode(arg, result); err != nil { return err } } for _, arg := range rce.Rargs { if err := v.validateNode(arg, result); err != nil { return err } } } // CoerceToDomain if ctd := node.GetCoerceToDomain(); ctd != nil { if err := v.validateNode(ctd.Arg, result); err != nil { return err } } // BooleanTest (IS TRUE, IS FALSE, etc.) if bt := node.GetBooleanTest(); bt != nil { if err := v.validateNode(bt.Arg, result); err != nil { return err } } // AIndices (array indices) if ai := node.GetAIndices(); ai != nil { if err := v.validateNode(ai.Lidx, result); err != nil { return err } if err := v.validateNode(ai.Uidx, result); err != nil { return err } } // AIndirection (array/field indirection) if aind := node.GetAIndirection(); aind != nil { if err := v.validateNode(aind.Arg, result); err != nil { return err } for _, ind := range aind.Indirection { if err := v.validateNode(ind, result); err != nil { return err } } } // CollateClause if cc := node.GetCollateClause(); cc != nil { if err := v.validateNode(cc.Arg, result); err != nil { return err } } // GroupingFunc if gf := node.GetGroupingFunc(); gf != nil { for _, arg := range gf.Args { if err := v.validateNode(arg, result); err != nil { return err } } } // JsonValueExpr if jve := node.GetJsonValueExpr(); jve != nil { if err := v.validateNode(jve.RawExpr, result); err != nil { return err } if err := v.validateNode(jve.FormattedExpr, result); err != nil { return err } } // JsonExpr if je := node.GetJsonExpr(); je != nil { if err := v.validateNode(je.FormattedExpr, result); err != nil { return err } if err := v.validateNode(je.PathSpec, result); err != nil { return err } for _, arg := range je.PassingValues { if err := v.validateNode(arg, result); err != nil { return err } } } // JsonIsPredicate if jip := node.GetJsonIsPredicate(); jip != nil { if err := v.validateNode(jip.Expr, result); err != nil { return err } } // JsonScalarExpr (JSON_SCALAR(...)) - PG17 SQL/JSON node. // Attack: SELECT JSON_SCALAR(pg_read_file('/etc/passwd')) FROM table if jse := node.GetJsonScalarExpr(); jse != nil { if err := v.validateNode(jse.Expr, result); err != nil { return err } } // JsonFuncExpr (JSON_VALUE / JSON_QUERY / JSON_EXISTS(...)) - PG17 SQL/JSON node. // Attack: SELECT JSON_VALUE(pg_read_file('/etc/passwd'), '$') FROM table if jfe := node.GetJsonFuncExpr(); jfe != nil { if err := v.validateJsonValueExpr(jfe.ContextItem, result); err != nil { return err } if err := v.validateNode(jfe.Pathspec, result); err != nil { return err } for _, arg := range jfe.Passing { if err := v.validateNode(arg, result); err != nil { return err } } if err := v.validateJsonBehavior(jfe.OnEmpty, result); err != nil { return err } if err := v.validateJsonBehavior(jfe.OnError, result); err != nil { return err } } // JsonParseExpr (JSON(...)) - PG17 SQL/JSON node. if jpe := node.GetJsonParseExpr(); jpe != nil { if err := v.validateJsonValueExpr(jpe.Expr, result); err != nil { return err } } // JsonSerializeExpr (JSON_SERIALIZE(...)) - PG17 SQL/JSON node. if jse := node.GetJsonSerializeExpr(); jse != nil { if err := v.validateJsonValueExpr(jse.Expr, result); err != nil { return err } } // JsonTable (JSON_TABLE(...)) - PG17 SQL/JSON node. if jt := node.GetJsonTable(); jt != nil { if err := v.validateJsonValueExpr(jt.ContextItem, result); err != nil { return err } for _, arg := range jt.Passing { if err := v.validateNode(arg, result); err != nil { return err } } for _, col := range jt.Columns { if err := v.validateNode(col, result); err != nil { return err } } if err := v.validateJsonBehavior(jt.OnError, result); err != nil { return err } } // XmlSerialize if xs := node.GetXmlSerialize(); xs != nil { if err := v.validateNode(xs.Expr, result); err != nil { return err } } // WindowDef if wd := node.GetWindowDef(); wd != nil { for _, part := range wd.PartitionClause { if err := v.validateNode(part, result); err != nil { return err } } for _, order := range wd.OrderClause { if err := v.validateNode(order, result); err != nil { return err } } if err := v.validateNode(wd.StartOffset, result); err != nil { return err } if err := v.validateNode(wd.EndOffset, result); err != nil { return err } } // SubPlan - BLOCK: This is an internal representation, should not appear in user queries if node.GetSubPlan() != nil { return fmt.Errorf("SubPlan nodes are not allowed") } // AlternativeSubPlan - BLOCK if node.GetAlternativeSubPlan() != nil { return fmt.Errorf("AlternativeSubPlan nodes are not allowed") } // ============================================================ // SECURITY: DEFAULT-DENY. // The recursive branches above validate every expression type we know how // to inspect. Any node type NOT in the recognized set below is REJECTED, // because an unhandled type can smuggle a dangerous FuncCall past the // blacklist (e.g. PG17's JSON_SCALAR wrapping pg_read_file). Whenever a new // recursive handler is added above, add its Node_* wrapper here as well. // ============================================================ switch node.Node.(type) { case // Recognized expression types (validated recursively above). *pg_query.Node_FuncCall, *pg_query.Node_FuncExpr, *pg_query.Node_ColumnRef, *pg_query.Node_TypeCast, *pg_query.Node_AExpr, *pg_query.Node_OpExpr, *pg_query.Node_BoolExpr, *pg_query.Node_NullTest, *pg_query.Node_BooleanTest, *pg_query.Node_CoalesceExpr, *pg_query.Node_CaseExpr, *pg_query.Node_CaseWhen, *pg_query.Node_ResTarget, *pg_query.Node_SortBy, *pg_query.Node_List, *pg_query.Node_AArrayExpr, *pg_query.Node_RowExpr, *pg_query.Node_MinMaxExpr, *pg_query.Node_NullIfExpr, *pg_query.Node_ScalarArrayOpExpr, *pg_query.Node_ArrayCoerceExpr, *pg_query.Node_CoerceViaIo, *pg_query.Node_CollateExpr, *pg_query.Node_CollateClause, *pg_query.Node_SubLink, *pg_query.Node_DistinctExpr, *pg_query.Node_XmlExpr, *pg_query.Node_XmlSerialize, *pg_query.Node_JsonConstructorExpr, *pg_query.Node_JsonValueExpr, *pg_query.Node_JsonExpr, *pg_query.Node_JsonIsPredicate, *pg_query.Node_JsonScalarExpr, *pg_query.Node_JsonFuncExpr, *pg_query.Node_JsonParseExpr, *pg_query.Node_JsonSerializeExpr, *pg_query.Node_JsonTable, *pg_query.Node_Aggref, *pg_query.Node_WindowFunc, *pg_query.Node_WindowDef, *pg_query.Node_GroupingFunc, *pg_query.Node_SubscriptingRef, *pg_query.Node_NamedArgExpr, *pg_query.Node_FieldSelect, *pg_query.Node_FieldStore, *pg_query.Node_RelabelType, *pg_query.Node_ConvertRowtypeExpr, *pg_query.Node_RowCompareExpr, *pg_query.Node_CoerceToDomain, *pg_query.Node_AIndices, *pg_query.Node_AIndirection, // Recognized safe leaf nodes (no child expressions to smuggle through). *pg_query.Node_AConst, *pg_query.Node_ParamRef, *pg_query.Node_SetToDefault, *pg_query.Node_CurrentOfExpr, *pg_query.Node_CaseTestExpr, *pg_query.Node_SqlvalueFunction, *pg_query.Node_AStar, *pg_query.Node_Integer, *pg_query.Node_Float, *pg_query.Node_Boolean, *pg_query.Node_String_, *pg_query.Node_BitString: return nil default: return fmt.Errorf("unsupported SQL expression type %T is not allowed", node.Node) } } // validateJsonValueExpr validates a JsonValueExpr, which appears as a concrete // (non-Node) field on several PG17 SQL/JSON expression nodes. Its RawExpr / // FormattedExpr children can hold arbitrary expressions (including FuncCalls), // so they must be recursed into. func (v *sqlValidator) validateJsonValueExpr(jve *pg_query.JsonValueExpr, result *SQLValidationResult) error { if jve == nil { return nil } if err := v.validateNode(jve.RawExpr, result); err != nil { return err } return v.validateNode(jve.FormattedExpr, result) } // validateJsonBehavior validates the ON EMPTY / ON ERROR behavior of a PG17 // SQL/JSON function. The DEFAULT branch can carry an arbitrary expression. func (v *sqlValidator) validateJsonBehavior(jb *pg_query.JsonBehavior, result *SQLValidationResult) error { if jb == nil { return nil } return v.validateNode(jb.Expr, result) } // validateFuncCall validates a function call func (v *sqlValidator) validateFuncCall(fc *pg_query.FuncCall, result *SQLValidationResult) error { // Get function name funcName := "" for _, namePart := range fc.Funcname { if s := namePart.GetString_(); s != nil { funcName = strings.ToLower(s.Sval) } } // Check for schema-qualified function calls if v.checkSchemaAccess && len(fc.Funcname) > 1 { schemaName := "" if s := fc.Funcname[0].GetString_(); s != nil { schemaName = strings.ToLower(s.Sval) } if schemaName != "" && schemaName != "pg_catalog" { return fmt.Errorf("schema-qualified function calls are not allowed: %s", schemaName) } } // Block dangerous function prefixes if v.checkDangerousFuncs { dangerousPrefixes := []string{ "pg_", // All pg_* functions (pg_read_file, pg_reload_conf, pg_stat_*, etc.) "lo_", // Large object functions (lo_import, lo_export, lo_from_bytea, lo_put, etc.) "dblink", // Database link functions "file_", // File functions "copy_", // Copy functions "binary_", // Binary functions } for _, prefix := range dangerousPrefixes { if strings.HasPrefix(funcName, prefix) { return fmt.Errorf("function '%s' is not allowed (dangerous prefix)", funcName) } } // Block specific dangerous functions - comprehensive list for RCE prevention dangerousFunctions := map[string]bool{ // Configuration and settings "current_setting": true, "set_config": true, // XML/XPath functions (XXE risks) "query_to_xml": true, "xpath": true, "xmlparse": true, "xmlroot": true, "xmlelement": true, "xmlforest": true, "xmlconcat": true, "xmlagg": true, "xmlpi": true, "xmlcomment": true, "xmlexists": true, "xml_is_well_formed": true, "xpath_exists": true, "table_to_xml": true, "cursor_to_xml": true, "database_to_xml": true, "schema_to_xml": true, // Transaction and system info "txid_current": true, "txid_current_snapshot": true, "txid_snapshot_xmin": true, "txid_snapshot_xmax": true, // Encoding functions (used in attack payloads) "encode": true, "decode": true, // Extension management "create_extension": true, // Copy operations "copy": true, "copy_to": true, "copy_from": true, "pg_copy_to": true, "pg_dump": true, "pg_dumpall": true, "pg_restore": true, "pg_basebackup": true, // Process and system functions "pg_terminate_backend": true, "pg_cancel_backend": true, "pg_rotate_logfile": true, // Advisory locks (can be abused for DoS) "pg_advisory_lock": true, "pg_advisory_unlock": true, "pg_advisory_lock_shared": true, "pg_advisory_unlock_shared": true, "pg_try_advisory_lock": true, "pg_try_advisory_lock_shared": true, // Backup and replication "pg_start_backup": true, "pg_stop_backup": true, "pg_switch_wal": true, "pg_create_restore_point": true, // Foreign data wrappers "postgres_fdw_handler": true, "file_fdw_handler": true, // Procedural languages (code execution) "plpgsql_call_handler": true, "plpython_call_handler": true, "plperl_call_handler": true, // System catalog modification "pg_catalog": true, "information_schema": true, // DuckDB file-access functions (data_analysis tool runs on DuckDB). // The user's data is already loaded into a session table, so these // arbitrary-path readers are never legitimately needed and would // allow reading any file on the app container. "read_text": true, "read_blob": true, "read_csv": true, "read_csv_auto": true, "read_parquet": true, "read_json": true, "read_json_auto": true, "read_ndjson": true, "read_ndjson_auto": true, "read_json_objects": true, "read_xlsx": true, "sniff_csv": true, "glob": true, "st_read": true, "st_read_meta": true, } if dangerousFunctions[funcName] { return fmt.Errorf("function '%s' is not allowed", funcName) } } // Check against whitelist if enabled if v.checkFunctionNames && !v.allowedFunctions[funcName] { return fmt.Errorf("function not allowed: %s", funcName) } // Validate function arguments recursively for _, arg := range fc.Args { if err := v.validateNode(arg, result); err != nil { return err } } return nil } // validateColumnRef validates a column reference func (v *sqlValidator) validateColumnRef(cr *pg_query.ColumnRef) error { if !v.checkSystemColumns { return nil } // Check for system column access for _, field := range cr.Fields { if s := field.GetString_(); s != nil { colName := strings.ToLower(s.Sval) // Block access to system columns systemColumns := []string{"xmin", "xmax", "cmin", "cmax", "ctid", "tableoid"} for _, sysCol := range systemColumns { if colName == sysCol { return fmt.Errorf("access to system column '%s' is not allowed", colName) } } // Block pg_ prefixed identifiers if strings.HasPrefix(colName, "pg_") { return fmt.Errorf("access to '%s' is not allowed", colName) } } } return nil } // getTypeName extracts the type name from a TypeName node func (v *sqlValidator) getTypeName(tn *pg_query.TypeName) string { var parts []string for _, name := range tn.Names { if s := name.GetString_(); s != nil { parts = append(parts, s.Sval) } } return strings.Join(parts, ".") }