MassAssignmentSecurity: add ability to specify your own sanitizer

Added an ability to specify your own behavior on mass assingment
protection, controlled by option:
ActiveModel::MassAssignmentSecurity.mass_assignment_sanitizer

7 files changed, 57 insertions, 39 deletions

diff --git a/activemodel/lib/active_model/mass_assignment_security.rb b/activemodel/lib/active_model/mass_assignment_security.rb
index 483b577681..cc30609f2b 100644
--- a/activemodel/lib/active_model/mass_assignment_security.rb
+++ b/activemodel/lib/active_model/mass_assignment_security.rb
@@ -1,5 +1,6 @@
 require 'active_support/core_ext/class/attribute.rb'
 require 'active_model/mass_assignment_security/permission_set'
+require 'active_model/mass_assignment_security/sanitizer'
 
 module ActiveModel
   # = Active Model Mass-Assignment Security
@@ -10,6 +11,7 @@ module ActiveModel
       class_attribute :_accessible_attributes
       class_attribute :_protected_attributes
       class_attribute :_active_authorizer
+      class_attribute :mass_assignment_sanitizer
     end
 
     # Mass assignment security provides an interface for protecting attributes
@@ -181,16 +183,14 @@ module ActiveModel
 
       def protected_attributes_configs
         self._protected_attributes ||= begin
-          default_black_list = BlackList.new(attributes_protected_by_default).tap do |w|
-            w.logger = self.logger if self.respond_to?(:logger)
-          end
+          default_black_list = BlackList.new(attributes_protected_by_default)
           Hash.new(default_black_list)
         end
       end
 
       def accessible_attributes_configs
         self._accessible_attributes ||= begin
-          default_white_list = WhiteList.new.tap { |w| w.logger = self.logger if self.respond_to?(:logger) }
+          default_white_list = WhiteList.new
           Hash.new(default_white_list)
         end
       end
@@ -199,7 +199,11 @@ module ActiveModel
   protected
 
     def sanitize_for_mass_assignment(attributes, role = :default)
-      mass_assignment_authorizer(role).sanitize(attributes)
+      (mass_assignment_sanitizer || default_mass_assignment_sanitizer).sanitize(attributes, mass_assignment_authorizer(role))
+    end
+
+    def default_mass_assignment_sanitizer
+      DefaultSanitizer.new(self.respond_to?(:logger) && self.logger)
     end
 
     def mass_assignment_authorizer(role = :default)
diff --git a/activemodel/lib/active_model/mass_assignment_security/permission_set.rb b/activemodel/lib/active_model/mass_assignment_security/permission_set.rb
index 9fcb94d48a..a1fcdf1a38 100644
--- a/activemodel/lib/active_model/mass_assignment_security/permission_set.rb
+++ b/activemodel/lib/active_model/mass_assignment_security/permission_set.rb
@@ -1,10 +1,8 @@
 require 'set'
-require 'active_model/mass_assignment_security/sanitizer'
 
 module ActiveModel
   module MassAssignmentSecurity
     class PermissionSet < Set
-      attr_accessor :logger
 
       def +(values)
         super(values.map(&:to_s))
@@ -14,6 +12,10 @@ module ActiveModel
         super(remove_multiparameter_id(key))
       end
 
+      def deny?(key)
+        raise NotImplementedError, "#deny?(key) suppose to be overwritten"
+      end
+
     protected
 
       def remove_multiparameter_id(key)
@@ -22,7 +24,6 @@ module ActiveModel
     end
 
     class WhiteList < PermissionSet
-      include Sanitizer
 
       def deny?(key)
         !include?(key)
@@ -30,7 +31,6 @@ module ActiveModel
     end
 
     class BlackList < PermissionSet
-      include Sanitizer
 
       def deny?(key)
         include?(key)
diff --git a/activemodel/lib/active_model/mass_assignment_security/sanitizer.rb b/activemodel/lib/active_model/mass_assignment_security/sanitizer.rb
index 150beb1ff2..5dbcf473bd 100644
--- a/activemodel/lib/active_model/mass_assignment_security/sanitizer.rb
+++ b/activemodel/lib/active_model/mass_assignment_security/sanitizer.rb
@@ -1,9 +1,9 @@
 module ActiveModel
   module MassAssignmentSecurity
-    module Sanitizer
+    class Sanitizer
       # Returns all attributes not denied by the authorizer.
-      def sanitize(attributes)
-        sanitized_attributes = attributes.reject { |key, value| deny?(key) }
+      def sanitize(attributes, authorizer)
+        sanitized_attributes = attributes.reject { |key, value| authorizer.deny?(key) }
         debug_protected_attribute_removal(attributes, sanitized_attributes)
         sanitized_attributes
       end
@@ -12,10 +12,24 @@ module ActiveModel
 
       def debug_protected_attribute_removal(attributes, sanitized_attributes)
         removed_keys = attributes.keys - sanitized_attributes.keys
-        warn!(removed_keys) if removed_keys.any?
+        process_removed_attributes(removed_keys) if removed_keys.any?
       end
+      
+      def process_removed_attributes(attrs)
+        raise NotImplementedError, "#process_removed_attributes(attrs) suppose to be overwritten"
+      end
+
+    end
+    class DefaultSanitizer < Sanitizer
 
