aws_advanced_ruby_driver_wrapper 1.0.0

data/lib/aws_advanced_ruby_driver_wrapper/plugins/kms_encryption/column_cipher.rb

@@ -0,0 +1,159 @@
+# frozen_string_literal: true
+
+#  Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#
+#  Licensed under the Apache License, Version 2.0 (the "License").
+#  You may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+#  http://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+
+require_relative 'encryption_service'
+require_relative 'errors'
+
+module AwsAdvancedRubyDriverWrapper
+  module Plugins
+    module Encryption
+      # Encrypts and decrypts single column values for the duration of one call.
+      #
+      # A cipher is created for each intercepted call and {#release}d when the call returns. That
+      # scope is what keeps plaintext data keys short lived: a key is decrypted at most once per
+      # column per call, however many values that call touches, and every copy is zeroed again on
+      # release.
+      #
+      # Values are returned as strings, matching what both drivers hand back for a text column.
+      class ColumnCipher
+        # @param key_manager [KeyManager]
+        # @param sql_runner [SqlRunner] used to read a bytea or blob column into binary data
+        def initialize(key_manager:, sql_runner:)
+          @key_manager = key_manager
+          @sql = sql_runner
+          @data_keys = {}
+          # Key material resolved by the id embedded in a value, cached for the cipher's lifetime so
+          # a result full of rows written under a rotated-away key costs one lookup, not one per row.
+          @key_metadata_by_id = {}
+        end
+
+        # Encrypts one value for the column the configuration describes.
+        #
+        # @param value [Object, nil] the plaintext value
+        # @param config [ColumnEncryptionConfig] the column's kms_encryption configuration
+        # @return [String, nil] the binary payload to store, nil when value is nil
+        # @raise [Errors::EncryptionError] if the value cannot be encrypted
+        def encrypt(value, config)
+          return nil if value.nil?
+
+          metadata = key_metadata!(config)
+          EncryptionService.encrypt(
+            value, data_key_for(metadata), hmac_key_for(metadata, config), config.algorithm, key_id: metadata.id
+          )
+        end
+
+        # Decrypts one value read from the column the configuration describes.
+        #
+        # The read fails closed: every value of a column configured for kms_encryption is expected to
+        # be an encrypted payload, so a value that cannot be verified and decrypted - too short to be
+        # a payload, an integrity tag that does not match, or a stored key with no HMAC key - is
+        # raised on rather than returned, so the wrapper never hands the application a value it cannot
+        # vouch for. A null column value reads back as nil.
+        #
+        # @param raw [Object, nil] the raw column value
+        # @param config [ColumnEncryptionConfig] the column's kms_encryption configuration
+        # @return [Object, nil] the decrypted value, or nil when the column value is null
+        # @raise [Errors::EncryptionError] if the value cannot be verified and decrypted, is not the
+        #   stored bytes of an encrypted column, or the column has no usable key material
+        def decrypt(raw, config)
+          return nil if raw.nil?
+
+          # A value read from an encrypted (binary) column is a String; nil, handled above, is the
+          # only other shape a NULL produces. Any other type is not the stored bytes of an encrypted
+          # column, so it is refused rather than returned unverified.
+          raise Errors::EncryptionError.decryption_failed("Cannot decrypt a #{raw.class} value") unless raw.is_a?(String)
+
+          # A column configured for kms_encryption must have key material; its absence is a
+          # misconfiguration worth failing on, as on the encrypt side.
+          key_metadata!(config)
+
+          bytes = @sql.read_binary(raw)
+          # Resolve the key the value was written with from the id in its payload, so a value written
+          # before a key rotation still decrypts. A value whose key cannot be resolved falls back to
+          # the current key, against which it then fails its integrity check and is refused.
+          metadata = key_metadata_for_value(bytes, config)
+          hmac_key = hmac_key_for(metadata, config)
+
+          EncryptionService.decrypt(bytes, data_key_for(metadata), hmac_key, config.algorithm, target_type: String)
+        end
+
+        # Zeroes every plaintext data key this cipher decrypted.
+        # @return [void]
+        def release
+          @data_keys.each_value { |data_key| EncryptionService.wipe(data_key) }
+          @data_keys.clear
+          nil
+        end
+
+        private
+
+        # The key a stored value was written with. The payload records the +key_storage+ id, so a
+        # value keeps decrypting after its column's key has been rotated: the current key is used
+        # when the ids match (the common case, and the only one that needs no extra lookup), and any
+        # other id is fetched from +key_storage+ and cached. A value with no embedded id (legacy
+        # data) or one whose key is gone falls back to the current key, against which it then fails
+        # its integrity check and is refused by {#decrypt}.
+        #
+        # @return [KeyMetadata, nil]
+        def key_metadata_for_value(bytes, config)
+          current = config.key_metadata
+          key_id = EncryptionService.key_id_from_payload(bytes)
+          return current if key_id.nil? || (current && current.id == key_id)
+
+          @key_metadata_by_id[key_id] ||= resolve_key_metadata(key_id) || current
+        end
+
+        # Looks a key up by the id embedded in a value. The lookup is opportunistic: a value that is
+        # not really an encrypted payload (legacy data, tampered bytes) can carry an arbitrary id, and
+        # a value written before this format carries none, so a lookup that finds nothing or fails must
+        # not surface as an error - the caller falls back to the current key, against which the value
+        # then fails its integrity check and is refused by {#decrypt}.
+        def resolve_key_metadata(key_id)
+          @key_manager.key_metadata_by_id(key_id)
+        rescue Errors::EncryptionPluginError
+          nil
+        end
+
+        def data_key_for(metadata)
+          @data_keys[metadata.encrypted_data_key] ||=
+            @key_manager.decrypt_data_key(metadata.encrypted_data_key, metadata.master_key_arn)
+        end
+
+        def hmac_key_for(metadata, config)
+          hmac_key = metadata.hmac_key
+          return hmac_key unless hmac_key.nil? || hmac_key.empty?
+
+          raise Errors::EncryptionError
+            .invalid_key('The stored key has no HMAC key')
+            .with_table(config.table_name)
+            .with_column(config.column_name)
+        end
+
+        def key_metadata!(config)
+          metadata = config.key_metadata
+          if metadata.nil?
+            raise Errors::EncryptionError
+              .invalid_key('The column has no key material')
+              .with_table(config.table_name)
+              .with_column(config.column_name)
+          end
+
+          metadata
+        end
+      end
+    end
+  end
+end

