Merge pull request #41 from FundingCircle/improve-batch-operations

Improve batch encryption and decryption

lib/vault/encrypted_model.rb

@@ -157,6 +157,21 @@ def vault_lazy_decrypt?
       def vault_lazy_decrypt!
         @vault_lazy_decrypt = true
       end
+
+      # works only with convergent encryption
+      def vault_persist_all(attribute, records, plaintexts)
+        options = __vault_attributes[attribute]
+
+        Vault::PerformInBatches.new(attribute, options).encrypt(records, plaintexts)
+      end
+
+      # works only with convergent encryption
+      # relevant only if lazy decryption is enabled
+      def vault_load_all(attribute, records)
+        options = __vault_attributes[attribute]
+
+        Vault::PerformInBatches.new(attribute, options).decrypt(records)
+      end
     end
 
     included do

lib/vault/perform_in_batches.rb

@@ -0,0 +1,55 @@
+module Vault
+  class PerformInBatches
+    def initialize(attribute, options)
+      @attribute  = attribute
+
+      @key        = options[:key]
+      @path       = options[:path]
+      @serializer = options[:serializer]
+      @column     = options[:encrypted_column]
+      @convergent = options[:convergent]
+    end
+
+    def encrypt(records, plaintexts)
+      raise 'Batch Operations work only with convergent attributes' unless @convergent
+
+      raw_plaintexts = serialize(plaintexts)
+
+      ciphertexts = Vault::Rails.batch_encrypt(path, key, raw_plaintexts, Vault.client)
+
+      records.each_with_index do |record, index|
+        record.send("#{column}=", ciphertexts[index])
+        record.save
+      end
+    end
+
+    def decrypt(records)
+      raise 'Batch Operations work only with convergent attributes' unless @convergent
+
+      ciphertexts = records.map { |record| record.send(column) }
+
+      raw_plaintexts = Vault::Rails.batch_decrypt(path, key, ciphertexts, Vault.client)
+      plaintexts = deserialize(raw_plaintexts)
+
+      records.each_with_index do |record, index|
+        record.instance_variable_set("@#{attribute}", plaintexts[index])
+      end
+    end
+
+    private
+
+    attr_reader :key, :path, :serializer, :column, :attribute
+
+    def serialize(plaintexts)
+      return plaintexts unless serializer
+
+      plaintexts.map { |plaintext| serializer.encode(plaintext) }
+    end
+
+    def deserialize(plaintexts)
+      return plaintexts unless serializer
+
+      plaintexts.map { |plaintext| serializer.decode(plaintext) }
+    end
+  end
+end

lib/vault/rails.rb

@@ -5,6 +5,7 @@
 
 require_relative 'encrypted_model'
 require_relative 'attribute_proxy'
+require_relative 'perform_in_batches'
 require_relative 'rails/configurable'
 require_relative 'rails/errors'
 require_relative 'rails/serializers/json_serializer'
@@ -201,7 +202,7 @@ def memory_encrypt(path, key, plaintext, _client, convergent)
 
       # Perform in-memory encryption. This is useful for testing and development.
       def memory_batch_encrypt(path, key, plaintexts, _client)
-        plaintexts.map { |plaintext| memory_encrypt(path, key, ciphertext, _client, true) }
+        plaintexts.map { |plaintext| memory_encrypt(path, key, plaintext, _client, true) }
       end
 
       # Perform in-memory decryption. This is useful for testing and development.
@@ -259,7 +260,9 @@ def vault_batch_encrypt(path, key, plaintexts, client)
           derived: true
         }
 
