feat: Add MCP query rules and digest statistics [WIP]

include/Discovery_Schema.h

@@ -5,6 +5,131 @@
 #include <string>
 #include <vector>
 #include <memory>
+#include <pthread.h>
+#include <unordered_map>
+#include "json.hpp"
+
+/**
+ * @brief MCP query rule structure
+ *
+ * Action is inferred from rule properties:
+ * - if error_msg != NULL    → block
+ * - if replace_pattern != NULL → rewrite
+ * - if timeout_ms > 0       → timeout
+ * - otherwise               → allow
+ *
+ * Note: 'hits' is only for in-memory tracking, not persisted to the table.
+ */
+struct MCP_Query_Rule {
+	int rule_id;
+	bool active;
+	char *username;
+	char *schemaname;
+	char *tool_name;
+	char *match_pattern;
+	bool negate_match_pattern;
+	int re_modifiers;         // bitmask: 1=CASELESS
+	int flagIN;
+	int flagOUT;
+	char *replace_pattern;
+	int timeout_ms;
+	char *error_msg;
+	char *ok_msg;
+	bool log;
+	bool apply;
+	char *comment;
+	uint64_t hits;            // in-memory only, not persisted to table
+	void* regex_engine;      // compiled regex (RE2)
+
+	MCP_Query_Rule() : rule_id(0), active(false), username(NULL), schemaname(NULL),
+	                   tool_name(NULL), match_pattern(NULL), negate_match_pattern(false),
+	                   re_modifiers(1), flagIN(0), flagOUT(0), replace_pattern(NULL),
+	                   timeout_ms(0), error_msg(NULL), ok_msg(NULL), log(false), apply(true),
+	                   comment(NULL), hits(0), regex_engine(NULL) {}
+};
+
+/**
+ * @brief MCP query digest statistics
+ */
+struct MCP_Query_Digest_Stats {
+	std::string tool_name;
+	int run_id;
+	uint64_t digest;
+	std::string digest_text;
+	unsigned int count_star;
+	time_t first_seen;
+	time_t last_seen;
+	unsigned long long sum_time;
+	unsigned long long min_time;
+	unsigned long long max_time;
+
+	MCP_Query_Digest_Stats() : run_id(-1), digest(0), count_star(0),
+	                           first_seen(0), last_seen(0),
+	                           sum_time(0), min_time(0), max_time(0) {}
+
+	void add_timing(unsigned long long duration_us, time_t timestamp) {
+		count_star++;
+		sum_time += duration_us;
+		if (duration_us < min_time || min_time == 0) min_time = duration_us;
+		if (duration_us > max_time) max_time = duration_us;
+		if (first_seen == 0) first_seen = timestamp;
+		last_seen = timestamp;
+	}
+};
+
+/**
+ * @brief MCP query processor output
+ *
+ * This structure collects all possible actions from matching MCP query rules.
+ * A single rule can perform multiple actions simultaneously (rewrite + timeout + block).
+ * Actions are inferred from rule properties:
+ * - if error_msg != NULL    → block
+ * - if replace_pattern != NULL → rewrite
+ * - if timeout_ms > 0       → timeout
+ * - if OK_msg != NULL       → return OK message
+ *
+ * The calling code checks these fields and performs the appropriate actions.
+ */
+struct MCP_Query_Processor_Output {
+	std::string *new_query;   // Rewritten query (caller must delete)
+	int timeout_ms;           // Query timeout in milliseconds (-1 = not set)
+	char *error_msg;          // Error message to return (NULL = not set)
+	char *OK_msg;             // OK message to return (NULL = not set)
+	int log;                  // Whether to log this query (-1 = not set, 0 = no, 1 = yes)
+	int next_query_flagIN;    // Flag for next query (-1 = not set)
+
+	void init() {
+		new_query = NULL;
+		timeout_ms = -1;
+		error_msg = NULL;
+		OK_msg = NULL;
+		log = -1;
+		next_query_flagIN = -1;
+	}
+
+	void destroy() {
+		if (new_query) {
+			delete new_query;
+			new_query = NULL;
+		}
+		if (error_msg) {
+			free(error_msg);
+			error_msg = NULL;
+		}
+		if (OK_msg) {
+			free(OK_msg);
+			OK_msg = NULL;
+		}
+	}
+
+	MCP_Query_Processor_Output() {
+		init();
+	}
+
+	~MCP_Query_Processor_Output() {
+		destroy();
+	}
+};
 
 /**
  * @brief Two-Phase Discovery Catalog Schema Manager
@@ -21,6 +146,15 @@ class Discovery_Schema {
 	SQLite3DB* db;
 	std::string db_path;
 
+	// MCP query rules management
+	std::vector<MCP_Query_Rule*> mcp_query_rules;
+	pthread_rwlock_t mcp_rules_lock;
+	volatile unsigned int mcp_rules_version;
+
+	// MCP query digest statistics
+	std::unordered_map<std::string, std::unordered_map<uint64_t, void*>> mcp_digest_umap;
+	pthread_rwlock_t mcp_digest_rwlock;
+
 	/**
 	 * @brief Initialize catalog schema with all tables
 	 * @return 0 on success, -1 on error
@@ -679,6 +813,72 @@ class Discovery_Schema {
 	 * @return Database file path
 	 */
 	std::string get_db_path() const { return db_path; }
