Merge pull request rails#51735 from heka1024/encryption-compressor

Introduce `compressor` option to `ActiveRecord::Encryption::Encryptor`

Author: eileencodes

activerecord/CHANGELOG.md

@@ -1,3 +1,31 @@
+*   `ActiveRecord::Encryption::Encryptor` now supports a `:compressor` option to customize the compression algorithm used.
+
+    ```ruby
+    module ZstdCompressor
+      def self.deflate(data)
+        Zstd.compress(data)
+      end
+
+      def self.inflate(data)
+        Zstd.decompress(data)
+      end
+    end
+
+    class User
+      encrypts :name, compressor: ZstdCompressor
+    end
+    ```
+
+    You disable compression by passing `compress: false`.
+
+    ```ruby
+    class User
+      encrypts :name, compress: false
+    end
+    ```
+
+    *heka1024*
+
 *   Add condensed `#inspect` for `ConnectionPool`, `AbstractAdapter`, and
     `DatabaseConfig`.
 

activerecord/lib/active_record/encryption/config.rb

@@ -8,7 +8,8 @@ module Encryption
     class Config
       attr_accessor :primary_key, :deterministic_key, :store_key_references, :key_derivation_salt, :hash_digest_class,
                     :support_unencrypted_data, :encrypt_fixtures, :validate_column_size, :add_to_filter_parameters,
-                    :excluded_from_filter_parameters, :extend_queries, :previous_schemes, :forced_encoding_for_deterministic_encryption
+                    :excluded_from_filter_parameters, :extend_queries, :previous_schemes, :forced_encoding_for_deterministic_encryption,
+                    :compressor
 
       def initialize
         set_defaults
@@ -55,6 +56,7 @@ def set_defaults
           self.previous_schemes = []
           self.forced_encoding_for_deterministic_encryption = Encoding::UTF_8
           self.hash_digest_class = OpenSSL::Digest::SHA1
+          self.compressor = Zlib
 
           # TODO: Setting to false for now as the implementation is a bit experimental
           self.extend_queries = false

activerecord/lib/active_record/encryption/encryptable_record.rb

@@ -46,11 +46,13 @@ module EncryptableRecord
         # * <tt>:previous</tt> - List of previous encryption schemes. When provided, they will be used in order when trying to read
         #   the attribute. Each entry of the list can contain the properties supported by #encrypts. Also, when deterministic
         #   encryption is used, they will be used to generate additional ciphertexts to check in the queries.
-        def encrypts(*names, key_provider: nil, key: nil, deterministic: false, support_unencrypted_data: nil, downcase: false, ignore_case: false, previous: [], **context_properties)
+        def encrypts(*names, key_provider: nil, key: nil, deterministic: false, support_unencrypted_data: nil, downcase: false, ignore_case: false, previous: [],
+                     compress: true, compressor: nil, **context_properties)
           self.encrypted_attributes ||= Set.new # not using :default because the instance would be shared across classes
 
           names.each do |name|
-            encrypt_attribute name, key_provider: key_provider, key: key, deterministic: deterministic, support_unencrypted_data: support_unencrypted_data, downcase: downcase, ignore_case: ignore_case, previous: previous, **context_properties
+            encrypt_attribute name, key_provider: key_provider, key: key, deterministic: deterministic, support_unencrypted_data: support_unencrypted_data, downcase: downcase, ignore_case: ignore_case, previous: previous,
+                              compress: compress, compressor: compressor, **context_properties
           end
         end
 
@@ -81,12 +83,13 @@ def global_previous_schemes_for(scheme)
             end
           end
 
-          def encrypt_attribute(name, key_provider: nil, key: nil, deterministic: false, support_unencrypted_data: nil, downcase: false, ignore_case: false, previous: [], **context_properties)
+          def encrypt_attribute(name, key_provider: nil, key: nil, deterministic: false, support_unencrypted_data: nil, downcase: false, ignore_case: false, previous: [],
+                                compress: true, compressor: nil, **context_properties)
             encrypted_attributes << name.to_sym
 
             decorate_attributes([name]) do |name, cast_type|
               scheme = scheme_for key_provider: key_provider, key: key, deterministic: deterministic, support_unencrypted_data: support_unencrypted_data, \