-        batch_input = plaintexts.map do |plaintext|
+        # Only present values can be encrypted by Vault. Empty values should be returned as they are.
+        non_empty_plaintexts = plaintexts.select { |plaintext| plaintext.present? }
+        batch_input = non_empty_plaintexts.map do |plaintext|
           {
             context: Base64.strict_encode64(Vault::Rails.convergent_encryption_context),
             plaintext: Base64.strict_encode64(plaintext)
@@ -269,7 +272,11 @@ def vault_batch_encrypt(path, key, plaintexts, client)
         options.merge!(batch_input: batch_input)
 
         secret = client.logical.write(route, options)
-        secret.data[:batch_results].map { |result| result[:ciphertext] }
+        vault_results = secret.data[:batch_results].map { |result| result[:ciphertext] }
+
+        plaintexts.map do |plaintext|
+          plaintext.present? ? vault_results.shift : plaintext
+        end
       end
 
       # Perform decryption using Vault. This will raise exceptions if Vault is
@@ -296,18 +303,22 @@ def vault_batch_decrypt(path, key, ciphertexts, client)
 
         route = File.join(path, 'decrypt', key)
 
-
-        batch_input = ciphertexts.map do |ciphertext|
+        # Only present values can be decrypted by Vault. Empty values should be returned as they are.
+        non_empty_ciphertexts = ciphertexts.select { |ciphertext| ciphertext.present? }
+        batch_input = non_empty_ciphertexts.map do |ciphertext|
           {
             context: Base64.strict_encode64(Vault::Rails.convergent_encryption_context),
             ciphertext: ciphertext
           }
         end
-
         options = { batch_input: batch_input }
 
         secret = client.logical.write(route, options)
-        secret.data[:batch_results].map { |result| Base64.strict_decode64(result[:plaintext]) }
+        vault_results = secret.data[:batch_results].map { |result| Base64.strict_decode64(result[:plaintext]) }
+
+        ciphertexts.map do |ciphertext|
+          ciphertext.present? ? vault_results.shift : ciphertext
+        end
       end
 
       # The symmetric key for the given params.
@@ -320,7 +331,7 @@ def memory_key_for(path, key)
       # newly encoded string.
       # @return [String]
       def force_encoding(str)
-        str.force_encoding(Vault::Rails.encoding).encode(Vault::Rails.encoding)
+        str.blank? ? str : str.force_encoding(Vault::Rails.encoding).encode(Vault::Rails.encoding)
       end
 
       private

spec/dummy/app/models/lazy_person.rb

@@ -29,4 +29,6 @@ class LazyPerson < ActiveRecord::Base
     decode: ->(raw) { raw && raw[3...-3] }
 
   vault_attribute :non_ascii
+
+  vault_attribute :passport_number, convergent: true
 end

spec/dummy/db/migrate/20181119142920_add_passport_number_to_people.rb

@@ -0,0 +1,5 @@
+class AddPassportNumberToPeople < ActiveRecord::Migration[5.0]
+  def change
+    add_column :people, :passport_number_encrypted, :string
+  end
+end

spec/dummy/db/schema.rb

@@ -10,7 +10,7 @@
 #
 # It's strongly recommended that you check this file into your version control system.
 
-ActiveRecord::Schema.define(version: 20181030234312) do
+ActiveRecord::Schema.define(version: 20181119142920) do
 
   create_table "people", force: :cascade do |t|
     t.string   "name"
@@ -20,15 +20,16 @@
     t.string   "business_card_encrypted"
     t.string   "favorite_color_encrypted"
     t.string   "non_ascii_encrypted"
-    t.datetime "created_at",               null: false
-    t.datetime "updated_at",               null: false
+    t.datetime "created_at",                null: false
+    t.datetime "updated_at",                null: false
     t.string   "email_encrypted"
     t.string   "county"
     t.string   "county_encrypted"
     t.string   "state"
     t.string   "state_encrypted"
     t.string   "date_of_birth"
     t.string   "date_of_birth_encrypted"
+    t.string   "passport_number_encrypted"
   end
 
 end

spec/integration/rails_spec.rb

@@ -481,4 +481,56 @@
       end
     end
   end
+
+  context 'batch encryption and decryption' do
+    before do
+      allow(Vault::Rails).to receive(:convergent_encryption_context).and_return('a' * 16).at_least(:once)
+    end
+
+    describe '.vault_load_all' do
+      it 'calls Vault just once' do
+        first_person = LazyPerson.create!(passport_number: '12345678')
+        second_person = LazyPerson.create!(passport_number: '12345679')
+
+        people = [first_person.reload, second_person.reload]
+        expect(Vault.logical).to receive(:write).once.and_call_original
+        LazyPerson.vault_load_all(:passport_number, people)
+
+        first_person.passport_number
+        second_person.passport_number
+      end
+
+      it 'loads the attribute of all records' do
+        first_person = LazyPerson.create!(passport_number: '12345678')
+        second_person = LazyPerson.create!(passport_number: '12345679')
+
+        first_person.reload
+        second_person.reload
+
+        LazyPerson.vault_load_all(:passport_number, [first_person, second_person])
+        expect(first_person.passport_number).to eq('12345678')
+        expect(second_person.passport_number).to eq('12345679')
+      end
+    end
+
+    describe '.vault_persist_all' do
+      it 'calls Vault just once' do
+        first_person = LazyPerson.new
+        second_person = LazyPerson.new
+
+        expect(Vault.logical).to receive(:write).once.and_call_original
+        LazyPerson.vault_persist_all(:passport_number, [first_person, second_person], %w(12345678 12345679))
+      end
+
+      it 'saves the attribute of all records' do
+        first_person = LazyPerson.new
+        second_person = LazyPerson.new
+
+        LazyPerson.vault_persist_all(:passport_number, [first_person, second_person], %w(12345678 12345679))
+
+        expect(first_person.reload.passport_number).to eq('12345678')
+        expect(second_person.reload.passport_number).to eq('12345679')
+      end
+    end
+  end
 end