query_guard 0.4.2 → 0.5.0

data/lib/query_guard/migrations/table_size_resolver.rb

@@ -0,0 +1,152 @@
+# frozen_string_literal: true
+
+module QueryGuard
+  module Migrations
+    # Extracts table names from Rails migration code
+    #
+    # Pragmatically identifies which tables a migration affects
+    # by parsing common Rails migration method calls.
+    #
+    # Detects: add_column, remove_column, change_column, create_table, drop_table,
+    # etc. on both active table operations and raw SQL.
+    class TableSizeResolver
+      # Extract all table names referenced in a migration
+      #
+      # @param migration_content [String] Migration file content
+      # @return [Array<String>] Table names (deduplicated)
+      def self.extract_table_names(migration_content)
+        tables = Set.new
+
+        lines = migration_content.lines
+
+        lines.each do |line|
+          next if line.strip.start_with?("#")
+          next if line.strip.empty?
+
+          # Try to capture table name: could be :symbol or "string" or 'string'
+          # add_column :users, :name, :string
+          if line.include?("add_column")
+            # Match: add_column :table_name or add_column "table_name" or add_column 'table_name'
+            if match = line.match(/add_column\s+[:"`']?(\w+)/)
+              tables << match[1]
+            end
+          end
+
+          # remove_column :posts, :old_field
+          if line.include?("remove_column")
+            if match = line.match(/remove_column\s+[:"`']?(\w+)/)
+              tables << match[1]
+            end
+          end
+
+          # change_column :users, :email, :text
+          if line.include?("change_column")
+            if match = line.match(/change_column\s+[:"`']?(\w+)/)
+              tables << match[1]
+            end
+          end
+
+          # rename_column :users, :old_name, :new_name
+          if line.include?("rename_column")
+            if match = line.match(/rename_column\s+[:"`']?(\w+)/)
+              tables << match[1]
+            end
+          end
+
+          # add_index :users, :email
+          if line.include?("add_index")
+            if match = line.match(/add_index\s+[:"`']?(\w+)/)
+              tables << match[1]
+            end
+          end
+
+          # remove_index :users, :email
+          if line.include?("remove_index")
+            if match = line.match(/remove_index\s+[:"`']?(\w+)/)
+              tables << match[1]
+            end
+          end
+
+          # drop_table :posts
+          if line.include?("drop_table")
+            if match = line.match(/drop_table\s+[:"`']?(\w+)/)
+              tables << match[1]
+            end
+          end
+
+          # create_table :accounts do ... end
+          if line.include?("create_table")
+            if match = line.match(/create_table\s+[:"`']?(\w+)/)
+              tables << match[1]
+            end
+          end
+
+          # Raw SQL UPDATE public.users SET ...
+          if line.upcase.include?("UPDATE")
+            if match = line.match(/UPDATE\s+(?:public\.)?(\w+)/i)
+              tables << match[1]
+            end
+          end
+
+          # Raw SQL DELETE FROM public.users
+          if line.upcase.include?("DELETE")
+            if match = line.match(/DELETE\s+FROM\s+(?:public\.)?(\w+)/i)
+              tables << match[1]
+            end
+          end
+
+          # Raw SQL INSERT INTO public.users
+          if line.upcase.include?("INSERT")
+            if match = line.match(/INSERT\s+INTO\s+(?:public\.)?(\w+)/i)
+              tables << match[1]
+            end
+          end
+
+          # Model.update_all(status: 'active') - convert model to table name
+          # or User.delete_all
+          if (line.include?("update_all") || line.include?("delete_all"))
+            if match = line.match(/(\b[A-Z]\w*)\s*\.\s*(?:update_all|delete_all)/)
+              model_name = match[1]
+              # Convert CamelCase to snake_case: User -> user, UserProfile -> user_profile
+              table_name = model_name
+                .gsub(/([A-Z])/, '_\1')
+                .downcase
+                .sub(/^_/, '')
+              # Pluralize simple table names: add 's' for most words
+              # This is a simple heuristic; full pluralization would require a library
+              table_name = table_name + 's' unless table_name.end_with?('s', 'x', 'z')
+              tables << table_name if table_name && !table_name.empty?
+            end
+          end
+        end
+
+        tables.to_a.sort
+      end
+
+      # Filter tables to only those likely affected by schema changes
+      #
+      # Excludes internal Rails tables like schema_migrations
+      #
+      # @param table_names [Array<String>] List of table names
+      # @return [Array<String>] Filtered list
+      def self.filter_schema_tables(table_names)
+        # Internal rails tables to ignore
+        internal_tables = %w[
+          schema_migrations ar_internal_metadata
+          delayed_jobs sidekiq_jobs
+        ]
+
+        table_names.reject { |name| internal_tables.include?(name) }
+      end
+
+      # Extract table names that schema changes would directly affect
+      #
+      # @param migration_content [String] Migration file content
+      # @return [Array<String>] Directly affected tables
+      def self.affected_tables(migration_content)
+        all_tables = extract_table_names(migration_content)
+        filter_schema_tables(all_tables)
+      end
+    end
+  end
+end

data/lib/query_guard/publish.rb

@@ -0,0 +1,38 @@
+# frozen_string_literal: true
+require "time"
+require_relative "client"
+
+module QueryGuard
+  class << self
+    def client
+      @client ||= Client.new(
+        base_url: (config.api_base_url || ENV["QUERYGUARD_API_URL"] || "http://localhost:4000"),
+        api_key:  (config.api_key      || ENV["QUERYGUARD_API_KEY"]),
+        project:  (config.project      || ENV["QUERYGUARD_PROJECT"] || "dev"),
+        env:      (config.env          || (defined?(Rails) ? Rails.env : "development"))
+      )
+    end
+
+    # thread-local queue for the current request
+    def buffer; Thread.current[:qg_buf] ||= [] end
+    def clear!; Thread.current[:qg_buf]  = []  end
+
+    def track!(attrs)
+      return unless enabled?
+      buffer << default_fields.merge(attrs)
+    end
+
+    def flush_now!
+      return if buffer.empty?
+      client.post("/api/v1/events", { events: buffer.dup })
+    ensure
+      clear!
+    end
+
+    private
+
+    def default_fields
+      { occurred_at: Time.now.utc.iso8601 }
+    end
+  end
+end

data/lib/query_guard/rspec.rb

@@ -0,0 +1,119 @@
+# frozen_string_literal: true
+
+require "rspec/expectations"
+require "query_guard/trace"
+
+# RSpec matcher for asserting query budgets.
+#
+# Usage:
+#   expect {
+#     User.where(active: true).to_a
+#   }.to_not exceed_query_budget(count: 5)
+#
+#   expect {
+#     User.all.to_a
+#   }.to_not exceed_query_budget(count: 10, duration_ms: 500)
+#
+# With a labeled budget:
+#   QueryGuard.config.budget.for("users#index", count: 5)
+#   
+#   expect {
+#     # code
+#   }.to_not exceed_query_budget("users#index")
+RSpec::Matchers.define :exceed_query_budget do |budget_key_or_limits = nil, **limits|
+  match do |block|
+    # Determine if we're checking a named budget or inline limits
+    if budget_key_or_limits.is_a?(String)
+      budget_key = budget_key_or_limits
+      budget_limits = QueryGuard.config&.budget&.budget_for(budget_key)
+      
+      if budget_limits.nil?
+        raise ArgumentError, "No budget defined for '#{budget_key}'. Define it with QueryGuard.config.budget.for('#{budget_key}', ...)"
+      end
+      
+      limits = budget_limits
+    elsif budget_key_or_limits.is_a?(Hash)
+      limits = budget_key_or_limits.merge(limits)
+    end
+
+    # Execute the block within a trace
+    _, report = QueryGuard::Trace.trace("rspec_matcher") do
+      block.call
+    end
+
+    @actual_count = report.query_count
+    @actual_duration = report.total_duration_ms
+    @limit_count = limits[:count]
+    @limit_duration = limits[:duration_ms]
+
+    # Check if any budget was exceeded
+    @exceeded = false
+    @violations = []
+
+    if @limit_count && @actual_count > @limit_count
+      @exceeded = true
+      @violations << "count: #{@actual_count} > #{@limit_count}"
+    end
+
+    if @limit_duration && @actual_duration > @limit_duration
+      @exceeded = true
+      @violations << "duration: #{@actual_duration.round(2)}ms > #{@limit_duration}ms"
+    end
+
+    @exceeded
+  end
+
+  failure_message do
+    "expected query budget not to be exceeded, but it was:\n  #{@violations.join("\n  ")}"
+  end
+
+  failure_message_when_negated do
+    parts = []
+    parts << "count: #{@actual_count} <= #{@limit_count}" if @limit_count
+    parts << "duration: #{@actual_duration.round(2)}ms <= #{@limit_duration}ms" if @limit_duration
+    "expected query budget to be exceeded, but it was within limits:\n  #{parts.join("\n  ")}"
+  end
+
+  description do
+    parts = []
+    parts << "count <= #{@limit_count}" if @limit_count
+    parts << "duration <= #{@limit_duration}ms" if @limit_duration
+    "not exceed query budget (#{parts.join(", ")})"
+  end
+
+  supports_block_expectations
+end
+
+module QueryGuard
+  module RSpec
+    # Convenience helper to check if code stays within a query budget
+    def within_query_budget(**limits, &block)
+      _, report = QueryGuard::Trace.trace("within_query_budget") do
+        block.call
+      end
+
+      violations = []
+
+      if limits[:count] && report.query_count > limits[:count]
+        violations << "Query count exceeded: #{report.query_count} > #{limits[:count]}"
+      end
+
+      if limits[:duration_ms] && report.total_duration_ms > limits[:duration_ms]
+        violations << "Duration exceeded: #{report.total_duration_ms.round(2)}ms > #{limits[:duration_ms]}ms"
+      end
+
+      if violations.any?
+        raise QueryGuard::Budget::Violation, violations.join("; ")
+      end
+
+      report
+    end
+  end
+end
+
+# Auto-include helpers when RSpec is available
+if defined?(::RSpec)
+  ::RSpec.configure do |config|
+    config.include QueryGuard::RSpec
+  end
+end

data/lib/query_guard/security.rb

@@ -0,0 +1,99 @@
+# frozen_string_literal: true
+require "digest"
+
+module QueryGuard
+  module Security
+    module_function
+
+    # Normalize SQL into a stable fingerprint:
+    # - collapse whitespace
+    # - replace quoted strings + numbers with ?
+    def fingerprint(sql)
+      s = sql.to_s.dup
+      s.gsub!(/\s+/, " ")
+      s.gsub!(/'(?:''|[^'])*'/, "?") # strings
+      s.gsub!(/\b\d+\b/, "?")        # integers
+      Digest::SHA1.hexdigest(s.strip.downcase)
+    end
+
+    def suspicious_sql_injection?(sql, patterns)
+      s = sql.to_s
+      patterns.any? { |re| re.match?(s) }
+    end
+
+    def possible_exfiltration_query?(sql)
+      s = sql.to_s.strip
+      return false unless s =~ /\ASELECT\b/i
+      # Heuristic: SELECT without WHERE and without LIMIT
+      no_where = !s.match?(/\bwhere\b/i)
+      no_limit = !s.match?(/\blimit\b/i)
+      no_where && no_limit
+    end
+
+    def post_request_checks!(env, stats, config)
+      return unless config.enable_security
+
+      actor = resolve_actor(env, config)
+      store = config.store || QueryGuard::Store.new
+
+      # --- Unusual query pattern (rate/variety) ---
+      if config.detect_unusual_query_pattern
+        bucket = Time.now.utc.strftime("%Y%m%d%H%M") # minute bucket
+        base = "qg:actor:#{actor}:#{bucket}"
+
+        total = store.incr("#{base}:queries", ttl: 120, by: stats[:count].to_i)
+        uniq_count = stats[:fingerprints]&.keys&.size.to_i
+        store.add_to_set("#{base}:uniqfp", stats[:request_id], ttl: 120) # keep request marker
+
+        uniq_fp_total = store.incr("#{base}:uniqfp_count", ttl: 120, by: uniq_count)
+
+        if total > config.max_queries_per_minute_per_actor
+          stats[:violations] << {
+            type: :unusual_query_rate,
+            actor: actor,
+            per_minute: total,
+            limit: config.max_queries_per_minute_per_actor
+          }
+        end
+
+        if uniq_fp_total > config.max_unique_query_fingerprints_per_minute_per_actor
+          stats[:violations] << {
+            type: :unusual_query_variety,
+            actor: actor,
+            unique_fingerprints_per_minute: uniq_fp_total,
+            limit: config.max_unique_query_fingerprints_per_minute_per_actor
+          }
+        end
+      end
+
+      # --- Data exfiltration (response size + endpoint hint) ---
+      if config.detect_data_exfiltration
+        bytes = stats[:response_bytes].to_i
+        path  = env["PATH_INFO"].to_s
+
+        if bytes > config.max_response_bytes_per_request
+          stats[:violations] << {
+            type: :data_exfiltration_large_response,
+            bytes: bytes,
+            limit: config.max_response_bytes_per_request,
+            path: path
+          }
+        end
+
+        if bytes > (config.max_response_bytes_per_request / 2) && path.match?(config.exfiltration_path_regex)
+          stats[:violations] << {
+            type: :data_exfiltration_suspected_export,
+            bytes: bytes,
+            path: path
+          }
+        end
+      end
+    end
+
+    def resolve_actor(env, config)
+      (config.actor_resolver && config.actor_resolver.call(env)) || "unknown"
+    rescue
+      "unknown"
+    end
+  end
+end

data/lib/query_guard/store.rb

@@ -0,0 +1,38 @@
+# frozen_string_literal: true
+require "active_support"
+require "active_support/cache"
+
+module QueryGuard
+  # Minimal store abstraction for rate counters.
+  # Default: in-process MemoryStore. In production you can swap with Redis-based adapter.
+  class Store
+    def initialize(cache: nil)
+      @cache = cache || ActiveSupport::Cache::MemoryStore.new(size: 8.megabytes)
+    end
+
+    # Increment integer key, expiring after ttl seconds
+    def incr(key, ttl: 60, by: 1)
+      val = (@cache.read(key) || 0).to_i + by
+      @cache.write(key, val, expires_in: ttl)
+      val
+    end
+
+    def read(key)
+      @cache.read(key)
+    end
+
+    # "Set" emulation: store a Hash of members => true with ttl
+    def add_to_set(key, member, ttl: 60)
+      h = @cache.read(key)
+      h = {} unless h.is_a?(Hash)
+      h[member] = true
+      @cache.write(key, h, expires_in: ttl)
+      h.size
+    end
+
+    def set_size(key)
+      h = @cache.read(key)
+      h.is_a?(Hash) ? h.size : 0
+    end
+  end
+end

data/lib/query_guard/subscriber.rb

@@ -1,5 +1,6 @@
 # frozen_string_literal: true
 require "active_support/notifications"
+require "query_guard/security"
 
 module QueryGuard
   module Subscriber
@@ -7,12 +8,14 @@ module QueryGuard
 
     def self.install!(config)
       return if @installed
-      @config = config
+
       @subscriber = ActiveSupport::Notifications.subscribe(SQL_EVENT) do |_, started, finished, _, payload|
+        context = Thread.current[:query_guard_context]
+        next unless context # only track inside our middleware window
+
         stats = Thread.current[:query_guard_stats]
-        next unless stats # only track inside our middleware window
+        next unless stats
 
-        # Skip schema and ignored
         name = payload[:name].to_s
         next if name == "SCHEMA"
 
@@ -20,29 +23,57 @@ module QueryGuard
         next if config.ignored_sql.any? { |r| r === sql }
 
         duration_ms = (finished - started) * 1000.0
+        
+        # Collect query into context (new behavior)
+        context.add_query(
+          sql: sql,
+          duration_ms: duration_ms,
+          name: name,
+          started_at: Time.at(started),
+          finished_at: Time.at(finished)
+        )
+
+        # Legacy: Also update Thread.current stats for backward compatibility
+        stats = Thread.current[:query_guard_stats] ||= { count: 0, total_duration_ms: 0.0, violations: [] }
         stats[:count] += 1
         stats[:total_duration_ms] += duration_ms
 
+        fp = QueryGuard::Security.fingerprint(sql)
+        stats[:fingerprints][fp] += 1
+
         if config.max_duration_ms_per_query && duration_ms > config.max_duration_ms_per_query
-          stats[:violations] << {
-            type: :slow_query,
-            duration_ms: duration_ms.round(2),
-            sql: sql
-          }
+          stats[:violations] << { type: :slow_query, duration_ms: duration_ms.round(2), sql: sql }
         end
 
         if config.block_select_star && sql =~ /\bSELECT\s+\*/i
           stats[:violations] << { type: :select_star, sql: sql }
         end
+
+        # --- SQL Injection detection ---
+        if config.enable_security && config.detect_sql_injection
+          if QueryGuard::Security.suspicious_sql_injection?(sql, config.sql_injection_patterns)
+            stats[:violations] << { type: :sql_injection_suspected, sql: sql }
+          end
+        end
+
+        # --- Data exfiltration query-shape heuristic ---
+        if config.enable_security && config.detect_data_exfiltration
+          if QueryGuard::Security.possible_exfiltration_query?(sql)
+            stats[:violations] << { type: :possible_data_exfiltration_query, sql: sql }
+          end
+        end
+
+        max = config.max_query_events_per_req || 200
+        if stats[:queries].length < max
+          stats[:queries] << {
+            sql: sql,
+            duration_ms: duration_ms.round(2),
+            occurred_at: Time.now.utc.iso8601
+          }
+        end
       end
-      @installed = true
-    end
 
-    def self.uninstall!
-      return unless @installed && @subscriber
-      ActiveSupport::Notifications.unsubscribe(@subscriber)
-      @installed = false
-      @subscriber = nil
+      @installed = true
     end
   end
 end