Implement IAM Authentication (#120)

aws · May 2, 2023 · 3547f49 · 3547f49
1 parent c8b1fde
commit 3547f49
Show file tree

Hide file tree

Showing 12 changed files with 549 additions and 40 deletions.
diff --git a/driver/CMakeLists.txt b/driver/CMakeLists.txt
@@ -57,15 +57,48 @@ WHILE(${DRIVER_INDEX} LESS ${DRIVERS_COUNT})
   SET(DRIVER_NAME "awsmysqlodbc${CONNECTOR_DRIVER_TYPE_SHORT}")
 
   SET(DRIVER_SRCS
-    aws_sdk_helper.cc base_metrics_holder.cc catalog.cc catalog_no_i_s.cc cluster_topology_info.cc
-    cluster_aware_hit_metrics_holder.cc cluster_aware_metrics_container.cc
-    cluster_aware_metrics.cc cluster_aware_time_metrics_holder.cc
-    connect.cc connection_handler.cc connection_proxy.cc cursor.cc desc.cc dll.cc driver.cc efm_proxy.cc
-    error.cc execute.cc failover_handler.cc
-    failover_reader_handler.cc failover_writer_handler.cc handle.cc host_info.cc info.cc
-    monitor.cc monitor_connection_context.cc monitor_service.cc monitor_thread_container.cc
-    my_prepared_stmt.cc my_stmt.cc mylog.cc mysql_proxy.cc options.cc parse.cc prepare.cc query_parsing.cc
-    results.cc topology_service.cc transact.cc utility.cc)
+    aws_sdk_helper.cc
+    base_metrics_holder.cc
+    catalog.cc
+    catalog_no_i_s.cc
+    cluster_topology_info.cc
+    cluster_aware_hit_metrics_holder.cc
+    cluster_aware_metrics_container.cc
+    cluster_aware_metrics.cc
+    cluster_aware_time_metrics_holder.cc
+    connect.cc
+    connection_handler.cc
+    connection_proxy.cc
+    cursor.cc
+    desc.cc
+    dll.cc
+    driver.cc
+    efm_proxy.cc
+    error.cc
+    execute.cc
+    failover_handler.cc
+    failover_reader_handler.cc
+    failover_writer_handler.cc
+    handle.cc
+    host_info.cc
+    iam_proxy.cc
+    info.cc
+    monitor.cc
+    monitor_connection_context.cc
+    monitor_service.cc
+    monitor_thread_container.cc
+    my_prepared_stmt.cc
+    my_stmt.cc
+    mylog.cc
+    mysql_proxy.cc
+    options.cc
+    parse.cc
+    prepare.cc
+    query_parsing.cc
+    results.cc
+    topology_service.cc
+    transact.cc
+    utility.cc)
 
   IF(UNICODE)
     SET(DRIVER_SRCS ${DRIVER_SRCS} unicode.cc)
@@ -82,10 +115,32 @@ WHILE(${DRIVER_INDEX} LESS ${DRIVERS_COUNT})
     CONFIGURE_FILE(${CMAKE_SOURCE_DIR}/driver/driver.def.cmake ${CMAKE_SOURCE_DIR}/driver/driver${CONNECTOR_DRIVER_TYPE_SHORT}.def @ONLY)
     CONFIGURE_FILE(${CMAKE_SOURCE_DIR}/driver/driver.rc.cmake ${CMAKE_SOURCE_DIR}/driver/driver${CONNECTOR_DRIVER_TYPE_SHORT}.rc @ONLY)
     SET(DRIVER_SRCS ${DRIVER_SRCS} driver${CONNECTOR_DRIVER_TYPE_SHORT}.def driver${CONNECTOR_DRIVER_TYPE_SHORT}.rc  
-                                   aws_sdk_helper.h base_metrics_holder.h catalog.h cluster_aware_hit_metrics_holder.h cluster_aware_metrics_container.h
-                                   cluster_aware_metrics.h cluster_aware_time_metrics_holder.h cluster_topology_info.h connection_handler.h connection_proxy.h
-                                   driver.h efm_proxy.h error.h failover.h host_info.h monitor.h monitor_connection_context.h monitor_service.h
-                                   monitor_thread_container.h mylog.h mysql_proxy.h myutil.h parse.h query_parsing.h topology_service.h
+                                   aws_sdk_helper.h                                   
+                                   base_metrics_holder.h
+                                   catalog.h
+                                   cluster_aware_hit_metrics_holder.h
+                                   cluster_aware_metrics_container.h
+                                   cluster_aware_metrics.h
+                                   cluster_aware_time_metrics_holder.h
+                                   cluster_topology_info.h
+                                   connection_handler.h
+                                   connection_proxy.h
+                                   driver.h
+                                   efm_proxy.h
+                                   error.h
+                                   failover.h
+                                   host_info.h
+                                   iam_proxy.h
+                                   monitor.h
+                                   monitor_connection_context.h
+                                   monitor_service.h
+                                   monitor_thread_container.h
+                                   mylog.h
+                                   mysql_proxy.h
+                                   myutil.h
+                                   parse.h
+                                   query_parsing.h
+                                   topology_service.h
                                    ../MYODBC_MYSQL.h ../MYODBC_CONF.h ../MYODBC_ODBC.h)
   ENDIF(WIN32)
 

