Move msgpack support to core and add activejob support

lib/kms_rails/active_job.rb

@@ -1,3 +1,4 @@
+require 'msgpack'
 require 'active_job'
 require 'kms_rails/core'
 
@@ -10,12 +11,12 @@ def included base
     end
 
     module ClassMethods
-      def kms_arg(field_number, key_id:, context_key: nil, context_value: nil)
-        kms_args([field_number], key_id: key_id, context_key: context_key, context_value: context_value)
+      def kms_arg(field_number, key_id:, msgpack: false, context_key: nil, context_value: nil)
+        kms_args([field_number], key_id: key_id, msgpack: msgpack, context_key: context_key, context_value: context_value)
       end
 
-      def kms_args(field_numbers, key_id:, context_key: nil, context_value: nil)
-        enc = Core.new(key_id: key_id, context_key: context_key, context_value: context_value)
+      def kms_args(field_numbers, key_id:, msgpack: false, context_key: nil, context_value: nil)
+        enc = Core.new(key_id: key_id, context_key: context_key, msgpack: msgpack, context_value: context_value)
 
         define_method 'serialize_arguments' do |args|
           args = args.dup

lib/kms_rails/active_record.rb

@@ -17,7 +17,7 @@ def kms_attr(field, key_id:, retain: false, msgpack: false, context_key: nil, co
         raise RuntimeError, "Field '#{real_field}' must exist to store encrypted data" unless self.column_names.include?(real_field)
         raise RuntimeError, "Field '#{field}' must not be a real column, '#{real_field}' is the real column" if self.column_names.include?(field)
 
-        enc = Core.new(key_id: key_id, context_key: context_key, context_value: context_value)
+        enc = Core.new(key_id: key_id, msgpack: msgpack, context_key: context_key, context_value: context_value)
 
         define_method "#{field}=" do |data|
           if data.nil? # Just set to nil if nil
@@ -27,8 +27,6 @@ def kms_attr(field, key_id:, retain: false, msgpack: false, context_key: nil, co
           end
 
           set_retained(field, data) if retain
-          data = data.to_msgpack if msgpack
-
           encrypted_data = enc.encrypt(data)
           data = nil
 
@@ -47,7 +45,6 @@ def kms_attr(field, key_id:, retain: false, msgpack: false, context_key: nil, co
             plaintext
           else
             plaintext = enc.decrypt(hash)
-            plaintext = MessagePack.unpack(plaintext) if msgpack
             set_retained(field, plaintext) if retain
             plaintext
           end

lib/kms_rails/core.rb

@@ -1,22 +1,25 @@
 require 'base64'
 require 'openssl'
+require 'msgpack'
 require 'aws-sdk'
 require 'kms_rails/configuration'
 
 module KmsRails
   class Core
     attr_reader :context_key, :context_value
 
-    def initialize(key_id:, context_key: nil, context_value: nil)
+    def initialize(key_id:, msgpack: false, context_key: nil, context_value: nil)
       @base_key_id = key_id
       @context_key = context_key
       @context_value = context_value
+      @msgpack = msgpack
     end
 
     def encrypt(data)
       return nil if data.nil?
 
       data_key = aws_generate_data_key(key_id)
+      data = data.to_msgpack if @msgpack
       encrypted = encrypt_attr(data, data_key.plaintext)
 
       self.class.shred_string(data_key.plaintext)
@@ -36,11 +39,15 @@ def encrypt64(data)
 
     def decrypt(data_obj)
       return nil if data_obj.nil?
-      decrypt_attr(
+
+      decrypted = decrypt_attr(
         data_obj['blob'], 
         aws_decrypt_key(data_obj['key']),
         data_obj['iv']
       )
+
+      decrypted = MessagePack.unpack(decrypted) if @msgpack
+      decrypted
     end
 
     def decrypt64(data_obj)

spec/kms_rails/active_job_spec.rb

@@ -66,8 +66,52 @@
       end
     end
 
-  end
+    context 'msgpack enabled' do
+      context '::kms_arg' do
+        subject { FirstArgMsgPackEncryptedJob }
+
+        it 'calls the encryption routine once' do
+          expect_any_instance_of(KmsRails::Aws::KMS::Client).to receive(:generate_data_key)
+            .once
+            .with(hash_including(key_id: 'alias/s', key_spec: 'AES_256'))
+            .and_call_original
+
+          subject.new({'a' => 'b', 'c' => 'd'}, 'bar', 'baz').serialize
+        end
+
+        it 'produces an encrypted argument' do
+          serialized = subject.new({'a' => 'b', 'c' => 'd'}, 'bar', 'baz').serialize['arguments']
+
+          expect(serialized.length).to eq(3)
+          expect(serialized[0].keys).to include('key', 'iv', 'blob')
+          expect(serialized[1]).to eq('bar')
+          expect(serialized[2]).to eq('baz')  
+        end
+      end
+
+      context '::kms_args' do
+        subject { SecondThirdArgMsgPackEncryptedJob }
+
+        it 'calls the encryption routine twice' do
+          expect_any_instance_of(KmsRails::Aws::KMS::Client).to receive(:generate_data_key)
+            .twice
+            .with(hash_including(key_id: 'alias/t', key_spec: 'AES_256'))
+            .and_call_original
 
+          subject.new('foo', {'a' => 'b', 'c' => 'd'}, {'q' => 'r', 's' => 't'}).serialize
+        end
+
+        it 'produces an encrypted argument' do
+          serialized = subject.new('foo', {'a' => 'b', 'c' => 'd'}, {'q' => 'r', 's' => 't'}).serialize['arguments']
+
+          expect(serialized.length).to eq(3)
+          expect(serialized[0]).to eq('foo')
+          expect(serialized[1].keys).to include('key', 'iv', 'blob')
+          expect(serialized[2].keys).to include('key', 'iv', 'blob')
+        end
+      end
+    end
+  end
 
   context 'deserialization' do
     context '1 argument' do 
@@ -118,5 +162,28 @@
         expect(job.perform_now).to eq(['baz', 'bar', 'foo'])
       end
     end
+
+    context 'msgpack enabled' do
+      let(:serialized) { {
+        'job_class'=>'FirstArgMsgPackEncryptedJob',
+        'job_id'=>'39413211-1db3-41d6-92c3-68500713061f',
+        'queue_name'=>'default',
+        'arguments'=>
+          [{'key'=>'NcWpobONtzmHoe7G2wH7v3fcgdqUbYUPehGXuHq5q0cgxMBzL3NhaWxhp5M=',
+            'iv'=>'r2ZIjyxAs3ZZAlLOdOFqdg==',
+            'blob'=>'Pif4qYHilkRiupLljJAyng==',
+            '_aj_symbol_keys'=>[]},
+           'bar',
+           'baz'],
+        'locale'=>:en
+      } }
+
+      subject { FirstArgMsgPackEncryptedJob }
+
+      it 'deserializes and decrypts arguments' do
+        job = subject.deserialize(serialized)
+        expect(job.perform_now).to eq(['baz', 'bar', {'a' => 'b', 'c' => 'd'}])
+      end
+    end
   end
 end

spec/support/test_job.rb

@@ -13,3 +13,19 @@ def perform(arg1, arg2, arg3)
     [arg3, arg2, arg1]
   end
 end
+
+class FirstArgMsgPackEncryptedJob < ActiveJob::Base
+  kms_arg 0, key_id: 's', msgpack: true
+
+  def perform(arg1, arg2, arg3) 
+    [arg3, arg2, arg1]
+  end
+end
+
+class SecondThirdArgMsgPackEncryptedJob < ActiveJob::Base
+  kms_args [1,2], key_id: 't', msgpack: true
+
+  def perform(arg1, arg2, arg3) 
+    [arg3, arg2, arg1]
+  end
+end