+
+	// ============================================================
+	// MCP QUERY RULES
+	// ============================================================
+
+	/**
+	 * @brief Load MCP query rules from SQLite
+	 */
+	void load_mcp_query_rules(SQLite3_result* resultset);
+
+	/**
+	 * @brief Evaluate MCP query rules for a tool invocation
+	 * @return MCP_Query_Processor_Output object populated with actions from matching rules
+	 *         Caller is responsible for destroying the returned object.
+	 */
+	MCP_Query_Processor_Output* evaluate_mcp_query_rules(
+		const std::string& tool_name,
+		const std::string& schemaname,
+		const nlohmann::json& arguments,
+		const std::string& original_query
+	);
+
+	/**
+	 * @brief Get current MCP query rules as resultset
+	 */
+	SQLite3_result* get_mcp_query_rules();
+
+	/**
+	 * @brief Get stats for MCP query rules (hits per rule)
+	 */
+	SQLite3_result* get_stats_mcp_query_rules();
+
+	// ============================================================
+	// MCP QUERY DIGEST
+	// ============================================================
+
+	/**
+	 * @brief Update MCP query digest statistics
+	 */
+	void update_mcp_query_digest(
+		const std::string& tool_name,
+		int run_id,
+		uint64_t digest,
+		const std::string& digest_text,
+		unsigned long long duration_us,
+		time_t timestamp
+	);
+
+	/**
+	 * @brief Get MCP query digest statistics
+	 * @param reset If true, reset stats after retrieval
+	 */
+	SQLite3_result* get_mcp_query_digest(bool reset = false);
+
+	/**
+	 * @brief Compute MCP query digest hash using SpookyHash
+	 */
+	static uint64_t compute_mcp_digest(
+		const std::string& tool_name,
+		const nlohmann::json& arguments
+	);
+
+	/**
+	 * @brief Fingerprint MCP query arguments (replace literals with ?)
+	 */
+	static std::string fingerprint_mcp_args(const nlohmann::json& arguments);
 };
 
 #endif /* CLASS_DISCOVERY_SCHEMA_H */

include/ProxySQL_Admin_Tables_Definitions.h

@@ -325,6 +325,95 @@
 #define STATS_SQLITE_TABLE_MCP_QUERY_TOOLS_COUNTERS "CREATE TABLE stats_mcp_query_tools_counters (tool VARCHAR NOT NULL , schema VARCHAR NOT NULL , count INT NOT NULL , first_seen INTEGER NOT NULL , last_seen INTEGER NOT NULL , sum_time INTEGER NOT NULL , min_time INTEGER NOT NULL , max_time INTEGER NOT NULL , PRIMARY KEY (tool, schema))"
 #define STATS_SQLITE_TABLE_MCP_QUERY_TOOLS_COUNTERS_RESET "CREATE TABLE stats_mcp_query_tools_counters_reset (tool VARCHAR NOT NULL , schema VARCHAR NOT NULL , count INT NOT NULL , first_seen INTEGER NOT NULL , last_seen INTEGER NOT NULL , sum_time INTEGER NOT NULL , min_time INTEGER NOT NULL , max_time INTEGER NOT NULL , PRIMARY KEY (tool, schema))"
 