diff --git a/driver/connect.cc b/driver/connect.cc
@@ -666,7 +666,9 @@ SQLRETURN DBC::connect(DataSource *dsrc, bool failover_enabled)
 #endif
 
 #if (MYSQL_VERSION_ID >= 50527 && MYSQL_VERSION_ID < 50600) || MYSQL_VERSION_ID >= 50607
-  if (dsrc->enable_cleartext_plugin)
+  // IAM authentication requires the plugin to be set.
+  if (dsrc->enable_cleartext_plugin || 
+      (dsrc->auth_mode8 && !myodbc_strcasecmp(AUTH_MODE_IAM, (const char*)dsrc->auth_mode8)))
   {
     connection_proxy->options(MYSQL_ENABLE_CLEARTEXT_PLUGIN, (char *)&on);
   }
@@ -823,20 +825,12 @@ SQLRETURN DBC::connect(DataSource *dsrc, bool failover_enabled)
     ds_set_strnattr(&dsrc->server8, (SQLCHAR*)host, strlen(host));
     dsrc->port = port;
 
-    const bool connect_result = dsrc->enable_dns_srv ?
-                            connection_proxy->real_connect_dns_srv(host,
-                              ds_get_utf8attr(dsrc->uid,      &dsrc->uid8),
-                              ds_get_utf8attr(dsrc->pwd,      &dsrc->pwd8),
-                              ds_get_utf8attr(dsrc->database, &dsrc->database8),
-                              flags)
-                            :
-                            connection_proxy->real_connect(host,
-                              ds_get_utf8attr(dsrc->uid,      &dsrc->uid8),
-                              ds_get_utf8attr(dsrc->pwd,      &dsrc->pwd8),
-                              ds_get_utf8attr(dsrc->database, &dsrc->database8),
-                              port, 
-                              ds_get_utf8attr(dsrc->socket,   &dsrc->socket8),
-                              flags);
+    const char* user = ds_get_utf8attr(dsrc->uid, &dsrc->uid8);
+    const char* password = ds_get_utf8attr(dsrc->pwd, &dsrc->pwd8);
+    const char* database = ds_get_utf8attr(dsrc->database, &dsrc->database8);
+    const char* socket = ds_get_utf8attr(dsrc->socket, &dsrc->socket8);
+
+    const bool connect_result = connection_proxy->connect(host, user, password, database, port, socket, flags);
     if (!connect_result)
     {
       unsigned int native_error= connection_proxy->error_code();

diff --git a/driver/connection_proxy.cc b/driver/connection_proxy.cc
@@ -45,6 +45,16 @@ CONNECTION_PROXY::~CONNECTION_PROXY() {
     }
 }
 
+bool CONNECTION_PROXY::connect(const char* host, const char* user, const char* password,
+    const char* database, unsigned int port, const char* socket, unsigned long flags) {
+
+    if (ds->enable_dns_srv) {
+        return this->real_connect_dns_srv(host, user, password, database, flags);
+    }
+
+    return this->real_connect(host, user, password, database, port, socket, flags);
+}
+
 void CONNECTION_PROXY::delete_ds() {
     next_proxy->delete_ds();
 }