-                downcase: downcase, ignore_case: ignore_case, previous: previous, **context_properties
+                downcase: downcase, ignore_case: ignore_case, previous: previous, compress: compress, compressor: compressor, **context_properties
 
               ActiveRecord::Encryption::EncryptedAttributeType.new(scheme: scheme, cast_type: cast_type, default: columns_hash[name.to_s]&.default)
             end

activerecord/lib/active_record/encryption/encryptor.rb

@@ -12,12 +12,20 @@ module Encryption
     # It interacts with a KeyProvider for getting the keys, and delegate to
     # ActiveRecord::Encryption::Cipher the actual encryption algorithm.
     class Encryptor
+      # The compressor to use for compressing the payload
+      attr_reader :compressor
+
       # === Options
       #
       # * <tt>:compress</tt> - Boolean indicating whether records should be compressed before encryption.
       #   Defaults to +true+.
-      def initialize(compress: true)
+      # * <tt>:compressor</tt> - The compressor to use.
+      #   1. If compressor is provided, it will be used.
+      #   2. If not, it will use ActiveRecord::Encryption.config.compressor which default value is +Zlib+.
+      #   If you want to use a custom compressor, it must respond to +deflate+ and +inflate+.
+      def initialize(compress: true, compressor: nil)
         @compress = compress
+        @compressor = compressor || ActiveRecord::Encryption.config.compressor
       end
 
       # Encrypts +clean_text+ and returns the encrypted result
@@ -78,6 +86,10 @@ def binary?
         serializer.binary?
       end
 
+      def compress? # :nodoc:
+        @compress
+      end
+
       private
         DECRYPT_ERRORS = [OpenSSL::Cipher::CipherError, Errors::EncryptedContentIntegrity, Errors::Decryption]
         ENCODING_ERRORS = [EncodingError, Errors::Encoding]
@@ -130,12 +142,8 @@ def compress_if_worth_it(string)
           end
         end
 
-        def compress?
-          @compress
-        end
-
         def compress(data)
-          Zlib::Deflate.deflate(data).tap do |compressed_data|
+          @compressor.deflate(data).tap do |compressed_data|
             compressed_data.force_encoding(data.encoding)
           end
         end
@@ -149,7 +157,7 @@ def uncompress_if_needed(data, compressed)
         end
 
         def uncompress(data)
-          Zlib::Inflate.inflate(data).tap do |uncompressed_data|
+          @compressor.inflate(data).tap do |uncompressed_data|
             uncompressed_data.force_encoding(data.encoding)
           end
         end

activerecord/lib/active_record/encryption/scheme.rb

@@ -11,7 +11,7 @@ class Scheme
       attr_accessor :previous_schemes
 
       def initialize(key_provider: nil, key: nil, deterministic: nil, support_unencrypted_data: nil, downcase: nil, ignore_case: nil,
-                     previous_schemes: nil, **context_properties)
+                     previous_schemes: nil, compress: true, compressor: nil, **context_properties)
         # Initializing all attributes to +nil+ as we want to allow a "not set" semantics so that we
         # can merge schemes without overriding values with defaults. See +#merge+
 
