aws_advanced_ruby_driver_wrapper 1.0.0 → 1.1.0

data/lib/aws_advanced_ruby_driver_wrapper/plugins/failover_plugin.rb

@@ -20,9 +20,11 @@ require_relative '../host/host_info'
 require_relative '../host/host_role'
 require_relative '../host/host_availability'
 require_relative '../logging'
+require_relative '../services/host_service'
 require_relative '../utils/rds_utils'
 require_relative '../utils/rds_url_type'
 require_relative '../utils/retry_util'
+require_relative '../utils/telemetry/telemetry_const'
 require_relative '../wrapper_property'
 require_relative 'failover_mode'
 
@@ -32,14 +34,25 @@ module AwsAdvancedRubyDriverWrapper
       include Logging
 
       ReaderFailoverResult = Data.define(:connection, :host_info)
+      FailoverCounters = Data.define(:triggered, :succeeded, :failed)
+
+      SUCCESSFUL_FAILOVER_ERRORS = [Errors::FailoverSuccessError, Errors::TransactionStateUnknownError].freeze
+
+      TELEMETRY_WRITER_FAILOVER = 'failover to writer host'
+      TELEMETRY_READER_FAILOVER = 'failover to replica'
 
       def initialize(service_container, props = ::Concurrent::Map.new)
         @service_container = service_container
         @wrapper_props = props
+        @telemetry_factory = service_container.telemetry_factory
+        @telemetry_failover_additional_top_trace = PropertyDefinition::TELEMETRY_FAILOVER_ADDITIONAL_TOP_TRACE.get_bool(props)
+        @writer_failover_counters = failover_counters('writer_failover')
+        @reader_failover_counters = failover_counters('reader_failover')
 
         @retry_util = Utils::RetryUtil.new(service_container)
         @failover_timeout = PropertyDefinition::FAILOVER_TIMEOUT_SEC.get_float(props)
         @reader_selector_strategy = PropertyDefinition::FAILOVER_READER_HOST_SELECTOR_STRATEGY.get(props)
+        Services::HostService.validate_strategy!(@reader_selector_strategy)
         @failover_mode = nil
         @rds_url_type = nil
         @closed_explicitly = false
@@ -235,9 +248,60 @@ module AwsAdvancedRubyDriverWrapper
         end
 
         if @failover_mode == FailoverMode::STRICT_WRITER
-          failover_writer
+          with_failover_telemetry(TELEMETRY_WRITER_FAILOVER) { count_failover(@writer_failover_counters) { failover_writer } }
         else
-          failover_reader
+          with_failover_telemetry(TELEMETRY_READER_FAILOVER) { count_failover(@reader_failover_counters) { failover_reader } }
+        end
+      end
+
+      def failover_counters(prefix)
+        FailoverCounters.new(
+          triggered: @telemetry_factory.create_counter("#{prefix}.triggered.count"),
+          succeeded: @telemetry_factory.create_counter("#{prefix}.completed.success.count"),
+          failed: @telemetry_factory.create_counter("#{prefix}.completed.failed.count")
+        )
+      end
+
+      # Records a failover as a span. Failover always ends in an error: a successful one ends in one
+      # of SUCCESSFUL_FAILOVER_ERRORS, which marks the span as succeeded, and any other error marks it
+      # as failed. With telemetry_failover_additional_top_trace enabled, the span is also submitted as
+      # a trace of its own.
+      def with_failover_telemetry(span_name)
+        telemetry_context = nil
+        @telemetry_factory.in_context(span_name, Utils::Telemetry::TraceLevel::NESTED) do |context|
+          telemetry_context = context
+          yield
+        rescue *SUCCESSFUL_FAILOVER_ERRORS => e
+          context.success = true
+          context.record_exception(e)
+          raise
+        rescue StandardError => e
+          context.success = false
+          context.record_exception(e)
+          raise
+        end
+      ensure
+        submit_failover_copy(telemetry_context) if telemetry_context && @telemetry_failover_additional_top_trace
+      end
+
+      # Runs while the failover's outcome is propagating, so an error here must not replace it.
+      def submit_failover_copy(telemetry_context)
+        @telemetry_factory.post_copy(telemetry_context, Utils::Telemetry::TraceLevel::FORCE_TOP_LEVEL)
+      rescue StandardError => e
+        logger.debug { "Unable to submit the failover trace as a trace of its own: #{e.message}" }
+      end
+
+      # Counts a failover as triggered, and then as succeeded or failed according to how it ends.
+      def count_failover(counters)
+        counters.triggered.inc
+        begin
+          yield
+        rescue *SUCCESSFUL_FAILOVER_ERRORS
+          counters.succeeded.inc
+          raise
+        rescue StandardError
+          counters.failed.inc
+          raise
         end
       end
 