@@ -78,6 +88,14 @@ unsigned int CONNECTION_PROXY::error_code() {
 }
 
 const char* CONNECTION_PROXY::error() {
+    if (has_custom_error_message) {		
+        // We disable this flag after fetching the custom message once
+        // so it does not obscure future proxy errors.
+        has_custom_error_message = false;
+
+        return this->custom_error_message.c_str();
+    }
+
     return next_proxy->error();
 }
 
@@ -393,3 +411,8 @@ void CONNECTION_PROXY::set_next_proxy(CONNECTION_PROXY* next_proxy) {
 MYSQL* CONNECTION_PROXY::move_mysql_connection() {
     return next_proxy ? next_proxy->move_mysql_connection() : nullptr;
 }
+
+void CONNECTION_PROXY::set_custom_error_message(const char* error_message) {
+    this->custom_error_message = error_message;
+    has_custom_error_message = true;
+}
diff --git a/driver/connection_proxy.h b/driver/connection_proxy.h
@@ -37,9 +37,19 @@ struct DataSource;
 
 class CONNECTION_PROXY {
 public:
+    CONNECTION_PROXY() = default;
     CONNECTION_PROXY(DBC* dbc, DataSource* ds);
     virtual ~CONNECTION_PROXY();
 
+    virtual bool connect(
+        const char* host,
+        const char* user,
+        const char* password,
+        const char* database,
+        unsigned int port,
+        const char* socket,
+        unsigned long flags);
+
     virtual void delete_ds();
     virtual uint64_t num_rows(MYSQL_RES* res);
     virtual unsigned int num_fields(MYSQL_RES* res);
@@ -163,10 +173,14 @@ class CONNECTION_PROXY {
 
     virtual MYSQL* move_mysql_connection();
 
+    void set_custom_error_message(const char* error_message);
+
 protected:
     DBC* dbc = nullptr;
     DataSource* ds = nullptr;
     CONNECTION_PROXY* next_proxy = nullptr;
+    bool has_custom_error_message = false;
+    std::string custom_error_message = "";
 };
 
 #endif /* __CONNECTION_PROXY__ */
diff --git a/driver/handle.cc b/driver/handle.cc
@@ -49,6 +49,7 @@
 
 #include "driver.h"
 #include "efm_proxy.h"
+#include "iam_proxy.h"
 #include "mysql_proxy.h"
 
 #include <mutex>
@@ -125,18 +126,18 @@ void DBC::init_proxy_chain(DataSource* dsrc)
         head = efm_proxy;
     }
 
-    ds_get_utf8attr(dsrc->auth_mode, &dsrc->auth_mode8);
-
-    if (!myodbc_strcasecmp(AUTH_MODE_IAM, reinterpret_cast<const char*>(dsrc->auth_mode8))) {
-        // CONNECTION_PROXY* iam_proxy = new IAM_PROXY(his, dsrc);
-        // iam_proxy->set_next_proxy(head);
-        // head = iam_proxy;
-    }
-
-    if (!myodbc_strcasecmp(AUTH_MODE_SECRETS_MANAGER, reinterpret_cast<const char*>(dsrc->auth_mode8))) {
-        // CONNECTION_PROXY* secrets_manager_proxy = new SECRETS_MANAGER_PROXY(his, dsrc);
-        // secrets_manager_proxy->set_next_proxy(head);
-        // head = secrets_manager_proxy;
+    if (dsrc->auth_mode) {
+        const char* auth_mode = ds_get_utf8attr(dsrc->auth_mode, &dsrc->auth_mode8);
+        if (!myodbc_strcasecmp(AUTH_MODE_IAM, auth_mode)) {
+            CONNECTION_PROXY* iam_proxy = new IAM_PROXY(this, dsrc);
+            iam_proxy->set_next_proxy(head);
+            head = iam_proxy;
+        }
+        else if (!myodbc_strcasecmp(AUTH_MODE_SECRETS_MANAGER, auth_mode)) {
+            // CONNECTION_PROXY* secrets_manager_proxy = new SECRETS_MANAGER_PROXY(his, dsrc);
+            // secrets_manager_proxy->set_next_proxy(head);
+            // head = secrets_manager_proxy;
+        }
     }
 
     this->connection_proxy = head;