aws_advanced_ruby_driver_wrapper 1.0.0

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

@@ -0,0 +1,281 @@
+# 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 '../../property_definition'
+require_relative 'audit_logger'
+require_relative 'data_key_cache'
+require_relative 'encryption_config'
+require_relative 'errors'
+require_relative 'key_manager'
+require_relative 'metadata_manager'
+require_relative 'schema_validator'
+require_relative 'sql_runner'
+
+module AwsAdvancedRubyDriverWrapper
+  module Plugins
+    module Encryption
+      # Wires up everything the kms_encryption plugin needs and owns its lifecycle.
+      #
+      # Construction is split in two, because the two halves become available at different times:
+      #
+      # * Everything that only needs properties (the configuration, the audit logger, the data key
+      #   cache, the KMS client) is built up front, so that a misconfiguration is reported when the
+      #   connection is opened rather than in the middle of a query.
+      # * Everything that needs to talk to the database (the metadata manager and the key manager,
+      #   both of which read the +encryption_metadata+ and +key_storage+ tables) is built the first
+      #   time a statement could touch an encrypted column, through {#ensure_initialized}. At that
+      #   point the connect pipeline has finished, so the plugin can open its own connections.
+      #
+      # The KMS client is created lazily as well, so that an application that never touches an
+      # encrypted column never has to reach KMS.
+      class KmsEncryptionUtility
+        include Logging
+
+        PLUGIN_NAME = 'KmsEncryptionPlugin'
+
+        # Logged at most once per process, so disabling the metadata cache does not warn on every
+        # connection.
+        @metadata_cache_warning_logged = false
+        class << self
+          attr_accessor :metadata_cache_warning_logged
+        end
+
+        attr_reader :config, :audit_logger, :data_key_cache
+
+        # @param service_container [Services::ServiceContainer]
+        # @param props [Concurrent::Map, Hash] the wrapper properties
+        # @param kms_client [Aws::KMS::Client, nil] a client to use instead of building one
+        # @raise [ArgumentError] if the kms_encryption properties are invalid, or no master keys are allowed
+        def initialize(service_container, props, kms_client: nil)
+          raise ArgumentError, 'service_container is required' if service_container.nil?
+
+          @service_container = service_container
+          @props = props
+          @config = EncryptionConfig.from_props(props)
+          require_master_key_allow_list!
+          @kms_client = kms_client
+          @lock = Mutex.new
+          @initialized = false
+          @closed = false
+
+          @audit_logger = AuditLogger.new(@config.audit_logging_enabled)
+          @data_key_cache = DataKeyCache.new(
+            max_size: @config.data_key_cache_max_size,
+            ttl_sec: @config.data_key_cache_expiration_sec,
+            enabled: @config.data_key_cache_enabled
+          )
+
+          logger.debug do
+            "Loaded the kms_encryption configuration: region=#{@config.kms_region}, " \
+              "schema=#{@config.metadata_schema}, metadata cache=#{@config.metadata_cache_enabled}, " \
+              "max retries=#{@config.key_management_max_retries}"
+          end
+
+          warn_if_metadata_cache_disabled
+        end
+
+        # @return [String] the name the plugin is known by
+        def plugin_name
+          PLUGIN_NAME
+        end
+
+        # @return [Boolean] whether the database backed components have been built
+        def initialized?
+          @lock.synchronize { @initialized }
+        end
+
+        # @return [Boolean] whether {#cleanup} has run
+        def closed?
+          @lock.synchronize { @closed }
+        end
+
+        # Builds the components that need a database connection, the first time one is needed.
+        #
+        # @return [void]
+        # @raise [Errors::MetadataError] if the initial metadata load fails
+        # @raise [Errors::EncryptionPluginError] if the components cannot be built
+        def ensure_initialized
+          return if @lock.synchronize { @initialized || @closed }
+
+          @lock.synchronize do
+            next if @initialized || @closed
+
+            build_database_components
+            @initialized = true
+          end
+          nil
+        end
+
+        # @return [MetadataManager, nil] nil until {#ensure_initialized} has run
+        def metadata_manager
+          @lock.synchronize { @metadata_manager }
+        end
+
+        # @return [KeyManager, nil] nil until {#ensure_initialized} has run
+        def key_manager
+          @lock.synchronize { @key_manager }
+        end
+
+        # @return [SqlRunner, nil] nil until {#ensure_initialized} has run
+        def sql_runner
+          @lock.synchronize { @sql_runner }
+        end
+
+        # Checks that the +encryption_metadata+ and +key_storage+ tables look the way the plugin
+        # expects. Nothing calls this automatically: it is meant for setup scripts and diagnostics.
+        #
+        # @return [SchemaValidator::ValidationResult]
+        # @raise [Errors::EncryptionPluginError] if the components cannot be built
+        def validate_schema
+          ensure_initialized
+          validator, metadata_manager = @lock.synchronize { [@schema_validator, @metadata_manager] }
+
+          metadata_manager.with_connection(operation: 'VALIDATE_SCHEMA') { |connection| validator.validate(connection) }
+        end
+
+        # The KMS client, created on first use.
+        #
+        # @return [Aws::KMS::Client]
+        def kms_client
+          @lock.synchronize { @kms_client ||= create_kms_client }
+        end
+
+        # @return [String] a description of how metadata is being read
+        def connection_mode_status
+          if initialized?
+            'The kms_encryption plugin is reading its metadata over its own short-lived connections'
+          else
+            'The kms_encryption plugin has not opened a metadata connection yet'
+          end
+        end
+
+        # Logs the connection mode, for troubleshooting.
+        # @return [void]
+        def log_current_status
+          logger.info("#{PLUGIN_NAME} status report")
+          logger.info(connection_mode_status)
+          nil
+        end
+
+        # Releases everything the plugin holds: the background metadata refresh, the cached data
+        # keys, and the KMS client.
+        #
+        # @return [void]
+        def cleanup
+          return if @lock.synchronize { @closed }
+
+          logger.debug("Cleaning up #{PLUGIN_NAME}")
+
+          metadata_manager, data_key_cache, kms_client = @lock.synchronize do
+            @closed = true
+            @initialized = false
+            [@metadata_manager, @data_key_cache, @kms_client]
+          end
+
+          quietly('stop the metadata refresh') { metadata_manager&.shutdown }
+          quietly('clear the data key cache') { data_key_cache&.shutdown }
+          quietly('close the KMS client') { kms_client.close if kms_client.respond_to?(:close) }
+
+          logger.debug("Finished cleaning up #{PLUGIN_NAME}")
+          nil
+        end
+
+        private
+
+        def require_master_key_allow_list!
+          return if @config.restricts_master_keys?
+
+          raise ArgumentError,
+                "#{PropertyDefinition::ENCRYPTION_ALLOWED_MASTER_KEY_ARNS.name} is required: list the KMS master key " \
+                'ARNs the kms_encryption plugin may use'
+        end
+
+        # Runs under @lock.
+        def build_database_components
+          @sql_runner = SqlRunner.new(@service_container.dialect_service.driver_dialect)
+          @audit_logger.log_connection_parameter_extraction(
+            strategy: 'ServiceContainer', connection_type: 'INDEPENDENT_CONNECTION'
+          )
+
+          @kms_client ||= create_kms_client
+          @key_manager = KeyManager.new(
+            kms_client: @kms_client,
+            service_container: @service_container,
+            sql_runner: @sql_runner,
+            config: @config,
+            data_key_cache: @data_key_cache,
+            audit_logger: @audit_logger
+          )
+          @metadata_manager = MetadataManager.new(
+            service_container: @service_container,
+            sql_runner: @sql_runner,
+            config: @config,
+            audit_logger: @audit_logger
+          )
+          @schema_validator = SchemaValidator.new(@config.metadata_schema, @sql_runner)
+
+          @metadata_manager.start
+          logger.debug('The kms_encryption plugin is ready to encrypt and decrypt column values')
+        end
+
+        # Warns once per process when the metadata cache is off, since that makes the plugin open a
+        # short-lived metadata connection for every statement that touches an encrypted column.
+        def warn_if_metadata_cache_disabled
+          return if @config.metadata_cache_enabled
+          return if self.class.metadata_cache_warning_logged
+
+          self.class.metadata_cache_warning_logged = true
+          logger.warn(
+            'The kms_encryption metadata cache is disabled (encryption_metadata_cache_enabled=false): the ' \
+            'plugin opens a short-lived metadata connection for every statement that touches an encrypted ' \
+            'column. Enable it in production and lower encryption_metadata_cache_refresh_interval_sec if ' \
+            'you need fresher metadata.'
+          )
+        end
+
+        def create_kms_client
+          ensure_sdk!
+          logger.debug { "Creating a KMS client for region #{@config.kms_region}" }
+
+          options = {
+            region: @config.kms_region,
+            credentials: PropertyDefinition::AWS_CREDENTIALS_PROVIDER.get(@props) ||
+                         Aws::CredentialProviderChain.new.resolve
+          }
+          options[:endpoint] = @config.kms_endpoint unless @config.kms_endpoint.nil?
+
+          Aws::KMS::Client.new(**options)
+        end
+
+        def ensure_sdk!
+          require 'aws-sdk-kms'
+        rescue LoadError
+          raise LoadError,
+                "The KMS kms_encryption plugin requires 'aws-sdk-kms'. " \
+                "Add it to your Gemfile: gem 'aws-sdk-kms'"
+        end
+
+        # Cleanup must release everything it can, so one failing step cannot stop the others.
+        def quietly(description)
+          yield
+        rescue StandardError => e
+          logger.warn("Failed to #{description} while cleaning up the kms_encryption plugin: #{e.message}")
+        end
+      end
+    end
+  end
+end

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

