Added client option - closes #37

Author: ankane

CHANGELOG.md

@@ -1,3 +1,7 @@
+## 1.8.0 (unreleased)
+
+- Added `client` option
+
 ## 1.7.0 (2025-04-03)
 
 - Dropped support for Ruby < 3.2 and Rails < 7.1

lib/kms_encrypted/box.rb

@@ -1,17 +1,18 @@
 module KmsEncrypted
   class Box
-    attr_reader :key_id, :version, :previous_versions
+    attr_reader :key_id, :version, :previous_versions, :client
 
-    def initialize(key_id: nil, version: nil, previous_versions: nil)
+    def initialize(key_id: nil, version: nil, previous_versions: nil, client: nil)
       @key_id = key_id || KmsEncrypted.key_id
       @version = version || 1
       @previous_versions = previous_versions || {}
+      @client = client
     end
 
     def encrypt(plaintext, context: nil)
       context = version_context(context, version)
       key_id = version_key_id(version)
-      ciphertext = KmsEncrypted::Client.new(key_id: key_id, data_key: true).encrypt(plaintext, context: context)
+      ciphertext = KmsEncrypted::Client.new(key_id: key_id, data_key: true, client: client).encrypt(plaintext, context: context)
       "v#{version}:#{encode64(ciphertext)}"
     end
 
@@ -43,11 +44,13 @@ def decrypt(ciphertext, context: nil)
       key_id ||= version_key_id(version)
       ciphertext = decode64(ciphertext)
       context = version_context(context, version)
+      client = version_client(version)
 
       KmsEncrypted::Client.new(
         key_id: key_id,
         data_key: true,
-        legacy_context: legacy_context
+        legacy_context: legacy_context,
+        client: client
       ).decrypt(ciphertext, context: context)
     end
 
@@ -68,6 +71,10 @@ def version_key_id(version)
       key_id
     end
 
+    def version_client(version)
+      previous_versions.dig(version, :client) || self.client
+    end
+
     def version_context(context, version)
       if context.respond_to?(:call)
         if context.arity == 0

lib/kms_encrypted/client.rb

@@ -2,10 +2,11 @@ module KmsEncrypted
   class Client
     attr_reader :key_id, :data_key
 
-    def initialize(key_id: nil, legacy_context: false, data_key: false)
+    def initialize(key_id: nil, legacy_context: false, data_key: false, client: nil)
       @key_id = key_id || KmsEncrypted.key_id
       @legacy_context = legacy_context
       @data_key = data_key
+      @service_client = client
     end
 
     def encrypt(plaintext, context: nil)
@@ -60,7 +61,7 @@ def client
             KmsEncrypted::Clients::Aws
           end
 
-        klass.new(key_id: key_id, legacy_context: @legacy_context)
+        klass.new(key_id: key_id, legacy_context: @legacy_context, client: @service_client)
       end
     end
   end

lib/kms_encrypted/clients/aws.rb

@@ -8,7 +8,7 @@ def encrypt(plaintext, context: nil)
         }
         options[:encryption_context] = generate_context(context) if context
 
-        KmsEncrypted.aws_client.encrypt(options).ciphertext_blob
+        client.encrypt(options).ciphertext_blob
       end
 
       def decrypt(ciphertext, context: nil)
@@ -18,14 +18,18 @@ def decrypt(ciphertext, context: nil)
         options[:encryption_context] = generate_context(context) if context
 
         begin
-          KmsEncrypted.aws_client.decrypt(options).plaintext
+          client.decrypt(options).plaintext
         rescue ::Aws::KMS::Errors::InvalidCiphertextException
           decryption_failed!
         end
       end
 
       private
 
+      def client
+        @client ||= KmsEncrypted.aws_client
+      end
+
       # make integers strings for convenience
       def generate_context(context)
         raise ArgumentError, "Context must be a hash" unless context.is_a?(Hash)

lib/kms_encrypted/clients/base.rb

@@ -3,9 +3,10 @@ module Clients
     class Base
       attr_reader :key_id
 
-      def initialize(key_id: nil, legacy_context: false)
+      def initialize(key_id: nil, legacy_context: false, client: nil)
         @key_id = key_id
         @legacy_context = legacy_context
+        @client = client
       end
 
       protected

lib/kms_encrypted/clients/google.rb

@@ -10,9 +10,9 @@ def encrypt(plaintext, context: nil)
         options[:additional_authenticated_data] = generate_context(context) if context
 
         # ensure namespace gets loaded