+// MCP query rules table - for firewall and query rewriting
+// Action is inferred from rule properties:
+//   - if error_msg is not NULL     → block
+//   - if replace_pattern is not NULL → rewrite
+//   - if timeout_ms > 0            → timeout
+//   - otherwise                    → allow
+#define ADMIN_SQLITE_TABLE_MCP_QUERY_RULES "CREATE TABLE mcp_query_rules (" \
+  "  rule_id INTEGER PRIMARY KEY AUTOINCREMENT NOT NULL ," \
+  "  active INT CHECK (active IN (0,1)) NOT NULL DEFAULT 0 ," \
+  "  username VARCHAR ," \
+  "  schemaname VARCHAR ," \
+  "  tool_name VARCHAR ," \
+  "  match_pattern VARCHAR ," \
+  "  negate_match_pattern INT CHECK (negate_match_pattern IN (0,1)) NOT NULL DEFAULT 0 ," \
+  "  re_modifiers VARCHAR DEFAULT 'CASELESS' ," \
+  "  flagIN INT NOT NULL DEFAULT 0 ," \
+  "  flagOUT INT CHECK (flagOUT >= 0) ," \
+  "  replace_pattern VARCHAR ," \
+  "  timeout_ms INT CHECK (timeout_ms >= 0) ," \
+  "  error_msg VARCHAR ," \
+  "  OK_msg VARCHAR ," \
+  "  log INT CHECK (log IN (0,1)) ," \
+  "  apply INT CHECK (apply IN (0,1)) NOT NULL DEFAULT 1 ," \
+  "  comment VARCHAR" \
+  ")"
+
+// MCP query rules runtime table - shows in-memory state of active rules
+// This table has the same schema as mcp_query_rules (no hits column).
+// The hits counter is only available in stats_mcp_query_rules table.
+// When this table is queried, it is automatically refreshed from the in-memory rules.
+#define ADMIN_SQLITE_TABLE_RUNTIME_MCP_QUERY_RULES "CREATE TABLE runtime_mcp_query_rules (" \
+  "  rule_id INTEGER PRIMARY KEY AUTOINCREMENT NOT NULL ," \
+  "  active INT CHECK (active IN (0,1)) NOT NULL DEFAULT 0 ," \
+  "  username VARCHAR ," \
+  "  schemaname VARCHAR ," \
+  "  tool_name VARCHAR ," \
+  "  match_pattern VARCHAR ," \
+  "  negate_match_pattern INT CHECK (negate_match_pattern IN (0,1)) NOT NULL DEFAULT 0 ," \
+  "  re_modifiers VARCHAR DEFAULT 'CASELESS' ," \
+  "  flagIN INT NOT NULL DEFAULT 0 ," \
+  "  flagOUT INT CHECK (flagOUT >= 0) ," \
+  "  replace_pattern VARCHAR ," \
+  "  timeout_ms INT CHECK (timeout_ms >= 0) ," \
+  "  error_msg VARCHAR ," \
+  "  OK_msg VARCHAR ," \
+  "  log INT CHECK (log IN (0,1)) ," \
+  "  apply INT CHECK (apply IN (0,1)) NOT NULL DEFAULT 1 ," \
+  "  comment VARCHAR" \
+  ")"
+
+// MCP query digest statistics table
+#define STATS_SQLITE_TABLE_MCP_QUERY_DIGEST "CREATE TABLE stats_mcp_query_digest (" \
+  "  tool_name VARCHAR NOT NULL ," \
+  "  run_id INT ," \
+  "  digest VARCHAR NOT NULL ," \
+  "  digest_text VARCHAR NOT NULL ," \
+  "  count_star INTEGER NOT NULL ," \
+  "  first_seen INTEGER NOT NULL ," \
+  "  last_seen INTEGER NOT NULL ," \
+  "  sum_time INTEGER NOT NULL ," \
+  "  min_time INTEGER NOT NULL ," \
+  "  max_time INTEGER NOT NULL ," \
+  "  PRIMARY KEY(tool_name, run_id, digest)" \
+  ")"
+
+// MCP query digest reset table
+#define STATS_SQLITE_TABLE_MCP_QUERY_DIGEST_RESET "CREATE TABLE stats_mcp_query_digest_reset (" \
+  "  tool_name VARCHAR NOT NULL ," \
+  "  run_id INT ," \
+  "  digest VARCHAR NOT NULL ," \
+  "  digest_text VARCHAR NOT NULL ," \
+  "  count_star INTEGER NOT NULL ," \
+  "  first_seen INTEGER NOT NULL ," \
+  "  last_seen INTEGER NOT NULL ," \
+  "  sum_time INTEGER NOT NULL ," \
+  "  min_time INTEGER NOT NULL ," \
+  "  max_time INTEGER NOT NULL ," \
+  "  PRIMARY KEY(tool_name, run_id, digest)" \
+  ")"
+
+// MCP query rules statistics table - shows hit counters for each rule
+// This table contains only rule_id and hits count.
+// It is automatically populated when stats_mcp_query_rules is queried.
+// The hits counter increments each time a rule matches during query processing.
+#define STATS_SQLITE_TABLE_MCP_QUERY_RULES "CREATE TABLE stats_mcp_query_rules (" \
+  "  rule_id INTEGER PRIMARY KEY NOT NULL ," \
+  "  hits INTEGER NOT NULL" \
+  ")"
+
 //#define STATS_SQLITE_TABLE_MEMORY_METRICS "CREATE TABLE stats_memory_metrics (Variable_Name VARCHAR NOT NULL PRIMARY KEY , Variable_Value VARCHAR NOT NULL)"
 
 

