LG-10342 Add the ability to specify a key_id for the KMS Client

app/services/encryption/kms_client.rb

@@ -25,6 +25,12 @@ class KmsClient
     end
     # rubocop:enable Layout/LineLength
 
+    attr_reader :kms_key_id
+
+    def initialize(kms_key_id: IdentityConfig.store.aws_kms_key_id)
+      @kms_key_id = kms_key_id
+    end
+
     def encrypt(plaintext, encryption_context)
       KmsLogger.log(:encrypt, encryption_context)
       return encrypt_kms(plaintext, encryption_context) if FeatureManagement.use_kms?
@@ -69,7 +75,7 @@ def encrypt_raw_kms(plaintext, encryption_context)
 
       KMS_CLIENT_POOL.with do |client|
         client.encrypt(
-          key_id: IdentityConfig.store.aws_kms_key_id,
+          key_id: kms_key_id,
           plaintext: plaintext,
           encryption_context: encryption_context,
         ).ciphertext_blob

config/application.yml.default

@@ -58,6 +58,8 @@ aws_http_retry_max_delay: 1
 aws_kms_key_id: alias/login-dot-gov-test-keymaker
 aws_kms_client_contextless_pool_size: 5
 aws_kms_client_multi_pool_size: 5
+aws_kms_multi_region_key_id: alias/login-dot-gov-keymaker-multi-region
+aws_kms_multi_region_write_enabled: false
 aws_logo_bucket: ''
 aws_region: 'us-west-2'
 backup_code_cost: '2000$8$1$'

lib/identity_config.rb

@@ -145,6 +145,8 @@ def self.build_store(config_map)
     config.add(:aws_kms_client_contextless_pool_size, type: :integer)
     config.add(:aws_kms_client_multi_pool_size, type: :integer)
     config.add(:aws_kms_key_id, type: :string)
+    config.add(:aws_kms_multi_region_key_id, type: :string)
+    config.add(:aws_kms_multi_region_write_enabled, type: :boolean)
     config.add(:aws_logo_bucket, type: :string)
     config.add(:aws_region, type: :string)
     config.add(:backup_code_cost, type: :string)

spec/services/encryption/kms_client_spec.rb

@@ -4,7 +4,7 @@
   before do
     stub_const(
       'Encryption::KmsClient::KMS_CLIENT_POOL',
-      FakeConnectionPool.new { Aws::KMS::Client.new(region: aws_region) },
+      FakeConnectionPool.new { aws_kms_client },
     )
 
     # rubocop:disable Layout/LineLength
@@ -36,6 +36,7 @@
     allow(IdentityConfig.store).to receive(:aws_kms_key_id).and_return(key_id)
   end
 
+  let(:aws_kms_client) { Aws::KMS::Client.new(region: aws_region) }
   let(:key_id) { 'key1' }
   let(:plaintext) { 'a' * 3000 + 'b' * 3000 + 'c' * 3000 }
   let(:encryption_context) { { 'context' => 'attribute-bundle', 'user_id' => '123-abc-456-def' } }
@@ -72,6 +73,33 @@
       end
     end
 
+    context 'with a KMS key ID specified' do
+      subject { described_class.new(kms_key_id: 'custom-key-id') }
+
+      before do
+        stub_mapped_aws_kms_client(
+          [
+            # rubocop:disable Layout/LineLength
+            { plaintext: 'a' * 3000, ciphertext: 'us-north-1:kms1', key_id: 'custom-key-id', region: 'us-north-1' },
+            { plaintext: 'b' * 3000, ciphertext: 'us-north-1:kms2', key_id: 'custom-key-id', region: 'us-north-1' },
+            { plaintext: 'c' * 3000, ciphertext: 'us-north-1:kms3', key_id: 'custom-key-id', region: 'us-north-1' },
+            # rubocop:enable Layout/LineLength
+          ],
+        )
+      end
+
+      it 'encrypts with the specified key ID' do
+        result = subject.encrypt(plaintext, encryption_context)
+
+        expect(result).to eq(kms_ciphertext)
+        expect(aws_kms_client.api_requests.count).to eq(3)
+
+        aws_kms_client.api_requests.each do |api_request|
+          expect(api_request[:params][:key_id]).to eq('custom-key-id')
+        end
+      end
+    end
+
     context 'with KMS disabled' do
       let(:kms_enabled) { false }
 

spec/support/aws_kms_client.rb

@@ -1,6 +1,9 @@
 module AwsKmsClientHelper
-  def stub_aws_kms_client(random_key = random_str, ciphered_key = random_str)
+  def stub_aws_kms_client(
+    random_key = random_str,
+    ciphered_key = random_str,
     aws_key_id = IdentityConfig.store.aws_kms_key_id
+  )
     Aws.config[:kms] = {
       stub_responses: {
         encrypt: { ciphertext_blob: ciphered_key, key_id: aws_key_id },