Add encryption context to KMS calls

app/models/profile.rb

@@ -34,29 +34,31 @@ def deactivate(reason)
   end
 
   def decrypt_pii(password)
-    Pii::Attributes.new_from_encrypted(
-      encrypted_pii,
-      password: password,
-    )
+    encryptor = Encryption::Encryptors::PiiEncryptor.new(password)
+    decrypted_json = encryptor.decrypt(encrypted_pii, user_uuid: user.uuid)
+    Pii::Attributes.new_from_json(decrypted_json)
   end
 
   def recover_pii(personal_key)
-    Pii::Attributes.new_from_encrypted(
-      encrypted_pii_recovery,
-      password: personal_key,
-    )
+    encryptor = Encryption::Encryptors::PiiEncryptor.new(personal_key)
+    decrypted_recovery_json = encryptor.decrypt(encrypted_pii_recovery, user_uuid: user.uuid)
+    Pii::Attributes.new_from_json(decrypted_recovery_json)
   end
 
   def encrypt_pii(pii, password)
     ssn = pii.ssn
     self.ssn_signature = Pii::Fingerprinter.fingerprint(ssn) if ssn
-    self.encrypted_pii = pii.encrypted(password)
+    encryptor = Encryption::Encryptors::PiiEncryptor.new(password)
+    self.encrypted_pii = encryptor.encrypt(pii.to_json, user_uuid: user.uuid)
     encrypt_recovery_pii(pii)
   end
 
   def encrypt_recovery_pii(pii)
     personal_key = personal_key_generator.create
-    self.encrypted_pii_recovery = pii.encrypted(personal_key_generator.normalize(personal_key))
+    encryptor = Encryption::Encryptors::PiiEncryptor.new(
+      personal_key_generator.normalize(personal_key),
+    )
+    self.encrypted_pii_recovery = encryptor.encrypt(pii.to_json, user_uuid: user.uuid)
     @personal_key = personal_key
   end
 

app/services/encryption/contextless_kms_client.rb