@@ -0,0 +1,332 @@
+# 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 '../../utils/conversion_utils'
+require_relative '../../utils/sql_encoding'
+require_relative 'column_encryption_config'
+require_relative 'connection_source'
+require_relative 'encryption_algorithm'
+require_relative 'errors'
+require_relative 'key_metadata'
+
+module AwsAdvancedRubyDriverWrapper
+  module Plugins
+    module Encryption
+      # Reads the +encryption_metadata+ table, which says which columns are encrypted and with
+      # which key, and keeps the answers in a cache.
+      #
+      # The cache is loaded once at startup and then refreshed on a background thread every
+      # +encryption_metadata_cache_refresh_interval_sec+, so that a column can be added to or
+      # removed from the kms_encryption configuration without restarting the application. Lookups
+      # fall back to querying the database directly when caching is disabled or the cache has
+      # expired.
+      class MetadataManager
+        include Logging
+        include Utils::ConversionUtils
+        include ConnectionSource
+
+        # How long {shutdown} waits for the refresh thread to finish.
+        SHUTDOWN_TIMEOUT_SEC = 5
+
+        # Exactly one of +connection+ or +service_container+ must be given (see {ConnectionSource}).
+        #
+        # @param sql_runner [SqlRunner]
+        # @param config [EncryptionConfig]
+        # @param connection [Object, nil] a caller-owned connection used for every operation
+        # @param service_container [Services::ServiceContainer, nil] opens a short-lived connection per operation
+        # @param audit_logger [AuditLogger, nil]
+        def initialize(sql_runner:, config:, connection: nil, service_container: nil, audit_logger: nil)
+          use_connection_source(connection: connection, service_container: service_container)
+          @sql = sql_runner
+          @config = config
+          @audit_logger = audit_logger
+          @cache = {}
+          @lock = Mutex.new
+          @last_refresh_time = nil
+          @running = false
+          @refresh_thread = nil
+        end
+
+        # Loads the cache and starts the background refresh.
+        #
+        # @return [void]
+        # @raise [Errors::MetadataError] if the initial load fails
+        def start
+          logger.debug('Initializing the kms_encryption metadata manager')
+          refresh if @config.metadata_cache_enabled
+          start_refresh_thread if @config.background_refresh_enabled?
+          nil
+        end
+
+        # Reloads every column configuration from the database and replaces the cache.
+        #
+        # @return [Integer] the number of cached column configurations
+        # @raise [Errors::MetadataError] if the load fails
+        def refresh
+          metadata = load_metadata
+
+          @lock.synchronize do
+            @cache = metadata
+            @last_refresh_time = Time.now
+          end
+
+          logger.debug { "Refreshed the kms_encryption metadata cache with #{metadata.size} column configuration(s)" }
+          @audit_logger&.log_metadata_operation(operation: 'refresh', success: true)
+          metadata.size
+        rescue Errors::MetadataError => e
+          @audit_logger&.log_metadata_operation(operation: 'refresh', success: false, error_message: e.message)
+          raise
+        end
+
+        # Reads every column configuration from the database, without touching the cache.
+        #
+        # @return [Hash{String => ColumnEncryptionConfig}] keyed by +"table.column"+
+        # @raise [Errors::MetadataError] if the load fails
+        def load_metadata
+          rows = with_connection(operation: 'LOAD_ENCRYPTION_METADATA') do |connection|
+            @sql.query(connection, load_metadata_sql)
+          end
+
+          rows.each_with_object({}) do |row, metadata|
+            column_config = to_column_config(row)
+            metadata[column_config.column_identifier] = column_config
+          end
+        rescue StandardError => e
+          raise e if e.is_a?(Errors::MetadataError)
+
+          raise Errors::MetadataError.load_failed("Failed to load kms_encryption metadata: #{e.message}")
+        end
+
+        # @param table_name [String, nil]
+        # @param column_name [String, nil]
+        # @return [Boolean] whether the column is configured for kms_encryption
+        # @raise [Errors::MetadataError] if the lookup has to hit the database and that fails
+        def column_encrypted?(table_name, column_name)
+          return false if table_name.nil? || column_name.nil?
+
+          if cache_usable?
+            hit = @lock.synchronize { @cache.key?(column_identifier(table_name, column_name)) }
+            return hit
+          end
+
+          column_encrypted_in_database?(table_name, column_name)
+        end
+
+        # @param table_name [String, nil]
+        # @param column_name [String, nil]
+        # @return [ColumnEncryptionConfig, nil] the column's configuration, nil when not encrypted
+        # @raise [Errors::MetadataError] if the lookup has to hit the database and that fails
+        def column_config(table_name, column_name)
+          return nil if table_name.nil? || column_name.nil?
+
+          return @lock.synchronize { @cache[column_identifier(table_name, column_name)] } if cache_usable?
+
+          column_config_from_database(table_name, column_name)
+        end
+
+        # Every encrypted column of one table.
+        #
+        # This is what a read is planned from: the statement says which tables it touches, and the
+        # configuration says which of their columns will come back encrypted.
+        #
+        # @param table_name [String, nil]
+        # @return [Array<ColumnEncryptionConfig>] empty when no column of the table is encrypted
+        # @raise [Errors::MetadataError] if the lookup has to hit the database and that fails
+        def table_configs(table_name)
+          return [] if table_name.nil?
+
+          return @lock.synchronize { cached_table_configs(table_name) } if cache_usable?
+
+          table_configs_from_database(table_name)
+        end
+
+        # @return [Time, nil] when the cache was last refreshed, nil when it never was
+        def last_refresh_time
+          @lock.synchronize { @last_refresh_time }
+        end
+
+        # @return [Integer] the number of cached column configurations
+        def cache_size
+          @lock.synchronize { @cache.size }
+        end
+
+        # Stops the background refresh and empties the cache.
+        # @return [void]
+        def shutdown
+          @running = false
+          thread = @refresh_thread
+          if thread
+            thread.wakeup if thread.alive?
+            thread.join(SHUTDOWN_TIMEOUT_SEC)
+            @refresh_thread = nil
+          end
+
+          @lock.synchronize do
+            @cache = {}
+            @last_refresh_time = nil
+          end
+          nil
+        end
+
+        private
+
+        def column_identifier(table_name, column_name)
+          "#{table_name}.#{column_name}"
+        end
+
+        # The cache can only answer a lookup when it is enabled, has been loaded, and has not
+        # expired.
+        def cache_usable?
+          return false unless @config.metadata_cache_enabled
+
+          last_refresh = @lock.synchronize { @last_refresh_time }
+          !last_refresh.nil? && (Time.now - last_refresh) < @config.metadata_cache_expiration_sec
+        end
+
+        def column_encrypted_in_database?(table_name, column_name)
+          row = with_connection(operation: 'CHECK_COLUMN_ENCRYPTED') do |connection|
+            @sql.query(connection, check_column_encrypted_sql, [table_name, column_name]).first
+          end
+          !row.nil?
+        rescue StandardError => e
+          raise lookup_error(e, table_name, column_name)
+        end
+
+        def column_config_from_database(table_name, column_name)
+          row = with_connection(operation: 'GET_COLUMN_CONFIG') do |connection|
+            @sql.query(connection, column_config_sql, [table_name, column_name]).first
+          end
+          row && to_column_config(row)
+        rescue StandardError => e
+          raise lookup_error(e, table_name, column_name)
+        end
+
+        # Callers hold {@lock}.
+        def cached_table_configs(table_name)
+          @cache.each_value.select { |config| config.table_name == table_name }
+        end
+
+        def table_configs_from_database(table_name)
+          rows = with_connection(operation: 'GET_TABLE_CONFIGS') do |connection|
+            @sql.query(connection, table_configs_sql, [table_name])
+          end
+          rows.map { |row| to_column_config(row) }
+        rescue StandardError => e
+          raise lookup_error(e, table_name, nil)
+        end
+
+        def lookup_error(error, table_name, column_name)
+          return error if error.is_a?(Errors::MetadataError)
+
+          Errors::MetadataError
+            .lookup_failed("Failed to load the kms_encryption configuration: #{error.message}")
+            .with_table(table_name)
+            .with_column(column_name)
+        end
+
+        # Builds a {ColumnEncryptionConfig}, and the {KeyMetadata} it points at, from one joined row.
+        #
+        # The row names its table and column in the connection's encoding, while the names they are
+        # looked up by are read from the application's SQL as UTF-8. They are kept as UTF-8 so the two
+        # compare equal whatever the connection's encoding is.
+        def to_column_config(row)
+          key_metadata = KeyMetadata.new(
+            id: row['key_id']&.to_i,
+            key_id: row['key_uuid'],
+            key_name: row['name'],
+            master_key_arn: row['master_key_arn'],
+            encrypted_data_key: row['encrypted_data_key'],
+            hmac_key: @sql.read_binary(row['hmac_key']),
+            key_spec: row['key_spec'] || KeyMetadata::DEFAULT_KEY_SPEC,
+            created_at: row['key_created_at'] && to_time(row['key_created_at']),
+            last_used_at: row['last_used_at'] && to_time(row['last_used_at'])
+          )
+
+          ColumnEncryptionConfig.new(
+            table_name: Utils::SqlEncoding.inspectable(row['table_name']),
+            column_name: Utils::SqlEncoding.inspectable(row['column_name']),
+            algorithm: row['encryption_algorithm'] || EncryptionAlgorithm::DEFAULT,
+            key_id: row['key_id']&.to_i,
+            key_metadata: key_metadata,
+            created_at: row['created_at'] && to_time(row['created_at']),
+            updated_at: row['updated_at'] && to_time(row['updated_at'])
+          )
+        end
+
+        def load_metadata_sql
+          "#{joined_columns_sql} ORDER BY em.table_name, em.column_name"
+        end
+
+        def column_config_sql
+          "#{joined_columns_sql} WHERE em.table_name #{eq} ? AND em.column_name #{eq} ?"
+        end
+
+        def table_configs_sql
+          "#{joined_columns_sql} WHERE em.table_name #{eq} ?"
+        end
+
+        def check_column_encrypted_sql
+          "SELECT 1 FROM #{@config.metadata_schema}.encryption_metadata " \
+            "WHERE table_name #{eq} ? AND column_name #{eq} ?"
+        end
+
+        # +ks.key_id+ is aliased because +em.key_id+ already occupies that name in the row: the
+        # metadata table's key_id is the integer foreign key into key_storage.id, while the key
+        # storage table's own key_id is the external identifier of the key.
+        #
+        # The join is an outer one so that a column whose key row is missing is still reported as
+        # configured for kms_encryption. An inner join would drop it, the column would look like an
+        # ordinary one, and a write would store the plaintext. Kept this way, its key metadata comes
+        # back empty, which fails validation and so fails the write instead.
+        def joined_columns_sql
+          schema = @config.metadata_schema
+          'SELECT em.table_name, em.column_name, em.encryption_algorithm, em.key_id, ' \
+            'em.created_at, em.updated_at, ' \
+            'ks.key_id AS key_uuid, ks.name, ks.master_key_arn, ks.encrypted_data_key, ks.hmac_key, ks.key_spec, ' \
+            'ks.created_at AS key_created_at, ks.last_used_at ' \
+            "FROM #{schema}.encryption_metadata em " \
+            "LEFT JOIN #{schema}.key_storage ks ON em.key_id #{eq} ks.id"
+        end
+
+        def eq
+          @sql.equals_operator
+        end
+
+        def start_refresh_thread
+          @running = true
+          interval = @config.metadata_cache_refresh_interval_sec
+
+          thread = Thread.new do
+            while @running
+              sleep(interval)
+              break unless @running
+
+              begin
+                refresh
+              rescue StandardError => e
+                logger.warn("Failed to refresh the kms_encryption metadata cache: #{e.message}")
+              end
+            end
+          end
+          thread.name = 'kms_encryption-metadata-refresh'
+          @refresh_thread = thread
+          logger.debug { "Started the kms_encryption metadata refresh thread with a #{interval}s interval" }
+        end
+      end
+    end
+  end
+end