-        client = KmsEncrypted.google_client
+        client = self.client
 
-        if defined?(::Google::Apis::CloudkmsV1::CloudKMSService) && KmsEncrypted.google_client.is_a?(::Google::Apis::CloudkmsV1::CloudKMSService)
+        if defined?(::Google::Apis::CloudkmsV1::CloudKMSService) && client.is_a?(::Google::Apis::CloudkmsV1::CloudKMSService)
           request = ::Google::Apis::CloudkmsV1::EncryptRequest.new(**options)
           response = client.encrypt_crypto_key(key_id, request)
           @last_key_version = response.name
@@ -32,9 +32,9 @@ def decrypt(ciphertext, context: nil)
         options[:additional_authenticated_data] = generate_context(context) if context
 
         # ensure namespace gets loaded
-        client = KmsEncrypted.google_client
+        client = self.client
 
-        if defined?(::Google::Apis::CloudkmsV1::CloudKMSService) && KmsEncrypted.google_client.is_a?(::Google::Apis::CloudkmsV1::CloudKMSService)
+        if defined?(::Google::Apis::CloudkmsV1::CloudKMSService) && client.is_a?(::Google::Apis::CloudkmsV1::CloudKMSService)
           request = ::Google::Apis::CloudkmsV1::DecryptRequest.new(**options)
           begin
             client.decrypt_crypto_key(key_id, request).plaintext
@@ -52,6 +52,12 @@ def decrypt(ciphertext, context: nil)
           end
         end
       end
+
+      private
+
+      def client
+        @client ||= KmsEncrypted.google_client
+      end
     end
   end
 end

lib/kms_encrypted/clients/vault.rb

@@ -7,7 +7,7 @@ def encrypt(plaintext, context: nil)
         }
         options[:context] = generate_context(context) if context
 
-        response = KmsEncrypted.vault_client.logical.write(
+        response = client.logical.write(
           "transit/encrypt/#{key_id.sub("vault/", "")}",
           options
         )
@@ -23,7 +23,7 @@ def decrypt(ciphertext, context: nil)
 
         response =
           begin
-            KmsEncrypted.vault_client.logical.write(
+            client.logical.write(
               "transit/decrypt/#{key_id.sub("vault/", "")}",
               options
             )
@@ -42,6 +42,10 @@ def decrypt(ciphertext, context: nil)
 
       private
 
+      def client
+        @client ||= KmsEncrypted.vault_client
+      end
+
       # turn hash into json
       def generate_context(context)
         Base64.encode64(super)

lib/kms_encrypted/database.rb

@@ -20,6 +20,10 @@ def previous_versions
       @previous_versions ||= evaluate_option(:previous_versions)
     end
 
+    def client
+      @client ||= evaluate_option(:client)
+    end
+
     def context(version)
       name = options[:name]
       context_method = name ? "kms_encryption_context_#{name}" : "kms_encryption_context"
@@ -36,7 +40,8 @@ def encrypt(plaintext)
       KmsEncrypted::Box.new(
         key_id: key_id,
         version: version,
-        previous_versions: previous_versions
+        previous_versions: previous_versions,
+        client: client
       ).encrypt(plaintext, context: context)
     end
 
@@ -49,7 +54,8 @@ def decrypt(ciphertext)
       KmsEncrypted::Box.new(
         key_id: key_id,
         version: version,
-        previous_versions: previous_versions
+        previous_versions: previous_versions,
+        client: client
       ).decrypt(ciphertext, context: context)
     end
 

lib/kms_encrypted/model.rb

@@ -1,6 +1,6 @@
 module KmsEncrypted
   module Model
-    def has_kms_key(name: nil, key_id: nil, eager_encrypt: false, version: 1, previous_versions: nil, upgrade_context: false)
+    def has_kms_key(name: nil, key_id: nil, eager_encrypt: false, version: 1, previous_versions: nil, upgrade_context: false, client: nil)
       key_id ||= KmsEncrypted.key_id
 
       key_method = name ? "kms_key_#{name}" : "kms_key"
@@ -28,7 +28,8 @@ def self.kms_keys
           name: name,
           version: version,
           previous_versions: previous_versions,
-          upgrade_context: upgrade_context
+          upgrade_context: upgrade_context,
+          client: client
         }
 
         if @kms_keys.size == 1