Normalize Lambda and OpenSearch client config

README.md

@@ -212,9 +212,19 @@ locally.
 - `AWS_ACCESS_KEY_ID`: AWS access key for OpenSearch and Lambda
 - `AWS_SECRET_ACCESS_KEY`: AWS secret key for OpenSearch and Lambda
 - `AWS_REGION`: AWS region for OpenSearch and Lambda services
+- `AWS_ROLE_ARN`: IAM role ARN to assume when using role-based AWS authentication.
+  Used by both OpenSearch (AOSS) and Lambda when `AWS_SESSION_TOKEN` is not set.
 - `AWS_SESSION_TOKEN`: (Optional) AWS session token for temporary credentials when using expiring AWS credentials.
   Use this with temporary AWS credentials for AWS-based OpenSearch access and Lambda.
-  For AOSS, when this is set, temporary credentials are used directly and `AWS_AOSS_ROLE_ARN` is not needed.
+  When this is set, temporary credentials are used directly and `AWS_ROLE_ARN` is not needed.
+
+### TIMDEX Semantic Builder Lambda Authentication
+
+Credential behavior for the Lambda client is:
+
+1. If `AWS_ROLE_ARN` is set, assume that role.
+2. Otherwise, if `AWS_ACCESS_KEY_ID` and `AWS_SECRET_ACCESS_KEY` are set, use those static credentials (and `AWS_SESSION_TOKEN` if present).
+3. If required credential env vars are missing, initialization fails with an explicit configuration error.
 
 ### AWS OpenSearch Service (Legacy)
 