@@ -0,0 +1,98 @@
+module Encryption
+  class ContextlessKmsClient
+    include Encodable
+
+    KEY_TYPE = {
+      KMS: 'KMSx',
+    }.freeze
+
+    def encrypt(plaintext)
+      return encrypt_kms(plaintext) if FeatureManagement.use_kms?
+      encrypt_local(plaintext)
+    end
+
+    def decrypt(ciphertext)
+      return decrypt_kms(ciphertext) if use_kms?(ciphertext)
+      decrypt_local(ciphertext)
+    end
+
+    def self.looks_like_kms?(ciphertext)
+      ciphertext.start_with?(KEY_TYPE[:KMS])
+    end
+
+    private
+
+    def use_kms?(ciphertext)
+      FeatureManagement.use_kms? && self.class.looks_like_kms?(ciphertext)
+    end
+
+    def encrypt_kms(plaintext)
+      if plaintext.bytesize > 4096
+        encrypt_in_chunks(plaintext)
+      else
+        KEY_TYPE[:KMS] + encrypt_raw_kms(plaintext)
+      end
+    end
+
+    # chunk plaintext into ~4096 byte chunks, but not less than 1024 bytes in a chunk if chunking.
+    # we do this by counting how many chunks we have and adding one.
+    # :reek:FeatureEnvy
+    def encrypt_in_chunks(plaintext)
+      plain_size = plaintext.bytesize
+      number_chunks = plain_size / 4096
+      chunk_size = plain_size / (1 + number_chunks)
+      ciphertext_set = plaintext.scan(/.{1,#{chunk_size}}/m).map(&method(:encrypt_raw_kms))
+      KEY_TYPE[:KMS] + ciphertext_set.map { |chunk| Base64.strict_encode64(chunk) }.to_json
+    end
+
+    def encrypt_raw_kms(plaintext)
+      raise ArgumentError, 'kms plaintext exceeds 4096 bytes' if plaintext.bytesize > 4096
+      aws_client.encrypt(
+        key_id: Figaro.env.aws_kms_key_id,
+        plaintext: plaintext,
+      ).ciphertext_blob
+    end
+
+    # :reek:DuplicateMethodCall
+    def decrypt_kms(ciphertext)
+      raw_ciphertext = ciphertext.sub(KEY_TYPE[:KMS], '')
+      if raw_ciphertext[0] == '[' && raw_ciphertext[-1] == ']'
+        decrypt_chunked_kms(raw_ciphertext)
+      else
+        decrypt_raw_kms(raw_ciphertext)
+      end
+    end
+
+    def decrypt_chunked_kms(raw_ciphertext)
+      ciphertext_set = JSON.parse(raw_ciphertext).map { |chunk| Base64.strict_decode64(chunk) }
+      ciphertext_set.map(&method(:decrypt_raw_kms)).join('')
+    rescue JSON::ParserError, ArgumentError
+      decrypt_raw_kms(raw_ciphertext)
+    end
+
+    def decrypt_raw_kms(raw_ciphertext)
+      aws_client.decrypt(ciphertext_blob: raw_ciphertext).plaintext
+    rescue Aws::KMS::Errors::InvalidCiphertextException
+      raise EncryptionError, 'Aws::KMS::Errors::InvalidCiphertextException'
+    end
+
+    def encrypt_local(plaintext)
+      encryptor.encrypt(plaintext, Figaro.env.password_pepper)
+    end
+
+    def decrypt_local(ciphertext)
+      encryptor.decrypt(ciphertext, Figaro.env.password_pepper)
+    end
+
+    def aws_client
+      @aws_client ||= Aws::KMS::Client.new(
+        instance_profile_credentials_timeout: 1, # defaults to 1 second
+        instance_profile_credentials_retries: 5, # defaults to 0 retries
+      )
+    end
+
+    def encryptor
+      @encryptor ||= Encryptors::AesEncryptor.new
+    end
+  end
+end

app/services/encryption/encryptors/pii_encryptor.rb

@@ -37,18 +37,22 @@ def initialize(password)
         @kms_client = KmsClient.new
       end
 
-      def encrypt(plaintext)
+      def encrypt(plaintext, user_uuid: nil)
         salt = SecureRandom.hex(32)
         cost = Figaro.env.scrypt_cost
         aes_encryption_key = scrypt_password_digest(salt: salt, cost: cost)
         aes_encrypted_ciphertext = aes_cipher.encrypt(plaintext, aes_encryption_key)
-        kms_encrypted_ciphertext = kms_client.encrypt(aes_encrypted_ciphertext)
+        kms_encrypted_ciphertext = kms_client.encrypt(
+          aes_encrypted_ciphertext, kms_encryption_context(user_uuid: user_uuid)
+        )
         Ciphertext.new(kms_encrypted_ciphertext, salt, cost).to_s
       end
 
-      def decrypt(ciphertext_string)
+      def decrypt(ciphertext_string, user_uuid: nil)
         ciphertext = Ciphertext.parse_from_string(ciphertext_string)
-        aes_encrypted_ciphertext = kms_client.decrypt(ciphertext.encrypted_data)
+        aes_encrypted_ciphertext = kms_client.decrypt(
+          ciphertext.encrypted_data, kms_encryption_context(user_uuid: user_uuid)
+        )
         aes_encryption_key = scrypt_password_digest(salt: ciphertext.salt, cost: ciphertext.cost)
         aes_cipher.decrypt(aes_encrypted_ciphertext, aes_encryption_key)
       end
@@ -57,6 +61,13 @@ def decrypt(ciphertext_string)
 
       attr_reader :password, :aes_cipher, :kms_client
 
+      def kms_encryption_context(user_uuid:)
+        {
+          'context' => 'pii-encryption',
+          'user_uuid' => user_uuid,
+        }
+      end
+
       def scrypt_password_digest(salt:, cost:)
         scrypt_salt = cost + OpenSSL::Digest::SHA256.hexdigest(salt)
         scrypted = SCrypt::Engine.hash_secret password, scrypt_salt, 32

app/services/encryption/encryptors/session_encryptor.rb

@@ -5,12 +5,12 @@ class SessionEncryptor
 
       def encrypt(plaintext)
         aes_ciphertext = AesEncryptor.new.encrypt(plaintext, aes_encryption_key)
-        kms_ciphertext = KmsClient.new.encrypt(aes_ciphertext)
+        kms_ciphertext = ContextlessKmsClient.new.encrypt(aes_ciphertext)
         encode(kms_ciphertext)
       end
 
       def decrypt(ciphertext)
-        aes_ciphertext = KmsClient.new.decrypt(decode(ciphertext))
+        aes_ciphertext = ContextlessKmsClient.new.decrypt(decode(ciphertext))
         aes_encryptor.decrypt(aes_ciphertext, aes_encryption_key)
       end
 

app/services/encryption/kms_client.rb

@@ -1,89 +1,125 @@
 require 'base64'
 
 module Encryption
-  class KmsClient
+  class KmsClient # rubocop:disable Metrics/ClassLength
     include Encodable
 
     KEY_TYPE = {
-      KMS: 'KMSx',
+      KMS: 'KMSc',
+      LOCAL_KEY: 'LOCc',
     }.freeze
 
-    def encrypt(plaintext)
-      return encrypt_kms(plaintext) if FeatureManagement.use_kms?
-      encrypt_local(plaintext)
+    def encrypt(plaintext, encryption_context)
+      return ContextlessKmsClient.new.encrypt(plaintext) unless FeatureManagement.use_kms_contexts?
+
+      return encrypt_kms(plaintext, encryption_context) if FeatureManagement.use_kms?
+      encrypt_local(plaintext, encryption_context)
     end
 
-    def decrypt(ciphertext)
-      return decrypt_kms(ciphertext) if use_kms?(ciphertext)
-      decrypt_local(ciphertext)
+    def decrypt(ciphertext, encryption_context)
+      return decrypt_contextless_kms(ciphertext) if self.class.looks_like_contextless?(ciphertext)
+      return decrypt_kms(ciphertext, encryption_context) if use_kms?(ciphertext)
+      decrypt_local(ciphertext, encryption_context)
     end
 
     def self.looks_like_kms?(ciphertext)
       ciphertext.start_with?(KEY_TYPE[:KMS])
     end
 
+    def self.looks_like_local_key?(ciphertext)
+      ciphertext.start_with?(KEY_TYPE[:LOCAL_KEY])
+    end
+
+    def self.looks_like_contextless?(ciphertext)
+      !looks_like_kms?(ciphertext) && !looks_like_local_key?(ciphertext)
+    end
+
     private
 
     def use_kms?(ciphertext)
       FeatureManagement.use_kms? && self.class.looks_like_kms?(ciphertext)
     end
 
-    def encrypt_kms(plaintext)
-      if plaintext.bytesize > 4096
-        encrypt_in_chunks(plaintext)
-      else
-        KEY_TYPE[:KMS] + encrypt_raw_kms(plaintext)
-      end
-    end
-
-    # chunk plaintext into ~4096 byte chunks, but not less than 1024 bytes in a chunk if chunking.
-    # we do this by counting how many chunks we have and adding one.
-    # :reek:FeatureEnvy
-    def encrypt_in_chunks(plaintext)
-      plain_size = plaintext.bytesize
-      number_chunks = plain_size / 4096
-      chunk_size = plain_size / (1 + number_chunks)
-      ciphertext_set = plaintext.scan(/.{1,#{chunk_size}}/m).map(&method(:encrypt_raw_kms))
-      KEY_TYPE[:KMS] + ciphertext_set.map { |chunk| Base64.strict_encode64(chunk) }.to_json
+    def encrypt_kms(plaintext, encryption_context)
+      KEY_TYPE[:KMS] + chunk_plaintext(plaintext).map do |chunk|
+        Base64.strict_encode64(
+          encrypt_raw_kms(chunk, encryption_context),
+        )
+      end.to_json
     end
 
-    def encrypt_raw_kms(plaintext)
+    def encrypt_raw_kms(plaintext, encryption_context)
       raise ArgumentError, 'kms plaintext exceeds 4096 bytes' if plaintext.bytesize > 4096
       aws_client.encrypt(
         key_id: Figaro.env.aws_kms_key_id,
         plaintext: plaintext,
+        encryption_context: encryption_context,
       ).ciphertext_blob
     end
 
-    # :reek:DuplicateMethodCall
-    def decrypt_kms(ciphertext)
-      raw_ciphertext = ciphertext.sub(KEY_TYPE[:KMS], '')
-      if raw_ciphertext[0] == '[' && raw_ciphertext[-1] == ']'
-        decrypt_chunked_kms(raw_ciphertext)
-      else
-        decrypt_raw_kms(raw_ciphertext)
-      end
-    end
-
-    def decrypt_chunked_kms(raw_ciphertext)
-      ciphertext_set = JSON.parse(raw_ciphertext).map { |chunk| Base64.strict_decode64(chunk) }
-      ciphertext_set.map(&method(:decrypt_raw_kms)).join('')
-    rescue JSON::ParserError, ArgumentError
-      decrypt_raw_kms(raw_ciphertext)
+    def decrypt_kms(ciphertext, encryption_context)
+      clipped_ciphertext = ciphertext.gsub(/\A#{KEY_TYPE[:KMS]}/, '')
+      ciphertext_chunks = JSON.parse(clipped_ciphertext)
+      ciphertext_chunks.map do |chunk|
+        decrypt_raw_kms(
+          Base64.strict_decode64(chunk),
+          encryption_context,
+        )
+      end.join('')
+    rescue JSON::ParserError, ArgumentError => error
+      raise EncryptionError, "Failed to parse KMS ciphertext: #{error}"
     end
 
-    def decrypt_raw_kms(raw_ciphertext)
-      aws_client.decrypt(ciphertext_blob: raw_ciphertext).plaintext
+    def decrypt_raw_kms(ciphertext, encryption_context)
+      aws_client.decrypt(
+        ciphertext_blob: ciphertext,
+        encryption_context: encryption_context,
+      ).plaintext
     rescue Aws::KMS::Errors::InvalidCiphertextException
       raise EncryptionError, 'Aws::KMS::Errors::InvalidCiphertextException'
     end
 
-    def encrypt_local(plaintext)
-      encryptor.encrypt(plaintext, Figaro.env.password_pepper)
+    def encrypt_local(plaintext, encryption_context)
+      KEY_TYPE[:LOCAL_KEY] + chunk_plaintext(plaintext).map do |chunk|
+        Base64.strict_encode64(
+          encryptor.encrypt(chunk, local_encryption_key(encryption_context)),
+        )
+      end.to_json
+    end
+
+    def decrypt_local(ciphertext, encryption_context)
+      clipped_ciphertext = ciphertext.gsub(/\A#{KEY_TYPE[:LOCAL_KEY]}/, '')
+      ciphertext_chunks = JSON.parse(clipped_ciphertext)
+      ciphertext_chunks.map do |chunk|
+        encryptor.decrypt(
+          Base64.strict_decode64(chunk),
+          local_encryption_key(encryption_context),
+        )
+      end.join('')
+    rescue JSON::ParserError, ArgumentError => error
+      raise EncryptionError, "Failed to parse local ciphertext: #{error}"
+    end
+
+    def local_encryption_key(encryption_context)
+      OpenSSL::HMAC.digest(
+        'sha256',
+        Figaro.env.password_pepper,
+        (encryption_context.keys + encryption_context.values).sort.join(''),
+      )
+    end
+
+    def decrypt_contextless_kms(ciphertext)
+      ContextlessKmsClient.new.decrypt(ciphertext)
     end
 
-    def decrypt_local(ciphertext)
-      encryptor.decrypt(ciphertext, Figaro.env.password_pepper)
+    # chunk plaintext into ~4096 byte chunks, but not less than 1024 bytes in a chunk if chunking.
+    # we do this by counting how many chunks we have and adding one.
+    # :reek:FeatureEnvy
+    def chunk_plaintext(plaintext)
+      plain_size = plaintext.bytesize
+      number_chunks = plain_size / 4096
+      chunk_size = plain_size / (1 + number_chunks)
+      plaintext.scan(/.{1,#{chunk_size}}/m)
     end
 
     def aws_client

app/services/encryption/user_access_key.rb

@@ -81,7 +81,7 @@ def transform_password_salt_to_scrypt_salt(salt)
     end
 
     def kms_client
-      KmsClient.new
+      ContextlessKmsClient.new
     end
 
     def split_scrypt_digest(digest)