@@ -414,7 +478,8 @@ module AwsAdvancedRubyDriverWrapper
           Host::HostRole::READER,
           @reader_selector_strategy
         )
-      rescue StandardError
+      rescue StandardError => e
+        logger.warn("Failed to select a reader using the '#{@reader_selector_strategy}' host selection strategy: #{e.message}")
         nil
       end
 

data/lib/aws_advanced_ruby_driver_wrapper/plugins/federated_auth/base_saml_auth_plugin.rb

@@ -0,0 +1,266 @@
+# 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 'concurrent'
+require_relative '../../errors'
+require_relative '../../property_definition'
+require_relative '../iam_auth_plugin'
+require_relative '../../utils/iam_auth_utils'
+require_relative '../../utils/rds_utils'
+require_relative '../../utils/rds_url_type'
+require_relative '../../utils/saml_utils'
+require_relative '../../utils/telemetry/telemetry_const'
+
+module AwsAdvancedRubyDriverWrapper
+  module Plugins
+    module FederatedAuth
+      # Shared connect flow for the SAML-based authentication plugins (AD FS and
+      # Okta).
+      #
+      # A subclass defines +self.fetch_saml_assertion+, +TELEMETRY_PREFIX+ (the prefix of its metric
+      # names) and +TELEMETRY_FETCH_SAML+ (the name of the span its SAML assertion fetch is recorded as).
+      class BaseSamlAuthPlugin
+        SUBSCRIBED_METHODS = Set['connect', 'internal_connect'].freeze
+        # Reuse the IAM plugin's cache name.
+        IAM_TOKEN_CACHE_NAME = AwsAdvancedRubyDriverWrapper::Plugins::IamAuthPlugin::IAM_TOKEN_CACHE_NAME
+
+        attr_reader :subscribed_methods
+
+        def initialize(service_container, props = ::Concurrent::Map.new)
+          ensure_aws_sdk!
+          @service_container = service_container
+          @wrapper_props = props
+          expiration = PropertyDefinition::IAM_EXPIRATION_SEC.get_float(props)
+          PropertyDefinition::IAM_EXPIRATION_SEC.validate!(expiration)
+          service_container.storage_service.register(IAM_TOKEN_CACHE_NAME, ttl: expiration)
+          @subscribed_methods = SUBSCRIBED_METHODS
+
+          # Each subclass names its metrics after itself, and its SAML assertion span after its IdP.
+          @telemetry_factory = service_container.telemetry_factory
+          @fetch_token_counter = @telemetry_factory.create_counter("#{self.class::TELEMETRY_PREFIX}.fetch_token.count")
+          @telemetry_factory.create_gauge(
+            "#{self.class::TELEMETRY_PREFIX}.token_cache.size",
+            &IamAuthPlugin.token_cache_size(service_container.storage_service)
+          )
+        end
+
+        def connect(host_info, driver_props, _is_initial_connection, pipeline_callable)
+          saml_connect(host_info, driver_props, pipeline_callable)
+        end
+
+        def internal_connect(host_info, driver_props, wrapper_props_override, _is_initial_connection, pipeline_callable)
+          saml_connect(host_info, driver_props, pipeline_callable, wrapper_props_override || @wrapper_props)
+        end
+
+        def self.clear_cache(storage_service)
+          storage_service.clear(IAM_TOKEN_CACHE_NAME)
+        end
+
+        # Exchanges a SAML assertion for temporary AWS credentials via STS
+        # AssumeRoleWithSAML.
+        #
+        # @return [Aws::Credentials]
+        def self.credentials_from_saml(saml_assertion, role_arn:, idp_arn:, region:)
+          require 'aws-sdk-core'
+          # AssumeRoleWithSAML is unsigned, so no AWS credentials are needed. Passing
+          # credentials: nil keeps the SDK from running the default credential chain (which
+          # on a credential-less, non-EC2 host wastes ~1s probing IMDS and logs a misleading
+          # "Error retrieving instance profile credentials").
+          client_opts = { credentials: nil }
+          client_opts[:region] = region if region && !region.to_s.empty?
+          resp = Aws::STS::Client.new(**client_opts).assume_role_with_saml(
+            role_arn:, principal_arn: idp_arn, saml_assertion:
+          )
+          creds = resp.credentials
+          Aws::Credentials.new(creds.access_key_id, creds.secret_access_key, creds.session_token)
+        end
+
+        private
+
+        def saml_connect(host_info, driver_props, pipeline_callable, wrapper_props_override = nil)
+          wrapper_props_override ||= @wrapper_props
+          saml_props = wrapper_props_override
+
+          validate_required!(saml_props)
+
+          # The connection's driver_props :user/:password double as the fallback IdP login when
+          # idp_username / idp_password are not set. We overwrite driver_props with db_user and the
+          # RDS token further down, and that hash is reused across reconnects (failover, retry,
+          # Blue/Green rerouting), so capture the fallback into the wrapper props once, up front.
+          # A later reconnect then still logs in to the IdP with the original credentials rather
+          # than db_user and a stale token.
+          persist_idp_credentials_fallback(saml_props, driver_props)
+          db_user = PropertyDefinition::DB_USER.get(saml_props)
+          idp_username = PropertyDefinition::IDP_USERNAME.get(saml_props)
+          idp_password = PropertyDefinition::IDP_PASSWORD.get(saml_props)
+
+          host = Utils::IamAuthUtils.resolve_host(
+            PropertyDefinition::IAM_HOST.get(wrapper_props_override), host_info
+          )
+          port = Utils::IamAuthUtils.resolve_port(
+            PropertyDefinition::IAM_PORT.get(wrapper_props_override),
+            host_info,
+            @service_container.dialect_service.db_dialect.default_port
+          )
+          rds_type = Utils::RdsUtils.identify_rds_type(host)
+
+          # Only a global writer cluster with no explicit iam_region needs AWS credentials to look
+          # its region up. In that one case the SAML exchange has to happen before the region lookup
+          # (and so before the token-cache check); every other case - including a global cluster
+          # whose iam_region is set - defers the IdP round-trip to the token-cache miss path, so a
+          # pool or failover storm does not trigger one IdP login per connection.
+          explicit_region = PropertyDefinition::IAM_REGION.get(wrapper_props_override)
+          preassumed_credentials = nil
+          if rds_type == Utils::RdsUrlType::RDS_GLOBAL_WRITER_CLUSTER && (explicit_region.nil? || explicit_region.to_s.empty?)
+            preassumed_credentials = obtain_credentials(saml_props, nil, idp_username, idp_password)
+          end
+
+          region = Utils::IamAuthUtils.region_for(
+            host:, props: wrapper_props_override, rds_type:,
+            credentials_provider: preassumed_credentials,
+            rds_client_func: preassumed_credentials ? -> { Aws::RDS::Client.new(credentials: preassumed_credentials) } : nil
+          )
+          unless region
+            raise Errors::SamlAuthError,
+                  "#{self.class.name}: unable to determine connection region. " \
+                  "If you are using a non-standard RDS URL, please set the 'iam_region' property."
+          end
+
+          fingerprint = Utils::SamlUtils.principal_fingerprint(
+            PropertyDefinition::IAM_ROLE_ARN.get(saml_props),
+            PropertyDefinition::IAM_IDP_ARN.get(saml_props)
+          )
+          cache_key = Utils::IamAuthUtils.cache_key(region, host, port, db_user, fingerprint)
+          entry = @service_container.storage_service.get(IAM_TOKEN_CACHE_NAME, cache_key)
+          expiration = PropertyDefinition::IAM_EXPIRATION_SEC.get_float(wrapper_props_override)
+          PropertyDefinition::IAM_EXPIRATION_SEC.validate!(expiration)
+
+          if Utils::IamAuthUtils.valid_entry?(entry)
+            token = entry.token
+            is_cached_token = true
+          else
+            token = fetch_and_cache_token(saml_props, region, host, port, db_user, cache_key, expiration,
+                                          idp_username, idp_password, preassumed_credentials)
+            is_cached_token = false
+          end
+
+          user_key = @service_container.dialect_service.driver_dialect.user_property_key
+          driver_props[user_key] = db_user
+          driver_props[:password] = token
+
+          begin
+            pipeline_callable.call
+          rescue StandardError => e
+            raise unless is_cached_token && @service_container.dialect_service.login_error?(e)
+
+            driver_props[:password] = fetch_and_cache_token(saml_props, region, host, port, db_user, cache_key,
+                                                            expiration, idp_username, idp_password, nil)
+            pipeline_callable.call
+          end
+        end
+
+        # Mints a fresh RDS token signed with SAML-derived credentials and caches it.
+        def fetch_and_cache_token(saml_props, region, host, port, db_user, cache_key, expiration,
+                                  idp_username, idp_password, credentials)
+          credentials ||= obtain_credentials(saml_props, region, idp_username, idp_password)
+          @fetch_token_counter.inc
+          token = Utils::IamAuthUtils.generate_token(
+            region:, hostname: host, port:, user: db_user, credentials_provider: credentials,
+            telemetry_factory: @telemetry_factory
+          )
+          @service_container.storage_service.set(
+            IAM_TOKEN_CACHE_NAME, cache_key, Utils::IamAuthUtils.build_token_entry(token, expiration)
+          )
+          token
+        end
+
+        # Captures the connection's user/password as the IdP login when idp_username /
+        # idp_password are unset, persisting them into the wrapper props so the value survives
+        # the driver_props overwrite later in the flow and any subsequent reconnect. Never stores
+        # a nil (Concurrent::Map rejects nil values), so a connection with no fallback simply
+        # keeps idp_username / idp_password unset.
+        def persist_idp_credentials_fallback(saml_props, driver_props)
+          if PropertyDefinition::IDP_USERNAME.get(saml_props).nil?
+            fallback_user = driver_props[:user] || driver_props[:username]
+            PropertyDefinition::IDP_USERNAME.set(saml_props, fallback_user) unless fallback_user.nil?
+          end
+          return unless PropertyDefinition::IDP_PASSWORD.get(saml_props).nil?
+
+          fallback_password = driver_props[:password]
+          PropertyDefinition::IDP_PASSWORD.set(saml_props, fallback_password) unless fallback_password.nil?
+        end
+
+        # Fetches a SAML assertion from the IdP and exchanges it for AWS credentials.
+        def obtain_credentials(saml_props, region, idp_username, idp_password)
+          assertion = @telemetry_factory.in_context(self.class::TELEMETRY_FETCH_SAML, Utils::Telemetry::TraceLevel::NESTED,
+                                                    record_failure: true) do
+            self.class.fetch_saml_assertion(**saml_fetch_options(saml_props, idp_username, idp_password))
+          end
+          self.class.credentials_from_saml(
+            assertion,
+            role_arn: PropertyDefinition::IAM_ROLE_ARN.get(saml_props),
+            idp_arn: PropertyDefinition::IAM_IDP_ARN.get(saml_props),
+            region:
+          )
+        rescue Errors::SamlAuthError
+          raise
+        rescue StandardError => e
+          # Surface IdP/STS failures (DNS, TLS, STS errors) as a SamlAuthError.
+          raise Errors::SamlAuthError, "#{self.class.name}: failed to obtain credentials via SAML: #{e.message}"
+        end
+
+        def validate_required!(saml_props)
+          missing = required_identity(saml_props).select { |_name, value| value.nil? || value.to_s.empty? }
+          return if missing.empty?
+
+          names = missing.map(&:first)
+          raise Errors::SamlAuthError,
+                "#{self.class.name}: #{names.join(', ')} #{names.length > 1 ? 'are' : 'is'} required"
+        end
+
+        # The properties that must be present for this IdP. Subclasses override to
+        # append their own (e.g. Okta's app_id).
+        def required_identity(saml_props)
+          [
+            ['db_user', PropertyDefinition::DB_USER.get(saml_props)],
+            ['idp_endpoint', PropertyDefinition::IDP_ENDPOINT.get(saml_props)],
+            ['iam_role_arn', PropertyDefinition::IAM_ROLE_ARN.get(saml_props)],
+            ['iam_idp_arn', PropertyDefinition::IAM_IDP_ARN.get(saml_props)]
+          ]
+        end
+
+        # The keyword arguments passed to .fetch_saml_assertion. Subclasses override.
+        def saml_fetch_options(saml_props, idp_username, idp_password)
+          {
+            idp_endpoint: PropertyDefinition::IDP_ENDPOINT.get(saml_props),
+            idp_username:,
+            idp_password:,
+            ssl_insecure: PropertyDefinition::SSL_INSECURE.get_bool(saml_props),
+            socket_timeout_sec: PropertyDefinition::HTTP_CLIENT_SOCKET_TIMEOUT_SEC.get_float(saml_props),
+            connect_timeout_sec: PropertyDefinition::HTTP_CLIENT_CONNECT_TIMEOUT_SEC.get_float(saml_props)
+          }
+        end
+
+        def ensure_aws_sdk!
+          require 'aws-sdk-rds'
+        rescue LoadError
+          raise LoadError,
+                "The SAML auth plugins require 'aws-sdk-rds'. Add it to your Gemfile: gem 'aws-sdk-rds'"
+        end
+      end
+    end
+  end
+end

