rails_agent_console 0.1.0

data/lib/rails_agent_console/query_validator.rb

@@ -0,0 +1,305 @@
+# frozen_string_literal: true
+
+require "ripper"
+require "set"
+
+require_relative "query_validator/model_suggestion"
+require_relative "query_validator/sensitive_columns"
+require_relative "query_validator/rules"
+
+module RailsAgentConsole
+  class QueryValidator
+    include Rules
+
+    def self.validate(code, write_mode: RailsAgentConsole.configuration.write_mode?)
+      new(code, write_mode: write_mode).validate
+    end
+
+    def initialize(code, write_mode: false)
+      @code = code.to_s
+      @write_mode = write_mode
+      @violations = []
+      @methods = []
+      @constants = []
+      @writes = []
+      @unknown = []
+      @unknown_constants = []
+      @bare = []
+      @blocked = false
+      @syntax_error = false
+      @locals = Set.new
+    end
+
+    def validate
+      if @code.strip.empty?
+        @violations << "empty code"
+        return result
+      end
+
+      sexp = Ripper.sexp(@code)
+      if sexp.nil?
+        @syntax_error = true
+        @violations << syntax_violation
+        return result
+      end
+
+      collect_locals(sexp)
+      inspect_node(sexp)
+      classify_methods
+      check_sensitive_columns
+      result
+    end
+
+    private
+
+    def syntax_violation
+      message = ParseError.message_for(@code)
+      message ? "the code is not valid Ruby: #{message}" : "the code is not valid Ruby"
+    end
+
+    def result
+      Result.new(code: @code, violations: @violations, writes: @writes,
+                 methods: @methods, constants: @constants, unknown: @unknown,
+                 unknown_constants: @unknown_constants, blocked: @blocked,
+                 syntax_error: @syntax_error)
+    end
+
+    def collect_locals(node, inside_params: false)
+      return unless node.is_a?(Array)
+
+      case node.first
+      when :var_field, :params, :block_var, :rest_param, :kwrest_param, :blockarg
+        collect_identifiers(node)
+        inside_params = true
+      end
+
+      node.each { |child| collect_locals(child, inside_params: inside_params) if child.is_a?(Array) }
+    end
+
+    def collect_identifiers(node)
+      return unless node.is_a?(Array)
+
+      @locals << node[1] if node.first == :@ident && node[1].is_a?(String)
+      node.each { |child| collect_identifiers(child) if child.is_a?(Array) }
+    end
+
+    def inspect_node(node)
+      return unless node.is_a?(Array)
+      return if handled_as_data?(node)
+      return if handled_as_declaration?(node)
+      return if handled_as_call?(node)
+
+      inspect_children(node)
+    end
+
+    def handled_as_data?(node)
+      case node.first
+      when :symbol_literal, :@label
+        true
+      when :@tstring_content
+        check_string(node[1])
+        true
+      when :@const
+        @constants << node[1]
+        true
+      else
+        false
+      end
+    end
+
+    def handled_as_declaration?(node)
+      case node.first
+      when :xstring_literal
+        block("shell command execution via backticks")
+      when :def, :defs, :class, :module, :sclass
+        block("defining methods or classes is not allowed in a generated query")
+      when :@ivar, :@gvar, :@cvar
+        block("access to variable #{node[1]}")
+      when :field
+        block("attribute assignment")
+        inspect_children(node)
+      when :super, :zsuper
+        block("super")
+      else
+        return false
+      end
+      true
+    end
+
+    def handled_as_call?(node)
+      case node.first
+      when :call
+        symbol_receiver(node)
+        record_message(node[3], receiver: node[1])
+        inspect_node(node[1])
+      when :command_call
+        record_message(node[3], receiver: node[1])
+        inspect_node(node[1])
+        inspect_node(node[4])
+      when :fcall
+        record_message(node[1])
+      when :command
+        record_message(node[1])
+        inspect_node(node[2])
+      when :vcall
+        record_bare(node[1])
+      else
+        return false
+      end
+      true
+    end
+
+    def symbol_receiver(node)
+      return unless node[1].is_a?(Array) && node[1].first == :symbol_literal
+      return unless node[3].is_a?(Array) && node[3].first == :@ident
+      return if %w[to_s to_sym to_proc].include?(node[3][1])
+
+      name = node[1].flatten.find { |part| part.is_a?(String) }
+      violate("`:#{name}.#{node[3][1]}` calls #{node[3][1]} on a Symbol, not on a table. " \
+              "A column of a joined table is written in SQL: where(\"#{name.pluralize}.#{node[3][1]} = ...\")")
+    end
+
+    def inspect_children(node)
+      node.each { |child| inspect_node(child) if child.is_a?(Array) }
+    end
+
+    def record_message(message, receiver: nil)
+      return @methods << "call" if message == :call
+      return unless message.is_a?(Array)
+
+      case message.first
+      when :@ident, :@op, :@kw
+        @methods << message[1] unless read_on_receiver?(message[1], receiver)
+      when :@const then @constants << message[1]
+      else inspect_node(message)
+      end
+    end
+
+    def read_on_receiver?(name, receiver)
+      value_construction?(name, receiver) || relation_load?(name, receiver)
+    end
+
+    def relation_load?(name, receiver)
+      return false unless name == "load" && !receiver.nil?
+
+      constant = constant_name(receiver)
+      constant.nil? || !UNSAFE_LOAD_RECEIVERS.include?(constant)
+    end
+
+    def value_construction?(name, receiver)
+      return false unless CONSTRUCTORS.include?(name)
+
+      constant = constant_name(receiver)
+      !constant.nil? && VALUE_CONSTANTS.include?(constant) && !model_constants.include?(constant)
+    end
+
+    def constant_name(node)
+      return unless node.is_a?(Array)
+
+      case node.first
+      when :var_ref, :top_const_ref then constant_name(node[1])
+      when :const_path_ref then constant_name(node[2])
+      when :@const then node[1]
+      end
+    end
+
+    def record_bare(message)
+      return unless message.is_a?(Array) && message.first == :@ident
+      return if @locals.include?(message[1])
+
+      @methods << message[1]
+      @bare << message[1]
+    end
+
+    def check_string(value)
+      block("raw SQL that modifies data: #{value.strip.inspect}") if value.to_s.match?(DANGEROUS_SQL)
+    end
+
+    def classify_methods
+      @methods.uniq.each do |name|
+        if FORBIDDEN.include?(name)
+          block("`#{name}` is never allowed from the agent console")
+        elsif WRITE.include?(name)
+          @writes << name
+          violate("`#{name}` writes to the database (read-only mode)") unless @write_mode
+        elsif allowed?(name)
+          next
+        else
+          @unknown << name
+          violate("`#{name}` is not on the read-only allowlist")
+        end
+      end
+
+      @constants.uniq.each do |name|
+        next if allowed_constant?(name)
+
+        if BLOCKED_CONSTANTS.include?(name)
+          block("constant `#{name}` is never allowed from the agent console")
+        else
+          @unknown_constants << name
+          violate("constant `#{name}` is not allowed#{ModelSuggestion.for(name)}")
+        end
+      end
+    end
+
+    def allowed?(name)
+      return allowed_bare?(name) if @bare.include?(name)
+
+      ALLOWED.include?(name) || schema_readers.include?(name) ||
+        extra_allowed_methods.include?(name) ||
+        (name.end_with?("?") && ALLOWED.include?(name.chomp("?")))
+    end
+
+    def allowed_bare?(name)
+      ALLOWED.include?(name) || extra_allowed_methods.include?(name) ||
+        (name.end_with?("?") && ALLOWED.include?(name.chomp("?")))
+    end
+
+    def schema_readers
+      @schema_readers ||= Set.new(
+        Schema.models.flat_map do |model|
+          model.column_names + model.column_names.map { |column| "#{column}?" } + model.association_names
+        end
+      )
+    rescue StandardError
+      Set.new
+    end
+
+    def allowed_constant?(name)
+      SAFE_CONSTANTS.include?(name) || model_constants.include?(name) ||
+        extra_allowed_constants.include?(name)
+    end
+
+    def model_constants
+      @model_constants ||= Set.new(Schema.model_names.flat_map { |model| model.split("::") })
+    rescue StandardError
+      Set.new
+    end
+
+    def extra_allowed_methods
+      Set.new(RailsAgentConsole.configuration.extra_allowed_methods.map(&:to_s))
+    end
+
+    def extra_allowed_constants
+      Set.new(RailsAgentConsole.configuration.extra_allowed_constants.map(&:to_s))
+    end
+
+    def check_sensitive_columns
+      SensitiveColumns.named_in(@code).each do |name|
+        block("`#{name}` holds a credential and is never read by the agent")
+      end
+    end
+
+    def violate(message)
+      @violations << message
+    end
+
+    def block(message)
+      @blocked = true
+      violate(message)
+    end
+  end
+end
+
+require_relative "query_validator/parse_error"
+require_relative "query_validator/result"