data/lib/aws_advanced_ruby_driver_wrapper/plugins/kms_encryption/column_encryption_config.rb

@@ -0,0 +1,61 @@
+# frozen_string_literal: true
+
+#  Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#
+#  Licensed under the Apache License, Version 2.0 (the "License").
+#  You may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+#  http://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+
+require_relative 'encryption_algorithm'
+require_relative 'key_metadata'
+
+module AwsAdvancedRubyDriverWrapper
+  module Plugins
+    module Encryption
+      # One row of the +encryption_metadata+ table joined with the key it points at: the
+      # algorithm to use for a single table column, and the key material to use with it.
+      ColumnEncryptionConfig = Data.define(
+        :table_name,
+        :column_name,
+        :algorithm,
+        :key_id,
+        :key_metadata,
+        :created_at,
+        :updated_at
+      )
+
+      class ColumnEncryptionConfig
+        # @param table_name [String]
+        # @param column_name [String]
+        # @param algorithm [String] an {EncryptionAlgorithm} name
+        # @param key_id [Integer, nil] the +encryption_metadata.key_id+ foreign key
+        # @param key_metadata [KeyMetadata, nil] the joined +key_storage+ row
+        # @param created_at [Time, nil]
+        # @param updated_at [Time, nil]
+        def initialize(table_name:, column_name:, algorithm: EncryptionAlgorithm::DEFAULT, key_id: nil,
+                       key_metadata: nil, created_at: nil, updated_at: nil)
+          super
+        end
+
+        # The cache key for this column, e.g. 'users.ssn'.
+        # @return [String]
+        def column_identifier
+          "#{table_name}.#{column_name}"
+        end
+
+        # @return [Boolean] true when the column has key material and a supported algorithm
+        def usable?
+          EncryptionAlgorithm.supported?(algorithm) && !key_metadata.nil? && key_metadata.valid?
+        end
+      end
+    end
+  end
+end