data/lib/aws_advanced_ruby_driver_wrapper/plugins/federated_auth/federated_auth_plugin.rb

@@ -0,0 +1,134 @@
+# 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 'cgi/escape'
+require 'uri'
+require_relative 'base_saml_auth_plugin'
+
+module AwsAdvancedRubyDriverWrapper
+  module Plugins
+    module FederatedAuth
+      # Federated authentication against Active Directory Federation Services
+      # (AD FS). Scrapes the IdP-initiated sign-on page for its form, posts the
+      # federated credentials, and extracts the SAML assertion from the response.
+      # Everything else (STS exchange, token minting, caching, retry) is handled by
+      # BaseSamlAuthPlugin.
+      class FederatedAuthPlugin < BaseSamlAuthPlugin
+        TELEMETRY_PREFIX = 'federated'
+        TELEMETRY_FETCH_SAML = 'Fetch ADFS SAML Assertion'
+        INPUT_TAG_PATTERN = /<input([^<>]*)>/
+        FORM_ACTION_PATTERN = /<form[^<>]*?\baction="([^"]+)"/
+
+        # Fetches a SAML assertion from AD FS.
+        #
+        # @return [String] the base64 SAML assertion
+        def self.fetch_saml_assertion(idp_endpoint:, idp_port:, idp_username:, idp_password:,
+                                      rp_identifier:, ssl_insecure:, socket_timeout_sec:, connect_timeout_sec:, **_rest)
+          # One cookie jar shared by the sign-on GET and the credential POST, so cookies AD FS
+          # (or a load balancer in front of it) sets on the sign-on page are returned with the POST.
+          cookies = {}
+          http = lambda do |method, url, body: nil, headers: {}|
+            Utils::SamlUtils.http_request(
+              method:, url:, body:, headers:, ssl_insecure:, cookies:,
+              open_timeout_sec: connect_timeout_sec, read_timeout_sec: socket_timeout_sec
+            )
+          end
+
+          sign_in_url = "https://#{idp_endpoint}:#{idp_port}/adfs/ls/IdpInitiatedSignOn.aspx?loginToRp=#{rp_identifier}"
+          sign_in_body = require_body!(http.call(:get, sign_in_url), 'sign-on page')
+
+          action = form_action(sign_in_body)
+          post_url = action&.start_with?('/') ? "https://#{idp_endpoint}:#{idp_port}#{action}" : sign_in_url
+
+          params = form_parameters(sign_in_body, idp_username, idp_password)
+          response = http.call(:post, post_url,
+                               body: URI.encode_www_form(params),
+                               headers: { 'Content-Type' => 'application/x-www-form-urlencoded' })
+          content = require_body!(response, 'sign-on POST')
+
+          match = Utils::SamlUtils::SAML_RESPONSE_PATTERN.match(content)
+          unless match
+            # AD FS answers a rejected login with the sign-on page again (no SAMLResponse); the
+            # body is not included as it may contain sensitive markup.
+            raise Errors::SamlAuthError,
+                  'FederatedAuthPlugin: no SAML assertion in the AD FS response - check idp_username, ' \
+                  'idp_password, and idp_endpoint'
+          end
+
+          match[1]
+        end
+
+        def self.require_body!(response, what)
+          return response.body if response.is_a?(Net::HTTPSuccess)
+
+          raise Errors::SamlAuthError,
+                "FederatedAuthPlugin: AD FS #{what} request failed (#{response.code} #{response.message})"
+        end
+
+        def self.form_action(body)
+          match = FORM_ACTION_PATTERN.match(body)
+          match && unescape_html(match[1])
+        end
+
+        # Builds the POST parameters from the sign-on form's input tags, filling in
+        # the federated username/password and preserving every other hidden field.
+        def self.form_parameters(body, idp_username, idp_password)
+          params = []
+          seen = {}
+          body.scan(INPUT_TAG_PATTERN).flatten.each do |tag|
+            name = value_by_key(tag, 'name')
+            next if name.empty?
+
+            lower = name.downcase
+            next if seen[lower]
+
+            seen[lower] = true
+            value = value_by_key(tag, 'value')
+
+            if lower.include?('username')
+              params << [name, idp_username]
+            elsif lower.include?('authmethod')
+              params << [name, value] unless value.empty?
+            elsif lower.include?('password')
+              params << [name, idp_password]
+            else
+              params << [name, value]
+            end
+          end
+          params
+        end
+
+        def self.value_by_key(input, key)
+          match = /#{Regexp.quote(key)}\s*=\s*"(.*?)"/.match(input)
+          match ? unescape_html(match[1]) : ''
+        end
+
+        def self.unescape_html(html)
+          CGI.unescapeHTML(html)
+        end
+
+        private
+
+        def saml_fetch_options(saml_props, idp_username, idp_password)
+          super.merge(
+            idp_port: PropertyDefinition::IDP_PORT.get(saml_props),
+            rp_identifier: PropertyDefinition::RP_IDENTIFIER.get(saml_props)
+          )
+        end
+      end
+    end
+  end
+end