-      def warn!(attrs)
+      attr_accessor :logger
+
+      def initialize(logger = nil)
+        self.logger = logger
+        super()
+      end
+      
+      def process_removed_attributes(attrs)
         self.logger.debug "WARNING: Can't mass-assign protected attributes: #{attrs.join(', ')}" if self.logger
       end
     end
diff --git a/activemodel/test/cases/mass_assignment_security/black_list_test.rb b/activemodel/test/cases/mass_assignment_security/black_list_test.rb
index ed168bc016..0ec7f8719c 100644
--- a/activemodel/test/cases/mass_assignment_security/black_list_test.rb
+++ b/activemodel/test/cases/mass_assignment_security/black_list_test.rb
@@ -16,13 +16,5 @@ class BlackListTest < ActiveModel::TestCase
     assert_equal false, @black_list.deny?('first_name')
   end
 
-  test "sanitize attributes" do
-    original_attributes = { 'first_name' => 'allowed', 'admin' => 'denied', 'admin(1)' => 'denied' }
-    attributes = @black_list.sanitize(original_attributes)
-
-    assert attributes.key?('first_name'), "Allowed key shouldn't be rejected"
-    assert !attributes.key?('admin'),     "Denied key should be rejected"
-    assert !attributes.key?('admin(1)'),  "Multi-parameter key should be detected"
-  end
 
 end
diff --git a/activemodel/test/cases/mass_assignment_security/sanitizer_test.rb b/activemodel/test/cases/mass_assignment_security/sanitizer_test.rb
index 9a73a5ad91..8547694c24 100644
--- a/activemodel/test/cases/mass_assignment_security/sanitizer_test.rb
+++ b/activemodel/test/cases/mass_assignment_security/sanitizer_test.rb
@@ -4,24 +4,21 @@ require 'active_support/core_ext/object/inclusion'
 
 class SanitizerTest < ActiveModel::TestCase
 
-  class SanitizingAuthorizer
-    include ActiveModel::MassAssignmentSecurity::Sanitizer
-
-    attr_accessor :logger
 
+  class Authorizer < ActiveModel::MassAssignmentSecurity::PermissionSet
     def deny?(key)
       key.in?(['admin'])
     end
-
   end
 
   def setup
-    @sanitizer = SanitizingAuthorizer.new
+    @sanitizer = ActiveModel::MassAssignmentSecurity::DefaultSanitizer.new
+    @authorizer = Authorizer.new
   end
 
   test "sanitize attributes" do
     original_attributes = { 'first_name' => 'allowed', 'admin' => 'denied' }
-    attributes = @sanitizer.sanitize(original_attributes)
+    attributes = @sanitizer.sanitize(original_attributes, @authorizer)
 
     assert attributes.key?('first_name'), "Allowed key shouldn't be rejected"
     assert !attributes.key?('admin'),     "Denied key should be rejected"
@@ -31,7 +28,7 @@ class SanitizerTest < ActiveModel::TestCase
     original_attributes = { 'first_name' => 'allowed', 'admin' => 'denied' }
     log = StringIO.new
     @sanitizer.logger = Logger.new(log)
-    @sanitizer.sanitize(original_attributes)
+    @sanitizer.sanitize(original_attributes, @authorizer)
     assert_match(/admin/, log.string, "Should log removed attributes: #{log.string}")
   end
 
diff --git a/activemodel/test/cases/mass_assignment_security/white_list_test.rb b/activemodel/test/cases/mass_assignment_security/white_list_test.rb
index aa3596ad2a..737b55492a 100644
--- a/activemodel/test/cases/mass_assignment_security/white_list_test.rb
+++ b/activemodel/test/cases/mass_assignment_security/white_list_test.rb
@@ -16,13 +16,4 @@ class WhiteListTest < ActiveModel::TestCase
     assert_equal true, @white_list.deny?('admin')
   end
 
-  test "sanitize attributes" do
-    original_attributes = { 'first_name' => 'allowed', 'admin' => 'denied', 'admin(1)' => 'denied' }
-    attributes = @white_list.sanitize(original_attributes)
-
-    assert attributes.key?('first_name'), "Allowed key shouldn't be rejected"
-    assert !attributes.key?('admin'),     "Denied key should be rejected"
-    assert !attributes.key?('admin(1)'),  "Multi-parameter key should be detected"
-  end
-
 end
diff --git a/activemodel/test/cases/mass_assignment_security_test.rb b/activemodel/test/cases/mass_assignment_security_test.rb
index 43a12eed61..a778240827 100644
--- a/activemodel/test/cases/mass_assignment_security_test.rb
+++ b/activemodel/test/cases/mass_assignment_security_test.rb
@@ -1,6 +1,15 @@
 require "cases/helper"
 require 'models/mass_assignment_specific'
 
+
+class CustomSanitizer < ActiveModel::MassAssignmentSecurity::Sanitizer
+
+  def process_removed_attributes(attrs)
+    raise StandardError
+  end
+
+end
+
 class MassAssignmentSecurityTest < ActiveModel::TestCase
 
   def test_attribute_protection
@@ -76,4 +85,15 @@ class MassAssignmentSecurityTest < ActiveModel::TestCase
     assert_equal sanitized, { }
   end
 
+  def test_custom_sanitizer
+    user = User.new
+    User.mass_assignment_sanitizer = CustomSanitizer.new
+    assert_raise StandardError do
+      user.sanitize_for_mass_assignment("admin" => true)
+    end
+  ensure
+    User.mass_assignment_sanitizer = nil
+
+  end
+
 end