@@ -24,8 +24,13 @@ def initialize(key_provider: nil, key: nil, deterministic: nil, support_unencryp
         @previous_schemes_param = previous_schemes
         @previous_schemes = Array.wrap(previous_schemes)
         @context_properties = context_properties
+        @compress = compress
+        @compressor = compressor
 
         validate_config!
+
+        @context_properties[:encryptor] = Encryptor.new(compress: @compress) unless @compress
+        @context_properties[:encryptor] = Encryptor.new(compressor: compressor) if compressor
       end
 
       def ignore_case?
@@ -78,6 +83,8 @@ def compatible_with?(other_scheme)
         def validate_config!
           raise Errors::Configuration, "ignore_case: can only be used with deterministic encryption" if @ignore_case && !@deterministic
           raise Errors::Configuration, "key_provider: and key: can't be used simultaneously" if @key_provider_param && @key
+          raise Errors::Configuration, "compressor: can't be used with compress: false" if !@compress && @compressor
+          raise Errors::Configuration, "compressor: can't be used with encryptor" if @compressor && @context_properties[:encryptor]
         end
 
         def key_provider_from_key

activerecord/test/cases/encryption/encryptable_record_test.rb

@@ -413,6 +413,11 @@ def name
     assert_equal json_bytes, EncryptedBookWithSerializedBinary.create!(logo: json_bytes).logo
   end
 
+  test "can compress data with custom compressor" do
+    name = "a" * 141
+    assert EncryptedBookWithCustomCompressor.create!(name: name).name.start_with?("[compressed]")
+  end
+
   private
     def build_derived_key_provider_with(hash_digest_class)
       ActiveRecord::Encryption.with_encryption_context(key_generator: ActiveRecord::Encryption::KeyGenerator.new(hash_digest_class: hash_digest_class)) do

activerecord/test/cases/encryption/encryptor_test.rb

@@ -88,6 +88,22 @@ class ActiveRecord::Encryption::EncryptorTest < ActiveRecord::EncryptionTestCase
     assert_equal Encoding::ISO_8859_1, decrypted_text.encoding
   end
 
+  test "accept a custom compressor" do
+    compressor = Module.new do
+      def self.deflate(data)
+        "compressed #{data}"
+      end
+
+      def self.inflate(data)
+        data.sub(/\Acompressed /, "")
+      end
+    end
+    @encryptor = ActiveRecord::Encryption::Encryptor.new(compressor: compressor)
+    content = SecureRandom.hex(5.kilobytes)
+
+    assert_encrypt_text content
+  end
+
   private
     def assert_encrypt_text(clean_text)
       encrypted_text = @encryptor.encrypt(clean_text)

activerecord/test/cases/encryption/scheme_test.rb

@@ -7,12 +7,36 @@ class ActiveRecord::Encryption::SchemeTest < ActiveRecord::EncryptionTestCase
   test "validates config options when using encrypted attributes" do
     assert_invalid_declaration deterministic: false, ignore_case: true
     assert_invalid_declaration key: "1234", key_provider: ActiveRecord::Encryption::DerivedSecretKeyProvider.new("my secret")
+    assert_invalid_declaration compress: false, compressor: Zlib
+    assert_invalid_declaration compressor: Zlib, encryptor: ActiveRecord::Encryption::Encryptor.new
 
     assert_valid_declaration deterministic: true
     assert_valid_declaration key: "1234"
     assert_valid_declaration key_provider: ActiveRecord::Encryption::DerivedSecretKeyProvider.new("my secret")
   end
 
+  test "should create a encryptor well when compressor is given" do
+    MyCompressor = Class.new do
+      def self.deflate(data)
+        "deflated #{data}"
+      end
+
+      def self.inflate(data)
+        data.sub("deflated ", "")
+      end
+    end
+
+    type = declare_encrypts_with compressor: MyCompressor
+
+    assert_equal MyCompressor, type.scheme.to_h[:encryptor].compressor
+  end
+
+  test "should create a encryptor well when compress is false" do
+    type = declare_encrypts_with compress: false
+
+    assert_not type.scheme.to_h[:encryptor].compress?
+  end
+
   private
     def assert_invalid_declaration(**options)
       assert_raises ActiveRecord::Encryption::Errors::Configuration do

activerecord/test/models/book_encrypted.rb

@@ -56,3 +56,19 @@ class EncryptedBookWithSerializedBinary < ActiveRecord::Base
   serialize :logo, coder: JSON
   encrypts :logo
 end
+
+class EncryptedBookWithCustomCompressor < ActiveRecord::Base
+  module CustomCompressor
+    def self.deflate(value)
+      "[compressed] #{value}"
+    end
+
+    def self.inflate(value)
+      value
+    end
+  end
+
+  self.table_name = "encrypted_books"
+
+  encrypts :name, compressor: CustomCompressor
+end

activerecord/test/schema/schema.rb

@@ -161,7 +161,7 @@
   create_table :encrypted_books, id: :integer, force: true do |t|
     t.references :author
     t.string :format
-    t.column :name, :string, default: "<untitled>"
+    t.column :name, :string, default: "<untitled>", limit: 1024
     t.column :original_name, :string
     t.column :logo, :binary