data/lib/aws_advanced_ruby_driver_wrapper/plugins/federated_auth/okta_auth_plugin.rb

@@ -0,0 +1,113 @@
+# 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 'cgi/escape'
+require 'json'
+require_relative 'base_saml_auth_plugin'
+
+module AwsAdvancedRubyDriverWrapper
+  module Plugins
+    module FederatedAuth
+      # Federated authentication against Okta. Authenticates to Okta's /authn API
+      # for a one-time session token, redeems it at the AWS app's SSO SAML endpoint
+      # for a SAML assertion, and then defers to BaseSamlAuthPlugin for the STS
+      # exchange, token minting, caching, and retry. Okta differs from AD FS only in
+      # configuration (it uses app_id, the AWS app configured on Okta) and in the
+      # shape of the IdP HTTP flow.
+      class OktaAuthPlugin < AwsAdvancedRubyDriverWrapper::Plugins::FederatedAuth::BaseSamlAuthPlugin
+        TELEMETRY_PREFIX = 'okta'
+        TELEMETRY_FETCH_SAML = 'Fetch OKTA SAML Assertion'
+        OKTA_AWS_APP_NAME = 'amazon_aws'
+        # Extracts the value attribute of the hidden SAMLResponse input regardless
+        # of attribute order.
+        SAML_INPUT_PATTERN = /<input[^>]*\bname="SAMLResponse"[^>]*>/i
+        VALUE_ATTR_PATTERN = /\bvalue="([^"]+)"/i
+
+        # Fetches a SAML assertion from Okta.
+        #
+        # @return [String] the base64 SAML assertion
+        def self.fetch_saml_assertion(idp_endpoint:, app_id:, idp_username:, idp_password:,
+                                      ssl_insecure:, socket_timeout_sec:, connect_timeout_sec:, **_rest)
+          http = lambda do |method, url, body: nil, headers: {}|
+            Utils::SamlUtils.http_request(
+              method:, url:, body:, headers:, ssl_insecure:,
+              open_timeout_sec: connect_timeout_sec, read_timeout_sec: socket_timeout_sec
+            )
+          end
+
+          session_token = session_token(http, idp_endpoint, idp_username, idp_password)
+          saml_url = "https://#{idp_endpoint}/app/#{OKTA_AWS_APP_NAME}/#{CGI.escape(app_id)}/sso/saml" \
+                     "?onetimetoken=#{CGI.escape(session_token)}"
+          response = http.call(:get, saml_url)
+          unless response.is_a?(Net::HTTPSuccess)
+            raise Errors::SamlAuthError,
+                  "OktaAuthPlugin: SAML request failed (#{response.code} #{response.message})"
+          end
+
+          tag = SAML_INPUT_PATTERN.match(response.body)
+          value = tag && VALUE_ATTR_PATTERN.match(tag[0])
+          raise Errors::SamlAuthError, 'OktaAuthPlugin: invalid SAML response from Okta' unless value
+
+          # The value attribute is HTML-entity-encoded (Okta encodes base64's trailing
+          # "=" padding, and "+" if present, as e.g. &#x3d; / &#x2b;); decode before
+          # handing the assertion to STS, matching FederatedAuthPlugin's AD FS flow.
+          CGI.unescapeHTML(value[1])
+        end
+
+        # Authenticates to Okta's /authn API and returns the one-time session token.
+        def self.session_token(http, idp_endpoint, idp_username, idp_password)
+          response = http.call(
+            :post, "https://#{idp_endpoint}/api/v1/authn",
+            body: JSON.generate(username: idp_username, password: idp_password),
+            headers: { 'Accept' => 'application/json', 'Content-Type' => 'application/json' }
+          )
+          unless response.is_a?(Net::HTTPSuccess)
+            raise Errors::SamlAuthError,
+                  "OktaAuthPlugin: Okta session token request failed (#{response.code} #{response.message})"
+          end
+
+          parsed = begin
+            JSON.parse(response.body)
+          rescue JSON::ParserError
+            nil
+          end
+          parsed = nil unless parsed.is_a?(Hash)
+          token = parsed && parsed['sessionToken']
+          if token.nil? || token.empty?
+            # Okta's `status` is a non-sensitive auth-state enum (e.g. MFA_REQUIRED,
+            # PASSWORD_EXPIRED, LOCKED_OUT); surface it for clarity. It never carries
+            # credentials or the session token itself.
+            status = parsed && parsed['status']
+            detail = status.is_a?(String) && !status.empty? ? " (status: #{status})" : ''
+            raise Errors::SamlAuthError, "OktaAuthPlugin: Okta did not return a session token#{detail}"
+          end
+
+          token
+        end
+
+        private
+
+        def required_identity(saml_props)
+          super + [['app_id', PropertyDefinition::APP_ID.get(saml_props)]]
+        end
+
+        def saml_fetch_options(saml_props, idp_username, idp_password)
+          super.merge(app_id: PropertyDefinition::APP_ID.get(saml_props))
+        end
+      end
+    end
+  end
+end