data/lib/aws_advanced_ruby_driver_wrapper/plugins/kms_encryption/connection_source.rb

@@ -0,0 +1,91 @@
+# frozen_string_literal: true
+
+#  Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#
+#  Licensed under the Apache License, Version 2.0 (the "License").
+#  You may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+#  http://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+
+require_relative '../../errors'
+
+module AwsAdvancedRubyDriverWrapper
+  module Plugins
+    module Encryption
+      # Gives the encryption components (KeyManager, MetadataManager, KeyManagementUtility) a
+      # connection to run a statement on. It is supplied one of two ways, and exactly one must be set:
+      #
+      # * +connection+ - a caller-owned connection, reused for every operation and never closed here.
+      #   This is how {KeyManagementUtility} wires things for a user with a connection in hand.
+      # * +service_container+ - the plugin's runtime source: each call opens a short-lived connection
+      #   of its own through the wrapper's connect pipeline and closes it afterward, so metadata and
+      #   key reads never ride the application's connection.
+      module ConnectionSource
+        # Records the connection source from a constructor. Exactly one argument must be non-nil.
+        #
+        # @raise [ArgumentError] unless exactly one of the two is given
+        def use_connection_source(connection:, service_container:)
+          if [connection, service_container].compact.size != 1
+            raise ArgumentError, 'provide exactly one of connection: or service_container:'
+          end
+
+          @connection = connection
+          @service_container = service_container
+        end
+
+        # Yields a usable connection. A supplied connection is yielded as-is and left open; otherwise a
+        # short-lived connection is opened for the call and closed afterward.
+        #
+        # @param operation [String, nil] operation name, for error context
+        # @yieldparam connection [Object] a pg or mysql2 connection
+        # @return [Object] whatever the block returns
+        # @raise [Errors::AwsError] if a runtime connection cannot be opened
+        def with_connection(operation: nil)
+          return yield(@connection) if @connection
+
+          connection = open_runtime_connection(operation)
+          begin
+            yield connection
+          ensure
+            close_runtime_connection(connection)
+          end
+        end
+
+        private
+
+        # Opens a short-lived connection through the connect pipeline, the way the plugin's runtime
+        # reads its metadata and keys (independent of the application's connection).
+        def open_runtime_connection(operation)
+          connection_service = @service_container.connection_service
+          connection = @service_container.plugin_manager.internal_connect(
+            connection_service.current_host_info,
+            connection_service.driver_props.dup,
+            connection_service.wrapper_props,
+            false
+          )
+          if connection.nil?
+            raise AwsAdvancedRubyDriverWrapper::Errors::AwsError,
+                  "The connect pipeline returned no connection for #{operation || 'a metadata query'}"
+          end
+
+          connection
+        end
+
+        def close_runtime_connection(connection)
+          return if connection.nil?
+
+          @service_container.dialect_service.driver_dialect.close_connection(connection)
+        rescue StandardError
+          nil
+        end
+      end
+    end
+  end
+end

data/lib/aws_advanced_ruby_driver_wrapper/plugins/kms_encryption/data_key_cache.rb

