CVE-2011-0448

@@ -251,6 +251,21 @@ def limited_update_conditions(where_sql, quoted_table_name, quoted_primary_key)
         "WHERE #{quoted_primary_key} IN (SELECT #{quoted_primary_key} FROM #{quoted_table_name} #{where_sql})"
       end
 
+      # Sanitizes the given LIMIT parameter in order to prevent SQL injection.
+      #
+      # +limit+ may be anything that can evaluate to a string via #to_s. It
+      # should look like an integer, or a comma-delimited list of integers.
+      #
+      # Returns the sanitized limit parameter, either as an integer, or as a
+      # string which contains a comma-delimited list of integers.
+      def sanitize_limit(limit)
+        if limit.to_s =~ /,/
+          Arel.sql limit.to_s.split(',').map{ |i| Integer(i) }.join(',')
+        else
+          Integer(limit)
+        end
+      end
+
       protected
         # Returns an array of record hashes with the column names as keys and
         # column values as values.
@@ -274,21 +289,6 @@ def delete_sql(sql, name = nil)
           update_sql(sql, name)
         end
 
-        # Sanitizes the given LIMIT parameter in order to prevent SQL injection.
-        #
-        # +limit+ may be anything that can evaluate to a string via #to_s. It
-        # should look like an integer, or a comma-delimited list of integers.
-        #
-        # Returns the sanitized limit parameter, either as an integer, or as a
-        # string which contains a comma-delimited list of integers.
-        def sanitize_limit(limit)
-          if limit.to_s =~ /,/
-            limit.to_s.split(',').map{ |i| i.to_i }.join(',')
-          else
-            limit.to_i
-          end
-        end
-
         # Send a rollback message to all records after they have been rolled back. If rollback
         # is false, only rollback records since the last save point.
         def rollback_transaction_records(rollback) #:nodoc

@@ -180,7 +180,7 @@ def build_arel
 
       arel = arel.having(*@having_values.uniq.reject{|h| h.blank?}) unless @having_values.empty?
 
-      arel = arel.take(@limit_value) if @limit_value
+      arel = arel.take(connection.sanitize_limit(@limit_value)) if @limit_value
       arel = arel.skip(@offset_value) if @offset_value
 
       arel = arel.group(*@group_values.uniq.reject{|g| g.blank?}) unless @group_values.empty?

@@ -141,16 +141,4 @@ def test_foreign_key_violations_are_translated_to_specific_exception
       end
     end
   end
-
-  def test_add_limit_offset_should_sanitize_sql_injection_for_limit_without_comas
-    sql_inject = "1 select * from schema"
-    assert_no_match(/schema/, @connection.add_limit_offset!("", :limit=>sql_inject))
-    assert_no_match(/schema/, @connection.add_limit_offset!("", :limit=>sql_inject, :offset=>7))
-  end
-
-  def test_add_limit_offset_should_sanitize_sql_injection_for_limit_with_comas
-    sql_inject = "1, 7 procedure help()"
-    assert_no_match(/procedure/, @connection.add_limit_offset!("", :limit=>sql_inject))
-    assert_no_match(/procedure/, @connection.add_limit_offset!("", :limit=>sql_inject, :offset=>7))
-  end
 end

@@ -48,6 +48,45 @@ class Boolean < ActiveRecord::Base; end
 class BasicsTest < ActiveRecord::TestCase
   fixtures :topics, :companies, :developers, :projects, :computers, :accounts, :minimalistics, 'warehouse-things', :authors, :categorizations, :categories, :posts
 
+  def test_limit_with_comma
+    assert_nothing_raised do
+      Topic.limit("1,2").all
+    end
+  end
+
+  def test_limit_without_comma
+    assert_nothing_raised do
+      assert_equal 1, Topic.limit("1").all.length
+    end
+
+    assert_nothing_raised do
+      assert_equal 1, Topic.limit(1).all.length
+    end
+  end
+
+  def test_invalid_limit
+    assert_raises(ArgumentError) do
+      Topic.limit("asdfadf").all
+    end
+  end
+
+  def test_limit_should_sanitize_sql_injection_for_limit_without_comas
+    assert_raises(ArgumentError) do
+      Topic.limit("1 select * from schema").all
+    end
+  end
+
+  def test_limit_should_sanitize_sql_injection_for_limit_with_comas
+    assert_raises(ArgumentError) do
+      Topic.limit("1, 7 procedure help()").all
+    end
+  end
+
+  def test_select_symbol
+    topic_ids = Topic.select(:id).map(&:id).sort
+    assert_equal Topic.find(:all).map(&:id).sort, topic_ids
+  end
+
   def test_table_exists
     assert !NonExistentTable.table_exists?
     assert Topic.table_exists?