data/lib/rails_agent_console/railtie.rb

@@ -0,0 +1,26 @@
+# frozen_string_literal: true
+
+require "rails/railtie"
+
+module RailsAgentConsole
+  class Railtie < ::Rails::Railtie
+    config.rails_agent_console = ActiveSupport::OrderedOptions.new
+
+    initializer "rails_agent_console.configure" do |app|
+      app.config.rails_agent_console.each do |key, value|
+        setter = "#{key}="
+        configuration = RailsAgentConsole.configuration
+        configuration.public_send(setter, value) if configuration.respond_to?(setter)
+      end
+    end
+
+    console do
+      require "irb"
+      require "rails_agent_console/console_installer"
+      require "rails_agent_console/console_session"
+
+      RailsAgentConsole::ConsoleInstaller.install!
+      RailsAgentConsole::ConsoleSession.start
+    end
+  end
+end

data/lib/rails_agent_console/rewriter/arguments.rb

@@ -0,0 +1,192 @@
+# frozen_string_literal: true
+
+module RailsAgentConsole
+  module Rewriter
+    class Arguments
+      ARGUMENTS = /\A\.\s*\w+[?!]?\s*\((.*)\)\s*\z/m
+      STRING_ARGUMENT = /\A\s*("(?:[^"\\]|\\.)*"|'(?:[^'\\]|\\.)*')\s*\z/
+      COLUMN_PAIR = /\A\s*(\w+):\s*([a-z_]\w*)\s*\z/
+      SYMBOL_PAIR = /\A\s*(\w+):\s*:(\w+)\s*\z/
+      SINGLE_PAIR = /\A\s*(\w+):\s*(.+?)\s*\z/m
+      HASH_KEY = /(?<![\w.:])(\w+):\s*\{/
+      ORDER_COLUMN = /\A\s*:?(\w+)(?::\s*:(asc|desc))?\s*\z/i
+      WEEKDAY = /\A\s*(\w+):\s*(.*\bwday\b.*?)\s*\z/m
+      WEEKDAY_TEST = /(?:Time\.(?:current|now)|Date\.(?:current|today))\.wday\s*==\s*([0-6])/
+      RAW_SQL_CALLS = %w[pluck pick order reorder].to_set.freeze
+      COMPARE_CALLS = %w[where not rewhere].to_set.freeze
+      DIRECTIONS = %w[asc desc].to_set.freeze
+      VALUES = %w[true false nil].to_set.freeze
+
+      NOTES = {
+        compare: "a hash value is a Ruby value, so comparing two columns is written in SQL",
+        weekday: "Time.current.wday is today's weekday, worked out once; the day of each row is EXTRACT(DOW ...)",
+        order: "ordering by a joined table's column is written as table.column",
+        joined_key: "the column belongs to a joined table, so the condition names that table",
+        table_key: "a condition names the joined table rather than the association, " \
+                   "so every reference to it agrees",
+        order_in_group: "a grouped query can only order by what it groups or aggregates, so it orders by MIN(...)",
+        raw_sql: "Rails only accepts an SQL expression in pluck and order through Arel.sql",
+        count: "counting a joined table counts its id",
+        nest: "an association of a joined model is joined through it",
+        bare: "an association is named with a symbol, not a bare word"
+      }.freeze
+      ASSOCIATION_CALLS = (JOINS.keys + %w[includes preload eager_load]).to_set.freeze
+      BARE_WORDS = /\A\s*[a-z_]\w*(?:\s*,\s*[a-z_]\w*)*\s*\z/
+
+      def initialize(chain, model)
+        @chain = chain
+        @model = model
+        @fixed = []
+      end
+
+      def rewrite
+        calls = nest_joins(@chain.calls)
+        tables = Joined.tables(calls, @model)
+        @grouped = calls.select { |call| call.name == "group" }.map(&:symbol_arguments)
+        calls = calls.map { |call| fix(call, tables) }
+        return if @fixed.empty?
+
+        [@chain.with_calls(calls), "#{@fixed.uniq.map { |key| NOTES[key] }.join("; ").sub(/\A./, &:upcase)}."]
+      end
+
+      private
+
+      def bare_association(call, inner)
+        return unless ASSOCIATION_CALLS.include?(call.name) && inner.match?(BARE_WORDS)
+
+        names = inner.split(",").map(&:strip)
+        return unless (names - @model.association_names).empty?
+
+        @fixed << :bare
+        names.map { |name| ":#{name}" }.join(", ")
+      end
+
+      def fix(call, tables)
+        original = call.text[ARGUMENTS, 1]
+        return call if original.nil?
+
+        inner = table_keys(call, original, tables)
+        replacement = bare_association(call, inner) || compare_columns(call, inner) || weekday(call, inner) ||
+                      joined_key(call, inner, tables) ||
+                      order_through(call, inner, tables) || order_in_group(call, inner) || wrap_sql(call, inner) ||
+                      count_joined(call, inner, tables) || (inner unless inner == original)
+        replacement ? Chain::Call.new(call.name, ".#{call.name}(#{replacement})") : call
+      end
+
+      def table_keys(call, inner, tables)
+        return inner unless COMPARE_CALLS.include?(call.name)
+
+        inner.gsub(HASH_KEY) do |key|
+          table = tables[Regexp.last_match(1)]
+          next key if table.nil? || table == Regexp.last_match(1)
+
+          @fixed << :table_key
+          "#{table}: {"
+        end
+      end
+
+      def compare_columns(call, inner)
+        left, right = inner.match(COLUMN_PAIR)&.captures
+        return unless COMPARE_CALLS.include?(call.name) && left && !VALUES.include?(right)
+        return unless ([left, right] - @model.column_names).empty?
+
+        fixed(:compare, %("#{@model.table}.#{left} = #{@model.table}.#{right}"))
+      end
+
+      def weekday(call, inner)
+        column, condition = inner.match(WEEKDAY)&.captures
+        return unless COMPARE_CALLS.include?(call.name) && column && @model.column_names.include?(column)
+        return unless condition.gsub(WEEKDAY_TEST, "").match?(/\A[\s()|]*\z/)
+
+        days = condition.scan(WEEKDAY_TEST).flatten.uniq.join(", ")
+        fixed(:weekday, %("EXTRACT(DOW FROM #{@model.table}.#{column}) IN (#{days})"))
+      end
+
+      def joined_key(call, inner, tables)
+        key, value = inner.match(SINGLE_PAIR)&.captures
+        return unless COMPARE_CALLS.include?(call.name) && key && !@model.column_names.include?(key)
+        return unless Chain::Scanner.single_argument?(value)
+
+        owners = tables.values.uniq.select { |table| Joined.model_for(table)&.column_names&.include?(key) }
+        fixed(:joined_key, "#{owners.first}: { #{key}: #{value} }") if owners.size == 1
+      end
+
+      def order_through(call, inner, tables)
+        name, column = inner.match(SYMBOL_PAIR)&.captures
+        return unless RAW_SQL_CALLS.include?(call.name) && name && !DIRECTIONS.include?(column.downcase)
+
+        table = tables[name] || (name if tables.value?(name))
+        return unless table && Joined.model_for(table)&.column_names&.include?(column)
+
+        fixed(:order, %("#{table}.#{column}"))
+      end
+
+      def order_in_group(call, inner)
+        column, direction = inner.match(ORDER_COLUMN)&.captures
+        return unless %w[order reorder].include?(call.name) && column && grouped_elsewhere?(column)
+
+        suffix = direction ? " #{direction.upcase}" : ""
+        fixed(:order_in_group, %(Arel.sql("MIN(#{@model.table}.#{column})#{suffix}")))
+      end
+
+      def grouped_elsewhere?(column)
+        return false if @grouped.empty? || @grouped.any?(&:nil?)
+
+        !@grouped.flatten.include?(column) && @model.column_names.include?(column)
+      end
+
+      def wrap_sql(call, inner)
+        literal = inner[STRING_ARGUMENT, 1]
+        return unless RAW_SQL_CALLS.include?(call.name) && literal&.include?("(")
+
+        fixed(:raw_sql, "Arel.sql(#{literal})")
+      end
+
+      def count_joined(call, inner, tables)
+        name = inner[/\A\s*:(\w+)\s*\z/, 1]
+        return unless call.name == "count" && name && !@model.column_names.include?(name)
+
+        table = joined_table(name.delete_suffix("_id"), tables)
+        table && fixed(:count, %("#{table}.id"))
+      end
+
+      def joined_table(name, tables)
+        tables[name] || tables.values.find { |table| [name, name.pluralize].include?(table) }
+      end
+
+      def nest_joins(calls)
+        joined = []
+        calls.map do |call|
+          names = JOINS.key?(call.name) && call.symbol_arguments
+          next call unless names
+
+          nested = names.map { |name| nest(name, joined) }
+          joined.concat(names)
+          nested == names ? call : nested_call(call, names, nested)
+        end
+      end
+
+      def nest(name, joined)
+        return name if Joined.target(@model, name) || @model.inbound.any? { |ref| ref[:table].to_s == name }
+
+        parent = joined.find do |candidate|
+          through = Joined.target(@model, candidate)
+          through && Joined.target(through, name)
+        end
+        parent ? "#{parent}: :#{name}" : name
+      end
+
+      def nested_call(call, names, nested)
+        @fixed << :nest
+        plain = nested.select.with_index { |item, index| item == names[index] }.map { |name| ":#{name}" }
+        pairs = nested.reject.with_index { |item, index| item == names[index] }
+        Chain::Call.new(call.name, ".#{call.name}(#{(plain + pairs).join(", ")})")
+      end
+
+      def fixed(key, replacement)
+        @fixed << key
+        replacement
+      end
+    end
+  end
+end

data/lib/rails_agent_console/rewriter/chain.rb

@@ -0,0 +1,126 @@
+# frozen_string_literal: true
+
+module RailsAgentConsole
+  module Rewriter
+    class Chain
+      RECEIVER = /\A[A-Za-z_]\w*(?:::[A-Z]\w*)*\z/
+      CALL_NAME = /\A\.\s*(\w+[?!]?)/
+
+      Call = Struct.new(:name, :text) do
+        def symbol_arguments
+          inner = text[/\A\.\s*\w+[?!]?\s*\((.*)\)\s*\z/m, 1]
+          return unless inner&.match?(/\A\s*:\w+(?:\s*,\s*:\w+)*\s*\z/)
+
+          inner.scan(/:(\w+)/).flatten
+        end
+      end
+
+      attr_reader :receiver, :calls
+
+      def self.parse(code)
+        text = code.to_s.strip
+        return if text.empty? || text.match?(/\bdo\b|;/)
+
+        pieces = Scanner.new(text).pieces
+        return if pieces.nil? || pieces.size < 2 || !pieces.first.match?(RECEIVER)
+
+        calls = pieces.drop(1).map { |piece| Call.new(piece[CALL_NAME, 1], piece) }
+        return unless calls.all? { |call| call.name && Scanner.call_shape?(call.text.sub(CALL_NAME, "")) }
+
+        new(pieces.first, calls)
+      end
+
+      def initialize(receiver, calls)
+        @receiver = receiver
+        @calls = calls
+      end
+
+      def with_calls(calls)
+        self.class.new(receiver, calls)
+      end
+
+      def to_s
+        receiver + calls.map(&:text).join
+      end
+
+      class Scanner
+        def self.call_shape?(rest)
+          new(rest).call_shape?
+        end
+
+        def self.single_argument?(text)
+          new(text).single_argument?
+        end
+
+        def initialize(text)
+          @text = text
+          @depth = 0
+          @quote = nil
+          @escaped = false
+        end
+
+        def pieces
+          cuts = []
+          @text.each_char.with_index do |char, index|
+            next if in_string?(char)
+            return nil unless balanced_after?(char)
+
+            cuts << index if char == "." && @depth.zero? && call_dot?(index)
+          end
+          return unless @depth.zero? && @quote.nil?
+
+          ([0] + cuts + [@text.length]).each_cons(2).map { |from, to| @text[from...to].strip }
+        end
+
+        def call_shape?
+          @text.each_char do |char|
+            return false if @quote.nil? && @depth.zero? && !char.match?(/[\s(){}]/)
+            next if in_string?(char)
+            return false unless balanced_after?(char)
+          end
+          @depth.zero? && @quote.nil?
+        end
+
+        def single_argument?
+          @text.each_char do |char|
+            next if in_string?(char)
+            return false unless balanced_after?(char)
+            return false if char == "," && @depth.zero?
+          end
+          @depth.zero? && @quote.nil?
+        end
+
+        private
+
+        def in_string?(char)
+          return string_opened?(char) if @quote.nil?
+
+          if @escaped then @escaped = false
+          elsif char == "\\" then @escaped = true
+          elsif char == @quote then @quote = nil
+          end
+          true
+        end
+
+        def string_opened?(char)
+          return false unless ["'", '"'].include?(char)
+
+          @quote = char
+          true
+        end
+
+        def balanced_after?(char)
+          @depth += 1 if "([{".include?(char)
+          @depth -= 1 if ")]}".include?(char)
+          @depth >= 0
+        end
+
+        def call_dot?(index)
+          before = index.positive? ? @text[index - 1] : ""
+          after = @text[index + 1].to_s
+          before != "." && after != "." && !before.match?(/\d/) && after.match?(/[\s\w]/)
+        end
+      end
+    end
+  end
+end

data/lib/rails_agent_console/rewriter/date_ranges.rb

@@ -0,0 +1,21 @@
+# frozen_string_literal: true
+
+module RailsAgentConsole
+  module Rewriter
+    module DateRanges
+      INCLUSIVE_END = /
+        (?<![.])\.\.\s*
+        (?<date>Time\.zone\.parse\((?<quote>["'])\d{4}-\d{2}-\d{2}\k<quote>\))
+        (?!\s*\.\s*end_of_day)
+      /x
+      NOTE = "A range that ended on a date stopped at midnight and missed that day, so it now runs to end_of_day."
+
+      module_function
+
+      def fix(code)
+        fixed = code.gsub(INCLUSIVE_END) { "..#{Regexp.last_match[:date]}.end_of_day" }
+        fixed == code ? [code, nil] : [fixed, NOTE]
+      end
+    end
+  end
+end