@@ -0,0 +1,220 @@
+# frozen_string_literal: true
+
+#  Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#
+#  Licensed under the Apache License, Version 2.0 (the "License").
+#  You may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+#  http://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+
+require_relative '../../logging'
+require_relative 'encryption_service'
+
+module AwsAdvancedRubyDriverWrapper
+  module Plugins
+    module Encryption
+      # An in-memory cache of plaintext data keys, so that reading an encrypted column does
+      # not need a KMS Decrypt call per statement.
+      #
+      # This deliberately does not use the shared StorageService: data keys need a bounded
+      # size with oldest-first eviction, their bytes have to be zeroed when they leave the
+      # cache, and every entry is handed out as a copy so that a caller wiping its own copy
+      # cannot corrupt the cached one.
+      class DataKeyCache
+        include Logging
+
+        MIN_CLEANUP_INTERVAL_SEC = 30
+        CLEANUP_INTERVAL_DIVISOR = 4
+
+        # A cached plaintext data key. The stored string is private to the cache.
+        CacheEntry = Data.define(:data_key, :created_at) do
+          def expired?(ttl_sec, now = Process.clock_gettime(Process::CLOCK_MONOTONIC))
+            now >= created_at + ttl_sec
+          end
+
+          # Redact the plaintext data key so it is never exposed if an instance is
+          # logged, interpolated, or rendered in a backtrace.
+          def inspect
+            "#<data CacheEntry data_key=#{AwsAdvancedRubyDriverWrapper::REDACTED.inspect}, created_at=#{created_at.inspect}>"
+          end
+          alias_method :to_s, :inspect
+
+          # `pp` / PrettyPrint does not call #inspect; route them through the
+          # redacted representation so `pp entry` cannot leak the data key.
+          def pretty_print(pp)
+            pp.text(inspect)
+          end
+        end
+
+        # A point-in-time snapshot of the cache counters.
+        CacheStats = Data.define(:size, :hits, :misses, :evictions) do
+          # @return [Float] the percentage of lookups that were served from the cache
+          def hit_rate
+            total = hits + misses
+            total.zero? ? 0.0 : (hits.to_f / total) * 100
+          end
+
+          def to_s
+            format('CacheStats{size=%<size>d, hits=%<hits>d, misses=%<misses>d, evictions=%<evictions>d, ' \
+                   'hit_rate=%<hit_rate>.2f%%}',
+                   size: size, hits: hits, misses: misses, evictions: evictions, hit_rate: hit_rate)
+          end
+        end
+
+        # @param max_size [Integer] the maximum number of data keys held at once
+        # @param ttl_sec [Numeric] how long a data key stays cached, in seconds
+        # @param enabled [Boolean] when false every lookup misses and nothing is stored
+        def initialize(max_size:, ttl_sec:, enabled: true)
+          @max_size = max_size
+          @ttl_sec = ttl_sec
+          @enabled = enabled
+          @cache = {}
+          @lock = Mutex.new
+          @hits = 0
+          @misses = 0
+          @evictions = 0
+          @running = enabled
+          @cleanup_thread = start_cleanup_thread if enabled
+        end
+
+        # @return [Boolean]
+        def enabled?
+          @enabled
+        end
+
+        # @param key [String] the cached data key, see {KeyManager#data_key_cache_key}
+        # @return [String, nil] a copy of the cached data key, or nil on a miss
+        def get(key)
+          return nil unless @enabled
+
+          @lock.synchronize do
+            entry = @cache[key]
+
+            if entry.nil?
+              @misses += 1
+              nil
+            elsif entry.expired?(@ttl_sec)
+              @cache.delete(key)
+              EncryptionService.wipe(entry.data_key)
+              @misses += 1
+              nil
+            else
+              @hits += 1
+              entry.data_key.dup
+            end
+          end
+        end
+
+        # Stores a copy of the given data key, evicting the oldest entry first when full.
+        #
+        # @param key [String]
+        # @param data_key [String] the plaintext data key; the caller keeps ownership of it
+        # @return [void]
+        def put(key, data_key)
+          return unless @enabled && data_key
+
+          @lock.synchronize do
+            existing = @cache.delete(key)
+            EncryptionService.wipe(existing.data_key) if existing
+            evict_oldest if @cache.size >= @max_size
+
+            @cache[key] = CacheEntry.new(data_key: data_key.dup, created_at: Process.clock_gettime(Process::CLOCK_MONOTONIC))
+          end
+        end
+
+        # @param key [String]
+        # @return [void]
+        def remove(key)
+          @lock.synchronize do
+            entry = @cache.delete(key)
+            EncryptionService.wipe(entry.data_key) if entry
+          end
+        end
+
+        # Zeroes and drops every cached data key.
+        # @return [void]
+        def clear
+          @lock.synchronize do
+            @cache.each_value { |entry| EncryptionService.wipe(entry.data_key) }
+            @cache.clear
+          end
+        end
+
+        # Drops every expired entry. Called on a schedule by the cleanup thread.
+        # @return [Integer] the number of entries removed
+        def remove_expired_entries
+          @lock.synchronize do
+            now = Process.clock_gettime(Process::CLOCK_MONOTONIC)
+            expired = @cache.select { |_, entry| entry.expired?(@ttl_sec, now) }
+            expired.each_key do |key|
+              EncryptionService.wipe(@cache.delete(key).data_key)
+            end
+            expired.size
+          end
+        end
+
+        # @return [Integer]
+        def size
+          @lock.synchronize { @cache.size }
+        end
+
+        # @return [CacheStats]
+        def stats
+          @lock.synchronize { CacheStats.new(size: @cache.size, hits: @hits, misses: @misses, evictions: @evictions) }
+        end
+
+        # Stops the cleanup thread and zeroes every cached key.
+        # @return [void]
+        def shutdown
+          @running = false
+          begin
+            @cleanup_thread&.wakeup
+          rescue ThreadError
+            nil
+          end
+          @cleanup_thread&.join(5)
+          clear
+        end
+
+        private
+
+        # Removes the entry that has been cached the longest. Callers already hold the lock.
+        def evict_oldest
+          oldest_key, oldest_entry = @cache.min_by { |_, entry| entry.created_at }
+          return if oldest_key.nil?
+
+          @cache.delete(oldest_key)
+          EncryptionService.wipe(oldest_entry.data_key)
+          @evictions += 1
+        end
+
+        def cleanup_interval_sec
+          [@ttl_sec.to_f / CLEANUP_INTERVAL_DIVISOR, MIN_CLEANUP_INTERVAL_SEC].max
+        end
+
+        def start_cleanup_thread
+          interval = cleanup_interval_sec
+          thread = Thread.new do
+            while @running
+              sleep(interval)
+              begin
+                remove_expired_entries if @running
+              rescue StandardError => e
+                logger.debug("DataKeyCache cleanup failed: #{e.message}")
+              end
+            end
+          end
+          thread.name = 'kms_encryption-data-key-cache-cleanup'
+          thread
+        end
+      end
+    end
+  end
+end