@@ -227,8 +237,8 @@ This is our legacy AWS OpenSearch Service Cluster. All production instances shou
 This is our upcoming configuration once migration is complete. This uses a different [authentication mechanism](https://github.com/awsdocs/amazon-opensearch-service-developer-guide/blob/master/doc_source/serverless-clients.md#ruby) than our legacy AWS OpenSearch Service.
 
 - `AWS_AOSS`: boolean. Set to `true` to enable AWS OpenSearch Serverless (AOSS).
-- `AWS_AOSS_ROLE_ARN`: AWS IAM role ARN to assume for AOSS authentication. **Required when** `AWS_AOSS=true` **and** `AWS_SESSION_TOKEN` is not set. This enables automatic credential refresh via role assumption.
-  When `AWS_SESSION_TOKEN` is present, temporary credentials are used directly and `AWS_AOSS_ROLE_ARN` is not needed. This is only used in local development. `AWS_AOSS_ROLE_ARN` is used in production.
+- `AWS_ROLE_ARN`: AWS IAM role ARN to assume for AOSS authentication. **Required when** `AWS_AOSS=true` **and** `AWS_SESSION_TOKEN` is not set. This enables automatic credential refresh via role assumption.
+  When `AWS_SESSION_TOKEN` is present, temporary credentials are used directly and `AWS_ROLE_ARN` is not needed.
 
 ### TIMDEX Semantic Builder Lambda
 

config/initializers/lambda.rb

@@ -1,16 +1,40 @@
 require 'aws-sdk-lambda'
+require 'aws_auth'
+require 'aws_config_validator'
+
+def validate_lambda_config!
+  AwsConfigValidator.validate_lambda_config
+end
+
+def lambda_credentials
+  if AwsAuth.role_arn_present?
+    Rails.logger.debug 'Configuring Lambda client with assumed role credentials'
+    return AwsAuth.assume_role_credentials(role_session_name: 'timdex-lambda')
+  end
+
+  if ENV['AWS_SESSION_TOKEN'].present?
+    Rails.logger.debug 'Configuring Lambda client with temporary static credentials (session token)'
+  else
+    Rails.logger.debug 'Configuring Lambda client with long-lived static credentials'
+  end
+
+  AwsAuth.static_credentials
+end
 
 def configure_lambda_client
-  options = {
-    region: ENV.fetch('AWS_REGION', 'us-east-1'),
-    access_key_id: ENV.fetch('AWS_ACCESS_KEY_ID'),
-    secret_access_key: ENV.fetch('AWS_SECRET_ACCESS_KEY')
-  }
-  options[:session_token] = ENV['AWS_SESSION_TOKEN'] if ENV['AWS_SESSION_TOKEN'].present?
+  validate_lambda_config!
+
+  Rails.logger.debug 'Configuring AWS Lambda client'
+
+  options = { region: ENV.fetch('AWS_REGION', 'us-east-1') }
+  options[:credentials] = lambda_credentials
 
   # AWS SDK sets this env in prod. However, we need to conditionally set it for tests so VCR can
   # intercept the requests with a fake URL.
-  options[:endpoint] = ENV['AWS_ENDPOINT_URL_LAMBDA'] if ENV['AWS_ENDPOINT_URL_LAMBDA'].present?
+  if ENV['AWS_ENDPOINT_URL_LAMBDA'].present?
+    Rails.logger.debug 'Using AWS_ENDPOINT_URL_LAMBDA override for Lambda client endpoint'
+    options[:endpoint] = ENV['AWS_ENDPOINT_URL_LAMBDA']
+  end
   Aws::Lambda::Client.new(options)
 end
 

config/initializers/opensearch.rb

@@ -1,7 +1,8 @@
 require 'faraday_middleware/aws_sigv4' if ENV['AWS_OPENSEARCH'] == 'true' && ENV.fetch('AWS_AOSS', 'false') == 'false'
 require 'opensearch-aws-sigv4'
 require 'aws-sigv4'
-require 'opensearch_config_validator'
+require 'aws_auth'
+require 'aws_config_validator'
 
 # Helper method to parse OPENSEARCH_LOG as a boolean
 # Environment variables are always strings, so 'false' is truthy
@@ -13,10 +14,10 @@ def opensearch_logging_enabled?
 # Priority is given to AWS AOSS, then AWS OpenSearch, and finally vanilla OpenSearch
 def configure_opensearch
   if ENV['AWS_AOSS'] == 'true'
-    OpensearchConfigValidator.validate_aws_aoss_config
+    AwsConfigValidator.validate_aws_aoss_config
     aws_aoss_client
   elsif ENV['AWS_OPENSEARCH'] == 'true'
-    OpensearchConfigValidator.validate_aws_os_config
+    AwsConfigValidator.validate_aws_os_config
     aws_os_client
   else
     os_client
@@ -44,10 +45,10 @@ def os_client
 # obtained by assuming a role.
 def aws_os_client
   OpenSearch::Client.new log: opensearch_logging_enabled?, url: ENV.fetch('OPENSEARCH_URL', nil) do |config|
-    Rails.logger.debug "Configuring Legacy AWS OpenSearch Service client"
+    Rails.logger.debug 'Configuring Legacy AWS OpenSearch Service client'
     # personal keys use expiring credentials with tokens
     if ENV['AWS_SESSION_TOKEN'].present?
-      Rails.logger.debug 'Using temporary credentials with session token'
+      Rails.logger.debug 'Using temporary credentials with session token for OpenSearch Service client'
       config.request :aws_sigv4,
                      service: 'es',
                      region: ENV.fetch('AWS_REGION', nil),
@@ -56,7 +57,7 @@ def aws_os_client
                      session_token: ENV['AWS_SESSION_TOKEN']
     # application keys don't use tokens
     else
-      Rails.logger.debug 'Using long-lived credentials without session token'
+      Rails.logger.debug 'Using long-lived credentials without session token for OpenSearch Service client'
       config.request :aws_sigv4,
                      service: 'es',
                      region: ENV.fetch('AWS_REGION', nil),
@@ -74,12 +75,22 @@ def aws_os_client
 # @note this configuration uses temporary credentials obtained by assuming a role or via the AWS console, unlike
 # AWS OpenSearch Service which can use long-lived access keys directly.
 def aws_aoss_client
-  Rails.logger.debug "Configuring AWS AOSS client"
+  Rails.logger.debug 'Configuring AWS OpenSearch Serverless (AOSS) client'
+
+  credentials_provider = if ENV.fetch('AWS_SESSION_TOKEN', false).present?
+                           Rails.logger.debug 'Using temporary credentials with session token for OpenSearch AOSS ' \
+                                              'client'
+                           AwsAuth.static_credentials
+                         else
+                           Rails.logger.debug 'Using long-lived credentials and assuming role for OpenSearch AOSS ' \
+                                              'client'
+                           AwsAuth.assume_role_credentials(role_session_name: 'timdex-opensearch')
+                         end
 
   signer = Aws::Sigv4::Signer.new(
     service: 'aoss',
-    region: ENV.fetch('AWS_REGION', nil),
-    credentials_provider: credentials
+    region: ENV.fetch('AWS_REGION', 'us-east-1'),
+    credentials_provider: credentials_provider
   )
 
   OpenSearch::Aws::Sigv4Client.new(
@@ -91,42 +102,4 @@ def aws_aoss_client
   )
 end
 
-def credentials
-  if ENV.fetch('AWS_SESSION_TOKEN', false).present?
-    Rails.logger.debug 'Using temporary credentials with session token'
-    temporary_credentials
-  else
-    Rails.logger.debug 'Using long-lived credentials and assuming role'
-    assume_role_credentials
-  end
-end
-
-# personal keys use expiring credentials with tokens, so we use them directly without assuming a role
-# application keys use long-lived credentials and assume a role to get temporary credentials for AOSS
-def temporary_credentials
-  Aws::Credentials.new(
-    ENV.fetch('AWS_ACCESS_KEY_ID', nil),
-    ENV.fetch('AWS_SECRET_ACCESS_KEY', nil),
-    ENV.fetch('AWS_SESSION_TOKEN', nil)
-  )
-end
-
-# AWS AOSS uses temporary credentials that are obtained by assuming a role. The
-# Aws::AssumeRoleCredentials class is used to get these temporary credentials. It requires the ARN of
-# the role to assume, a session name, and a client for the AWS Security Token Service (STS) which is
-# used to perform the AssumeRole operation. It uses the AWS region and access keys from the
-# environment variables to create the STS client. When the session token expires, the
-# Aws::AssumeRoleCredentials will automatically refresh the credentials by calling AssumeRole again.
-def assume_role_credentials
-  Aws::AssumeRoleCredentials.new(
-    role_arn: ENV.fetch('AWS_AOSS_ROLE_ARN', nil),
-    role_session_name: 'timdex-opensearch',
-    client: Aws::STS::Client.new(
-      region: ENV.fetch('AWS_REGION', nil),
-      access_key_id: ENV.fetch('AWS_ACCESS_KEY_ID', nil),
-      secret_access_key: ENV.fetch('AWS_SECRET_ACCESS_KEY', nil)
-    )
-  )
-end
-
 Timdex::OSClient = configure_opensearch

lib/aws_auth.rb

@@ -0,0 +1,55 @@
+# AwsAuth centralizes AWS credential and role-assumption helpers shared by
+# OpenSearch and Lambda initialization.
+module AwsAuth
+  module_function
+
+  def validate_required_vars!(required_vars, error_prefix:)
+    missing_vars = required_vars.select { |_key, value| value.blank? }.keys
+
+    return unless missing_vars.any?
+
+    raise "#{error_prefix}: These required environment variables are not set: #{missing_vars.join(', ')}"
+  end
+
+  def validate_base_aws_config!(error_prefix:)
+    validate_required_vars!(
+      {
+        'AWS_REGION' => ENV.fetch('AWS_REGION', nil),
+        'AWS_ACCESS_KEY_ID' => ENV.fetch('AWS_ACCESS_KEY_ID', nil),
+        'AWS_SECRET_ACCESS_KEY' => ENV.fetch('AWS_SECRET_ACCESS_KEY', nil)
+      },
+      error_prefix: error_prefix
+    )
+  end
+
+  def role_arn_present?
+    ENV['AWS_ROLE_ARN'].present?
+  end
+
+  def static_credentials
+    Aws::Credentials.new(
+      ENV.fetch('AWS_ACCESS_KEY_ID', nil),
+      ENV.fetch('AWS_SECRET_ACCESS_KEY', nil),
+      ENV.fetch('AWS_SESSION_TOKEN', nil)
+    )
+  end
+
+  def sts_client
+    options = {
+      region: ENV.fetch('AWS_REGION', 'us-east-1'),
+      access_key_id: ENV.fetch('AWS_ACCESS_KEY_ID', nil),
+      secret_access_key: ENV.fetch('AWS_SECRET_ACCESS_KEY', nil)
+    }
+    options[:session_token] = ENV['AWS_SESSION_TOKEN'] if ENV['AWS_SESSION_TOKEN'].present?
+
+    Aws::STS::Client.new(options)
+  end
+
+  def assume_role_credentials(role_session_name:)
+    Aws::AssumeRoleCredentials.new(
+      role_arn: ENV.fetch('AWS_ROLE_ARN', nil),
+      role_session_name: role_session_name,
+      client: sts_client
+    )
+  end
+end

lib/aws_config_validator.rb

@@ -0,0 +1,34 @@
+require 'aws_auth'
+
+# AwsConfigValidator validates AWS-related environment configuration for
+# OpenSearch and Lambda initialization.
+class AwsConfigValidator
+  class << self
+    def validate_lambda_config
+      AwsAuth.validate_base_aws_config!(error_prefix: 'AWS Lambda Config Error')
+    end
+
+    def validate_aws_aoss_config
+      required_vars = {
+        'OPENSEARCH_URL' => ENV.fetch('OPENSEARCH_URL', nil),
+        'AWS_REGION' => ENV.fetch('AWS_REGION', nil),
+        'AWS_ACCESS_KEY_ID' => ENV.fetch('AWS_ACCESS_KEY_ID', nil),
+        'AWS_SECRET_ACCESS_KEY' => ENV.fetch('AWS_SECRET_ACCESS_KEY', nil)
+      }
+
+      # Required only when AWS_SESSION_TOKEN is not present (using role assumption)
+      required_vars['AWS_ROLE_ARN'] = ENV.fetch('AWS_ROLE_ARN', nil) if ENV['AWS_SESSION_TOKEN'].blank?
+
+      AwsAuth.validate_required_vars!(required_vars, error_prefix: 'AWS AOSS Config Error')
+    end
+
+    def validate_aws_os_config
+      AwsAuth.validate_required_vars!({
+                                        'OPENSEARCH_URL' => ENV.fetch('OPENSEARCH_URL', nil),
+                                        'AWS_REGION' => ENV.fetch('AWS_REGION', nil),
+                                        'AWS_ACCESS_KEY_ID' => ENV.fetch('AWS_ACCESS_KEY_ID', nil),
+                                        'AWS_SECRET_ACCESS_KEY' => ENV.fetch('AWS_SECRET_ACCESS_KEY', nil)
+                                      }, error_prefix: 'AWS OpenSearch Config Error')
+    end
+  end
+end