Initial pass adding AWS IAM Authentication #1263

README.md

@@ -284,8 +284,10 @@ Mysql2::Client.new(
   :get_server_public_key = true/false,
   :default_file = '/path/to/my.cfg',
   :default_group = 'my.cfg section',
-  :default_auth = 'authentication_windows_client'
-  :init_command => sql
+  :default_auth = 'authentication_windows_client',
+  :init_command => sql,
+  :use_iam_authentication => true/false,
+  :host_region,
   )
 ```
 
@@ -348,6 +350,32 @@ When secure_auth is enabled, the server will refuse a connection if the account
 The MySQL 5.6.5 client library may also refuse to attempt a connection if provided an older format password.
 To bypass this restriction in the client, pass the option `:secure_auth => false` to Mysql2::Client.new().
 
+### AWS IAM Authentication
+
+You may use AWS IAM Authentication instead of setting a password in
+the configuration.  A temporary token used in place of the password
+will be fetched as necessary and used for connections until it
+expires.  The value for :host_region will either use the one provided,
+or if not provided, the environment variable AWS_REGION.
+
+You must add the `aws-sdk-rds` gem to your bundle to use this functionality.
+
+| `:use_iam_authentication` | true                                 |
+| ---                       | ---                                  |
+| `:username`               | The database username configured to use IAM Authentication |
+| `:host`                   | The database host |
+| `:port`                   | The database port |
+| `:host_region`            | An AWS region name, e.g. `us-east-1` |
+
+As prerequisites, you must enable IAM authentication on the RDS
+instance, create an IAM policy, attach the policy to the target IAM
+user or role, create the database user set to use the AWS
+Authentication Plugin, and then run your ruby code using that IAM user or
+role.  See
+[AWS documentation](https://docs.aws.amazon.com/AmazonRDS/latest/UserGuide/UsingWithRDS.IAMDBAuth.Connecting.html)
+for details on these steps.
+
+
 ### Flags option parsing
 
 The `:flags` parameter accepts an integer, a string, or an array. The integer

lib/mysql2.rb

@@ -38,6 +38,7 @@
 require 'mysql2/client'
 require 'mysql2/field'
 require 'mysql2/statement'
+require 'mysql2/aws_iam_auth'
 
 # = Mysql2
 #

lib/mysql2/aws_iam_auth.rb

@@ -0,0 +1,67 @@
+require 'singleton'
+
+module Mysql2
+  # Generates and caches AWS IAM Authentication tokens to use in place of MySQL user passwords
+  class AwsIamAuth
+    include Singleton
+    attr_reader :mutex
+    attr_accessor :passwords
+
+    # Tokens are valid for up to 15 minutes.
+    # We will assume ours expire in 14 minutes to be safe.
+    TOKEN_EXPIRES_IN = (60 * 14) # 14 minutes
+
+    def initialize
+      begin
+        require 'aws-sdk-rds'
+      rescue LoadError
+        raise LoadError, "gem aws-sdk-rds was not found.  Please add this gem to your bundle to use AWS IAM Authentication."
+      end
+
+      @mutex = Mutex.new
+      # Key identifies a unique set of authentication parameters
+      # Value is a Hash
+      # :password is the token value
+      # :expires_at is (just before) the token was generated plus 14 minutes
+      @passwords = {}
+      instance_credentials = Aws::InstanceProfileCredentials.new
+      @generator = Aws::RDS::AuthTokenGenerator.new(:credentials => instance_credentials)
+    end
+
+    def password(user, host, port, opts)
+      params = to_params(user, host, port, opts)
+      key = key_from_params(params)
+      passwd = nil
+      AwsIamAuth.instance.mutex.synchronize do
+        begin
+          passwd = @passwords[key][:password] if @passwords.dig(key, :password) && Time.now.utc < @passwords.dig(key, :expires_at)
+        rescue KeyError
+          passwd = nil
+        end
+      end
+      return passwd unless passwd.nil?
+
+      AwsIamAuth.instance.mutex.synchronize do
+        @passwords[key] = {}
+        @passwords[key][:expires_at] = Time.now.utc + TOKEN_EXPIRES_IN
+        @passwords[key][:password] = password_from_iam(params)
+      end
+    end
+
+    def password_from_iam(params)
+      @generator.auth_token(params)
+    end
+
+    def to_params(user, host, port, opts)
+      params = {}
+      params[:region] = opts[:host_region] || ENV['AWS_REGION']
+      params[:endpoint] = "#{host}:#{port}"
+      params[:user_name] = user
+      params
+    end
+
+    def key_from_params(params)
+      "#{params[:user_name]}/#{params[:endpoint]}/#{params[:region]}"
+    end
+  end
+end

lib/mysql2/client.rb

@@ -94,6 +94,11 @@ def initialize(opts = {})
       socket = socket.to_s unless socket.nil?
       conn_attrs = parse_connect_attrs(opts[:connect_attrs])
 
+      if opts[:use_iam_authentication]
+        aws = Mysql2::AwsIamAuth.instance
+        pass = aws.password(user, host, port, opts)
+      end
+
       connect user, pass, host, port, database, socket, flags, conn_attrs
     end