data/lib/aws_advanced_ruby_driver_wrapper/plugins/kms_encryption/encryption_algorithm.rb

@@ -0,0 +1,75 @@
+# frozen_string_literal: true
+
+#  Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#
+#  Licensed under the Apache License, Version 2.0 (the "License").
+#  You may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+#  http://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+
+require_relative 'errors'
+
+module AwsAdvancedRubyDriverWrapper
+  module Plugins
+    module Encryption
+      # The symmetric algorithms the plugin can use to encrypt column values. The names are
+      # the ones stored in the +encryption_metadata.encryption_algorithm+ column.
+      module EncryptionAlgorithm
+        AES_256_GCM = 'AES-256-GCM'
+        AES_128_GCM = 'AES-128-GCM'
+        DEFAULT = AES_256_GCM
+
+        # Expected data key length in bytes for each algorithm.
+        KEY_LENGTHS = {
+          AES_256_GCM => 32,
+          AES_128_GCM => 16
+        }.freeze
+
+        # OpenSSL cipher names for each algorithm.
+        CIPHERS = {
+          AES_256_GCM => 'aes-256-gcm',
+          AES_128_GCM => 'aes-128-gcm'
+        }.freeze
+
+        ALL = KEY_LENGTHS.keys.freeze
+
+        module_function
+
+        # @param name [String] an algorithm name, e.g. 'AES-256-GCM'
+        # @return [Integer] the required data key length in bytes
+        # @raise [Errors::EncryptionError] if the algorithm is not supported
+        def key_length(name)
+          KEY_LENGTHS.fetch(name) { raise unsupported(name) }
+        end
+
+        # @param name [String] an algorithm name, e.g. 'AES-256-GCM'
+        # @return [String] the OpenSSL cipher name
+        # @raise [Errors::EncryptionError] if the algorithm is not supported
+        def cipher_name(name)
+          CIPHERS.fetch(name) { raise unsupported(name) }
+        end
+
+        # @param name [String, nil]
+        # @return [Boolean]
+        def supported?(name)
+          KEY_LENGTHS.key?(name)
+        end
+
+        # @param name [String, nil]
+        # @return [Errors::EncryptionError]
+        def unsupported(name)
+          Errors::EncryptionError
+            .invalid_algorithm("Unsupported kms_encryption algorithm: #{name.inspect}. Supported algorithms: #{ALL.join(', ')}")
+            .with_algorithm(name)
+        end
+      end
+    end
+  end
+end