include/proxysql_admin.h

@@ -643,6 +643,10 @@ class ProxySQL_Admin {
 	void save_mysql_firewall_whitelist_rules_from_runtime(bool, SQLite3_result *);
 	void save_mysql_firewall_whitelist_sqli_fingerprints_from_runtime(bool, SQLite3_result *);
 
+	// MCP query rules
+	char* load_mcp_query_rules_to_runtime();
+	void save_mcp_query_rules_from_runtime(bool _runtime = false);
+
 	char* load_pgsql_firewall_to_runtime();
 
 	void load_scheduler_to_runtime();
@@ -700,6 +704,8 @@ class ProxySQL_Admin {
 	void stats___mysql_gtid_executed();
 	void stats___mysql_client_host_cache(bool reset);
 	void stats___mcp_query_tools_counters(bool reset);
+	void stats___mcp_query_digest(bool reset);
+	void stats___mcp_query_rules();
 
 	// Update prometheus metrics
 	void p_stats___memory_metrics();

lib/Admin_Bootstrap.cpp

@@ -810,6 +810,12 @@ bool ProxySQL_Admin::init(const bootstrap_info_t& bootstrap_info) {
 	insert_into_tables_defs(tables_defs_admin, "pgsql_firewall_whitelist_sqli_fingerprints", ADMIN_SQLITE_TABLE_PGSQL_FIREWALL_WHITELIST_SQLI_FINGERPRINTS);
 	insert_into_tables_defs(tables_defs_admin, "runtime_pgsql_firewall_whitelist_sqli_fingerprints", ADMIN_SQLITE_TABLE_RUNTIME_PGSQL_FIREWALL_WHITELIST_SQLI_FINGERPRINTS);
 
+	// MCP query rules
+	insert_into_tables_defs(tables_defs_admin, "mcp_query_rules", ADMIN_SQLITE_TABLE_MCP_QUERY_RULES);
+	insert_into_tables_defs(tables_defs_admin, "runtime_mcp_query_rules", ADMIN_SQLITE_TABLE_RUNTIME_MCP_QUERY_RULES);
+
+	insert_into_tables_defs(tables_defs_config, "mcp_query_rules", ADMIN_SQLITE_TABLE_MCP_QUERY_RULES);
+
 	insert_into_tables_defs(tables_defs_config, "pgsql_servers", ADMIN_SQLITE_TABLE_PGSQL_SERVERS);
 	insert_into_tables_defs(tables_defs_config, "pgsql_users", ADMIN_SQLITE_TABLE_PGSQL_USERS);
 	insert_into_tables_defs(tables_defs_config, "pgsql_ldap_mapping", ADMIN_SQLITE_TABLE_PGSQL_LDAP_MAPPING);
@@ -902,6 +908,11 @@ bool ProxySQL_Admin::init(const bootstrap_info_t& bootstrap_info) {
 	insert_into_tables_defs(tables_defs_stats,"stats_mcp_query_tools_counters", STATS_SQLITE_TABLE_MCP_QUERY_TOOLS_COUNTERS);
 	insert_into_tables_defs(tables_defs_stats,"stats_mcp_query_tools_counters_reset", STATS_SQLITE_TABLE_MCP_QUERY_TOOLS_COUNTERS_RESET);
 
+	// MCP query digest stats
+	insert_into_tables_defs(tables_defs_stats,"stats_mcp_query_digest", STATS_SQLITE_TABLE_MCP_QUERY_DIGEST);
+	insert_into_tables_defs(tables_defs_stats,"stats_mcp_query_digest_reset", STATS_SQLITE_TABLE_MCP_QUERY_DIGEST_RESET);
+	insert_into_tables_defs(tables_defs_stats,"stats_mcp_query_rules", STATS_SQLITE_TABLE_MCP_QUERY_RULES); // Reuse same schema for stats
+
 	// init ldap here
 	init_ldap();