diff --git a/.github/workflows/CI-unit-tests-tsan.yml b/.github/workflows/CI-unit-tests-tsan.yml index ad4d238ef2..0f6971114b 100644 --- a/.github/workflows/CI-unit-tests-tsan.yml +++ b/.github/workflows/CI-unit-tests-tsan.yml @@ -1,8 +1,8 @@ name: CI-unit-tests-tsan run-name: '${{ github.event.workflow_run && github.event.workflow_run.head_branch || github.ref_name }} ${{ github.workflow }} ${{ github.event.workflow_run && github.event.workflow_run.head_sha || github.sha }}' -# Builds the mysqlx + plugin-chassis unit tests with WITHTSAN=1 and -# runs them under ThreadSanitizer to catch race conditions. Closes +# Builds the mysqlx, plugin-chassis, and provider-neutral IAM concurrency tests +# with WITHTSAN=1 and runs them under ThreadSanitizer to catch race conditions. Closes # Phase 2 of issue #5675. # # Architecture (vs. the earlier closed PR #5720): @@ -120,7 +120,7 @@ jobs: run: | make ubuntu24-tap - - name: Run mysqlx-tsan-g1 TAP group inside Docker + - name: Run concurrency TAP group inside Docker # Run the TSAN-instrumented unit tests INSIDE the same Docker # image used for the build — same source mount at /opt/proxysql, # same toolchain, same libstdc++/libgcc that the binaries linked diff --git a/README.md b/README.md index 714ce1d1b0..8c35e38f82 100644 --- a/README.md +++ b/README.md @@ -76,6 +76,8 @@ tar xzf proxysql--linux-amd64.tar.gz The archive contains `bin/proxysql`, a sample `etc/proxysql.cnf`, the `systemd/` units, and helper tools. The v4.0 build additionally ships the runtime plugins under `lib/proxysql/` (`ProxySQL_MySQLX_Plugin.so`, `ProxySQL_GenAI_Plugin.so`). +See [AWS locality-aware backend selection](doc/aws-locality-awareness.md) for +the optional external-provider contract and MySQL configuration controls. Alternatively you can also use the available repositories: diff --git a/deps/Makefile b/deps/Makefile index b05fcb108c..6bac398aca 100644 --- a/deps/Makefile +++ b/deps/Makefile @@ -2,7 +2,11 @@ PROXYSQL_PATH := $(shell while [ ! -f ./src/proxysql_global.cpp ]; do cd ..; done; pwd) +.DEFAULT_GOAL := default + include $(PROXYSQL_PATH)/include/makefiles_vars.mk +include $(PROXYSQL_PATH)/common_mk/openssl_flags.mk +include $(PROXYSQL_PATH)/common_mk/openssl_version_check.mk # to compile libmariadb_client with support for valgrind enabled, run: @@ -80,10 +84,6 @@ default: $(targets) ### deps targets -include $(PROXYSQL_PATH)/common_mk/openssl_flags.mk -include $(PROXYSQL_PATH)/common_mk/openssl_version_check.mk - - libinjection/libinjection/src/libinjection.a: cd libinjection && rm -rf libinjection-*/ || true cd libinjection && tar -zxf libinjection-3.10.0.tar.gz @@ -277,6 +277,7 @@ endif cd mariadb-client-library/mariadb_client && patch -p0 < ../ma_password.c.patch # cd mariadb-client-library/mariadb_client && patch libmariadb/ma_secure.c < ../ma_secure.c.patch cd mariadb-client-library/mariadb_client && patch -p0 < ../mysql.h.patch + cd mariadb-client-library/mariadb_client && patch -p0 < ../tls_server_name.patch cd mariadb-client-library/mariadb_client && patch -p0 < ../ma_priv.h.patch cd mariadb-client-library/mariadb_client && patch -p0 < ../ma_alloc.c.patch cd mariadb-client-library/mariadb_client && patch -p0 < ../ma_charset.c.patch diff --git a/deps/mariadb-client-library/mariadb_lib.c.patch b/deps/mariadb-client-library/mariadb_lib.c.patch index ab3bbaf2c3..96db50b99e 100644 --- a/deps/mariadb-client-library/mariadb_lib.c.patch +++ b/deps/mariadb-client-library/mariadb_lib.c.patch @@ -69,7 +69,7 @@ index e8db51a0..684aff1a 100644 mysql->stmts= NULL; } } -@@ -2458,6 +2475,42 @@ mysql_close(MYSQL *mysql) +@@ -2458,6 +2475,58 @@ mysql_close(MYSQL *mysql) return; } @@ -85,6 +85,22 @@ index e8db51a0..684aff1a 100644 + mysql->options.reconnect=0; + end_server(mysql); + } ++ else if (mysql->options.extension && ++ mysql->options.extension->async_context) ++ { ++ struct mysql_async_context *ctxt= ++ mysql->options.extension->async_context; ++ if (ctxt->pending_gai_res) ++ { ++ freeaddrinfo(ctxt->pending_gai_res); ++ ctxt->pending_gai_res= 0; ++ } ++ if (ctxt->pvio) ++ { ++ ma_pvio_close(ctxt->pvio); ++ ctxt->pvio= 0; ++ } ++ } + } + mysql_close_memory(mysql); + mysql_close_options(mysql); diff --git a/deps/mariadb-client-library/sslkeylogfile.patch b/deps/mariadb-client-library/sslkeylogfile.patch index 76766679dc..7712e8b132 100644 --- a/deps/mariadb-client-library/sslkeylogfile.patch +++ b/deps/mariadb-client-library/sslkeylogfile.patch @@ -14,13 +14,14 @@ diff --git include/mysql.h include/mysql.h index 9ee86227..c07717c5 100644 --- include/mysql.h +++ include/mysql.h -@@ -257,7 +257,8 @@ extern const char *SQLSTATE_UNKNOWN; +@@ -257,8 +257,9 @@ extern const char *SQLSTATE_UNKNOWN; MARIADB_OPT_RESTRICTED_AUTH, MARIADB_OPT_RPL_REGISTER_REPLICA, MARIADB_OPT_STATUS_CALLBACK, -- MARIADB_OPT_SERVER_PLUGINS -+ MARIADB_OPT_SERVER_PLUGINS, -+ MARIADB_OPT_SSL_KEYLOG_CALLBACK + MARIADB_OPT_SERVER_PLUGINS, +- MARIADB_OPT_TLS_SERVER_NAME = MARIADB_OPT_SERVER_PLUGINS + 2 ++ MARIADB_OPT_SSL_KEYLOG_CALLBACK = MARIADB_OPT_SERVER_PLUGINS + 1, ++ MARIADB_OPT_TLS_SERVER_NAME = MARIADB_OPT_SERVER_PLUGINS + 2 }; enum mariadb_value { diff --git a/deps/mariadb-client-library/tls_server_name.patch b/deps/mariadb-client-library/tls_server_name.patch new file mode 100644 index 0000000000..a8cdd9ba29 --- /dev/null +++ b/deps/mariadb-client-library/tls_server_name.patch @@ -0,0 +1,199 @@ +diff --git include/ma_common.h include/ma_common.h +--- include/ma_common.h ++++ include/ma_common.h +@@ -73,6 +73,7 @@ struct st_mysql_options_extension { + unsigned int tls_cipher_strength; + char *tls_version; + my_bool read_only; ++ char *tls_server_name; + char *connection_handler; + my_bool (*set_option)(MYSQL *mysql, const char *config_option, const char *config_value); + MA_HASHTBL userdata; +diff --git include/ma_tls.h include/ma_tls.h +--- include/ma_tls.h ++++ include/ma_tls.h +@@ -148,6 +148,7 @@ int ma_pvio_tls_get_protocol_version_id(MARIADB_TLS *ctls); + void ma_tls_set_connection(MYSQL *mysql); + + /* Function prototypes */ ++const char *ma_tls_get_server_name(MYSQL *mysql); + MARIADB_TLS *ma_pvio_tls_init(MYSQL *mysql); + my_bool ma_pvio_tls_connect(MARIADB_TLS *ctls); + ssize_t ma_pvio_tls_read(MARIADB_TLS *ctls, const uchar *buffer, size_t length); +diff --git include/mysql.h include/mysql.h +--- include/mysql.h ++++ include/mysql.h +@@ -260,6 +260,7 @@ extern const char *SQLSTATE_UNKNOWN; +- MARIADB_OPT_SERVER_PLUGINS ++ MARIADB_OPT_SERVER_PLUGINS, ++ MARIADB_OPT_TLS_SERVER_NAME = MARIADB_OPT_SERVER_PLUGINS + 2 + }; + + enum mariadb_value { + MARIADB_CHARSET_ID, + MARIADB_CHARSET_NAME, +diff --git libmariadb/ma_tls.c libmariadb/ma_tls.c +--- libmariadb/ma_tls.c ++++ libmariadb/ma_tls.c +@@ -51,6 +51,14 @@ + my_bool ma_tls_initialized= FALSE; + unsigned int mariadb_deinitialize_ssl= 1; + ++const char *ma_tls_get_server_name(MYSQL *mysql) ++{ ++ if (mysql->options.extension && mysql->options.extension->tls_server_name && ++ mysql->options.extension->tls_server_name[0]) ++ return mysql->options.extension->tls_server_name; ++ return mysql->host; ++} ++ + const char *tls_protocol_version[]= + {"SSLv3", "TLSv1.0", "TLSv1.1", "TLSv1.2", "TLSv1.3", "Unknown"}; + +diff --git libmariadb/mariadb_lib.c libmariadb/mariadb_lib.c +--- libmariadb/mariadb_lib.c ++++ libmariadb/mariadb_lib.c +@@ -2304,6 +2304,7 @@ static void mysql_close_options(MYSQL *mysql) + free(mysql->options.extension->tls_fp_list); + free(mysql->options.extension->tls_pw); + free(mysql->options.extension->tls_version); ++ free(mysql->options.extension->tls_server_name); + free(mysql->options.extension->url); + free(mysql->options.extension->connection_handler); + free(mysql->options.extension->proxy_header); +@@ -3768,6 +3769,9 @@ mysql_optionsv(MYSQL *mysql,enum mysql_option option, ...) + case MYSQL_OPT_TLS_VERSION: + OPT_SET_EXTENDED_VALUE_STR(&mysql->options, tls_version, (char *)arg1); + break; ++ case MARIADB_OPT_TLS_SERVER_NAME: ++ OPT_SET_EXTENDED_VALUE_STR(&mysql->options, tls_server_name, (char *)arg1); ++ break; + case MARIADB_OPT_IO_WAIT: + CHECK_OPT_EXTENSION_SET(&mysql->options); + mysql->options.extension->io_wait = (int(*)(my_socket, my_bool, int))arg1; +@@ -3927,6 +3931,9 @@ mysql_get_optionv(MYSQL *mysql, enum mysql_option option, void *arg, ...) + case MYSQL_OPT_TLS_VERSION: + *((char **)arg)= mysql->options.extension ? mysql->options.extension->tls_version : NULL; + break; ++ case MARIADB_OPT_TLS_SERVER_NAME: ++ *((char **)arg)= mysql->options.extension ? mysql->options.extension->tls_server_name : NULL; ++ break; + case MYSQL_OPT_CONNECT_ATTRS: + /* mysql_get_optionsv(mysql, MYSQL_OPT_CONNECT_ATTRS, keys, vals, elements) */ + { +diff --git libmariadb/secure/gnutls.c libmariadb/secure/gnutls.c +--- libmariadb/secure/gnutls.c ++++ libmariadb/secure/gnutls.c +@@ -1176,6 +1176,7 @@ my_bool ma_tls_connect(MARIADB_TLS *ctls) + MYSQL *mysql= (MYSQL *)gnutls_session_get_ptr(ssl); + MARIADB_PVIO *pvio; + int ret; ++ const char *server_name; + + if (!mysql) + return 1; +@@ -1198,6 +1199,11 @@ my_bool ma_tls_connect(MARIADB_TLS *ctls) + gnutls_transport_set_int(ssl, mysql_get_socket(mysql)); + #endif + ++ server_name= ma_tls_get_server_name(mysql); ++ if (server_name && gnutls_server_name_set(ssl, GNUTLS_NAME_DNS, ++ server_name, strlen(server_name)) < 0) ++ return 1; ++ + do { + ret = gnutls_handshake(ssl); + } while (ret < 0 && gnutls_error_is_fatal(ret) == 0); +@@ -1359,7 +1365,7 @@ static int my_verify_callback(gnutls_session_t ssl) + + if ((mysql->options.extension->tls_verify_server_cert)) + { +- const char *hostname= mysql->host; ++ const char *hostname= ma_tls_get_server_name(mysql); + + if (gnutls_certificate_verify_peers3 (ssl, hostname, &status) < 0) + return GNUTLS_E_CERTIFICATE_ERROR; +diff --git libmariadb/secure/ma_schannel.c libmariadb/secure/ma_schannel.c +--- libmariadb/secure/ma_schannel.c ++++ libmariadb/secure/ma_schannel.c +@@ -300,7 +300,7 @@ SECURITY_STATUS ma_schannel_client_handshake(MARIADB_TLS *ctls) + + sRet = InitializeSecurityContext(&sctx->CredHdl, + NULL, +- pvio->mysql->host, ++ (SEC_CHAR *)ma_tls_get_server_name(pvio->mysql), + SFlags, + 0, + SECURITY_NATIVE_DREP, +@@ -514,7 +514,7 @@ my_bool ma_schannel_verify_certs(MARIADB_TLS *ctls, BOOL verify_server_name) + pServerCert, + store, + crl_file != 0 || crl_path != 0, +- mysql->host, ++ ma_tls_get_server_name(mysql), + verify_server_name, + errmsg, sizeof(errmsg)); + +diff --git libmariadb/secure/openssl.c libmariadb/secure/openssl.c +--- libmariadb/secure/openssl.c ++++ libmariadb/secure/openssl.c +@@ -463,6 +463,7 @@ my_bool ma_tls_connect(MARIADB_TLS *ctls) + MYSQL *mysql; + MARIADB_PVIO *pvio; + int rc; ++ const char *server_name; + #ifdef OPENSSL_USE_BIOMETHOD + BIO_METHOD *bio_method= NULL; + BIO *bio; +@@ -486,6 +487,13 @@ my_bool ma_tls_connect(MARIADB_TLS *ctls) + SSL_set_fd(ssl, (int)mysql_get_socket(mysql)); + #endif + ++ server_name= ma_tls_get_server_name(mysql); ++ if (server_name && !SSL_set_tlsext_host_name(ssl, server_name)) ++ { ++ ma_tls_set_error(mysql); ++ return 1; ++ } ++ + while (try_connect && (rc= SSL_connect(ssl)) == -1) + { + switch((SSL_get_error(ssl, rc))) { +@@ -659,6 +667,7 @@ int ma_tls_verify_server_cert(MARIADB_TLS *ctls) + MYSQL *mysql; + SSL *ssl; + MARIADB_PVIO *pvio; ++ const char *server_name; + #if !defined(HAVE_OPENSSL_CHECK_HOST) + X509_NAME *x509sn; + int cn_pos; +@@ -673,7 +682,8 @@ int ma_tls_verify_server_cert(MARIADB_TLS *ctls) + mysql= (MYSQL *)SSL_get_app_data(ssl); + pvio= mysql->net.pvio; + +- if (!mysql->host) ++ server_name= ma_tls_get_server_name(mysql); ++ if (!server_name) + { + pvio->set_error(mysql, CR_SSL_CONNECTION_ERROR, SQLSTATE_UNKNOWN, + ER(CR_SSL_CONNECTION_ERROR), "Invalid (empty) hostname"); +@@ -687,8 +697,8 @@ int ma_tls_verify_server_cert(MARIADB_TLS *ctls) + return 1; + } + #ifdef HAVE_OPENSSL_CHECK_HOST +- if (X509_check_host(cert, mysql->host, strlen(mysql->host), 0, 0) != 1 +- && X509_check_ip_asc(cert, mysql->host, 0) != 1) ++ if (X509_check_host(cert, server_name, strlen(server_name), 0, 0) != 1 ++ && X509_check_ip_asc(cert, server_name, 0) != 1) + goto error; + #else + x509sn= X509_get_subject_name(cert); +@@ -708,7 +718,7 @@ int ma_tls_verify_server_cert(MARIADB_TLS *ctls) + if ((size_t)ASN1_STRING_length(cn_asn1) != strlen(cn_str)) + goto error; + +- if (strcmp(cn_str, mysql->host)) ++ if (strcmp(cn_str, server_name)) + goto error; + #endif + X509_free(cert); diff --git a/doc/PLUGIN_API.md b/doc/PLUGIN_API.md index 6adca93746..80ec1cc79f 100644 --- a/doc/PLUGIN_API.md +++ b/doc/PLUGIN_API.md @@ -103,7 +103,7 @@ All types are defined in `include/ProxySQL_Plugin.h`: ```cpp struct ProxySQL_PluginDescriptor { const char *name; // Human-readable plugin name - uint32_t abi_version; // PROXYSQL_PLUGIN_ABI_VERSION (1, 2, or 3) + uint32_t abi_version; // PROXYSQL_PLUGIN_ABI_VERSION (1 through 8) proxysql_plugin_init_cb init; // bool (*)(ProxySQL_PluginServices *) proxysql_plugin_start_cb start; // bool (*)() proxysql_plugin_stop_cb stop; // bool (*)() @@ -115,7 +115,7 @@ struct ProxySQL_PluginDescriptor { | Field | Type | Description | |--------------------|---------------|-----------------------------------------------------------| | `name` | `const char*` | Plugin identifier, used in logging. | -| `abi_version` | `uint32_t` | Set from `PROXYSQL_PLUGIN_ABI_VERSION`. Value `1` = pre-chassis descriptor (six fields). Value `2` = adds `register_schemas` (four-phase lifecycle). Value `3` = same descriptor layout as `2`; `ProxySQL_PluginServices` adds a tail-appended `register_runtime_view`. A v3/v3.1 ProxySQL core rejects `abi_version > 1`; the current PROXYSQL40 core accepts `[1, 3]`. | +| `abi_version` | `uint32_t` | Set from `PROXYSQL_PLUGIN_ABI_VERSION`. Value `1` is the pre-chassis six-field descriptor; ABI 2 adds `register_schemas`; ABI 3 adds `register_runtime_view`; ABI 4 adds the view's `db_kind`; ABI 5 adds IAM provider install/limits; ABI 6 adds metadata-provider install; ABI 7 adds the MySQL locality projection callback; ABI 8 adds IAM provider rollback. A v3/v3.1 core rejects `abi_version > 1`; the current PROXYSQL40 core accepts `[1, 8]`. | | `init` | callback | Phase D — called with live services; register tables and commands here (or finish context setup if `register_schemas` already did it). | | `start` | callback | Phase E — start threads, open sockets, load config. | | `stop` | callback | Called on shutdown. Pairs with `init`, not `start`: if `init` returned true and `start` later failed, `stop` is still called so the plugin can release resources it allocated in `init`. | @@ -128,16 +128,16 @@ Return `true` on success, `false` on failure. A `false` return from #### ABI version -`include/ProxySQL_Plugin.h` exposes `PROXYSQL_PLUGIN_ABI_VERSION` (3 under +`include/ProxySQL_Plugin.h` exposes `PROXYSQL_PLUGIN_ABI_VERSION` (8 under PROXYSQL40, undefined in pre-chassis builds — the descriptor is then a legacy six-field struct with `abi_version = 1`). Plugins MUST assign `abi_version` from this macro rather than hard-coding a literal; the core's loader uses it to detect layout skew and reject plugins built -for an unsupported ABI. ABI 3 keeps the descriptor layout identical to -ABI 2 — the only addition is a tail-appended `register_runtime_view` -field on `ProxySQL_PluginServices` — so plugins that compile against -ABI 2 still load on the current core; the trailing services field is -simply invisible to them. See `ProxySQL_Plugin.h` for the exact rules. +for an unsupported ABI. ABIs 3 through 8 keep the descriptor layout +identical to ABI 2; each adds only tail fields to service or view structs. +Plugins compiled against an older ABI therefore still load on the current +core, where the trailing fields remain invisible to them. See +`ProxySQL_Plugin.h` for the exact rules. ### The Entry Point @@ -169,9 +169,44 @@ struct ProxySQL_PluginServices { proxysql_plugin_register_command_alias_cb register_command_alias; // ABI 3 tail extension: proxysql_plugin_register_runtime_view_cb register_runtime_view; + // ABI 5 tail extensions, live only while init() is running: + proxysql_plugin_install_aws_iam_token_source_cb install_aws_iam_token_source; + proxysql_plugin_get_aws_iam_limits_cb get_aws_iam_limits; + // ABI 6 tail extension, live only while init() is running: + proxysql_plugin_install_aws_metadata_provider_cb install_aws_metadata_provider; + // ABI 7 tail extension, live during register_schemas() and init(): + proxysql_plugin_refresh_mysql_aws_locality_stats_cb refresh_mysql_aws_locality_stats; + // ABI 8 tail extension, live only while init() is running: + proxysql_plugin_uninstall_aws_iam_token_source_cb uninstall_aws_iam_token_source; }; ``` +`install_aws_iam_token_source` allows an optional external provider to supply +IAM database-authentication tokens. The provider passes a newly allocated +source, a destroy callback, and an extra `dlopen()` handle. Core retains that +handle until all session leases drain, then destroys the source and closes the +module. A plugin must not call this callback outside `init()` or retain the +service pointer after initialization. Without an installed provider, IAM +authentication remains available as a policy but fails closed. + +`install_aws_metadata_provider` installs an optional external asynchronous +metadata provider for MySQL locality selection. Core owns the retained lease +registry: new leases are rejected during shutdown, active manager/callback +leases drain, the provider's `shutdown()` and destroy callback run, and the +extra module handle closes last. Without an installed provider, locality uses +configured weights and reports the fixed `provider_unavailable` category. + +`refresh_mysql_aws_locality_stats` projects the current MySQL-owned immutable +locality snapshot into a schema supplied by the calling plugin. It performs no +provider or network I/O and is live during `register_schemas()` so an external +provider can register its stats table and runtime-view callback together. +Public core deliberately registers no locality table on its own. + +`uninstall_aws_iam_token_source` lets the same plugin roll back its IAM source +when a later initialization step fails. It rejects a null or different source, +stops new leases, drains active leases, destroys the source, and releases the +retained module handle before returning. + ### Service Callbacks #### `register_table` @@ -542,7 +577,7 @@ void register_stats_table(ProxySQL_PluginServices& services, - **No dependency resolution**: Plugins are loaded in the order listed in `proxysql.cnf`. If one plugin depends on another, the dependency must be listed first. -- **ABI version range**: The current core accepts `abi_version` values in `[1, 3]`. Newly built plugins should set `abi_version = PROXYSQL_PLUGIN_ABI_VERSION`. +- **ABI version range**: The current core accepts `abi_version` values in `[1, 8]`. Newly built plugins should set `abi_version = PROXYSQL_PLUGIN_ABI_VERSION`. - **Compiler coupling**: Plugins must match the ProxySQL core's C++ compiler and standard library due to `std::string` in `ProxySQL_PluginCommandResult`. diff --git a/doc/README.md b/doc/README.md index 9da8eeafe6..648186b9a6 100644 --- a/doc/README.md +++ b/doc/README.md @@ -51,4 +51,4 @@ Detailed technical documentation for specific components: --- -> This approach provides the benefits of AI-generated assistance while being transparent about limitations and reducing verification overhead. \ No newline at end of file +> This approach provides the benefits of AI-generated assistance while being transparent about limitations and reducing verification overhead. diff --git a/doc/aws-locality-awareness.md b/doc/aws-locality-awareness.md new file mode 100644 index 0000000000..fa16f1fc50 --- /dev/null +++ b/doc/aws-locality-awareness.md @@ -0,0 +1,114 @@ +# AWS locality-aware MySQL backend selection + +ProxySQL 4.0 can apply temporary locality multipliers while selecting eligible +Amazon RDS and Aurora MySQL backends. The MySQL module owns the configuration, +policy validation, immutable selection snapshot, and effective-weight +calculation. Locality never changes `mysql_servers.weight`, +`runtime_mysql_servers.weight`, saved configuration, or ProxySQL Cluster +checksums. + +Locality metadata is supplied asynchronously by an optional compatible +external provider. Without a provider, or when metadata is unavailable or too +old, selection stays neutral and uses the configured server weights. + +## Hostgroup policy + +Add `aws.locality_awareness` to the existing +`mysql_hostgroup_attributes.hostgroup_settings` JSON. Both multipliers are +required: + +```sql +INSERT INTO mysql_hostgroup_attributes(hostgroup_id, hostgroup_settings) +VALUES ( + 10, + '{ + "aws": { + "locality_awareness": { + "same_region_multiplier": 2.0, + "same_az_multiplier": 4.0, + "refresh_interval_seconds": 300, + "stale_ttl_seconds": 1800 + } + } + }' +); + +LOAD MYSQL SERVERS TO RUNTIME; +SAVE MYSQL SERVERS TO DISK; +``` + +The accepted values are: + +```text +1.0 <= same_region_multiplier <= same_az_multiplier <= 10.0 + +refresh_interval_seconds default: 300 +stale_ttl_seconds default: 1800 + +30 <= refresh_interval_seconds <= 86400 +refresh_interval_seconds <= stale_ttl_seconds <= 604800 +``` + +Multipliers are JSON numbers. An invalid `aws.locality_awareness` object +disables locality bias for that hostgroup when servers are loaded. Diagnostics +identify the rejected field and hostgroup without logging the supplied value. + +## Master switch + +The process-wide MySQL variable defaults to `false`: + +```sql +SET mysql-aws_locality_awareness = true; +LOAD MYSQL VARIABLES TO RUNTIME; +SAVE MYSQL VARIABLES TO DISK; +``` + +Disabling the variable immediately restores configured-weight selection, +cancels or supersedes outstanding locality requests, and stops new refresh +scheduling. It does not change existing backend connections or server rows. + +## Selection contract + +For each selection attempt, after the normal health, capacity, lag, GTID, +backoff, and session-compatibility checks, ProxySQL calculates: + +```text +remote or unknown configured_weight +same Region, different AZ int(configured_weight * same_region_multiplier) +same AZ int(configured_weight * same_az_multiplier) +``` + +Conversion to an integer truncates toward zero. The tiers are not cumulative, +and configured weight zero stays zero. Global hostgroup selection and +thread-local idle-connection reuse use the same immutable snapshot and never +perform provider or network work on the selection path. + +ProxySQL recognizes eligible RDS and Aurora endpoint shapes to form neutral +metadata requests. A compatible provider is responsible for authoritative +endpoint discovery and normalized Region, Availability Zone, and account +metadata. Custom CNAMEs, arbitrary MySQL hosts, proxy endpoints, malformed +names, and endpoints the provider cannot confirm remain neutral. + +## Provider absence and failures + +The public provider interface uses retained leases so shutdown rejects new +work, drains active callbacks, joins the locality manager worker, and only then +permits provider destruction and module unload. Provider absence is exposed as +the fixed `provider_unavailable` category. + +Metadata states are `pending`, `fresh`, `stale`, `expired`, `error`, and +`disabled`. Only `fresh` and unexpired `stale` values activate a multiplier. +All other states use multiplier `1.0`, so effective weights equal configured +weights. A failed refresh can retain the last successful value through the +bounded stale TTL; it becomes neutral after expiry. + +An external provider may register a read-only runtime diagnostics table using +the plugin table/view services and the MySQL-owned projection callback. Public +core does not register `stats_mysql_aws_locality`; without a provider that +registers it, the table does not exist. + +## Rollback + +Set `mysql-aws_locality_awareness` to `false` and load MySQL variables to +runtime. To remove a policy, delete the `aws.locality_awareness` object from +that hostgroup's `hostgroup_settings`, then load MySQL servers to runtime. diff --git a/include/Aws_Iam_Provider.h b/include/Aws_Iam_Provider.h new file mode 100644 index 0000000000..44cd0af632 --- /dev/null +++ b/include/Aws_Iam_Provider.h @@ -0,0 +1,203 @@ +#ifndef AWS_IAM_PROVIDER_H +#define AWS_IAM_PROVIDER_H + +#include "Aws_Iam_Types.h" + +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace prometheus { class Registry; } + +class SecureString { +public: + // A custom cleanser must remain callable for the SecureString's full + // lifetime. Provider DSOs must use the process-lifetime default cleanser + // for values that can outlive the provider module. + using CleanseFn = void (*)(void*, size_t); + + SecureString() = default; + explicit SecureString(std::string_view value, CleanseFn cleanse = OPENSSL_cleanse) + : size_(value.size()), cleanse_(cleanse ? cleanse : OPENSSL_cleanse) { + if (size_ != 0) { + bytes_.reset(new unsigned char[size_ + 1]); + std::memcpy(bytes_.get(), value.data(), size_); + bytes_[size_] = 0; + } + } + SecureString(SecureString&& other) noexcept + : bytes_(std::move(other.bytes_)), size_(other.size_), cleanse_(other.cleanse_) { + other.size_ = 0; + other.cleanse_ = OPENSSL_cleanse; + } + SecureString& operator=(SecureString&& other) noexcept { + if (this != &other) { + clear(); + bytes_ = std::move(other.bytes_); + size_ = other.size_; + cleanse_ = other.cleanse_; + other.size_ = 0; + other.cleanse_ = OPENSSL_cleanse; + } + return *this; + } + SecureString(const SecureString&) = delete; + SecureString& operator=(const SecureString&) = delete; + ~SecureString() { clear(); } + + SecureString clone() const { + return empty() ? SecureString() : + SecureString(std::string_view(c_str(), size_), cleanse_); + } + const char* c_str() const { + return bytes_ ? reinterpret_cast(bytes_.get()) : ""; + } + size_t size() const { return size_; } + bool empty() const { return size_ == 0; } + void clear() { + if (bytes_) { + cleanse_(bytes_.get(), size_); + bytes_[size_] = 0; + bytes_.reset(); + } + size_ = 0; + } + +private: + std::unique_ptr bytes_; + size_t size_ { 0 }; + CleanseFn cleanse_ { OPENSSL_cleanse }; +}; + +enum class AwsIamStatus : uint8_t { + OK, SUPPORT_NOT_COMPILED, INVALID_CONFIG, PROVIDER_ERROR, + CREDENTIAL_PROVIDER_ERROR, QUEUE_FULL, WAITER_LIMIT, TIMEOUT, + CANCELED, SHUTDOWN, +}; + +struct AwsIamRedactedFailure { + std::string category; + std::string aws_error_code; + std::string request_id; +}; + +struct AwsIamTokenResult { + AwsIamStatus status { AwsIamStatus::PROVIDER_ERROR }; + SecureString token; + std::chrono::steady_clock::time_point expires_at {}; + uint64_t generation { 0 }; + AwsIamRedactedFailure failure; +}; + +struct AwsIamCompletion { + uint64_t opaque_id { 0 }; + AwsIamTokenResult result; +}; + +class AwsIamCompletionSink { +public: + virtual void post(AwsIamCompletion&&) = 0; + virtual ~AwsIamCompletionSink() = default; +}; + +struct AwsIamRequestHandle { uint64_t value { 0 }; }; + +struct AwsIamStatsSnapshot { + uint64_t token_requests { 0 }; + uint64_t token_cache_hits { 0 }; + uint64_t token_refresh_successes { 0 }; + uint64_t token_refresh_failures { 0 }; + uint64_t credential_provider_failures { 0 }; + uint64_t queue_rejections { 0 }; + uint64_t backend_connection_successes { 0 }; + uint64_t backend_connection_failures { 0 }; + uint64_t token_cache_entries { 0 }; + uint64_t in_flight_generations { 0 }; + uint64_t queued_generations { 0 }; + uint64_t waiting_sessions { 0 }; +}; + +struct AwsIamNamedStat { + const char *name; + uint64_t value; +}; + +using AwsIamNamedStats = std::array; + +AwsIamNamedStats aws_iam_stats_mysql_global_rows(const AwsIamStatsSnapshot&); +void initialize_aws_iam_prometheus_metrics(prometheus::Registry&); +void update_aws_iam_prometheus_metrics(const AwsIamStatsSnapshot&); + +class AwsIamTokenSource { +public: + virtual bool support_compiled() const { return true; } + virtual AwsIamRequestHandle request(const AwsIamTokenKey&, uint64_t opaque_id, + std::weak_ptr) = 0; + virtual AwsIamTokenResult request_blocking(const AwsIamTokenKey&, + std::chrono::steady_clock::time_point deadline) = 0; + virtual void cancel(AwsIamRequestHandle) = 0; + virtual void invalidate(const AwsIamTokenKey&, uint64_t generation) = 0; + virtual void record_backend_connection(bool success) = 0; + virtual void record_waiting_session(bool waiting) = 0; + virtual AwsIamStatsSnapshot snapshot() const = 0; + virtual ~AwsIamTokenSource() = default; +}; + +struct AwsIamRuntimeConfig { + size_t max_total_waiters; + size_t max_waiters_per_key; +}; + +using AwsIamTokenSourceDestroyFn = void (*)(AwsIamTokenSource *); + +class AwsIamTokenSourceLease { +public: + AwsIamTokenSourceLease() = default; + ~AwsIamTokenSourceLease(); + AwsIamTokenSourceLease(AwsIamTokenSourceLease&& other) noexcept; + AwsIamTokenSourceLease& operator=(AwsIamTokenSourceLease&& other) noexcept; + + AwsIamTokenSource *get() const { return source_; } + AwsIamTokenSource *operator->() const { return source_; } + explicit operator bool() const { return source_ != nullptr; } + + AwsIamTokenSourceLease(const AwsIamTokenSourceLease&) = delete; + AwsIamTokenSourceLease& operator=(const AwsIamTokenSourceLease&) = delete; + +private: + explicit AwsIamTokenSourceLease(AwsIamTokenSource *source) : source_(source) {} + void release(); + AwsIamTokenSource *source_ { nullptr }; + + friend AwsIamTokenSourceLease acquire_global_aws_iam_token_source(); +}; + +std::unique_ptr create_aws_iam_token_source( + const AwsIamRuntimeConfig& config); + +void publish_global_aws_iam_token_source(AwsIamTokenSource *source); +AwsIamTokenSourceLease acquire_global_aws_iam_token_source(); +void shutdown_global_aws_iam_token_source(); + +// Transfers ownership of a provider created by a dynamically loaded plugin. +// `module_handle` is an extra dlopen() reference retained by core until every +// source lease has drained and `destroy` has run. +// `destroy` runs while core still owns the retirement claim. It must not call +// uninstall_global_aws_iam_token_source() or +// shutdown_global_aws_iam_token_source(), because either call would wait for +// that same claim and deadlock. +bool install_global_aws_iam_token_source( + AwsIamTokenSource *source, AwsIamTokenSourceDestroyFn destroy, void *module_handle); +bool uninstall_global_aws_iam_token_source(AwsIamTokenSource *expected_source); + +extern AwsIamTokenSource* GloAwsIamTokenSource; + +#endif diff --git a/include/Aws_Iam_Types.h b/include/Aws_Iam_Types.h new file mode 100644 index 0000000000..2b3df54ac8 --- /dev/null +++ b/include/Aws_Iam_Types.h @@ -0,0 +1,19 @@ +#ifndef AWS_IAM_TYPES_H +#define AWS_IAM_TYPES_H + +#include +#include + +struct AwsIamTokenKey { + std::string endpoint; + uint16_t port; + std::string region; + std::string database_user; + + bool operator==(const AwsIamTokenKey& other) const { + return endpoint == other.endpoint && port == other.port && + region == other.region && database_user == other.database_user; + } +}; + +#endif // AWS_IAM_TYPES_H diff --git a/include/Aws_Locality_Manager.h b/include/Aws_Locality_Manager.h new file mode 100644 index 0000000000..cf7b0f18e2 --- /dev/null +++ b/include/Aws_Locality_Manager.h @@ -0,0 +1,141 @@ +#ifndef __CLASS_AWS_LOCALITY_MANAGER_H +#define __CLASS_AWS_LOCALITY_MANAGER_H + +#include "Aws_Locality_Types.h" +#include "json_fwd.hpp" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +AwsLocalityPolicy parse_aws_locality_policy( + const nlohmann::json& policy_json, + uint32_t hostgroup_id, + AwsLocalityPolicyError& error); + +AwsEndpointCandidate recognize_rds_endpoint( + uint32_t hostgroup_id, + std::string_view hostname, + uint16_t port); + +AwsLocalityClass classify_aws_locality( + const AwsLocalLocation& local, + const AwsBackendLocation& backend); + +uint64_t aws_locality_effective_weight( + int64_t configured_weight, + double multiplier); + +uint64_t aws_locality_saturating_add(uint64_t lhs, uint64_t rhs); +size_t aws_locality_weighted_index( + const uint64_t* weights, + size_t count, + uint64_t random_value); + +using AwsMetadataProviderDestroyFn = void (*)(AwsMetadataProvider*); + +class AwsMetadataProviderLease { +public: + AwsMetadataProviderLease() = default; + ~AwsMetadataProviderLease(); + AwsMetadataProviderLease(AwsMetadataProviderLease&& other) noexcept; + AwsMetadataProviderLease& operator=(AwsMetadataProviderLease&& other) noexcept; + + AwsMetadataProvider* get() const { return provider_; } + AwsMetadataProvider* operator->() const { return provider_; } + explicit operator bool() const { return provider_ != nullptr; } + + AwsMetadataProviderLease(const AwsMetadataProviderLease&) = delete; + AwsMetadataProviderLease& operator=(const AwsMetadataProviderLease&) = delete; + +private: + explicit AwsMetadataProviderLease(AwsMetadataProvider* provider) + : provider_(provider) {} + void release(); + AwsMetadataProvider* provider_ { nullptr }; + + friend AwsMetadataProviderLease acquire_global_aws_metadata_provider(); +}; + +bool install_global_aws_metadata_provider( + AwsMetadataProvider* provider, + AwsMetadataProviderDestroyFn destroy, + void* module_handle); +AwsMetadataProviderLease acquire_global_aws_metadata_provider(); +void shutdown_global_aws_metadata_provider(); + +struct AwsLocalitySnapshotEntry { + uint32_t hostgroup_id { 0 }; + std::string hostname; + uint16_t port { 0 }; + AwsEndpointType endpoint_type { AwsEndpointType::unknown }; + int64_t configured_weight { 0 }; + AwsLocalLocation local; + AwsBackendLocation backend; + AwsLocalityClass locality { AwsLocalityClass::unknown }; + double multiplier { 1.0 }; + AwsLocalityMetadataStatus status { AwsLocalityMetadataStatus::disabled }; + int64_t last_success_timestamp { 0 }; + int64_t last_attempt_timestamp { 0 }; + std::string failure_category; +}; + +struct AwsLocalitySnapshot { + uint64_t generation { 0 }; + bool enabled { false }; + std::unordered_multimap entries; + std::unordered_set hostgroups; + + const AwsLocalitySnapshotEntry* find( + uint32_t hostgroup_id, + std::string_view hostname, + uint16_t port) const; + uint64_t effective_weight( + uint32_t hostgroup_id, + std::string_view hostname, + uint16_t port, + int64_t configured_weight) const; + bool has_hostgroup(uint32_t hostgroup_id) const { + return hostgroups.find(hostgroup_id) != hostgroups.end(); + } +}; + +struct AwsLocalityManagerConfig { + using SteadyClock = std::function; + using WallClock = std::function; + + AwsLocalityManagerConfig(); + SteadyClock steady_clock; + WallClock wall_clock; + std::chrono::milliseconds request_timeout { std::chrono::seconds(5) }; + std::chrono::milliseconds disable_wait_timeout { std::chrono::milliseconds(250) }; + std::function before_completion; +}; + +class MySQLAwsLocalityManager { +public: + explicit MySQLAwsLocalityManager(AwsLocalityManagerConfig config = {}); + ~MySQLAwsLocalityManager(); + MySQLAwsLocalityManager(const MySQLAwsLocalityManager&) = delete; + MySQLAwsLocalityManager& operator=(const MySQLAwsLocalityManager&) = delete; + + void configure(std::vector hostgroups); + void set_enabled(bool enabled); + void request_refresh(); + std::shared_ptr snapshot() const; + std::vector diagnostic_rows() const; + void shutdown(); + +private: + class Impl; + std::unique_ptr impl_; +}; + +#endif // __CLASS_AWS_LOCALITY_MANAGER_H diff --git a/include/Aws_Locality_Types.h b/include/Aws_Locality_Types.h new file mode 100644 index 0000000000..47ff5a91f0 --- /dev/null +++ b/include/Aws_Locality_Types.h @@ -0,0 +1,169 @@ +#ifndef __CLASS_AWS_LOCALITY_TYPES_H +#define __CLASS_AWS_LOCALITY_TYPES_H + +#include +#include +#include +#include +#include +#include +#include +#include + +inline std::string aws_locality_normalized_hostname(std::string_view input) { + if (input.empty()) return {}; + std::string hostname(input); + if (hostname.back() == '.') hostname.pop_back(); + if (hostname.empty() || hostname.back() == '.') return {}; + for (char& character : hostname) { + const unsigned char value = static_cast(character); + if (!(std::isalnum(value) || character == '-' || character == '.')) return {}; + character = static_cast(std::tolower(value)); + } + return hostname; +} + +enum class AwsEndpointType : uint8_t { + unknown, + instance, + cluster, + reader, + custom, +}; + +enum class AwsLocalityClass : uint8_t { + unknown, + remote, + same_region, + same_az, +}; + +enum class AwsLocalityMetadataStatus : uint8_t { + disabled, + pending, + fresh, + stale, + expired, + error, +}; + +struct AwsLocalityPolicy { + bool valid { false }; + double same_region_multiplier { 1.0 }; + double same_az_multiplier { 1.0 }; + uint32_t refresh_interval_seconds { 300 }; + uint32_t stale_ttl_seconds { 1800 }; +}; + +struct AwsLocalityPolicyError { + uint32_t hostgroup_id { 0 }; + std::string field; +}; + +struct AwsEndpointCandidate { + bool recognized { false }; + uint32_t hostgroup_id { 0 }; + std::string hostname; + uint16_t port { 0 }; + std::string region; + std::string partition; +}; + +struct AwsLocalLocation { + std::string region; + std::string availability_zone; + std::string account_id; +}; + +struct AwsBackendLocation { + AwsEndpointType endpoint_type { AwsEndpointType::unknown }; + std::string region; + std::string availability_zone; + std::string account_id; +}; + +enum class AwsMetadataRequestKind : uint8_t { + local_location, + rds_region, +}; + +enum class AwsMetadataStatus : uint8_t { + ok, + provider_unavailable, + access_denied, + throttled, + imds_unavailable, + timeout, + cancelled, + invalid_response, + shutdown, +}; + +struct AwsMetadataRequestHandle { + uint64_t value { 0 }; +}; + +struct AwsMetadataRequest { + AwsMetadataRequestKind kind { AwsMetadataRequestKind::local_location }; + uint64_t opaque_id { 0 }; + uint64_t generation { 0 }; + std::string region; + std::string partition; + std::vector endpoints; + std::chrono::steady_clock::time_point deadline {}; +}; + +struct AwsMetadataEndpoint { + std::string hostname; + uint16_t port { 0 }; + AwsEndpointType endpoint_type { AwsEndpointType::unknown }; + std::string region; + std::string availability_zone; + std::string account_id; +}; + +struct AwsMetadataResult { + AwsMetadataStatus status { AwsMetadataStatus::provider_unavailable }; + AwsLocalLocation local; + std::vector endpoints; + std::string failure_category; +}; + +struct AwsMetadataCompletion { + uint64_t opaque_id { 0 }; + uint64_t generation { 0 }; + AwsMetadataResult result; +}; + +class AwsMetadataCompletionSink { +public: + virtual void post(AwsMetadataCompletion&& completion) = 0; + virtual ~AwsMetadataCompletionSink() = default; +}; + +class AwsMetadataProvider { +public: + virtual AwsMetadataRequestHandle request( + const AwsMetadataRequest& request, + std::weak_ptr sink) = 0; + virtual void cancel(AwsMetadataRequestHandle handle) = 0; + virtual void shutdown() = 0; + virtual ~AwsMetadataProvider() = default; +}; + +struct AwsLocalityBackendConfig { + AwsEndpointCandidate endpoint; + int64_t configured_weight { 0 }; + + AwsLocalityBackendConfig() = default; + AwsLocalityBackendConfig(AwsEndpointCandidate value, int64_t weight = 0) + : endpoint(std::move(value)), configured_weight(weight) {} +}; + +struct AwsLocalityHostgroupConfig { + uint32_t hostgroup_id { 0 }; + AwsLocalityPolicy policy; + std::vector backends; +}; + +#endif // __CLASS_AWS_LOCALITY_TYPES_H diff --git a/include/Base_HostGroups_Manager.h b/include/Base_HostGroups_Manager.h index 9a47dcc68e..43d778b965 100644 --- a/include/Base_HostGroups_Manager.h +++ b/include/Base_HostGroups_Manager.h @@ -19,6 +19,9 @@ class MetricsCollector; #include "proxysql.h" #include "cpp.h" #include "GTID_Server_Data.h" +#ifdef PROXYSQL40 +#include "Aws_Locality_Types.h" +#endif #include @@ -310,6 +313,10 @@ class BaseHGC { // MySQL Host Group Container char * init_connect; char * comment; char * ignore_session_variables_text; // this is the original version (text format) of ignore_session_variables + char * aws_iam_region; +#ifdef PROXYSQL40 + AwsLocalityPolicy aws_locality_policy; +#endif uint32_t max_num_online_servers; uint32_t throttle_connections_per_sec; int32_t monitor_slave_lag_when_null; diff --git a/include/MySQL_Backend_Auth.h b/include/MySQL_Backend_Auth.h new file mode 100644 index 0000000000..f82b249433 --- /dev/null +++ b/include/MySQL_Backend_Auth.h @@ -0,0 +1,67 @@ +#ifndef MYSQL_BACKEND_AUTH_H +#define MYSQL_BACKEND_AUTH_H + +#include +#include +#include + +#include "Aws_Iam_Types.h" + +class MySQL_Authentication; + +enum class MySQLBackendAuthType : uint8_t { + PASSWORD, + AWS_IAM, + INVALID, +}; + +struct MySQLBackendAuthPolicy { + MySQLBackendAuthType type{MySQLBackendAuthType::INVALID}; + std::string database_user; + std::string failure_code; + bool ignored_password{false}; +}; + +MySQLBackendAuthPolicy parse_mysql_backend_auth_policy( + std::string_view database_user, + std::string_view attributes, + bool backend_password_is_nonempty); + +MySQLBackendAuthPolicy resolve_mysql_backend_auth_policy( + MySQL_Authentication& authentication, + const char* mapped_backend_username); + +const char* mysql_backend_auth_type_name(MySQLBackendAuthType type); + +enum class AwsIamConnectionConfigStatus : uint8_t { + OK, + SUPPORT_NOT_COMPILED, + MISSING_REGION, + REGION_ENDPOINT_MISMATCH, + INVALID_ENDPOINT, + UNIX_SOCKET_NOT_ALLOWED, + TLS_REQUIRED, + CA_TRUST_REQUIRED, +}; + +struct AwsIamConnectionConfigInput { + std::string database_user; + std::string configured_endpoint; + uint16_t port; + std::string region; + bool use_ssl; + std::string ssl_ca; + std::string ssl_capath; + bool support_compiled; +}; + +struct AwsIamConnectionConfigResult { + AwsIamConnectionConfigStatus status; + AwsIamTokenKey key; + std::string failure_code; +}; + +AwsIamConnectionConfigResult validate_mysql_aws_iam_connection( + const AwsIamConnectionConfigInput& input); + +#endif // MYSQL_BACKEND_AUTH_H diff --git a/include/MySQL_HostGroups_Manager.h b/include/MySQL_HostGroups_Manager.h index 1d88201a0b..de85449b96 100644 --- a/include/MySQL_HostGroups_Manager.h +++ b/include/MySQL_HostGroups_Manager.h @@ -1,6 +1,10 @@ #ifndef PROXYSQL_MYSQL_HOSTGROUPS_MANAGER_H #define PROXYSQL_MYSQL_HOSTGROUPS_MANAGER_H #include "proxysql.h" +#include "MySQL_Backend_Auth.h" +#ifdef PROXYSQL40 +#include "Aws_Locality_Manager.h" +#endif #include "cpp.h" #include "proxysql_gtid.h" @@ -185,8 +189,9 @@ class MySrvConnList { conns->remove_index_fast((unsigned int)i); } MySQL_Connection *remove(int); - MySQL_Connection * get_random_MyConn(MySQL_Session *sess, bool ff); - void get_random_MyConn_inner_search(unsigned int start, unsigned int end, unsigned int& conn_found_idx, unsigned int& connection_quality_level, unsigned int& number_of_matching_session_variables, const MySQL_Connection * client_conn); + MySQL_Connection * get_random_MyConn( + MySQL_Session *sess, bool ff, MySQLBackendAuthType requested_type); + void get_random_MyConn_inner_search(unsigned int start, unsigned int end, unsigned int& conn_found_idx, unsigned int& connection_quality_level, unsigned int& number_of_matching_session_variables, const MySQL_Connection * client_conn, MySQLBackendAuthType requested_type); unsigned int conns_length() { return conns->len; } void drop_all_connections(); void mark_connections_unhealthy(); @@ -623,6 +628,9 @@ class MySQL_HostGroups_Manager : public Base_HostGroups_Manager { * present, distinguishing between 'READER' and 'WRITER' hostgroups. */ std::unordered_map> hostgroup_server_mapping; +#ifdef PROXYSQL40 + std::unique_ptr aws_locality_manager_; +#endif /** * @brief Holds the previous computed checksum for 'mysql_servers'. * @details Used to check if the servers checksums has changed during 'commit', if a change is detected, @@ -883,6 +891,17 @@ class MySQL_HostGroups_Manager : public Base_HostGroups_Manager { MySQL_HostGroups_Manager(); ~MySQL_HostGroups_Manager(); void init(); +#ifdef PROXYSQL40 + void refresh_aws_locality_configuration(); + void set_aws_locality_awareness_enabled(bool enabled); + void refresh_aws_locality_stats(SQLite3DB* statsdb) const; + static bool project_aws_locality_stats( + SQLite3DB* statsdb, + const std::vector& rows); + MySQLAwsLocalityManager* aws_locality_manager() const { + return aws_locality_manager_.get(); + } +#endif #if 0 void wrlock(); void wrunlock(); @@ -1034,7 +1053,7 @@ class MySQL_HostGroups_Manager : public Base_HostGroups_Manager { */ int remove_server_in_hg(uint32_t hid, const string& addr, uint16_t port); - MySQL_Connection * get_MyConn_from_pool(unsigned int hid, MySQL_Session *sess, bool ff, char * gtid_uuid, uint64_t gtid_trxid, int max_lag_ms); + MySQL_Connection * get_MyConn_from_pool(unsigned int hid, MySQL_Session *sess, bool ff, char * gtid_uuid, uint64_t gtid_trxid, int max_lag_ms, MySQLBackendAuthType requested_type = MySQLBackendAuthType::PASSWORD); void drop_all_idle_connections(); int get_multiple_idle_connections(int, unsigned long long, MySQL_Connection **, int); diff --git a/include/MySQL_Session.h b/include/MySQL_Session.h index 011d11a180..73cabb2c41 100644 --- a/include/MySQL_Session.h +++ b/include/MySQL_Session.h @@ -7,7 +7,9 @@ #ifndef PROXYSQL_MYSQL_SESSION_H #define PROXYSQL_MYSQL_SESSION_H +#include #include +#include #include #include @@ -16,6 +18,7 @@ #include "MySQL_Variables.h" #include "MySQL_User_Variables.h" #include "Base_Session.h" +#include "Aws_Iam_Provider.h" #ifndef PROXYJSON #define PROXYJSON @@ -243,6 +246,10 @@ class MySQL_Session: public Base_Session #endif // IDLE_THREADS #include +#include #include +#include #include +#include #include +#include #include "prometheus_helpers.h" @@ -41,6 +47,71 @@ extern class MySQL_Variables mysql_variables; +/** + * Thread-safe, independently-lived delivery boundary between IAM provider + * threads and one MySQL worker. It deliberately owns no session or + * connection pointer: producers can only enqueue an opaque completion and + * wake the worker's existing control pipe. + */ +class AwsIamWorkerInbox final : public AwsIamCompletionSink { +private: + struct BoundedCompletions { + explicit BoundedCompletions(size_t maximum) : maximum(maximum) {} + size_t maximum; + std::deque values; + }; + + std::mutex mutex_; + BoundedCompletions completions_; + bool closed_ { false }; + int wake_fd_ { -1 }; + +public: + explicit AwsIamWorkerInbox(int worker_write_fd, size_t capacity = 1024) + : completions_(capacity), wake_fd_(::dup(worker_write_fd)) { + if (wake_fd_ < 0) closed_ = true; + } + + ~AwsIamWorkerInbox() override { close(); } + AwsIamWorkerInbox(const AwsIamWorkerInbox&) = delete; + AwsIamWorkerInbox& operator=(const AwsIamWorkerInbox&) = delete; + + void post(AwsIamCompletion&& completion) override { + std::lock_guard guard(mutex_); + if (closed_ || completions_.values.size() >= completions_.maximum) return; + const bool wake_worker = completions_.values.empty(); + completions_.values.emplace_back(std::move(completion)); + if (wake_worker) { + const unsigned char byte = 0; + ssize_t ignored = ::write(wake_fd_, &byte, sizeof(byte)); + (void)ignored; + } + } + + std::deque drain() { + std::lock_guard guard(mutex_); + std::deque drained; + drained.swap(completions_.values); + return drained; + } + + bool available() { + std::lock_guard guard(mutex_); + return !closed_ && wake_fd_ >= 0; + } + + void close() { + std::lock_guard guard(mutex_); + if (closed_) return; + closed_ = true; + completions_.values.clear(); + if (wake_fd_ >= 0) { + ::close(wake_fd_); + wake_fd_ = -1; + } + } +}; + #ifdef PROXYSQL31 class MySQL_Caching_Sha2_RSA; #endif @@ -134,6 +205,15 @@ class __attribute__((aligned(64))) MySQL_Thread : public Base_Thread PtrArray *cached_connections; unsigned int push_local_counter; // round-robin counter for bounded local caching: cache 1-in-N where N = mysql_threads +#ifdef PROXYSQL40 + struct AwsLocalityCachedCandidate { + MySrvC* parent; + unsigned int cached_index; + }; + // Capacity grows when a connection enters the local cache, never while a + // query is choosing a backend from that cache. + std::vector aws_locality_candidates; +#endif #ifdef IDLE_THREADS struct epoll_event events[MY_EPOLL_THREAD_MAXEVENTS]; @@ -199,6 +279,9 @@ class __attribute__((aligned(64))) MySQL_Thread : public Base_Thread #endif // IDLE_THREADS int pipefd[2]; + std::shared_ptr aws_iam_inbox; + std::unordered_map aws_iam_waiters; + uint64_t next_aws_iam_waiter_id { 1 }; // int shutdown; kill_queue_t kq; @@ -230,6 +313,12 @@ class __attribute__((aligned(64))) MySQL_Thread : public Base_Thread ~MySQL_Thread(); //MySQL_Session * create_new_session_and_client_data_stream(int _fd); bool init(); + uint64_t register_aws_iam_waiter(MySQL_Session *session); + void cancel_aws_iam_waiter(uint64_t opaque_id); + void drain_aws_iam_completions(); + std::weak_ptr aws_iam_completion_sink() const { + return aws_iam_inbox; + } void run___get_multiple_idle_connections(int& num_idles); void run___cleanup_mirror_queue(); //void ProcessAllMyDS_BeforePoll(); @@ -251,7 +340,7 @@ class __attribute__((aligned(64))) MySQL_Thread : public Base_Thread void unregister_session_connection_handler(int idx, bool _new=false); void listener_handle_new_connection(MySQL_Data_Stream *myds, unsigned int n); void Get_Memory_Stats(); - MySQL_Connection * get_MyConn_local(unsigned int, MySQL_Session *sess, char *gtid_uuid, uint64_t gtid_trxid, int max_lag_ms); + MySQL_Connection * get_MyConn_local(unsigned int, MySQL_Session *sess, char *gtid_uuid, uint64_t gtid_trxid, int max_lag_ms, MySQLBackendAuthType requested_type = MySQLBackendAuthType::PASSWORD); void push_MyConn_local(MySQL_Connection *); void return_local_connections(); void Scan_Sessions_to_Kill(PtrArray *mysess); @@ -565,6 +654,9 @@ class MySQL_Threads_Handler bool passthrough_auth_empty_password; bool passthrough_auth_unknown_users; bool passthrough_auth_require_tls; +#ifdef PROXYSQL40 + bool aws_locality_awareness; +#endif int passthrough_default_hg; int passthrough_auth_cache_ttl_s; int passthrough_auth_max_inflight_probes; diff --git a/include/ProxySQL_Plugin.h b/include/ProxySQL_Plugin.h index a16cb7dee4..1b6cd36408 100644 --- a/include/ProxySQL_Plugin.h +++ b/include/ProxySQL_Plugin.h @@ -8,11 +8,14 @@ // abi_version values it doesn't understand. #ifdef PROXYSQL40 +#include #include #include class SQLite3DB; class SQLite3_result; +class AwsIamTokenSource; +class AwsMetadataProvider; namespace prometheus { class Registry; } // Descriptor ABI version the plugin was compiled for. Plugins set @@ -36,8 +39,16 @@ namespace prometheus { class Registry; } // struct with {table_name, refresh, opaque} automatically get // db_kind = admin_db (value 0) via zero-initialization of the // trailing field — matching the pre-ABI-4 behaviour. -constexpr unsigned int PROXYSQL_PLUGIN_ABI_VERSION = 4u; -constexpr unsigned int PROXYSQL_PLUGIN_ABI_VERSION_MAX = 4u; +// ABI 5: ProxySQL_PluginServices gains AWS IAM provider installation and +// sizing callbacks. They are live only during normal plugin init. +// ABI 6: ProxySQL_PluginServices gains the general AWS metadata-provider +// installation callback used by locality discovery. +// ABI 7: ProxySQL_PluginServices gains the MySQL-owned AWS-locality stats +// projection callback used by the AWS plugin's runtime view. +// ABI 8: ProxySQL_PluginServices gains an IAM-provider uninstall callback +// so a plugin can roll back a partially successful init. +constexpr unsigned int PROXYSQL_PLUGIN_ABI_VERSION = 8u; +constexpr unsigned int PROXYSQL_PLUGIN_ABI_VERSION_MAX = 8u; enum class ProxySQL_PluginDBKind : uint8_t { admin_db = 0, @@ -229,6 +240,24 @@ struct ProxySQL_PluginRuntimeView { using proxysql_plugin_register_runtime_view_cb = bool (*)(const ProxySQL_PluginRuntimeView &); + +// ABI-5 extension for optional external IAM database-authentication providers. +// The source is owned by core after successful installation; `module_handle` +// is a retained dlopen() reference released only after all session leases drain. +using proxysql_plugin_install_aws_iam_token_source_cb = + bool (*)(AwsIamTokenSource *, void (*)(AwsIamTokenSource *), void *module_handle); + +using proxysql_plugin_uninstall_aws_iam_token_source_cb = + bool (*)(AwsIamTokenSource *expected_source); + +using proxysql_plugin_get_aws_iam_limits_cb = + void (*)(size_t *max_total_waiters, size_t *max_waiters_per_key); + +using proxysql_plugin_install_aws_metadata_provider_cb = + bool (*)(AwsMetadataProvider *, void (*)(AwsMetadataProvider *), void *module_handle); + +using proxysql_plugin_refresh_mysql_aws_locality_stats_cb = + void (*)(SQLite3DB *); #endif /* PROXYSQL40 */ // Services provided to plugins across the four-phase lifecycle. @@ -292,6 +321,16 @@ struct ProxySQL_PluginServices { // at the same point they register their tables, so the callback // is wired in both phases. proxysql_plugin_register_runtime_view_cb register_runtime_view; + // ABI-5 extension. Both are null outside plugin init(). + proxysql_plugin_install_aws_iam_token_source_cb install_aws_iam_token_source; + proxysql_plugin_get_aws_iam_limits_cb get_aws_iam_limits; + // ABI-6 extension. Null outside normal plugin init(). + proxysql_plugin_install_aws_metadata_provider_cb install_aws_metadata_provider; + // ABI-7 extension. Live in Phase B and normal init; it performs no I/O. + proxysql_plugin_refresh_mysql_aws_locality_stats_cb refresh_mysql_aws_locality_stats; + // ABI-8 extension. Live only during normal plugin init and intended for + // rollback of the same plugin's partially installed IAM provider. + proxysql_plugin_uninstall_aws_iam_token_source_cb uninstall_aws_iam_token_source; #endif /* PROXYSQL40 */ }; diff --git a/include/mysql_connection.h b/include/mysql_connection.h index df7e5b34ef..6c3422662a 100644 --- a/include/mysql_connection.h +++ b/include/mysql_connection.h @@ -28,6 +28,8 @@ #define STATUS_MYSQL_CONNECTION_HAS_WARNINGS 0x00001000 #include "Servers_SslParams.h" +#include "Aws_Iam_Provider.h" +#include "MySQL_Backend_Auth.h" #ifdef PROXYSQLED25519 #include "MySQL_Ed25519.h" @@ -40,6 +42,13 @@ class Variable { void fill_client_internal_session(nlohmann::json &j, int idx); }; +struct MySQLAwsIamIdentity { + AwsIamTokenKey key; + uint64_t token_generation{0}; + uint8_t fresh_token_retries{0}; + SecureString handshake_token; +}; + enum charset_action { UNKNOWN, NAMES, @@ -82,6 +91,11 @@ unsigned int mysql_user_variable_replay_error_code(unsigned int backend_error_co class MySQL_Connection { private: + MySQLBackendAuthType backend_auth_type_{MySQLBackendAuthType::PASSWORD}; + bool rowless_passthrough_authorized_{false}; + std::unique_ptr aws_iam_identity_; + bool aws_iam_connector_secret_active_{false}; + bool aws_iam_async_connect_pending_{false}; void update_warning_count_from_connection(); void update_warning_count_from_statement(); bool is_expired(unsigned long long timeout); @@ -206,6 +220,15 @@ class MySQL_Connection { MySQL_Connection(); ~MySQL_Connection(); + void set_backend_auth_type(MySQLBackendAuthType); + MySQLBackendAuthType backend_auth_type() const; + /** Record that this PASSWORD connection authenticated via rowless pass-through. */ + void set_rowless_passthrough_authorized(bool); + /** Check whether the current resolved policy permits password reset/reuse. */ + bool can_reset_for_backend_auth_policy(const MySQLBackendAuthPolicy&) const; + void attach_aws_iam_token(const AwsIamTokenKey&, AwsIamTokenResult&&); + void clear_aws_iam_handshake_secret(); + bool has_aws_iam_handshake_secret() const; bool set_autocommit(bool); bool set_no_backslash_escapes(bool); unsigned int set_charset(unsigned int, enum charset_action); @@ -329,7 +352,12 @@ class MySQL_Connection { bool match_ff_req_options(const MySQL_Connection *c); bool match_tracked_options(const MySQL_Connection *c); - bool requires_CHANGE_USER(const MySQL_Connection *client_conn); + bool backend_auth_compatible( + const char *requested_username, + MySQLBackendAuthType requested_type) const; + bool requires_CHANGE_USER( + const MySQL_Connection *client_conn, + MySQLBackendAuthType requested_type = MySQLBackendAuthType::PASSWORD) const; unsigned int number_of_matching_session_variables(const MySQL_Connection *client_conn, unsigned int& not_matching); unsigned long get_mysql_thread_id() { return mysql ? mysql->thread_id : 0; } static void set_ssl_params(MYSQL *mysql, MySQLServers_SslParams *ssl_params); diff --git a/include/proxysql_structs.h b/include/proxysql_structs.h index 07f899c094..ac7ff1cd10 100644 --- a/include/proxysql_structs.h +++ b/include/proxysql_structs.h @@ -295,6 +295,7 @@ enum session_status { CONNECTING_SERVER, LDAP_AUTH_CLIENT, AUTHENTICATING_BACKEND_FOR_CLIENT, + WAITING_AWS_IAM_TOKEN, PINGING_SERVER, WAITING_CLIENT_DATA, WAITING_SERVER_DATA, @@ -1298,6 +1299,9 @@ __thread bool mysql_thread___passthrough_auth_enabled; __thread bool mysql_thread___passthrough_auth_empty_password; __thread bool mysql_thread___passthrough_auth_unknown_users; __thread bool mysql_thread___passthrough_auth_require_tls; +#ifdef PROXYSQL40 +__thread bool mysql_thread___aws_locality_awareness; +#endif __thread int mysql_thread___passthrough_default_hg; __thread int mysql_thread___passthrough_auth_cache_ttl_s; __thread int mysql_thread___passthrough_auth_max_inflight_probes; @@ -1655,6 +1659,9 @@ extern __thread bool mysql_thread___passthrough_auth_enabled; extern __thread bool mysql_thread___passthrough_auth_empty_password; extern __thread bool mysql_thread___passthrough_auth_unknown_users; extern __thread bool mysql_thread___passthrough_auth_require_tls; +#ifdef PROXYSQL40 +extern __thread bool mysql_thread___aws_locality_awareness; +#endif extern __thread int mysql_thread___passthrough_default_hg; extern __thread int mysql_thread___passthrough_auth_cache_ttl_s; extern __thread int mysql_thread___passthrough_auth_max_inflight_probes; diff --git a/lib/Admin_FlushVariables.cpp b/lib/Admin_FlushVariables.cpp index 6a00f448d9..d195106c7e 100644 --- a/lib/Admin_FlushVariables.cpp +++ b/lib/Admin_FlushVariables.cpp @@ -615,8 +615,19 @@ FlushVariableStats ProxySQL_Admin::flush_mysql_variables___database_to_runtime(S ASSERT_SQLITE_OK(rc, db); } } +#ifdef PROXYSQL40 + const bool aws_locality_awareness_enabled = + GloMTH->get_variable_int("aws_locality_awareness") != 0; +#endif GloMTH->wrunlock(); +#ifdef PROXYSQL40 + if (MyHGM != nullptr) { + MyHGM->set_aws_locality_awareness_enabled( + aws_locality_awareness_enabled); + } +#endif + { // NOTE: 'GloMTH->wrunlock()' should have been called before this point to avoid possible // deadlocks. See issue #3847. diff --git a/lib/Aws_Iam_Provider.cpp b/lib/Aws_Iam_Provider.cpp new file mode 100644 index 0000000000..b4e489cda3 --- /dev/null +++ b/lib/Aws_Iam_Provider.cpp @@ -0,0 +1,362 @@ +#include "Aws_Iam_Provider.h" +#include "proxysql.h" + +#include "prometheus/counter.h" +#include "prometheus/family.h" +#include "prometheus/gauge.h" +#include "prometheus/registry.h" + +#include +#include +#include +#include +#include +#include + +namespace { + +constexpr std::array kAwsIamCounterNames {{ + "proxysql_mysql_aws_iam_token_requests_total", + "proxysql_mysql_aws_iam_token_cache_hits_total", + "proxysql_mysql_aws_iam_token_refresh_successes_total", + "proxysql_mysql_aws_iam_token_refresh_failures_total", + "proxysql_mysql_aws_iam_credential_provider_failures_total", + "proxysql_mysql_aws_iam_queue_rejections_total", + "proxysql_mysql_aws_iam_backend_connection_successes_total", + "proxysql_mysql_aws_iam_backend_connection_failures_total", +}}; + +constexpr std::array kAwsIamGaugeNames {{ + "proxysql_mysql_aws_iam_token_cache_entries", + "proxysql_mysql_aws_iam_in_flight_generations", + "proxysql_mysql_aws_iam_queued_generations", + "proxysql_mysql_aws_iam_waiting_sessions", +}}; + +struct AwsIamPrometheusState { + std::mutex mutex; + prometheus::Registry *registry { nullptr }; + std::array counters {}; + std::array gauges {}; +}; + +AwsIamPrometheusState& aws_iam_prometheus_state() { + static AwsIamPrometheusState state; + return state; +} + +std::array aws_iam_counter_values(const AwsIamStatsSnapshot& stats) { + return {{ + stats.token_requests, + stats.token_cache_hits, + stats.token_refresh_successes, + stats.token_refresh_failures, + stats.credential_provider_failures, + stats.queue_rejections, + stats.backend_connection_successes, + stats.backend_connection_failures, + }}; +} + +std::array aws_iam_gauge_values(const AwsIamStatsSnapshot& stats) { + return {{ + stats.token_cache_entries, + stats.in_flight_generations, + stats.queued_generations, + stats.waiting_sessions, + }}; +} + +std::mutex global_source_mutex; +std::condition_variable global_source_cv; +AwsIamTokenSource *leased_global_source = nullptr; +size_t global_source_leases = 0; +bool global_source_accepting = false; +bool global_source_retirement_active = false; +// Count queued as well as active retirement callers so a replacement cannot +// slip between two waiters and be claimed by the lagging caller. +size_t global_source_retirement_requests = 0; + +// A plugin supplies an extra dlopen() reference with its source. Core keeps +// that reference until every session lease has drained, then invokes the +// plugin's destroy callback before dlclose(). This prevents a source vtable +// from pointing at an already-unmapped plugin during shutdown. +struct AwsIamOwnedSource { + AwsIamTokenSource *source { nullptr }; + AwsIamTokenSourceDestroyFn destroy { nullptr }; + void *module_handle { nullptr }; + + AwsIamOwnedSource() = default; + AwsIamOwnedSource(const AwsIamOwnedSource&) = delete; + AwsIamOwnedSource& operator=(const AwsIamOwnedSource&) = delete; + AwsIamOwnedSource(AwsIamOwnedSource&& other) noexcept + : source(other.source), destroy(other.destroy), module_handle(other.module_handle) { + other.source = nullptr; + other.destroy = nullptr; + other.module_handle = nullptr; + } + AwsIamOwnedSource& operator=(AwsIamOwnedSource&& other) noexcept { + if (this != &other) { + reset(); + source = other.source; + destroy = other.destroy; + module_handle = other.module_handle; + other.source = nullptr; + other.destroy = nullptr; + other.module_handle = nullptr; + } + return *this; + } + + void reset() noexcept { + if (source != nullptr) { + if (destroy != nullptr) { + destroy(source); + } else { + delete source; + } + } + if (module_handle != nullptr) dlclose(module_handle); + source = nullptr; + destroy = nullptr; + module_handle = nullptr; + } +}; + +AwsIamOwnedSource installed_source; + +class AwsIamNotCompiledTokenSource final : public AwsIamTokenSource { +public: + bool support_compiled() const override { return false; } + + AwsIamRequestHandle request(const AwsIamTokenKey&, uint64_t opaque_id, + std::weak_ptr sink) override { + AwsIamRequestHandle handle { next_handle_.fetch_add(1, std::memory_order_relaxed) }; + if (auto live_sink = sink.lock()) { + AwsIamCompletion completion; + completion.opaque_id = opaque_id; + completion.result.status = AwsIamStatus::SUPPORT_NOT_COMPILED; + completion.result.failure.category = "support_not_compiled"; + live_sink->post(std::move(completion)); + } + return handle; + } + + AwsIamTokenResult request_blocking(const AwsIamTokenKey&, + std::chrono::steady_clock::time_point) override { + AwsIamTokenResult result; + result.status = AwsIamStatus::SUPPORT_NOT_COMPILED; + result.failure.category = "support_not_compiled"; + return result; + } + + void cancel(AwsIamRequestHandle) override {} + void invalidate(const AwsIamTokenKey&, uint64_t) override {} + void record_backend_connection(bool) override {} + void record_waiting_session(bool) override {} + AwsIamStatsSnapshot snapshot() const override { return {}; } + +private: + std::atomic next_handle_ { 1 }; +}; + +} // namespace + +AwsIamNamedStats aws_iam_stats_mysql_global_rows(const AwsIamStatsSnapshot& stats) { + return {{ + { "AwsIam_Token_requests", stats.token_requests }, + { "AwsIam_Token_cache_hits", stats.token_cache_hits }, + { "AwsIam_Token_refresh_successes", stats.token_refresh_successes }, + { "AwsIam_Token_refresh_failures", stats.token_refresh_failures }, + { "AwsIam_Credential_provider_failures", stats.credential_provider_failures }, + { "AwsIam_Queue_rejections", stats.queue_rejections }, + { "AwsIam_Backend_connection_successes", stats.backend_connection_successes }, + { "AwsIam_Backend_connection_failures", stats.backend_connection_failures }, + { "AwsIam_Token_cache_entries", stats.token_cache_entries }, + { "AwsIam_In_flight_generations", stats.in_flight_generations }, + { "AwsIam_Queued_generations", stats.queued_generations }, + { "AwsIam_Waiting_sessions", stats.waiting_sessions }, + }}; +} + +void initialize_aws_iam_prometheus_metrics(prometheus::Registry& registry) { + AwsIamPrometheusState& state = aws_iam_prometheus_state(); + std::lock_guard lock(state.mutex); + if (state.registry == ®istry) return; + // A ProxySQL process owns one registry for its lifetime. Supporting a new + // registry here also keeps isolated tests deterministic after fresh setup. + state.registry = ®istry; + state.counters.fill(nullptr); + state.gauges.fill(nullptr); + for (size_t i = 0; i < kAwsIamCounterNames.size(); ++i) { + auto& family = prometheus::BuildCounter() + .Name(kAwsIamCounterNames[i]) + .Help("ProxySQL AWS IAM backend authentication counter.") + .Register(registry); + state.counters[i] = std::addressof(family.Add({})); + } + for (size_t i = 0; i < kAwsIamGaugeNames.size(); ++i) { + auto& family = prometheus::BuildGauge() + .Name(kAwsIamGaugeNames[i]) + .Help("ProxySQL AWS IAM backend authentication gauge.") + .Register(registry); + state.gauges[i] = std::addressof(family.Add({})); + } +} + +void update_aws_iam_prometheus_metrics(const AwsIamStatsSnapshot& stats) { + AwsIamPrometheusState& state = aws_iam_prometheus_state(); + std::lock_guard lock(state.mutex); + if (state.registry == nullptr) return; + const auto counters = aws_iam_counter_values(stats); + const auto gauges = aws_iam_gauge_values(stats); + for (size_t i = 0; i < counters.size(); ++i) { + const double current = state.counters[i]->Value(); + if (static_cast(counters[i]) > current) { + state.counters[i]->Increment(static_cast(counters[i]) - current); + } + } + for (size_t i = 0; i < gauges.size(); ++i) { + state.gauges[i]->Set(static_cast(gauges[i])); + } +} + +void AwsIamTokenSourceLease::release() { + if (source_ == nullptr) return; + { + std::lock_guard lock(global_source_mutex); + if (global_source_leases != 0) --global_source_leases; + } + source_ = nullptr; + global_source_cv.notify_all(); +} + +AwsIamTokenSourceLease::~AwsIamTokenSourceLease() { release(); } + +AwsIamTokenSourceLease::AwsIamTokenSourceLease( + AwsIamTokenSourceLease&& other) noexcept : source_(other.source_) { + other.source_ = nullptr; +} + +AwsIamTokenSourceLease& AwsIamTokenSourceLease::operator=( + AwsIamTokenSourceLease&& other) noexcept { + if (this != &other) { + release(); + source_ = other.source_; + other.source_ = nullptr; + } + return *this; +} + +void publish_global_aws_iam_token_source(AwsIamTokenSource *source) { + std::lock_guard lock(global_source_mutex); + if (global_source_retirement_requests != 0) { + if (source != nullptr) { + proxy_warning("Refusing to publish an AWS IAM token source during retirement\n"); + } + return; + } + if (installed_source.source != nullptr && source != installed_source.source) { + proxy_warning("Refusing to replace the plugin-owned AWS IAM token source\n"); + return; + } + if (source != nullptr && GloVars.prometheus_registry != nullptr) { + initialize_aws_iam_prometheus_metrics(*GloVars.prometheus_registry); + update_aws_iam_prometheus_metrics(source->snapshot()); + } + leased_global_source = source; + global_source_accepting = source != nullptr; + GloAwsIamTokenSource = source; +} + +bool install_global_aws_iam_token_source( + AwsIamTokenSource *source, AwsIamTokenSourceDestroyFn destroy, void *module_handle) { + if (source == nullptr || destroy == nullptr || module_handle == nullptr) return false; + + std::lock_guard lock(global_source_mutex); + if (global_source_retirement_requests != 0 || global_source_accepting || + leased_global_source != nullptr || + installed_source.source != nullptr) { + return false; + } + if (GloVars.prometheus_registry != nullptr) { + initialize_aws_iam_prometheus_metrics(*GloVars.prometheus_registry); + update_aws_iam_prometheus_metrics(source->snapshot()); + } + installed_source.source = source; + installed_source.destroy = destroy; + installed_source.module_handle = module_handle; + leased_global_source = source; + global_source_accepting = true; + GloAwsIamTokenSource = source; + return true; +} + +bool uninstall_global_aws_iam_token_source(AwsIamTokenSource *expected_source) { + AwsIamOwnedSource retired_source; + { + std::unique_lock lock(global_source_mutex); + ++global_source_retirement_requests; + global_source_cv.wait(lock, [] { + return !global_source_retirement_active; + }); + if (expected_source == nullptr || installed_source.source != expected_source || + leased_global_source != expected_source) { + --global_source_retirement_requests; + lock.unlock(); + global_source_cv.notify_all(); + return false; + } + global_source_retirement_active = true; + global_source_accepting = false; + GloAwsIamTokenSource = nullptr; + global_source_cv.wait(lock, [] { return global_source_leases == 0; }); + leased_global_source = nullptr; + retired_source = std::move(installed_source); + } + retired_source.reset(); + { + std::lock_guard lock(global_source_mutex); + global_source_retirement_active = false; + --global_source_retirement_requests; + } + global_source_cv.notify_all(); + return true; +} + +AwsIamTokenSourceLease acquire_global_aws_iam_token_source() { + std::lock_guard lock(global_source_mutex); + if (!global_source_accepting || leased_global_source == nullptr) return {}; + ++global_source_leases; + return AwsIamTokenSourceLease(leased_global_source); +} + +void shutdown_global_aws_iam_token_source() { + AwsIamOwnedSource retired_source; + { + std::unique_lock lock(global_source_mutex); + ++global_source_retirement_requests; + global_source_cv.wait(lock, [] { + return !global_source_retirement_active; + }); + global_source_retirement_active = true; + global_source_accepting = false; + GloAwsIamTokenSource = nullptr; + global_source_cv.wait(lock, [] { return global_source_leases == 0; }); + leased_global_source = nullptr; + retired_source = std::move(installed_source); + } + retired_source.reset(); + { + std::lock_guard lock(global_source_mutex); + global_source_retirement_active = false; + --global_source_retirement_requests; + } + global_source_cv.notify_all(); +} + +std::unique_ptr create_aws_iam_token_source( + const AwsIamRuntimeConfig& config) { + (void)config; + return std::make_unique(); +} diff --git a/lib/Aws_Locality_Manager.cpp b/lib/Aws_Locality_Manager.cpp new file mode 100644 index 0000000000..ee3f417736 --- /dev/null +++ b/lib/Aws_Locality_Manager.cpp @@ -0,0 +1,1149 @@ +#include "Aws_Locality_Manager.h" + +#include "json.hpp" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +using nlohmann::json; + +namespace { + +AwsLocalityPolicy invalid_policy( + uint32_t hostgroup_id, + const char* field, + AwsLocalityPolicyError& error) { + error.hostgroup_id = hostgroup_id; + error.field = field; + return {}; +} + +bool read_multiplier(const json& object, const char* field, double& value) { + const auto it = object.find(field); + if (it == object.end() || !it->is_number()) { + return false; + } + + value = it->get(); + return std::isfinite(value) && value >= 1.0 && value <= 10.0; +} + +bool read_seconds( + const json& object, + const char* field, + uint32_t default_value, + uint32_t minimum, + uint32_t maximum, + uint32_t& value) { + const auto it = object.find(field); + if (it == object.end()) { + value = default_value; + return value >= minimum && value <= maximum; + } + if (!it->is_number_unsigned() && !it->is_number_integer()) { + return false; + } + + const int64_t parsed = it->get(); + if (parsed < static_cast(minimum) || + parsed > static_cast(maximum)) { + return false; + } + value = static_cast(parsed); + return true; +} + +bool ends_with(const std::string& value, const std::string& suffix) { + return value.size() >= suffix.size() && + value.compare(value.size() - suffix.size(), suffix.size(), suffix) == 0; +} + +bool valid_region(const std::string& region) { + if (region.empty() || region.front() == '-' || region.back() == '-') { + return false; + } + + unsigned int hyphens = 0; + for (const char character : region) { + if (character == '-') { + ++hyphens; + } else if (!std::islower(static_cast(character)) && + !std::isdigit(static_cast(character))) { + return false; + } + } + return hyphens >= 2 && + std::isdigit(static_cast(region.back())); +} + +bool is_rds_proxy_endpoint_prefix(const std::string& prefix) { + // RDS Proxy endpoints have the canonical shape + // .proxy-..rds.amazonaws.com. A normal + // DB identifier is allowed to begin with "proxy-", so only the reserved + // generated-ID label after a proxy name identifies this endpoint type. + const size_t separator = prefix.rfind('.'); + return separator != std::string::npos && separator != 0 && + prefix.size() - separator - 1 > 6 && + prefix.compare(separator + 1, 6, "proxy-") == 0; +} + +} // namespace + +AwsLocalityPolicy parse_aws_locality_policy( + const json& policy_json, + uint32_t hostgroup_id, + AwsLocalityPolicyError& error) { + error = {}; + if (!policy_json.is_object()) { + return invalid_policy(hostgroup_id, "locality_awareness", error); + } + + AwsLocalityPolicy policy; + if (!read_multiplier(policy_json, "same_region_multiplier", + policy.same_region_multiplier)) { + return invalid_policy(hostgroup_id, "same_region_multiplier", error); + } + if (!read_multiplier(policy_json, "same_az_multiplier", + policy.same_az_multiplier) || + policy.same_az_multiplier < policy.same_region_multiplier) { + return invalid_policy(hostgroup_id, "same_az_multiplier", error); + } + if (!read_seconds(policy_json, "refresh_interval_seconds", 300, 30, 86400, + policy.refresh_interval_seconds)) { + return invalid_policy(hostgroup_id, "refresh_interval_seconds", error); + } + if (!read_seconds(policy_json, "stale_ttl_seconds", 1800, + policy.refresh_interval_seconds, 604800, + policy.stale_ttl_seconds)) { + return invalid_policy(hostgroup_id, "stale_ttl_seconds", error); + } + + policy.valid = true; + return policy; +} + +AwsEndpointCandidate recognize_rds_endpoint( + uint32_t hostgroup_id, + std::string_view hostname_input, + uint16_t port) { + AwsEndpointCandidate result; + result.hostgroup_id = hostgroup_id; + result.port = port; + result.hostname = aws_locality_normalized_hostname(hostname_input); + if (result.hostname.empty()) { + return result; + } + + const std::string china_suffix = ".rds.amazonaws.com.cn"; + const std::string standard_suffix = ".rds.amazonaws.com"; + const std::string* suffix = nullptr; + if (ends_with(result.hostname, china_suffix)) { + suffix = &china_suffix; + } else if (ends_with(result.hostname, standard_suffix)) { + suffix = &standard_suffix; + } else { + return result; + } + + const std::string before_suffix = result.hostname.substr( + 0, result.hostname.size() - suffix->size()); + const size_t region_separator = before_suffix.rfind('.'); + if (region_separator == std::string::npos || region_separator == 0 || + region_separator + 1 == before_suffix.size()) { + return result; + } + + const std::string endpoint_prefix = before_suffix.substr(0, region_separator); + result.region = before_suffix.substr(region_separator + 1); + if (is_rds_proxy_endpoint_prefix(endpoint_prefix) || !valid_region(result.region)) { + result.region.clear(); + return result; + } + + if (result.region.compare(0, 3, "cn-") == 0) { + result.partition = "aws-cn"; + } else if (result.region.compare(0, 7, "us-gov-") == 0) { + result.partition = "aws-us-gov"; + } else { + result.partition = "aws"; + } + result.recognized = true; + return result; +} + +AwsLocalityClass classify_aws_locality( + const AwsLocalLocation& local, + const AwsBackendLocation& backend) { + if (local.region.empty() || backend.region.empty()) { + return AwsLocalityClass::unknown; + } + if (local.region != backend.region) { + return AwsLocalityClass::remote; + } + if (backend.endpoint_type == AwsEndpointType::instance && + !local.availability_zone.empty() && + local.availability_zone == backend.availability_zone && + !local.account_id.empty() && + local.account_id == backend.account_id) { + return AwsLocalityClass::same_az; + } + return AwsLocalityClass::same_region; +} + +uint64_t aws_locality_effective_weight( + int64_t configured_weight, + double multiplier) { + if (configured_weight <= 0 || !std::isfinite(multiplier) || multiplier <= 0.0) { + return 0; + } + + const long double product = static_cast(configured_weight) * + static_cast(multiplier); + const long double maximum = static_cast( + std::numeric_limits::max()); + if (product >= maximum) { + return std::numeric_limits::max(); + } + return static_cast(product); +} + +uint64_t aws_locality_saturating_add(uint64_t lhs, uint64_t rhs) { + const uint64_t maximum = std::numeric_limits::max(); + return maximum - lhs < rhs ? maximum : lhs + rhs; +} + +size_t aws_locality_weighted_index( + const uint64_t* weights, + size_t count, + uint64_t random_value) { + if (weights == nullptr || count == 0) { + return count; + } + + uint64_t total = 0; + for (size_t i = 0; i < count; ++i) { + total = aws_locality_saturating_add(total, weights[i]); + } + if (total == 0) { + return count; + } + + const uint64_t target = random_value % total; + uint64_t cumulative = 0; + for (size_t i = 0; i < count; ++i) { + cumulative = aws_locality_saturating_add(cumulative, weights[i]); + if (target < cumulative) { + return i; + } + } + return count; +} + +namespace { + +std::mutex metadata_provider_mutex; +std::condition_variable metadata_provider_cv; +AwsMetadataProvider* leased_metadata_provider = nullptr; +AwsMetadataProviderDestroyFn metadata_provider_destroy = nullptr; +void* metadata_provider_module = nullptr; +size_t metadata_provider_leases = 0; +bool metadata_provider_accepting = false; + +bool dns_identity_length(std::string_view input, size_t& length) { + if (input.empty()) return false; + length = input.size(); + if (input[length - 1] == '.') --length; + if (length == 0 || input[length - 1] == '.') return false; + for (size_t index = 0; index < length; ++index) { + const unsigned char value = static_cast(input[index]); + if (!(std::isalnum(value) || input[index] == '-' || input[index] == '.')) { + return false; + } + } + return true; +} + +uint64_t identity_hash( + uint32_t hostgroup_id, + std::string_view hostname, + uint16_t port) { + uint64_t hash = 14695981039346656037ULL; + auto append = [&](unsigned char value) { + hash ^= value; + hash *= 1099511628211ULL; + }; + for (unsigned int shift = 0; shift < 32; shift += 8) { + append(static_cast(hostgroup_id >> shift)); + } + append(static_cast(port)); + append(static_cast(port >> 8)); + size_t length = 0; + const bool valid_dns = dns_identity_length(hostname, length); + append(valid_dns ? 1 : 0); + if (!valid_dns) length = hostname.size(); + for (size_t index = 0; index < length; ++index) { + const unsigned char value = static_cast(hostname[index]); + append(valid_dns ? static_cast(std::tolower(value)) : value); + } + return hash; +} + +bool same_hostname_identity(std::string_view lhs, std::string_view rhs) { + size_t lhs_length = 0; + size_t rhs_length = 0; + const bool lhs_dns = dns_identity_length(lhs, lhs_length); + const bool rhs_dns = dns_identity_length(rhs, rhs_length); + if (lhs_dns != rhs_dns) return false; + if (!lhs_dns) return lhs == rhs; + if (lhs_length != rhs_length) return false; + for (size_t index = 0; index < lhs_length; ++index) { + if (std::tolower(static_cast(lhs[index])) != + std::tolower(static_cast(rhs[index]))) { + return false; + } + } + return true; +} + +std::string endpoint_key(std::string_view hostname_input, uint16_t port) { + const std::string hostname = aws_locality_normalized_hostname(hostname_input); + return (hostname.empty() ? "raw:" + std::string(hostname_input) + : "dns:" + hostname) + "\n" + std::to_string(port); +} + +const char* failure_category(AwsMetadataStatus status) { + switch (status) { + case AwsMetadataStatus::ok: return ""; + case AwsMetadataStatus::provider_unavailable: return "provider_unavailable"; + case AwsMetadataStatus::access_denied: return "access_denied"; + case AwsMetadataStatus::throttled: return "throttled"; + case AwsMetadataStatus::imds_unavailable: return "imds_unavailable"; + case AwsMetadataStatus::timeout: return "timeout"; + case AwsMetadataStatus::cancelled: return "cancelled"; + case AwsMetadataStatus::invalid_response: return "invalid_response"; + case AwsMetadataStatus::shutdown: return "cancelled"; + } + return "invalid_response"; +} + +int64_t wall_seconds(std::chrono::system_clock::time_point value) { + return std::chrono::duration_cast( + value.time_since_epoch()).count(); +} + +} // namespace + +void AwsMetadataProviderLease::release() { + if (provider_ == nullptr) { + return; + } + { + std::lock_guard lock(metadata_provider_mutex); + if (metadata_provider_leases != 0) { + --metadata_provider_leases; + } + } + provider_ = nullptr; + metadata_provider_cv.notify_all(); +} + +AwsMetadataProviderLease::~AwsMetadataProviderLease() { + release(); +} + +AwsMetadataProviderLease::AwsMetadataProviderLease( + AwsMetadataProviderLease&& other) noexcept + : provider_(other.provider_) { + other.provider_ = nullptr; +} + +AwsMetadataProviderLease& AwsMetadataProviderLease::operator=( + AwsMetadataProviderLease&& other) noexcept { + if (this != &other) { + release(); + provider_ = other.provider_; + other.provider_ = nullptr; + } + return *this; +} + +bool install_global_aws_metadata_provider( + AwsMetadataProvider* provider, + AwsMetadataProviderDestroyFn destroy, + void* module_handle) { + if (provider == nullptr || destroy == nullptr) { + return false; + } + + std::lock_guard lock(metadata_provider_mutex); + if (metadata_provider_accepting || leased_metadata_provider != nullptr || + metadata_provider_destroy != nullptr) { + return false; + } + leased_metadata_provider = provider; + metadata_provider_destroy = destroy; + metadata_provider_module = module_handle; + metadata_provider_accepting = true; + return true; +} + +AwsMetadataProviderLease acquire_global_aws_metadata_provider() { + std::lock_guard lock(metadata_provider_mutex); + if (!metadata_provider_accepting || leased_metadata_provider == nullptr) { + return {}; + } + ++metadata_provider_leases; + return AwsMetadataProviderLease(leased_metadata_provider); +} + +void shutdown_global_aws_metadata_provider() { + AwsMetadataProvider* provider = nullptr; + AwsMetadataProviderDestroyFn destroy = nullptr; + void* module_handle = nullptr; + { + std::unique_lock lock(metadata_provider_mutex); + metadata_provider_accepting = false; + metadata_provider_cv.wait(lock, [] { + return metadata_provider_leases == 0; + }); + provider = leased_metadata_provider; + destroy = metadata_provider_destroy; + module_handle = metadata_provider_module; + leased_metadata_provider = nullptr; + metadata_provider_destroy = nullptr; + metadata_provider_module = nullptr; + } + + if (provider != nullptr) { + provider->shutdown(); + destroy(provider); + } + if (module_handle != nullptr) { + dlclose(module_handle); + } +} + +const AwsLocalitySnapshotEntry* AwsLocalitySnapshot::find( + uint32_t hostgroup_id, + std::string_view hostname, + uint16_t port) const { + const auto range = entries.equal_range(identity_hash(hostgroup_id, hostname, port)); + for (auto it = range.first; it != range.second; ++it) { + const auto& entry = it->second; + if (entry.hostgroup_id == hostgroup_id && entry.port == port && + same_hostname_identity(entry.hostname, hostname)) { + return &entry; + } + } + return nullptr; +} + +uint64_t AwsLocalitySnapshot::effective_weight( + uint32_t hostgroup_id, + std::string_view hostname, + uint16_t port, + int64_t configured_weight) const { + const auto* entry = find(hostgroup_id, hostname, port); + return aws_locality_effective_weight( + configured_weight, entry == nullptr ? 1.0 : entry->multiplier); +} + +AwsLocalityManagerConfig::AwsLocalityManagerConfig() + : steady_clock([] { return std::chrono::steady_clock::now(); }), + wall_clock([] { return std::chrono::system_clock::now(); }) {} + +class MySQLAwsLocalityManager::Impl { +public: + explicit Impl(AwsLocalityManagerConfig config) + : config_(std::move(config)), sink_(std::make_shared(this)) { + auto initial = std::make_shared(); + std::atomic_store_explicit( + &published_, std::shared_ptr(std::move(initial)), + std::memory_order_release); + } + + ~Impl() noexcept { + try { + shutdown(); + } catch (...) { + // Destructors must not let allocation/system exceptions from the + // final snapshot publication escape. Ensure the scheduler thread + // cannot make std::thread's destructor terminate the process. + try { + std::thread worker; + { + std::lock_guard lock(mutex_); + stopping_ = true; + enabled_ = false; + cancel_requested_ = true; + cv_.notify_all(); + worker = std::move(worker_); + } + if (worker.joinable()) worker.join(); + sink_->detach(); + } catch (...) { + std::terminate(); + } + } + } + + void configure(std::vector hostgroups) { + std::lock_guard lock(mutex_); + if (stopping_) { + return; + } + for (auto& hostgroup : hostgroups) { + for (auto& backend : hostgroup.backends) { + backend.endpoint.hostgroup_id = hostgroup.hostgroup_id; + } + } + hostgroups_.clear(); + for (auto& hostgroup : hostgroups) { + if (hostgroup.policy.valid) { + hostgroups_.push_back(std::move(hostgroup)); + } + } + ++generation_; + cancel_requested_ = true; + force_refresh_ = enabled_ && !hostgroups_.empty(); + if (force_refresh_) { + ensure_worker_locked(); + } + publish_locked(); + cv_.notify_all(); + } + + void set_enabled(bool enabled) { + std::unique_lock lock(mutex_); + if (stopping_) { + return; + } + enabled_ = enabled; + if (enabled_ && !hostgroups_.empty()) { + ensure_worker_locked(); + force_refresh_ = true; + } else { + cancel_requested_ = true; + disable_acknowledged_ = !worker_.joinable(); + } + publish_locked(); + cv_.notify_all(); + if (!enabled_ && worker_.joinable()) { + cv_.wait_for(lock, config_.disable_wait_timeout, [&] { + return disable_acknowledged_ || stopping_; + }); + } + } + + void request_refresh() { + std::lock_guard lock(mutex_); + if (stopping_ || !enabled_ || hostgroups_.empty()) { + return; + } + force_refresh_ = true; + publish_locked(); + cv_.notify_all(); + } + + std::shared_ptr snapshot() const { + return std::atomic_load_explicit(&published_, std::memory_order_acquire); + } + + std::vector diagnostic_rows() const { + const auto current = snapshot(); + std::vector rows; + rows.reserve(current->entries.size()); + for (const auto& item : current->entries) { + rows.push_back(item.second); + } + std::sort(rows.begin(), rows.end(), [](const auto& lhs, const auto& rhs) { + if (lhs.hostgroup_id != rhs.hostgroup_id) { + return lhs.hostgroup_id < rhs.hostgroup_id; + } + if (lhs.hostname != rhs.hostname) { + return lhs.hostname < rhs.hostname; + } + return lhs.port < rhs.port; + }); + return rows; + } + + void shutdown() { + std::thread worker; + { + std::lock_guard lock(mutex_); + if (shutdown_complete_) { + return; + } + stopping_ = true; + enabled_ = false; + cancel_requested_ = true; + publish_locked(); + cv_.notify_all(); + worker = std::move(worker_); + } + if (worker.joinable()) { + worker.join(); + } + sink_->detach(); + { + std::lock_guard lock(mutex_); + shutdown_complete_ = true; + publish_locked(); + } + } + +private: + struct EndpointRecord { + bool has_value { false }; + AwsBackendLocation value; + std::chrono::steady_clock::time_point success_steady {}; + int64_t success_wall { 0 }; + int64_t attempt_wall { 0 }; + std::string error; + }; + + struct LocalRecord { + bool has_value { false }; + AwsLocalLocation value; + std::chrono::steady_clock::time_point success_steady {}; + int64_t success_wall { 0 }; + int64_t attempt_wall { 0 }; + std::string error; + }; + + struct InFlight { + AwsMetadataRequest request; + AwsMetadataRequestHandle handle; + }; + + class CompletionSink final : public AwsMetadataCompletionSink { + public: + explicit CompletionSink(Impl* owner) : owner_(owner) {} + + void post(AwsMetadataCompletion&& completion) override { + Impl* owner = nullptr; + { + std::lock_guard lock(mutex_); + if (owner_ == nullptr) { + return; + } + owner = owner_; + ++active_; + } + owner->on_completion(std::move(completion)); + { + std::lock_guard lock(mutex_); + --active_; + cv_.notify_all(); + } + } + + void detach() { + std::unique_lock lock(mutex_); + owner_ = nullptr; + cv_.wait(lock, [&] { return active_ == 0; }); + } + + private: + std::mutex mutex_; + std::condition_variable cv_; + Impl* owner_ { nullptr }; + size_t active_ { 0 }; + }; + + void ensure_worker_locked() { + if (!worker_.joinable()) { + worker_ = std::thread([this] { worker_loop(); }); + } + } + + uint32_t minimum_refresh_seconds_locked() const { + uint32_t result = 86400; + for (const auto& hostgroup : hostgroups_) { + result = std::min(result, hostgroup.policy.refresh_interval_seconds); + } + return result; + } + + std::vector build_cycle_locked( + std::chrono::steady_clock::time_point now) { + std::vector requests; + AwsMetadataRequest local; + local.kind = AwsMetadataRequestKind::local_location; + local.opaque_id = next_opaque_id_++; + local.generation = generation_; + local.deadline = now + config_.request_timeout; + requests.push_back(local); + + std::map> regions; + for (const auto& hostgroup : hostgroups_) { + for (const auto& backend : hostgroup.backends) { + const auto& endpoint = backend.endpoint; + if (!endpoint.recognized) { + continue; + } + regions[endpoint.region][endpoint_key(endpoint.hostname, endpoint.port)] = endpoint; + } + } + for (const auto& region : regions) { + AwsMetadataRequest request; + request.kind = AwsMetadataRequestKind::rds_region; + request.opaque_id = next_opaque_id_++; + request.generation = generation_; + request.region = region.first; + request.deadline = now + config_.request_timeout; + for (const auto& item : region.second) { + request.endpoints.push_back(item.second); + if (request.partition.empty()) { + request.partition = item.second.partition; + } + } + requests.push_back(std::move(request)); + } + + const int64_t attempt = wall_seconds(config_.wall_clock()); + local_.attempt_wall = attempt; + ++local_in_flight_; + for (const auto& request : requests) { + in_flight_.emplace(request.opaque_id, InFlight {request, {}}); + if (request.kind == AwsMetadataRequestKind::rds_region) { + ++region_in_flight_[request.region]; + for (const auto& endpoint : request.endpoints) { + endpoint_cache_[endpoint_key(endpoint.hostname, endpoint.port)].attempt_wall = attempt; + } + } + } + publish_locked(); + return requests; + } + + void cancel_all_requests(AwsMetadataProvider* provider) { + std::vector handles; + { + std::lock_guard lock(mutex_); + for (const auto& item : in_flight_) { + if (item.second.handle.value != 0) { + handles.push_back(item.second.handle); + } + } + in_flight_.clear(); + local_in_flight_ = 0; + region_in_flight_.clear(); + publish_locked(); + } + if (provider != nullptr) { + for (const auto handle : handles) { + provider->cancel(handle); + } + } + } + + void provider_unavailable(uint64_t generation) { + std::lock_guard lock(mutex_); + if (generation != generation_ || stopping_) { + return; + } + const int64_t attempt = wall_seconds(config_.wall_clock()); + local_.attempt_wall = attempt; + local_.error = "provider_unavailable"; + for (const auto& hostgroup : hostgroups_) { + for (const auto& backend : hostgroup.backends) { + auto& record = endpoint_cache_[endpoint_key( + backend.endpoint.hostname, backend.endpoint.port)]; + record.attempt_wall = attempt; + record.error = "provider_unavailable"; + } + } + in_flight_.clear(); + local_in_flight_ = 0; + region_in_flight_.clear(); + publish_locked(); + } + + void worker_loop() { + AwsMetadataProviderLease provider_lease; + bool cycle_started = false; + auto next_due = config_.steady_clock(); + + for (;;) { + std::vector requests; + bool cancel = false; + bool release_provider = false; + bool stop_now = false; + bool acknowledge_disable = false; + uint64_t dispatch_generation = 0; + { + std::unique_lock lock(mutex_); + while (!stopping_) { + if (cancel_requested_ || force_refresh_) { + break; + } + if (!enabled_ || hostgroups_.empty()) { + cv_.wait(lock, [&] { + return stopping_ || cancel_requested_ || force_refresh_ || + (enabled_ && !hostgroups_.empty()); + }); + continue; + } + const auto now = config_.steady_clock(); + if (!cycle_started || now >= next_due) { + break; + } + const auto delay = next_due - now; + cv_.wait_for(lock, delay); + } + + if (stopping_) { + cancel = true; + release_provider = true; + } else { + cancel = cancel_requested_; + cancel_requested_ = false; + if (cancel) { + // Cancel the previous generation before constructing new work. + // Otherwise the cancellation sweep can consume requests that + // were created in this same scheduler iteration. + if (!enabled_) { + release_provider = true; + acknowledge_disable = true; + } + } else if (!enabled_ || hostgroups_.empty()) { + release_provider = true; + acknowledge_disable = !enabled_; + force_refresh_ = false; + publish_locked(); + } else { + const auto now = config_.steady_clock(); + if (force_refresh_ || !cycle_started || now >= next_due) { + force_refresh_ = false; + requests = build_cycle_locked(now); + dispatch_generation = generation_; + cycle_started = true; + next_due = now + std::chrono::seconds( + minimum_refresh_seconds_locked()); + } + } + } + stop_now = stopping_; + } + + if (cancel) { + cancel_all_requests(provider_lease.get()); + } + if (release_provider) { + provider_lease = {}; + } + if (acknowledge_disable) { + std::lock_guard lock(mutex_); + disable_acknowledged_ = true; + cv_.notify_all(); + } + if (stop_now) { + break; + } + if (requests.empty()) { + continue; + } + + if (!provider_lease) { + provider_lease = acquire_global_aws_metadata_provider(); + } + if (!provider_lease) { + provider_unavailable(dispatch_generation); + continue; + } + + for (const auto& request : requests) { + AwsMetadataRequestHandle handle; + try { + handle = provider_lease->request(request, sink_); + } catch (...) { + AwsMetadataCompletion completion; + completion.opaque_id = request.opaque_id; + completion.generation = request.generation; + completion.result.status = AwsMetadataStatus::provider_unavailable; + on_completion(std::move(completion)); + continue; + } + std::lock_guard lock(mutex_); + const auto it = in_flight_.find(request.opaque_id); + if (it != in_flight_.end()) { + it->second.handle = handle; + } + } + } + } + + void on_completion(AwsMetadataCompletion&& completion) { + if (config_.before_completion) { + config_.before_completion(); + } + std::lock_guard lock(mutex_); + const auto pending = in_flight_.find(completion.opaque_id); + if (pending == in_flight_.end()) { + return; + } + const AwsMetadataRequest request = pending->second.request; + in_flight_.erase(pending); + if (request.kind == AwsMetadataRequestKind::local_location) { + if (local_in_flight_ != 0) { + --local_in_flight_; + } + } else { + auto region = region_in_flight_.find(request.region); + if (region != region_in_flight_.end() && region->second != 0) { + --region->second; + if (region->second == 0) { + region_in_flight_.erase(region); + } + } + } + if (completion.generation != request.generation || + completion.generation != generation_ || stopping_) { + publish_locked(); + return; + } + + const auto now_steady = config_.steady_clock(); + const int64_t now_wall = wall_seconds(config_.wall_clock()); + if (request.kind == AwsMetadataRequestKind::local_location) { + local_.attempt_wall = now_wall; + if (completion.result.status == AwsMetadataStatus::ok && + !completion.result.local.region.empty()) { + local_.has_value = true; + local_.value = std::move(completion.result.local); + local_.success_steady = now_steady; + local_.success_wall = now_wall; + local_.error.clear(); + } else { + local_.error = completion.result.status == AwsMetadataStatus::ok + ? "invalid_response" : failure_category(completion.result.status); + } + } else { + apply_region_completion_locked(request, completion.result, + now_steady, now_wall); + } + publish_locked(); + } + + void apply_region_completion_locked( + const AwsMetadataRequest& request, + const AwsMetadataResult& result, + std::chrono::steady_clock::time_point now_steady, + int64_t now_wall) { + if (result.status != AwsMetadataStatus::ok) { + for (const auto& endpoint : request.endpoints) { + auto& record = endpoint_cache_[endpoint_key(endpoint.hostname, endpoint.port)]; + record.attempt_wall = now_wall; + record.error = failure_category(result.status); + } + return; + } + + std::unordered_map returned; + for (const auto& endpoint : result.endpoints) { + const std::string hostname = aws_locality_normalized_hostname(endpoint.hostname); + if (!hostname.empty() && endpoint.region == request.region) { + returned[endpoint_key(hostname, endpoint.port)] = &endpoint; + if (endpoint.port == 0) { + returned[endpoint_key(hostname, 0)] = &endpoint; + } + } + } + + for (const auto& endpoint : request.endpoints) { + auto& record = endpoint_cache_[endpoint_key(endpoint.hostname, endpoint.port)]; + record.attempt_wall = now_wall; + const auto exact = returned.find(endpoint_key(endpoint.hostname, endpoint.port)); + const auto no_port = returned.find(endpoint_key(endpoint.hostname, 0)); + const AwsMetadataEndpoint* match = exact != returned.end() + ? exact->second : (no_port != returned.end() ? no_port->second : nullptr); + if (match == nullptr || match->endpoint_type == AwsEndpointType::unknown) { + record.error = "endpoint_not_found"; + continue; + } + record.has_value = true; + record.value = {match->endpoint_type, match->region, + match->availability_zone, match->account_id}; + record.success_steady = now_steady; + record.success_wall = now_wall; + record.error.clear(); + } + } + + AwsLocalitySnapshotEntry build_entry_locked( + const AwsLocalityHostgroupConfig& hostgroup, + const AwsLocalityBackendConfig& backend_config, + std::chrono::steady_clock::time_point now) const { + AwsLocalitySnapshotEntry entry; + entry.hostgroup_id = hostgroup.hostgroup_id; + entry.hostname = backend_config.endpoint.hostname; + entry.port = backend_config.endpoint.port; + entry.configured_weight = backend_config.configured_weight; + entry.local = local_.value; + + const auto endpoint = endpoint_cache_.find(endpoint_key( + backend_config.endpoint.hostname, backend_config.endpoint.port)); + if (endpoint != endpoint_cache_.end()) { + entry.backend = endpoint->second.value; + entry.endpoint_type = endpoint->second.value.endpoint_type; + entry.last_attempt_timestamp = std::max( + local_.attempt_wall, endpoint->second.attempt_wall); + if (local_.success_wall != 0 && endpoint->second.success_wall != 0) { + entry.last_success_timestamp = std::min( + local_.success_wall, endpoint->second.success_wall); + } + } else { + entry.last_attempt_timestamp = local_.attempt_wall; + } + + if (!enabled_) { + entry.status = AwsLocalityMetadataStatus::disabled; + return entry; + } + if (!backend_config.endpoint.recognized) { + entry.status = AwsLocalityMetadataStatus::error; + entry.failure_category = "invalid_response"; + return entry; + } + + // A missing provider is different from a transient provider error: no + // concrete locality authority remains installed. Do not continue to + // apply a multiplier derived from a previous provider's cached result. + const bool provider_unavailable = local_.error == "provider_unavailable" || + (endpoint != endpoint_cache_.end() && + endpoint->second.error == "provider_unavailable"); + if (provider_unavailable) { + entry.status = AwsLocalityMetadataStatus::error; + entry.failure_category = "provider_unavailable"; + return entry; + } + + const bool endpoint_pending = region_in_flight_.find( + backend_config.endpoint.region) != region_in_flight_.end(); + if (!local_.has_value || endpoint == endpoint_cache_.end() || + !endpoint->second.has_value) { + if ((!local_.has_value && local_in_flight_ != 0) || + (endpoint != endpoint_cache_.end() && !endpoint->second.has_value && + endpoint_pending)) { + entry.status = AwsLocalityMetadataStatus::pending; + } else { + entry.status = AwsLocalityMetadataStatus::error; + } + if (endpoint != endpoint_cache_.end() && !endpoint->second.error.empty()) { + entry.failure_category = endpoint->second.error; + } else if (!local_.error.empty()) { + entry.failure_category = local_.error; + } + return entry; + } + + const auto local_age = now >= local_.success_steady + ? now - local_.success_steady : std::chrono::steady_clock::duration::zero(); + const auto endpoint_age = now >= endpoint->second.success_steady + ? now - endpoint->second.success_steady : std::chrono::steady_clock::duration::zero(); + const auto age = std::max(local_age, endpoint_age); + const auto refresh = std::chrono::seconds(hostgroup.policy.refresh_interval_seconds); + const auto stale_ttl = std::chrono::seconds(hostgroup.policy.stale_ttl_seconds); + if (age > stale_ttl) { + entry.status = AwsLocalityMetadataStatus::expired; + return entry; + } + const bool refresh_failed = !local_.error.empty() || !endpoint->second.error.empty(); + entry.status = age > refresh || refresh_failed + ? AwsLocalityMetadataStatus::stale : AwsLocalityMetadataStatus::fresh; + entry.failure_category = !endpoint->second.error.empty() + ? endpoint->second.error : local_.error; + entry.locality = classify_aws_locality(local_.value, endpoint->second.value); + if (entry.locality == AwsLocalityClass::same_az) { + entry.multiplier = hostgroup.policy.same_az_multiplier; + } else if (entry.locality == AwsLocalityClass::same_region) { + entry.multiplier = hostgroup.policy.same_region_multiplier; + } + return entry; + } + + void publish_locked() { + auto next = std::make_shared(); + next->generation = generation_; + next->enabled = enabled_; + const auto now = config_.steady_clock(); + for (const auto& hostgroup : hostgroups_) { + next->hostgroups.insert(hostgroup.hostgroup_id); + for (const auto& backend : hostgroup.backends) { + auto entry = build_entry_locked(hostgroup, backend, now); + next->entries.emplace(identity_hash(entry.hostgroup_id, + entry.hostname, entry.port), std::move(entry)); + } + } + std::atomic_store_explicit( + &published_, std::shared_ptr(std::move(next)), + std::memory_order_release); + } + + AwsLocalityManagerConfig config_; + mutable std::mutex mutex_; + std::condition_variable cv_; + std::vector hostgroups_; + uint64_t generation_ { 0 }; + uint64_t next_opaque_id_ { 1 }; + bool enabled_ { false }; + bool stopping_ { false }; + bool shutdown_complete_ { false }; + bool cancel_requested_ { false }; + bool force_refresh_ { false }; + bool disable_acknowledged_ { true }; + std::thread worker_; + std::shared_ptr sink_; + std::shared_ptr published_; + LocalRecord local_; + std::unordered_map endpoint_cache_; + std::unordered_map in_flight_; + size_t local_in_flight_ { 0 }; + std::unordered_map region_in_flight_; +}; + +MySQLAwsLocalityManager::MySQLAwsLocalityManager(AwsLocalityManagerConfig config) + : impl_(new Impl(std::move(config))) {} + +MySQLAwsLocalityManager::~MySQLAwsLocalityManager() = default; + +void MySQLAwsLocalityManager::configure( + std::vector hostgroups) { + impl_->configure(std::move(hostgroups)); +} + +void MySQLAwsLocalityManager::set_enabled(bool enabled) { + impl_->set_enabled(enabled); +} + +void MySQLAwsLocalityManager::request_refresh() { + impl_->request_refresh(); +} + +std::shared_ptr MySQLAwsLocalityManager::snapshot() const { + return impl_->snapshot(); +} + +std::vector MySQLAwsLocalityManager::diagnostic_rows() const { + return impl_->diagnostic_rows(); +} + +void MySQLAwsLocalityManager::shutdown() { + impl_->shutdown(); +} diff --git a/lib/BaseHGC.cpp b/lib/BaseHGC.cpp index 5f9ed952b5..f9f55ed1cf 100644 --- a/lib/BaseHGC.cpp +++ b/lib/BaseHGC.cpp @@ -62,6 +62,7 @@ void BaseHGC::reset_attributes() { attributes.init_connect = NULL; attributes.comment = NULL; attributes.ignore_session_variables_text = NULL; + attributes.aws_iam_region = NULL; } attributes.initialized = true; attributes.configured = false; @@ -80,6 +81,11 @@ void BaseHGC::reset_attributes() { attributes.comment = NULL; free(attributes.ignore_session_variables_text); attributes.ignore_session_variables_text = NULL; + free(attributes.aws_iam_region); + attributes.aws_iam_region = NULL; +#ifdef PROXYSQL40 + attributes.aws_locality_policy = {}; +#endif if (attributes.ignore_session_variables_json) { delete attributes.ignore_session_variables_json; attributes.ignore_session_variables_json = NULL; diff --git a/lib/Base_Session.cpp b/lib/Base_Session.cpp index 7e6e980cac..2d12ac7e67 100644 --- a/lib/Base_Session.cpp +++ b/lib/Base_Session.cpp @@ -28,6 +28,7 @@ template Base_Session::find_backend(int); template PgSQL_Backend * Base_Session::find_backend(int); +template MySQL_Backend * Base_Session::create_backend(int, MySQL_Data_Stream *); template MySQL_Backend * Base_Session::find_or_create_backend(int, MySQL_Data_Stream *); template PgSQL_Backend * Base_Session::find_or_create_backend(int, PgSQL_Data_Stream *); diff --git a/lib/Makefile b/lib/Makefile index 63bcc297c4..f7871d82e7 100644 --- a/lib/Makefile +++ b/lib/Makefile @@ -93,7 +93,7 @@ MYCXXFLAGS := $(STDCPP) $(MYCFLAGS) $(PSQLCH) $(PSQL40) $(PSQL31) $(PSQLFFTO) $( default: libproxysql.a .PHONY: default -_OBJ_CXX := ProxySQL_GloVars.oo network.oo debug.oo configfile.oo Query_Cache.oo SpookyV2.oo MySQL_Authentication.oo MySQL_Passthrough_Auth_Cache.oo gen_utils.oo sqlite3db.oo mysql_connection.oo MySQL_HostGroups_Manager.oo mysql_data_stream.oo MySQL_Thread.oo MySQL_Session.oo MySQL_Protocol.oo mysql_backend.oo Query_Processor.oo MySQL_Query_Processor.oo PgSQL_Query_Processor.oo ProxySQL_Admin.oo ProxySQL_Config.oo ProxySQL_Restapi.oo MySQL_Monitor.oo MySQL_Logger.oo log_utils.oo thread.oo MySQL_PreparedStatement.oo ProxySQL_Cluster.oo ClickHouse_Authentication.oo ClickHouse_Server.oo ProxySQL_Statistics.oo Chart_bundle_js.oo ProxySQL_HTTP_Server.oo ProxySQL_RESTAPI_Server.oo font-awesome.min.css.oo main-bundle.min.css.oo MySQL_Variables.oo MySQL_User_Variables.oo c_tokenizer.oo proxysql_utils.oo proxysql_coredump.oo proxysql_sslkeylog.oo \ +_OBJ_CXX := ProxySQL_GloVars.oo network.oo debug.oo configfile.oo Query_Cache.oo SpookyV2.oo MySQL_Authentication.oo MySQL_Backend_Auth.oo Aws_Iam_Provider.oo Aws_Locality_Manager.oo MySQL_Passthrough_Auth_Cache.oo gen_utils.oo sqlite3db.oo mysql_connection.oo MySQL_HostGroups_Manager.oo mysql_data_stream.oo MySQL_Thread.oo MySQL_Session.oo MySQL_Protocol.oo mysql_backend.oo Query_Processor.oo MySQL_Query_Processor.oo PgSQL_Query_Processor.oo ProxySQL_Admin.oo ProxySQL_Config.oo ProxySQL_Restapi.oo MySQL_Monitor.oo MySQL_Logger.oo log_utils.oo thread.oo MySQL_PreparedStatement.oo ProxySQL_Cluster.oo ClickHouse_Authentication.oo ClickHouse_Server.oo ProxySQL_Statistics.oo Chart_bundle_js.oo ProxySQL_HTTP_Server.oo ProxySQL_RESTAPI_Server.oo font-awesome.min.css.oo main-bundle.min.css.oo MySQL_Variables.oo MySQL_User_Variables.oo c_tokenizer.oo proxysql_utils.oo proxysql_coredump.oo proxysql_sslkeylog.oo \ sha256crypt.oo \ ProxySQL_PluginManager.oo \ BaseSrvList.oo BaseHGC.oo Base_HostGroups_Manager.oo \ diff --git a/lib/MyHGC.cpp b/lib/MyHGC.cpp index 7f3e1d4dbb..df54889d7e 100644 --- a/lib/MyHGC.cpp +++ b/lib/MyHGC.cpp @@ -31,9 +31,32 @@ MySrvC *MyHGC::get_random_MySrvC(char * gtid_uuid, uint64_t gtid_trxid, int max_ MySrvC **mysrvcCandidates = mysrvcCandidates_static; unsigned int num_candidates = 0; bool max_connections_reached = false; + bool use_aws_locality = false; +#ifdef PROXYSQL40 + std::shared_ptr aws_locality_snapshot; + if (mysql_thread___aws_locality_awareness && MyHGM != nullptr && + MyHGM->aws_locality_manager() != nullptr) { + aws_locality_snapshot = MyHGM->aws_locality_manager()->snapshot(); + use_aws_locality = aws_locality_snapshot != nullptr && + aws_locality_snapshot->enabled && + aws_locality_snapshot->has_hostgroup(hid); + } +#endif if (l>32) { mysrvcCandidates = (MySrvC **)malloc(sizeof(MySrvC *)*l); } + auto candidate_weight_sum = [&]() -> uint64_t { +#ifdef PROXYSQL40 + if (use_aws_locality) { + // Locality selection has a uniform fallback for an all-zero + // configured-weight set. These availability checks therefore only + // need to know whether an eligible candidate exists; defer snapshot + // lookups until the actual lottery below. + return num_candidates; + } +#endif + return sum; + }; if (l) { //int j=0; for (j=0; j32) { free(mysrvcCandidates); @@ -275,7 +298,7 @@ MySrvC *MyHGC::get_random_MySrvC(char * gtid_uuid, uint64_t gtid_trxid, int max_ unsigned int New_sum=sum; - if (New_sum==0) { + if (candidate_weight_sum()==0) { proxy_debug(PROXY_DEBUG_MYSQL_CONNPOOL, 7, "Returning MySrvC NULL because no backend ONLINE or with weight\n"); if (l>32) { free(mysrvcCandidates); @@ -327,6 +350,59 @@ MySrvC *MyHGC::get_random_MySrvC(char * gtid_uuid, uint64_t gtid_trxid, int max_ } } +#ifdef PROXYSQL40 + if (use_aws_locality) { + uint64_t total_weight = 0; + for (j = 0; j < num_candidates; ++j) { + mysrvc = mysrvcCandidates[j]; + total_weight = aws_locality_saturating_add(total_weight, + aws_locality_snapshot->effective_weight( + hid, mysrvc->address, mysrvc->port, mysrvc->weight)); + } + const uint64_t random_value = + (static_cast(rand_fast()) << 32) | + static_cast(rand_fast()); + size_t selected = num_candidates; + if (total_weight == 0 && num_candidates != 0) { + selected = random_value % num_candidates; + } else if (total_weight != 0) { + const uint64_t target = random_value % total_weight; + uint64_t cumulative = 0; + for (j = 0; j < num_candidates; ++j) { + mysrvc = mysrvcCandidates[j]; + cumulative = aws_locality_saturating_add(cumulative, + aws_locality_snapshot->effective_weight( + hid, mysrvc->address, mysrvc->port, mysrvc->weight)); + if (target < cumulative) { + selected = j; + break; + } + } + } + if (selected < num_candidates) { + mysrvc = mysrvcCandidates[selected]; + proxy_debug(PROXY_DEBUG_MYSQL_CONNPOOL, 7, + "Returning MySrvC %p, server %s:%d with AWS locality weighting\n", + mysrvc, mysrvc->address, mysrvc->port); + if (l>32) { + free(mysrvcCandidates); + } +#ifdef TEST_AURORA + array_mysrvc_cands += num_candidates; +#endif // TEST_AURORA + return mysrvc; + } + proxy_debug(PROXY_DEBUG_MYSQL_CONNPOOL, 7, + "Returning MySrvC NULL because no AWS locality candidate is eligible\n"); + if (l>32) { + free(mysrvcCandidates); + } +#ifdef TEST_AURORA + array_mysrvc_cands += num_candidates; +#endif // TEST_AURORA + return NULL; + } +#endif unsigned int k; k=rand_fast()%New_sum; diff --git a/lib/MySQL_Authentication.cpp b/lib/MySQL_Authentication.cpp index 58ae1ceb4d..fcd1c3071a 100644 --- a/lib/MySQL_Authentication.cpp +++ b/lib/MySQL_Authentication.cpp @@ -8,6 +8,7 @@ #include "proxysql_atomic.h" #include "MySQL_Authentication.hpp" +#include "MySQL_Backend_Auth.h" #include @@ -48,6 +49,12 @@ void cleanse_and_free_password(char*& password) { } } +const char* invalid_attributes_storage(enum cred_username_type usertype) { + // Backend policies must stay invalid after runtime loading so they cannot + // silently fall back to password authentication. + return usertype == USERNAME_BACKEND ? "null" : ""; +} + } // namespace void free_account_details(account_details_t& ad) { @@ -269,8 +276,8 @@ bool MySQL_Authentication::add(char * username, char * password, enum cred_usern } } catch(nlohmann::json::exception& e) { - ad->attributes=strdup(""); - proxy_error("Invalid attributes for user %s: %s\n", username, attributes); + ad->attributes=strdup(invalid_attributes_storage(usertype)); + proxy_error("Invalid attributes for user %s; ignoring invalid JSON\n", username); } } else { ad->attributes=strdup(attributes); // default, empty string @@ -293,8 +300,8 @@ bool MySQL_Authentication::add(char * username, char * password, enum cred_usern ad->attributes=strdup(attributes); } catch(nlohmann::json::exception& e) { - ad->attributes=strdup(""); - proxy_error("Invalid attributes for user %s: %s\n", username, attributes); + ad->attributes=strdup(invalid_attributes_storage(usertype)); + proxy_error("Invalid attributes for user %s; ignoring invalid JSON\n", username); } } else { ad->attributes=strdup(attributes); // default, empty string @@ -330,11 +337,26 @@ bool MySQL_Authentication::add(char * username, char * password, enum cred_usern cg.bt_map.insert(std::make_pair(hash1,ad)); cg.cred_array->add(ad); } + bool warn_for_iam_password = false; + if (usertype == USERNAME_BACKEND) { + const MySQLBackendAuthPolicy policy = parse_mysql_backend_auth_policy( + username, + ad->attributes != nullptr ? ad->attributes : "", + ad->password != nullptr && ad->password[0] != '\0'); + warn_for_iam_password = + policy.type == MySQLBackendAuthType::AWS_IAM && policy.ignored_password; + } #ifdef PROXYSQL_AUTH_PTHREAD_MUTEX pthread_rwlock_unlock(&cg.lock); #else spin_wrunlock(&cg.lock); #endif + if (warn_for_iam_password) { + proxy_warning( + "mysql_users backend entry for '%s' uses aws_iam authentication; " + "clear the unused backend password\n", + username); + } return true; }; diff --git a/lib/MySQL_Backend_Auth.cpp b/lib/MySQL_Backend_Auth.cpp new file mode 100644 index 0000000000..ac43b44a86 --- /dev/null +++ b/lib/MySQL_Backend_Auth.cpp @@ -0,0 +1,201 @@ +#include "../deps/json/json.hpp" + +#include "MySQL_Backend_Auth.h" +#include "MySQL_Authentication.hpp" + +#include + +namespace { + +MySQLBackendAuthPolicy invalid_policy(std::string_view database_user, const char* failure_code) { + MySQLBackendAuthPolicy policy; + policy.database_user = database_user; + policy.failure_code = failure_code; + return policy; +} + +AwsIamConnectionConfigResult invalid_connection_config( + AwsIamConnectionConfigStatus status, const char* failure_code) { + AwsIamConnectionConfigResult result; + result.status = status; + result.failure_code = failure_code; + return result; +} + +bool is_dns_label(const std::string& label) { + if (label.empty() || label.size() > 63 || label.front() == '-' || label.back() == '-') { + return false; + } + for (const unsigned char c : label) { + if (!((c >= 'a' && c <= 'z') || (c >= 'A' && c <= 'Z') || + (c >= '0' && c <= '9') || c == '-')) { + return false; + } + } + return true; +} + +bool split_rds_endpoint(const std::string& endpoint, std::string& endpoint_region) { + if (endpoint.empty() || endpoint.size() > 253 || endpoint.back() == '.') { + return false; + } + + std::vector labels; + size_t begin = 0; + while (begin < endpoint.size()) { + const size_t end = endpoint.find('.', begin); + const std::string label = endpoint.substr(begin, end == std::string::npos ? end : end - begin); + if (!is_dns_label(label)) { + return false; + } + labels.push_back(label); + if (end == std::string::npos) { + break; + } + begin = end + 1; + } + + static const std::vector> suffixes { + { "rds", "amazonaws", "com" }, + { "rds", "amazonaws", "com", "cn" }, + { "rds", "c2s", "ic", "gov" }, + { "rds", "sc2s", "sgov", "gov" }, + }; + for (const auto& suffix : suffixes) { + if (labels.size() <= suffix.size() + 1) { + continue; + } + const size_t suffix_start = labels.size() - suffix.size(); + bool suffix_matches = true; + for (size_t i = 0; i < suffix.size(); ++i) { + if (labels[suffix_start + i] != suffix[i]) { + suffix_matches = false; + break; + } + } + if (suffix_matches) { + endpoint_region = labels[suffix_start - 1]; + return true; + } + } + return false; +} + +} // namespace + +const char* mysql_backend_auth_type_name(MySQLBackendAuthType type) { + switch (type) { + case MySQLBackendAuthType::PASSWORD: + return "password"; + case MySQLBackendAuthType::AWS_IAM: + return "aws_iam"; + case MySQLBackendAuthType::INVALID: + return "invalid"; + } + return "invalid"; +} + +MySQLBackendAuthPolicy parse_mysql_backend_auth_policy( + std::string_view database_user, + std::string_view attributes, + bool backend_password_is_nonempty) +{ + MySQLBackendAuthPolicy policy; + policy.database_user = database_user; + + if (attributes.empty()) { + policy.type = MySQLBackendAuthType::PASSWORD; + return policy; + } + + const nlohmann::json parsed = nlohmann::json::parse( + attributes.data(), attributes.data() + attributes.size(), nullptr, false); + if (parsed.is_discarded() || !parsed.is_object()) { + return invalid_policy(database_user, "attributes_not_object"); + } + + const auto backend_auth = parsed.find("backend_auth"); + if (backend_auth == parsed.end()) { + policy.type = MySQLBackendAuthType::PASSWORD; + return policy; + } + if (!backend_auth->is_object()) { + return invalid_policy(database_user, "backend_auth_not_object"); + } + + const auto type = backend_auth->find("type"); + if (type == backend_auth->end()) { + return invalid_policy(database_user, "type_missing"); + } + if (!type->is_string()) { + return invalid_policy(database_user, "type_not_string"); + } + if (type->get() != "aws_iam") { + return invalid_policy(database_user, "type_unsupported"); + } + + policy.type = MySQLBackendAuthType::AWS_IAM; + policy.ignored_password = backend_password_is_nonempty; + return policy; +} + +MySQLBackendAuthPolicy resolve_mysql_backend_auth_policy( + MySQL_Authentication& authentication, + const char* mapped_backend_username) +{ + const std::string_view database_user = + mapped_backend_username != nullptr ? mapped_backend_username : ""; + if (mapped_backend_username == nullptr) { + return invalid_policy(database_user, "backend_user_not_found"); + } + + account_details_t account = authentication.lookup( + const_cast(mapped_backend_username), USERNAME_BACKEND, { false, false, true }); + if (account.password == nullptr) { + free_account_details(account); + return invalid_policy(database_user, "backend_user_not_found"); + } + + const MySQLBackendAuthPolicy policy = parse_mysql_backend_auth_policy( + database_user, + account.attributes != nullptr ? account.attributes : "", + account.password[0] != '\0'); + free_account_details(account); + return policy; +} + +AwsIamConnectionConfigResult validate_mysql_aws_iam_connection( + const AwsIamConnectionConfigInput& input) +{ + if (!input.support_compiled) { + return invalid_connection_config(AwsIamConnectionConfigStatus::SUPPORT_NOT_COMPILED, + "support_not_compiled"); + } + if (input.region.empty()) { + return invalid_connection_config(AwsIamConnectionConfigStatus::MISSING_REGION, "missing_region"); + } + if (input.port == 0) { + return invalid_connection_config(AwsIamConnectionConfigStatus::UNIX_SOCKET_NOT_ALLOWED, + "unix_socket_not_allowed"); + } + + std::string endpoint_region; + if (!split_rds_endpoint(input.configured_endpoint, endpoint_region)) { + return invalid_connection_config(AwsIamConnectionConfigStatus::INVALID_ENDPOINT, "invalid_endpoint"); + } + if (endpoint_region != input.region) { + return invalid_connection_config(AwsIamConnectionConfigStatus::REGION_ENDPOINT_MISMATCH, + "region_endpoint_mismatch"); + } + if (!input.use_ssl) { + return invalid_connection_config(AwsIamConnectionConfigStatus::TLS_REQUIRED, "tls_required"); + } + if (input.ssl_ca.empty() && input.ssl_capath.empty()) { + return invalid_connection_config(AwsIamConnectionConfigStatus::CA_TRUST_REQUIRED, "ca_trust_required"); + } + + AwsIamConnectionConfigResult result; + result.status = AwsIamConnectionConfigStatus::OK; + result.key = { input.configured_endpoint, input.port, input.region, input.database_user }; + return result; +} diff --git a/lib/MySQL_HostGroups_Manager.cpp b/lib/MySQL_HostGroups_Manager.cpp index f7a9824841..29fc23d865 100644 --- a/lib/MySQL_HostGroups_Manager.cpp +++ b/lib/MySQL_HostGroups_Manager.cpp @@ -29,11 +29,14 @@ using json = nlohmann::json; #include "ev.h" #include +#include #include #include using std::function; +extern MySQL_Authentication *GloMyAuth; + #define SAFE_SQLITE3_STEP(_stmt) do {\ do {\ @@ -129,8 +132,7 @@ T j_get_srv_default_int_val( } -//static void * HGCU_thread_run() { -static void * HGCU_thread_run() { +void * HGCU_thread_run() { PtrArray *conn_array=new PtrArray(); set_thread_name("MyHGCU", GloVars.set_thread_name); while(1) { @@ -150,15 +152,34 @@ static void * HGCU_thread_run() { } conn_array->add(myconn); } + for (unsigned int i = 0; i < conn_array->len;) { + myconn = (MySQL_Connection *)conn_array->index(i); + const char *backend_username = myconn->userinfo != nullptr + ? myconn->userinfo->username : nullptr; + const MySQLBackendAuthPolicy policy = GloMyAuth != nullptr + ? resolve_mysql_backend_auth_policy(*GloMyAuth, backend_username) + : MySQLBackendAuthPolicy {}; + const bool reset_allowed = + myconn->backend_auth_type() == MySQLBackendAuthType::PASSWORD && + (GloMyAuth == nullptr || + myconn->can_reset_for_backend_auth_policy(policy)); + if (!reset_allowed) { + conn_array->remove_index_fast(i); + myconn->send_quit = false; + MyHGM->destroy_MyConn_from_pool(myconn); + continue; + } + ++i; + } unsigned int l=conn_array->len; int *errs=(int *)malloc(sizeof(int)*l); int *statuses=(int *)malloc(sizeof(int)*l); my_bool *ret=(my_bool *)malloc(sizeof(my_bool)*l); int i; for (i=0;i<(int)l;i++) { + myconn=(MySQL_Connection *)conn_array->index(i); myconn->reset(); MyHGM->increase_reset_counter(); - myconn=(MySQL_Connection *)conn_array->index(i); if (myconn->mysql->net.pvio && myconn->mysql->net.fd && myconn->mysql->net.buff) { MySQL_Connection_userinfo *userinfo = myconn->userinfo; char *auth_password = NULL; @@ -669,6 +690,9 @@ hg_metrics_map = std::make_tuple( ); MySQL_HostGroups_Manager::MySQL_HostGroups_Manager() { +#ifdef PROXYSQL40 + aws_locality_manager_ = std::make_unique(); +#endif status.client_connections=0; status.client_connections_prim_pass=0; status.client_connections_addl_pass=0; @@ -781,6 +805,11 @@ void MySQL_HostGroups_Manager::init() { } void MySQL_HostGroups_Manager::shutdown() { +#ifdef PROXYSQL40 + if (aws_locality_manager_) { + aws_locality_manager_->shutdown(); + } +#endif queue.add(NULL); HGCU_thread->join(); delete HGCU_thread; @@ -790,6 +819,11 @@ void MySQL_HostGroups_Manager::shutdown() { } MySQL_HostGroups_Manager::~MySQL_HostGroups_Manager() { +#ifdef PROXYSQL40 + if (aws_locality_manager_) { + aws_locality_manager_->shutdown(); + } +#endif while (MyHostGroups->len) { MyHGC *myhgc=(MyHGC *)MyHostGroups->remove_index_fast(0); delete myhgc; @@ -809,6 +843,140 @@ MySQL_HostGroups_Manager::~MySQL_HostGroups_Manager() { pthread_mutex_destroy(&lock); } +#ifdef PROXYSQL40 +void MySQL_HostGroups_Manager::refresh_aws_locality_configuration() { + std::vector hostgroups; + + wrlock(); + for (unsigned int i = 0; i < MyHostGroups->len; ++i) { + MyHGC* hostgroup = static_cast(MyHostGroups->index(i)); + if (!hostgroup->attributes.aws_locality_policy.valid) { + continue; + } + + AwsLocalityHostgroupConfig config; + config.hostgroup_id = hostgroup->hid; + config.policy = hostgroup->attributes.aws_locality_policy; + config.backends.reserve(hostgroup->mysrvs->servers->len); + for (unsigned int j = 0; j < hostgroup->mysrvs->servers->len; ++j) { + MySrvC* server = static_cast(hostgroup->mysrvs->servers->index(j)); + config.backends.emplace_back( + recognize_rds_endpoint(hostgroup->hid, server->address, server->port), + server->weight); + } + hostgroups.emplace_back(std::move(config)); + } + wrunlock(); + + if (aws_locality_manager_) { + aws_locality_manager_->configure(std::move(hostgroups)); + } +} + +void MySQL_HostGroups_Manager::set_aws_locality_awareness_enabled(bool enabled) { + if (aws_locality_manager_) { + aws_locality_manager_->set_enabled(enabled); + } +} + +namespace { + +const char* aws_endpoint_type_name(AwsEndpointType type) { + switch (type) { + case AwsEndpointType::instance: return "instance"; + case AwsEndpointType::cluster: return "cluster"; + case AwsEndpointType::reader: return "reader"; + case AwsEndpointType::custom: return "custom"; + case AwsEndpointType::unknown: return "unknown"; + } + return "unknown"; +} + +const char* aws_locality_name(AwsLocalityClass locality) { + switch (locality) { + case AwsLocalityClass::remote: return "remote"; + case AwsLocalityClass::same_region: return "same_region"; + case AwsLocalityClass::same_az: return "same_az"; + case AwsLocalityClass::unknown: return "unknown"; + } + return "unknown"; +} + +const char* aws_metadata_status_name(AwsLocalityMetadataStatus status) { + switch (status) { + case AwsLocalityMetadataStatus::disabled: return "disabled"; + case AwsLocalityMetadataStatus::pending: return "pending"; + case AwsLocalityMetadataStatus::fresh: return "fresh"; + case AwsLocalityMetadataStatus::stale: return "stale"; + case AwsLocalityMetadataStatus::expired: return "expired"; + case AwsLocalityMetadataStatus::error: return "error"; + } + return "error"; +} + +const char* aws_account_match_name(const AwsLocalitySnapshotEntry& row) { + if (row.local.account_id.empty() || row.backend.account_id.empty()) { + return "unknown"; + } + return row.local.account_id == row.backend.account_id ? "same" : "different"; +} + +bool aws_locality_status_is_active(AwsLocalityMetadataStatus status) { + return status == AwsLocalityMetadataStatus::fresh || + status == AwsLocalityMetadataStatus::stale; +} + +std::mutex aws_locality_stats_projection_mutex; + +} // namespace + +bool MySQL_HostGroups_Manager::project_aws_locality_stats( + SQLite3DB* statsdb, + const std::vector& rows) { + std::lock_guard projection_lock(aws_locality_stats_projection_mutex); + if (statsdb == nullptr || !statsdb->execute("BEGIN")) return false; + bool success = statsdb->execute("DELETE FROM stats_mysql_aws_locality"); + for (const auto& row : rows) { + if (!success) break; + const double active_multiplier = aws_locality_status_is_active(row.status) + ? row.multiplier : 1.0; + const uint64_t effective_weight = aws_locality_effective_weight( + row.configured_weight, active_multiplier); + char* query = sqlite3_mprintf( + "INSERT INTO stats_mysql_aws_locality (" + "hostgroup_id,hostname,port,endpoint_type,configured_weight," + "effective_weight,local_region,local_az,backend_region,backend_az," + "account_match,locality,active_multiplier,metadata_status," + "last_success_timestamp,last_attempt_timestamp,last_error_category) " + "VALUES (%u,'%q',%u,'%q',%lld,%llu,'%q','%q','%q','%q','%q','%q'," + "%.17g,'%q',%lld,%lld,'%q')", + row.hostgroup_id, row.hostname.c_str(), static_cast(row.port), + aws_endpoint_type_name(row.endpoint_type), + static_cast(row.configured_weight), + static_cast(effective_weight), + row.local.region.c_str(), row.local.availability_zone.c_str(), + row.backend.region.c_str(), row.backend.availability_zone.c_str(), + aws_account_match_name(row), aws_locality_name(row.locality), + active_multiplier, aws_metadata_status_name(row.status), + static_cast(row.last_success_timestamp), + static_cast(row.last_attempt_timestamp), + row.failure_category.c_str()); + success = query != nullptr && statsdb->execute(query); + sqlite3_free(query); + } + if (success) success = statsdb->execute("COMMIT"); + if (!success) statsdb->execute("ROLLBACK"); + return success; +} + +void MySQL_HostGroups_Manager::refresh_aws_locality_stats(SQLite3DB* statsdb) const { + const std::vector rows = aws_locality_manager_ + ? aws_locality_manager_->diagnostic_rows() + : std::vector(); + project_aws_locality_stats(statsdb, rows); +} +#endif + void MySQL_HostGroups_Manager::p_update_mysql_error_counter(p_mysql_error_type err_type, unsigned int hid, char* address, uint16_t port, unsigned int code) { p_hg_dyn_counter::metric metric = p_hg_dyn_counter::mysql_error; if (err_type == p_mysql_error_type::proxysql) { @@ -1620,6 +1788,9 @@ bool MySQL_HostGroups_Manager::commit( update_aws_rds_bgd_hosts_monitor_resultset(); wrunlock(); +#ifdef PROXYSQL40 + refresh_aws_locality_configuration(); +#endif unsigned long long curtime2=monotonic_time(); curtime1 = curtime1/1000; curtime2 = curtime2/1000; @@ -2533,7 +2704,11 @@ void MySQL_HostGroups_Manager::unshun_server_all_hostgroups(const char * address * @note This method locks the connection pool to ensure thread safety during access. It releases the lock once * the operation is completed. */ -MySQL_Connection * MySQL_HostGroups_Manager::get_MyConn_from_pool(unsigned int _hid, MySQL_Session *sess, bool ff, char * gtid_uuid, uint64_t gtid_trxid, int max_lag_ms) { +MySQL_Connection * MySQL_HostGroups_Manager::get_MyConn_from_pool( + unsigned int _hid, MySQL_Session *sess, bool ff, char *gtid_uuid, + uint64_t gtid_trxid, int max_lag_ms, + MySQLBackendAuthType requested_type) +{ MySQL_Connection * conn = nullptr; // Pointer to hold the retrieved MySQL_Connection // Acquire a write lock to access the connection pool @@ -2551,7 +2726,7 @@ MySQL_Connection * MySQL_HostGroups_Manager::get_MyConn_from_pool(unsigned int _ mysrvc = myhgc->get_random_MySrvC(gtid_uuid, gtid_trxid, max_lag_ms, sess); if (mysrvc) { // a MySrvC exists. If not, we return NULL = no targets // Attempt to get a random MySQL_Connection from the server's free connection pool - conn=mysrvc->ConnectionsFree->get_random_MyConn(sess, ff); + conn=mysrvc->ConnectionsFree->get_random_MyConn(sess, ff, requested_type); // If a connection is obtained, mark it as used and update connection pool statistics if (conn) { @@ -2583,17 +2758,21 @@ void MySQL_HostGroups_Manager::destroy_MyConn_from_pool(MySQL_Connection *c, boo bool to_del=true; // the default, legacy behavior MySrvC *mysrvc=(MySrvC *)c->parent; - if (c->healthy && mysrvc->get_status() == MYSQL_SERVER_STATUS_ONLINE && c->send_quit && queue.size() < __sync_fetch_and_add(&GloMTH->variables.connpoll_reset_queue_length, 0)) { + if (c->healthy && mysrvc->get_status() == MYSQL_SERVER_STATUS_ONLINE && + c->send_quit && + queue.size() < __sync_fetch_and_add(&GloMTH->variables.connpoll_reset_queue_length, 0)) { if (c->async_state_machine==ASYNC_IDLE) { - // overall, the backend seems healthy and so it is the connection. Try to reset it - int myerr=mysql_errno(c->mysql); - if (myerr >= 2000 && myerr < 3000) { - // client library error . We must not try to save the connection - proxy_debug(PROXY_DEBUG_MYSQL_CONNPOOL, 7, "Not trying to reset MySQL_Connection %p, server %s:%d . Error code %d\n", c, mysrvc->address, mysrvc->port, myerr); - } else { - proxy_debug(PROXY_DEBUG_MYSQL_CONNPOOL, 7, "Trying to reset MySQL_Connection %p, server %s:%d\n", c, mysrvc->address, mysrvc->port); - to_del=false; - queue.add(c); + if (c->backend_auth_type() != MySQLBackendAuthType::AWS_IAM) { + // overall, the backend seems healthy and so it is the connection. Try to reset it + int myerr=mysql_errno(c->mysql); + if (myerr >= 2000 && myerr < 3000) { + // client library error . We must not try to save the connection + proxy_debug(PROXY_DEBUG_MYSQL_CONNPOOL, 7, "Not trying to reset MySQL_Connection %p, server %s:%d . Error code %d\n", c, mysrvc->address, mysrvc->port, myerr); + } else { + proxy_debug(PROXY_DEBUG_MYSQL_CONNPOOL, 7, "Trying to reset MySQL_Connection %p, server %s:%d\n", c, mysrvc->address, mysrvc->port); + to_del=false; + queue.add(c); + } } } else { // the connection seems health, but we are trying to destroy it @@ -2607,15 +2786,30 @@ void MySQL_HostGroups_Manager::destroy_MyConn_from_pool(MySQL_Connection *c, boo default: if (c->mysql->thread_id) { MySQL_Connection_userinfo *ui=c->userinfo; - char *auth_password=NULL; - if (ui->password) { - if (ui->password[0]=='*') { // we don't have the real password, let's pass sha1 - auth_password=ui->sha1_pass; - } else { - auth_password=ui->password; + KillArgs *ka = nullptr; + if (c->backend_auth_type() == MySQLBackendAuthType::AWS_IAM) { + const char *region = mysrvc->myhgc != nullptr && + mysrvc->myhgc->attributes.aws_iam_region != nullptr + ? mysrvc->myhgc->attributes.aws_iam_region : ""; + ka = new KillArgs( + ui->username, nullptr, mysrvc->address, mysrvc->port, + mysrvc->myhgc->hid, c->mysql->thread_id, + KILL_CONNECTION, mysrvc->use_ssl, nullptr, + c->connected_host_details.ip, + MySQLBackendAuthType::AWS_IAM, mysrvc->address, + region, ui->username, + std::chrono::steady_clock::now() + std::chrono::seconds(5)); + } else { + char *auth_password=NULL; + if (ui->password) { + if (ui->password[0]=='*') { // we don't have the real password, let's pass sha1 + auth_password=ui->sha1_pass; + } else { + auth_password=ui->password; + } } + ka = new KillArgs(ui->username, auth_password, c->parent->address, c->parent->port, c->parent->myhgc->hid, c->mysql->thread_id, KILL_CONNECTION, c->parent->use_ssl, NULL, c->connected_host_details.ip); } - KillArgs *ka = new KillArgs(ui->username, auth_password, c->parent->address, c->parent->port, c->parent->myhgc->hid, c->mysql->thread_id, KILL_CONNECTION, c->parent->use_ssl, NULL, c->connected_host_details.ip); pthread_attr_t attr; pthread_attr_init(&attr); pthread_attr_setdetachstate(&attr, PTHREAD_CREATE_DETACHED); @@ -6196,11 +6390,37 @@ bool AWS_Aurora_Info::update(int r, int _port, char *_end_addr, int maxl, int al */ void init_myhgc_hostgroup_settings(const char* hostgroup_settings, MyHGC* myhgc) { const uint32_t hid = myhgc->hid; + free(myhgc->attributes.aws_iam_region); + myhgc->attributes.aws_iam_region = NULL; +#ifdef PROXYSQL40 + myhgc->attributes.aws_locality_policy = {}; +#endif if (hostgroup_settings[0] != '\0') { try { nlohmann::json j = nlohmann::json::parse(hostgroup_settings); +#ifdef PROXYSQL40 + const auto aws = j.find("aws"); + if (aws != j.end()) { + if (!aws->is_object()) { + proxy_error("Invalid AWS locality policy field 'aws' for hostgroup %u. Value rejected.\n", hid); + } else { + const auto locality = aws->find("locality_awareness"); + if (locality != aws->end()) { + AwsLocalityPolicyError error; + myhgc->attributes.aws_locality_policy = parse_aws_locality_policy( + *locality, hid, error); + if (!myhgc->attributes.aws_locality_policy.valid) { + proxy_error( + "Invalid AWS locality policy field '%s' for hostgroup %u. Value rejected.\n", + error.field.c_str(), hid); + } + } + } + } +#endif + const auto handle_warnings_check = [](int8_t handle_warnings) -> bool { return handle_warnings == 0 || handle_warnings == 1; }; const int8_t handle_warnings = j_get_srv_default_int_val(j, hid, "handle_warnings", handle_warnings_check); myhgc->attributes.handle_warnings = handle_warnings; @@ -6214,12 +6434,28 @@ void init_myhgc_hostgroup_settings(const char* hostgroup_settings, MyHGC* myhgc) { return (default_query_timeout >= 1000 && default_query_timeout <= 20*24*3600*1000); }; const int32_t default_query_timeout = j_get_srv_default_int_val(j, hid, "default_query_timeout", default_query_timeout_check); myhgc->attributes.default_query_timeout = default_query_timeout; + + const auto aws_iam_region = j.find("aws_iam_region"); + if (aws_iam_region != j.end()) { + if (!aws_iam_region->is_string()) { + proxy_error("Invalid 'aws_iam_region' value for hostgroup %d. Value rejected.\n", hid); + } else { + const std::string region = aws_iam_region->get(); + const bool valid_region = !region.empty() && std::all_of(region.begin(), region.end(), + [](unsigned char c) { + return (c >= 'A' && c <= 'Z') || (c >= 'a' && c <= 'z') || + (c >= '0' && c <= '9') || c == '-'; + }); + if (valid_region) { + myhgc->attributes.aws_iam_region = strdup(region.c_str()); + } else { + proxy_error("Invalid 'aws_iam_region' value for hostgroup %d. Value rejected.\n", hid); + } + } + } } - catch (const json::exception& e) { - proxy_error( - "JSON parsing for 'mysql_hostgroup_attributes.hostgroup_settings' for hostgroup %d failed with exception `%s`.\n", - hid, e.what() - ); + catch (const json::exception&) { + proxy_error("hostgroup_settings_parse_failed for hostgroup %d. Value rejected.\n", hid); } } } diff --git a/lib/MySQL_Session.cpp b/lib/MySQL_Session.cpp index 432dad8519..288d891f72 100644 --- a/lib/MySQL_Session.cpp +++ b/lib/MySQL_Session.cpp @@ -25,6 +25,8 @@ using json = nlohmann::json; #include "MySQL_Logger.hpp" #include "StatCounters.h" #include "MySQL_Authentication.hpp" +#include "MySQL_Backend_Auth.h" +#include "Aws_Iam_Provider.h" #include "MySQL_Passthrough_Auth_Cache.h" #include "MySQL_LDAP_Authentication.hpp" #include "MySQL_Protocol.h" @@ -77,6 +79,11 @@ using json = nlohmann::json; #define SHOW_STATUS_LIKE_SSL_VERSION "SHOW STATUS LIKE 'Ssl_version" #define SHOW_STATUS_LIKE_SSL_VERSION_LEN 29 +static void record_aws_iam_backend_connection(bool success) { + AwsIamTokenSourceLease source = acquire_global_aws_iam_token_source(); + if (source) source->record_backend_connection(success); +} + #define EXPMARIA using std::function; @@ -161,6 +168,34 @@ extern MySQL_STMT_Manager_v14 *GloMyStmt; extern SQLite3_Server *GloSQLite3Server; +static bool session_authorizes_rowless_passthrough( + const MySQL_Session *session, const char *backend_username, + const MySQLBackendAuthPolicy& policy) +{ + return policy.type == MySQLBackendAuthType::INVALID && session != nullptr && + session->passthrough_credential && backend_username != nullptr && + backend_username[0] != '\0' && + policy.failure_code == "backend_user_not_found"; +} + +static MySQLBackendAuthPolicy resolved_backend_auth_policy_for_session( + const MySQL_Session *session) +{ + const char *backend_username = + session != nullptr && session->client_myds != nullptr && + session->client_myds->myconn != nullptr && + session->client_myds->myconn->userinfo != nullptr + ? session->client_myds->myconn->userinfo->username : nullptr; + MySQLBackendAuthPolicy policy = + resolve_mysql_backend_auth_policy(*GloMyAuth, backend_username); + if (session_authorizes_rowless_passthrough( + session, backend_username, policy)) { + policy.type = MySQLBackendAuthType::PASSWORD; + policy.failure_code.clear(); + } + return policy; +} + #ifdef PROXYSQLCLICKHOUSE extern ClickHouse_Authentication *GloClickHouseAuth; extern ClickHouse_Server *GloClickHouseServer; @@ -209,9 +244,21 @@ KillArgs::KillArgs(char* u, char* p, char* h, unsigned int P, unsigned int _hid, } } -KillArgs::KillArgs(char* u, char* p, char* h, unsigned int P, unsigned int _hid, unsigned long i, int kt, int _use_ssl, MySQL_Thread *_mt, char *ip) { +KillArgs::KillArgs(char* u, char* p, char* h, unsigned int P, + unsigned int _hid, unsigned long i, int kt, int _use_ssl, + MySQL_Thread *_mt, char *ip) : + KillArgs(u, p, h, P, _hid, i, kt, _use_ssl, _mt, ip, + MySQLBackendAuthType::PASSWORD, h, "", u, + std::chrono::steady_clock::now() + std::chrono::seconds(5)) {} + +KillArgs::KillArgs(char* u, char* p, char* h, unsigned int P, + unsigned int _hid, unsigned long i, int kt, int _use_ssl, + MySQL_Thread *_mt, char *ip, MySQLBackendAuthType auth_type, + const char *endpoint, const char *aws_region, const char *db_user, + std::chrono::steady_clock::time_point deadline) { username=u ? strdup(u) : nullptr; - password=p ? strdup(p) : nullptr; + password=auth_type == MySQLBackendAuthType::AWS_IAM + ? nullptr : (p ? strdup(p) : nullptr); hostname=h ? strdup(h) : nullptr; ip_addr = NULL; if (ip) @@ -222,6 +269,11 @@ KillArgs::KillArgs(char* u, char* p, char* h, unsigned int P, unsigned int _hid, kill_type=kt; use_ssl=_use_ssl; mt=_mt; + backend_auth_type=auth_type; + configured_endpoint=endpoint != nullptr ? endpoint : ""; + region=aws_region != nullptr ? aws_region : ""; + database_user=db_user != nullptr ? db_user : ""; + token_deadline=deadline; } KillArgs::~KillArgs() { @@ -241,6 +293,13 @@ const char* KillArgs::get_host_address() const { return host_address; } +static void cleanse_iam_connector_password(MYSQL *mysql) { + if (mysql == nullptr || mysql->passwd == nullptr) return; + OPENSSL_cleanse(mysql->passwd, strlen(mysql->passwd)); + free(mysql->passwd); + mysql->passwd = nullptr; +} + /** * @brief Thread function to kill a query or connection on a MySQL server. @@ -253,6 +312,12 @@ const char* KillArgs::get_host_address() const { */ void* kill_query_thread(void *arg) { KillArgs *ka=(KillArgs *)arg; + const bool iam_mode = + ka->backend_auth_type == MySQLBackendAuthType::AWS_IAM; + AwsIamTokenResult iam_result; + AwsIamTokenSourceLease iam_source; + const char *connect_user = ka->username; + const char *connect_password = ka->password; //! It initializes a new MySQL_Thread object to handle MySQL-related operations. std::unique_ptr mysql_thr(new MySQL_Thread()); set_thread_name("KillQuery", GloVars.set_thread_name); @@ -276,11 +341,88 @@ void* kill_query_thread(void *arg) { //! If SSL is enabled and port information is available, it retrieves SSL parameters for the server from MyHGM and configures the MySQL connection accordingly. if (ka->use_ssl && ka->port) { - ssl_params = MyHGM->get_Server_SSL_Params(ka->hostname, ka->port, ka->username); + const char *ssl_endpoint = iam_mode && !ka->configured_endpoint.empty() + ? ka->configured_endpoint.c_str() : ka->hostname; + const char *ssl_user = iam_mode && !ka->database_user.empty() + ? ka->database_user.c_str() : ka->username; + ssl_params = MyHGM->get_Server_SSL_Params( + const_cast(ssl_endpoint), ka->port, + const_cast(ssl_user)); MySQL_Connection::set_ssl_params(mysql,ssl_params); mysql_options(mysql, MARIADB_OPT_SSL_KEYLOG_CALLBACK, (void*)proxysql_keylog_write_line_callback); } + if (iam_mode) { + iam_source = acquire_global_aws_iam_token_source(); + AwsIamConnectionConfigInput input; + input.database_user = ka->database_user; + input.configured_endpoint = ka->configured_endpoint; + input.port = ka->port; + input.region = ka->region; + input.use_ssl = ka->use_ssl != 0; + input.ssl_ca = ssl_params != nullptr + ? ssl_params->ssl_ca + : (mysql_thread___ssl_p2s_ca != nullptr ? mysql_thread___ssl_p2s_ca : ""); + input.ssl_capath = ssl_params != nullptr + ? ssl_params->ssl_capath + : (mysql_thread___ssl_p2s_capath != nullptr ? mysql_thread___ssl_p2s_capath : ""); + input.support_compiled = iam_source && iam_source->support_compiled(); + const AwsIamConnectionConfigResult config = + validate_mysql_aws_iam_connection(input); + if (config.status != AwsIamConnectionConfigStatus::OK || + !iam_source) { + proxy_error( + "AWS IAM kill helper failure user='%s' hostgroup=%u endpoint='%s'" + " region='%s' category='%s' code='' request_id=''\n", + ka->database_user.c_str(), ka->hid, + ka->configured_endpoint.c_str(), ka->region.c_str(), + !iam_source + ? "token_source_unavailable" : config.failure_code.c_str()); + goto __exit_kill_query_thread; + } + + iam_result = iam_source->request_blocking( + config.key, ka->token_deadline); + if (iam_result.status != AwsIamStatus::OK || iam_result.token.empty()) { + proxy_error( + "AWS IAM kill helper failure user='%s' hostgroup=%u endpoint='%s'" + " region='%s' category='%s' code='%s' request_id='%s'\n", + ka->database_user.c_str(), ka->hid, + ka->configured_endpoint.c_str(), ka->region.c_str(), + iam_result.failure.category.empty() + ? "token_request_failed" : iam_result.failure.category.c_str(), + iam_result.failure.aws_error_code.c_str(), + iam_result.failure.request_id.c_str()); + goto __exit_kill_query_thread; + } + const auto remaining = ka->token_deadline - + std::chrono::steady_clock::now(); + const auto connect_timeout_seconds = + std::chrono::duration_cast(remaining).count(); + if (connect_timeout_seconds <= 0) { + proxy_error( + "AWS IAM kill helper failure user='%s' hostgroup=%u endpoint='%s'" + " region='%s' category='helper_deadline_exceeded' code='' request_id=''\n", + ka->database_user.c_str(), ka->hid, + ka->configured_endpoint.c_str(), ka->region.c_str()); + goto __exit_kill_query_thread; + } + + my_bool enabled = 1; + my_bool reconnect = 0; + const unsigned int connect_timeout = + static_cast(connect_timeout_seconds); + mysql_options(mysql, MYSQL_OPT_CONNECT_TIMEOUT, &connect_timeout); + mysql_options(mysql, MYSQL_OPT_SSL_ENFORCE, &enabled); + mysql_options(mysql, MYSQL_OPT_SSL_VERIFY_SERVER_CERT, &enabled); + mysql_options(mysql, MYSQL_ENABLE_CLEARTEXT_PLUGIN, &enabled); + mysql_options(mysql, MYSQL_OPT_RECONNECT, &reconnect); + mysql_options(mysql, MARIADB_OPT_TLS_SERVER_NAME, + config.key.endpoint.c_str()); + connect_user = ka->database_user.c_str(); + connect_password = iam_result.token.c_str(); + } + MYSQL *ret; //! Depending on the type of operation (kill_type), constructs a KILL command string (buf) to terminate the specified query or connection. @@ -301,7 +443,8 @@ void* kill_query_thread(void *arg) { default: break; } - ret=mysql_real_connect(mysql, ka->get_host_address(), ka->username, ka->password, NULL, ka->port, NULL, 0); + ret=mysql_real_connect(mysql, ka->get_host_address(), connect_user, + connect_password, NULL, ka->port, NULL, 0); } else { switch (ka->kill_type) { case KILL_QUERY: @@ -313,11 +456,22 @@ void* kill_query_thread(void *arg) { default: break; } - ret=mysql_real_connect(mysql,"localhost",ka->username,ka->password,NULL,0,ka->hostname,0); + ret=mysql_real_connect(mysql,"localhost",connect_user,connect_password,NULL,0,ka->hostname,0); + } + if (iam_mode) { + cleanse_iam_connector_password(mysql); + iam_result.token.clear(); + iam_source->record_backend_connection(ret != nullptr); } if (!ret) { int myerr = mysql_errno(mysql); - if (ssl_params != NULL && myerr == 2026) { + if (iam_mode) { + proxy_error( + "AWS IAM kill helper failure user='%s' hostgroup=%u endpoint='%s'" + " region='%s' category='backend_connect' code='' request_id=''\n", + ka->database_user.c_str(), ka->hid, + ka->configured_endpoint.c_str(), ka->region.c_str()); + } else if (ssl_params != NULL && myerr == 2026) { proxy_error("Failed to connect to server %s:%d to run KILL %s %lu. SSL Params: %s , %s , %s , %s , %s , %s , %s , %s\n", ka->hostname, ka->port, ( ka->kill_type==KILL_QUERY ? "QUERY" : "CONNECTION" ) , ka->id, ssl_params->ssl_ca.c_str() , ssl_params->ssl_cert.c_str() , ssl_params->ssl_key.c_str() , ssl_params->ssl_capath.c_str() , @@ -349,6 +503,10 @@ void* kill_query_thread(void *arg) { mysql_query(mysql,buf); __exit_kill_query_thread: //! clean-up + if (iam_mode) { + cleanse_iam_connector_password(mysql); + iam_result.token.clear(); + } if (mysql) mysql_close(mysql); delete ka; @@ -621,6 +779,9 @@ bool Query_Info::is_select_NOT_for_update() { void MySQL_Session::set_status(enum session_status e) { + if (e != WAITING_AWS_IAM_TOKEN && aws_iam_waiter_id != 0) { + cancel_aws_iam_wait(); + } if (e==session_status___NONE) { if (mybe) { if (mybe->server_myds) { @@ -634,6 +795,180 @@ void MySQL_Session::set_status(enum session_status e) { status=e; } +void MySQL_Session::accept_aws_iam_completion( + uint64_t opaque_id, AwsIamTokenResult&& result) { + // The worker registry has already consumed opaque_id. Re-check the + // session-side gate because a timeout/state transition may have won just + // before this drain. + if (status != WAITING_AWS_IAM_TOKEN || aws_iam_waiter_id != opaque_id || + aws_iam_connection == nullptr) { + return; + } + aws_iam_completion = std::move(result); + aws_iam_completion_ready = true; + pause_until = 0; + to_process = 1; +} + +void MySQL_Session::cancel_aws_iam_wait() { + if (!aws_iam_completion_ready && aws_iam_token_source_lease && + aws_iam_request_handle.value != 0) { + aws_iam_token_source_lease->cancel(aws_iam_request_handle); + } + if (aws_iam_waiter_id != 0 && thread != nullptr) { + thread->cancel_aws_iam_waiter(aws_iam_waiter_id); + } + if (aws_iam_waiting_session_counted && aws_iam_token_source_lease) { + aws_iam_token_source_lease->record_waiting_session(false); + } + aws_iam_waiting_session_counted = false; + + if (aws_iam_connection != nullptr && mybe != nullptr && + mybe->server_myds != nullptr && + mybe->server_myds->myconn == aws_iam_connection) { + mybe->server_myds->destroy_MySQL_Connection_From_Pool(false); + } + + aws_iam_completion.token.clear(); + aws_iam_completion = AwsIamTokenResult {}; + aws_iam_token_key = AwsIamTokenKey {}; + aws_iam_request_handle = {}; + aws_iam_connection = nullptr; + aws_iam_waiter_id = 0; + aws_iam_deadline_us = 0; + aws_iam_completion_ready = false; + aws_iam_connect_token_key = AwsIamTokenKey {}; + aws_iam_connect_token_generation = 0; + aws_iam_fresh_token_retry_attempted = false; + pause_until = 0; + aws_iam_token_source_lease = AwsIamTokenSourceLease {}; +} + +void MySQL_Session::fail_aws_iam_backend( + const char *failure_code, const AwsIamRedactedFailure *provider_failure) { + const std::string database_user = aws_iam_token_key.database_user; + const std::string endpoint = aws_iam_token_key.endpoint; + const std::string region = aws_iam_token_key.region; + const char *category = provider_failure != nullptr && + !provider_failure->category.empty() + ? provider_failure->category.c_str() + : (failure_code != nullptr ? failure_code : "unknown"); + const char *aws_code = provider_failure != nullptr + ? provider_failure->aws_error_code.c_str() : ""; + const char *request_id = provider_failure != nullptr + ? provider_failure->request_id.c_str() : ""; + + proxy_error( + "AWS IAM backend token failure user='%s' hostgroup=%d endpoint='%s' region='%s'" + " category='%s' code='%s' request_id='%s'\n", + database_user.c_str(), current_hostgroup, endpoint.c_str(), region.c_str(), + category, aws_code, request_id); + + MySQL_Data_Stream *backend = mybe != nullptr ? mybe->server_myds : nullptr; + cancel_aws_iam_wait(); + while (!previous_status.empty()) previous_status.pop(); + + static const char generic_error[] = "Unable to connect to backend"; + if (client_myds != nullptr) { + client_myds->setDSS_STATE_QUERY_SENT_NET(); + client_myds->myprot.generate_pkt_ERR( + true, NULL, NULL, client_myds->pkt_sid + 1, 9002, + (char *)"HY000", generic_error, true); + client_myds->pkt_sid++; + } + RequestEnd(backend, 9002, generic_error); + if (backend != nullptr) backend->max_connect_time = 0; + // RequestEnd deliberately leaves fast-forward sessions in their current + // state. Make every IAM failure terminal explicitly, matching the existing + // CONNECTING_SERVER failure disposition and preventing waiter re-entry. + set_status(WAITING_CLIENT_DATA); +} + +void MySQL_Session::fail_invalid_backend_auth_policy( + MySQL_Data_Stream *backend, const char *database_user, + const char *failure_code) +{ + proxy_error( + "Invalid backend authentication policy user='%s' hostgroup=%d category='%s'\n", + database_user != nullptr ? database_user : "", current_hostgroup, + failure_code != nullptr ? failure_code : "invalid_policy"); + + static const char generic_error[] = "Unable to connect to backend"; + if (client_myds != nullptr) { + client_myds->setDSS_STATE_QUERY_SENT_NET(); + client_myds->myprot.generate_pkt_ERR( + true, NULL, NULL, client_myds->pkt_sid + 1, 9002, + (char *)"HY000", generic_error, true); + client_myds->pkt_sid++; + } + RequestEnd(backend, 9002, generic_error); + while (!previous_status.empty()) previous_status.pop(); + if (backend != nullptr && backend->myconn != nullptr) { + backend->destroy_MySQL_Connection_From_Pool(false); + } + if (backend != nullptr) backend->max_connect_time = 0; + set_status(WAITING_CLIENT_DATA); +} + +int MySQL_Session::handler_again___status_WAITING_AWS_IAM_TOKEN() { + if (status != WAITING_AWS_IAM_TOKEN || aws_iam_waiter_id == 0 || + aws_iam_connection == nullptr || mybe == nullptr || + mybe->server_myds == nullptr || + mybe->server_myds->myconn != aws_iam_connection) { + fail_aws_iam_backend("invalid_wait_state"); + return 0; + } + + const unsigned long long backend_deadline = mybe->server_myds->max_connect_time; + if ((backend_deadline != 0 && thread->curtime >= backend_deadline) || + (aws_iam_deadline_us != 0 && thread->curtime >= aws_iam_deadline_us)) { + fail_aws_iam_backend( + backend_deadline != 0 && backend_deadline <= aws_iam_deadline_us + ? "backend_deadline" : "token_timeout"); + return 0; + } + + if (!aws_iam_completion_ready) return 0; + if (aws_iam_completion.status != AwsIamStatus::OK || + aws_iam_completion.token.empty()) { + AwsIamRedactedFailure failure = aws_iam_completion.failure; + fail_aws_iam_backend("token_request_failed", &failure); + return 0; + } + if (previous_status.empty() || previous_status.top() != CONNECTING_SERVER) { + fail_aws_iam_backend("invalid_resume_state"); + return 0; + } + + MySQL_Connection *connection = aws_iam_connection; + AwsIamTokenKey key = std::move(aws_iam_token_key); + AwsIamTokenResult completion = std::move(aws_iam_completion); + aws_iam_connect_token_key = key; + aws_iam_connect_token_generation = completion.generation; + previous_status.pop(); + if (aws_iam_waiting_session_counted && aws_iam_token_source_lease) { + aws_iam_token_source_lease->record_waiting_session(false); + } + aws_iam_waiting_session_counted = false; + aws_iam_request_handle = {}; + aws_iam_connection = nullptr; + aws_iam_waiter_id = 0; + aws_iam_deadline_us = 0; + aws_iam_completion_ready = false; + aws_iam_completion = AwsIamTokenResult {}; + aws_iam_token_key = AwsIamTokenKey {}; + pause_until = 0; + set_status(CONNECTING_SERVER); + aws_iam_token_source_lease = AwsIamTokenSourceLease {}; + + connection->attach_aws_iam_token(key, std::move(completion)); + connection->handler(0); + mybe->server_myds->fd = connection->fd; + mybe->server_myds->DSS = STATE_MARIADB_CONNECTING; + connection->reusable = true; + return 0; +} + /** * @brief Constructs a new MySQL session object. */ @@ -725,6 +1060,7 @@ MySQL_Session::MySQL_Session() { * @brief Resets the MySQL session to its initial state. */ void MySQL_Session::reset() { + cancel_aws_iam_wait(); pending_user_variable_set.reset(); current_query_user_variable_safe = false; current_query_user_variable_unsafe_fallback = false; @@ -1969,11 +2305,18 @@ int MySQL_Session::handler_again___status_AUTHENTICATING_BACKEND_FOR_CLIENT() { // the backend conn before returning it, mirroring the client-side guard // above (NULL-safe: len==0 falls back to mysql_thread___default_schema). if (mybe && mybe->server_myds && mybe->server_myds->myconn) { - MySQL_Connection_userinfo *bui = mybe->server_myds->myconn->userinfo; + MySQL_Connection *backend_conn = mybe->server_myds->myconn; + MySQL_Connection_userinfo *bui = backend_conn->userinfo; if (bui && bui->schemaname == NULL) { bui->set_schemaname( default_schema, default_schema ? strlen(default_schema) : 0); } + const char *backend_username = bui != nullptr ? bui->username : nullptr; + const MySQLBackendAuthPolicy policy = + resolve_mysql_backend_auth_policy(*GloMyAuth, backend_username); + backend_conn->set_rowless_passthrough_authorized( + session_authorizes_rowless_passthrough( + this, backend_username, policy)); mybe->server_myds->return_MySQL_Connection_To_Pool(); } @@ -2106,7 +2449,8 @@ int MySQL_Session::handler_again___status_AUTHENTICATING_BACKEND_FOR_CLIENT() { // fresh connection has fd == -1, so connect_start runs // mysql_real_connect_start with the borrowed credential below. MySQL_Connection *mc = MyHGM->get_MyConn_from_pool( - mybe->hostgroup_id, this, true /*ff*/, NULL, 0, -1); + mybe->hostgroup_id, this, true /*ff*/, NULL, 0, -1, + MySQLBackendAuthType::PASSWORD); if (mc == NULL) { // Pool throttle fired or no backend. Pass-through does not retry // (a credential verdict requires a reachable backend; retrying just @@ -2168,6 +2512,19 @@ int MySQL_Session::handler_again___status_RESETTING_CONNECTION() { if (myds->mypolls==NULL) { thread->mypolls.add(POLLIN|POLLOUT, myds->fd, myds, thread->curtime); } + const char *backend_username = myconn->userinfo != nullptr + ? myconn->userinfo->username : nullptr; + const MySQLBackendAuthPolicy policy = + resolve_mysql_backend_auth_policy(*GloMyAuth, backend_username); + if (!myconn->can_reset_for_backend_auth_policy(policy)) { + myds->destroy_MySQL_Connection_From_Pool(false); + myds->fd = 0; + delete mybe->server_myds; + mybe->server_myds = NULL; + while (!previous_status.empty()) previous_status.pop(); + set_status(session_status___NONE); + return -1; + } myds->DSS=STATE_MARIADB_QUERY; // we recreate local_stmts : see issue #752 delete myconn->local_stmts; @@ -2270,7 +2627,32 @@ void MySQL_Session::handler_again___new_thread_to_kill_connection() { } } - KillArgs *ka = new KillArgs(ui->username, auth_password, myds->myconn->parent->address, myds->myconn->parent->port, myds->myconn->parent->myhgc->hid, myds->myconn->mysql->thread_id, KILL_QUERY, myds->myconn->parent->use_ssl, thread, myds->myconn->connected_host_details.ip); + MySQL_Connection *connection = myds->myconn; + KillArgs *ka = nullptr; + if (connection->backend_auth_type() == MySQLBackendAuthType::AWS_IAM) { + const char *database_user = connection->userinfo != nullptr && + connection->userinfo->username != nullptr + ? connection->userinfo->username : ""; + const char *region = connection->parent->myhgc != nullptr && + connection->parent->myhgc->attributes.aws_iam_region != nullptr + ? connection->parent->myhgc->attributes.aws_iam_region : ""; + ka = new KillArgs( + const_cast(database_user), nullptr, + connection->parent->address, connection->parent->port, + connection->parent->myhgc->hid, connection->mysql->thread_id, + KILL_QUERY, connection->parent->use_ssl, thread, + connection->connected_host_details.ip, + MySQLBackendAuthType::AWS_IAM, connection->parent->address, + region, database_user, + std::chrono::steady_clock::now() + std::chrono::seconds(5)); + } else { + ka = new KillArgs( + ui->username, auth_password, connection->parent->address, + connection->parent->port, connection->parent->myhgc->hid, + connection->mysql->thread_id, KILL_QUERY, + connection->parent->use_ssl, thread, + connection->connected_host_details.ip); + } pthread_attr_t attr; pthread_attr_init(&attr); pthread_attr_setdetachstate(&attr, PTHREAD_CREATE_DETACHED); @@ -2791,6 +3173,14 @@ bool MySQL_Session::handler_again___verify_backend_autocommit() { bool MySQL_Session::handler_again___verify_backend_user_schema() { MySQL_Data_Stream *myds=mybe->server_myds; + const MySQLBackendAuthPolicy requested_policy = + resolved_backend_auth_policy_for_session(this); + if (requested_policy.type == MySQLBackendAuthType::INVALID) { + fail_invalid_backend_auth_policy( + myds, requested_policy.database_user.c_str(), + requested_policy.failure_code.c_str()); + return true; + } proxy_debug(PROXY_DEBUG_MYSQL_CONNECTION, 5, "Session %p , client: %s , backend: %s\n", this, client_myds->myconn->userinfo->username, mybe->server_myds->myconn->userinfo->username); proxy_debug(PROXY_DEBUG_MYSQL_CONNECTION, 5, "Session %p , client: %s , backend: %s\n", this, client_myds->myconn->userinfo->schemaname, mybe->server_myds->myconn->userinfo->schemaname); if (client_myds->myconn->userinfo->hash!=mybe->server_myds->myconn->userinfo->hash) { @@ -2807,7 +3197,9 @@ bool MySQL_Session::handler_again___verify_backend_user_schema() { } } // if we reach here, the username is the same - if (myds->myconn->requires_CHANGE_USER(client_myds->myconn)) { + const MySQLBackendAuthType requested_type = requested_policy.type; + if (myds->myconn->requires_CHANGE_USER( + client_myds->myconn, requested_type)) { // if we reach here, even if the username is the same, // the backend connection has some session variable set // that the client never asked for @@ -3511,6 +3903,35 @@ bool MySQL_Session::handler_again___status_CONNECTING_SERVER(int *_rc) { } if (mybe->server_myds->max_connect_time) { if (thread->curtime >= mybe->server_myds->max_connect_time) { + MySQL_Connection *timed_out_connection = mybe->server_myds->myconn; + if (timed_out_connection != nullptr && + timed_out_connection->backend_auth_type() == MySQLBackendAuthType::AWS_IAM) { + record_aws_iam_backend_connection(false); + const char *database_user = timed_out_connection->userinfo != nullptr && + timed_out_connection->userinfo->username != nullptr + ? timed_out_connection->userinfo->username : ""; + const char *region = timed_out_connection->parent->myhgc != nullptr && + timed_out_connection->parent->myhgc->attributes.aws_iam_region != nullptr + ? timed_out_connection->parent->myhgc->attributes.aws_iam_region : ""; + proxy_error( + "AWS IAM backend connection failure user='%s' hostgroup=%u endpoint='%s'" + " region='%s' category='backend_connect' code='timeout' request_id=''\n", + database_user, timed_out_connection->parent->myhgc->hid, + timed_out_connection->parent->address, region); + timed_out_connection->connect_cont(MYSQL_WAIT_TIMEOUT); + static const char generic_error[] = "Unable to connect to backend"; + client_myds->setDSS_STATE_QUERY_SENT_NET(); + client_myds->myprot.generate_pkt_ERR( + true, NULL, NULL, 1, 9002, (char *)"HY000", generic_error, true); + RequestEnd(mybe->server_myds, 9002, generic_error); + while (!previous_status.empty()) previous_status.pop(); + mybe->server_myds->destroy_MySQL_Connection_From_Pool(false); + aws_iam_connect_token_key = AwsIamTokenKey {}; + aws_iam_connect_token_generation = 0; + aws_iam_fresh_token_retry_attempted = false; + mybe->server_myds->max_connect_time = 0; + NEXT_IMMEDIATE_NEW(WAITING_CLIENT_DATA); + } if (mirror) { PROXY_TRACE(); } @@ -3557,6 +3978,9 @@ bool MySQL_Session::handler_again___status_CONNECTING_SERVER(int *_rc) { } if (mybe->server_myds->myconn==NULL) { handler___client_DSS_QUERY_SENT___server_DSS_NOT_INITIALIZED__get_connection(); + if (status != CONNECTING_SERVER) { + return true; + } } if (mybe->server_myds->myconn==NULL) { if (mirror) { @@ -3583,6 +4007,12 @@ bool MySQL_Session::handler_again___status_CONNECTING_SERVER(int *_rc) { } enum session_status st=status; if (mybe->server_myds->myconn->async_state_machine==ASYNC_IDLE) { + if (mybe->server_myds->myconn->backend_auth_type() == + MySQLBackendAuthType::AWS_IAM) { + aws_iam_connect_token_key = AwsIamTokenKey {}; + aws_iam_connect_token_generation = 0; + aws_iam_fresh_token_retry_attempted = false; + } if (handle_session_track_capabilities() == false) { pause_until = thread->curtime + mysql_thread___connect_retries_delay*1000; return false; @@ -3610,6 +4040,12 @@ bool MySQL_Session::handler_again___status_CONNECTING_SERVER(int *_rc) { } switch (rc) { case 0: + if (myconn->backend_auth_type() == MySQLBackendAuthType::AWS_IAM) { + record_aws_iam_backend_connection(true); + aws_iam_connect_token_key = AwsIamTokenKey {}; + aws_iam_connect_token_generation = 0; + aws_iam_fresh_token_retry_attempted = false; + } myds->myds_type=MYDS_BACKEND; myds->DSS=STATE_MARIADB_GENERIC; status=WAITING_CLIENT_DATA; @@ -3714,7 +4150,66 @@ bool MySQL_Session::handler_again___status_CONNECTING_SERVER(int *_rc) { current_query_user_variable_safe = false; current_query_user_variable_unsafe_fallback = false; current_query_user_variable_context_change = false; - MyHGM->p_update_mysql_error_counter(p_mysql_error_type::mysql, myconn->parent->myhgc->hid, myconn->parent->address, myconn->parent->port, mysql_errno(myconn->mysql)); + const unsigned int connect_errno = mysql_errno(myconn->mysql); + MyHGM->p_update_mysql_error_counter(p_mysql_error_type::mysql, myconn->parent->myhgc->hid, myconn->parent->address, myconn->parent->port, connect_errno); + if (myconn->backend_auth_type() == MySQLBackendAuthType::AWS_IAM) { + record_aws_iam_backend_connection(false); + const AwsIamTokenKey failed_key = aws_iam_connect_token_key; + const uint64_t failed_generation = + aws_iam_connect_token_generation; + const unsigned int failed_hostgroup = myconn->parent->myhgc->hid; + AwsIamTokenSourceLease invalidation_source; + if (connect_errno == ER_ACCESS_DENIED_ERROR && + !aws_iam_fresh_token_retry_attempted && + failed_generation != 0 && + !failed_key.endpoint.empty() && failed_key.port != 0 && + !failed_key.region.empty() && + !failed_key.database_user.empty()) { + invalidation_source = acquire_global_aws_iam_token_source(); + } + if (invalidation_source) { + proxy_error( + "AWS IAM backend connection failure user='%s' hostgroup=%u endpoint='%s'" + " region='%s' category='backend_auth_rejected' code='' request_id=''\n", + failed_key.database_user.c_str(), failed_hostgroup, + failed_key.endpoint.c_str(), failed_key.region.c_str()); + invalidation_source->invalidate( + failed_key, failed_generation); + myds->connect_retries_on_failure = 0; + myds->destroy_MySQL_Connection_From_Pool(false); + aws_iam_connect_token_key = AwsIamTokenKey {}; + aws_iam_connect_token_generation = 0; + aws_iam_fresh_token_retry_attempted = true; + NEXT_IMMEDIATE_NEW(CONNECTING_SERVER); + } + if (connect_errno == ER_ACCESS_DENIED_ERROR && + aws_iam_fresh_token_retry_attempted) { + proxy_error( + "AWS IAM backend connection failure user='%s' hostgroup=%u endpoint='%s'" + " region='%s' category='backend_auth_rejected' code='' request_id=''" + " hint='verify_system_clock_for_sigv4_clock_skew'\n", + failed_key.database_user.c_str(), failed_hostgroup, + failed_key.endpoint.c_str(), failed_key.region.c_str()); + } else { + proxy_error( + "AWS IAM backend connection failure user='%s' hostgroup=%u endpoint='%s'" + " region='%s' category='backend_connect' code='' request_id=''\n", + failed_key.database_user.c_str(), failed_hostgroup, + failed_key.endpoint.c_str(), failed_key.region.c_str()); + } + static const char generic_error[] = "Unable to connect to backend"; + client_myds->setDSS_STATE_QUERY_SENT_NET(); + client_myds->myprot.generate_pkt_ERR( + true, NULL, NULL, 1, 9002, (char *)"HY000", generic_error, true); + RequestEnd(myds, 9002, generic_error); + while (!previous_status.empty()) previous_status.pop(); + myds->destroy_MySQL_Connection_From_Pool(false); + aws_iam_connect_token_key = AwsIamTokenKey {}; + aws_iam_connect_token_generation = 0; + aws_iam_fresh_token_retry_attempted = false; + myds->max_connect_time = 0; + NEXT_IMMEDIATE_NEW(WAITING_CLIENT_DATA); + } /* * Pass-through divert (spec §6.4). * @@ -3892,6 +4387,22 @@ bool MySQL_Session::handler_again___status_CHANGING_USER_SERVER(int *_rc) { if (myds->mypolls==NULL) { thread->mypolls.add(POLLIN|POLLOUT, mybe->server_myds->fd, mybe->server_myds, thread->curtime); } + const MySQLBackendAuthPolicy requested_policy = + resolved_backend_auth_policy_for_session(this); + const MySQLBackendAuthType requested_type = requested_policy.type; + if (requested_type == MySQLBackendAuthType::INVALID) { + fail_invalid_backend_auth_policy( + myds, requested_policy.database_user.c_str(), + requested_policy.failure_code.c_str()); + NEXT_IMMEDIATE_NEW(WAITING_CLIENT_DATA); + } + if (myconn->backend_auth_type() == MySQLBackendAuthType::AWS_IAM || + requested_type == MySQLBackendAuthType::AWS_IAM) { + myds->destroy_MySQL_Connection_From_Pool(false); + myds->fd = 0; + myds->DSS = STATE_NOT_INITIALIZED; + NEXT_IMMEDIATE_NEW(CONNECTING_SERVER); + } // we recreate local_stmts : see issue #752 delete myconn->local_stmts; myconn->local_stmts=new MySQL_STMTs_local_v14(false); // false by default, it is a backend @@ -6096,6 +6607,14 @@ int MySQL_Session::handler() { } } break; + case WAITING_AWS_IAM_TOKEN: + { + handler_again___status_WAITING_AWS_IAM_TOKEN(); + if (status != WAITING_AWS_IAM_TOKEN) { + goto handler_again; + } + } + break; case PINGING_SERVER: { int rc=handler_again___status_PINGING_SERVER(); @@ -8796,6 +9315,31 @@ void MySQL_Session::handler___status_WAITING_CLIENT_DATA___STATE_SLEEP___MYSQL_C void MySQL_Session::handler___client_DSS_QUERY_SENT___server_DSS_NOT_INITIALIZED__get_connection() { // Get a MySQL Connection + const char *backend_username = + client_myds != nullptr && client_myds->myconn != nullptr && + client_myds->myconn->userinfo != nullptr + ? client_myds->myconn->userinfo->username : nullptr; + MySQLBackendAuthPolicy backend_auth_policy = + resolve_mysql_backend_auth_policy(*GloMyAuth, backend_username); + const bool force_fresh_iam_connection = + backend_auth_policy.type == MySQLBackendAuthType::AWS_IAM && + aws_iam_fresh_token_retry_attempted; + // A pass-through credential has already been authorized by a successful + // backend probe (or its cache). Unknown-user pass-through intentionally + // has no USERNAME_BACKEND row, so retain its established password mode. + // Every other invalid policy, including malformed IAM attributes, remains + // fail-closed and an IAM row can never fall back to password mode. + const bool rowless_passthrough = session_authorizes_rowless_passthrough( + this, backend_username, backend_auth_policy); + if (rowless_passthrough) { + backend_auth_policy.type = MySQLBackendAuthType::PASSWORD; + backend_auth_policy.failure_code.clear(); + } + if (backend_auth_policy.type == MySQLBackendAuthType::INVALID) { + aws_iam_token_key.database_user = backend_auth_policy.database_user; + fail_aws_iam_backend(backend_auth_policy.failure_code.c_str()); + return; + } MySQL_Connection *mc=NULL; MySQL_Backend * _gtid_from_backend = NULL; @@ -8815,7 +9359,9 @@ void MySQL_Session::handler___client_DSS_QUERY_SENT___server_DSS_NOT_INITIALIZED } } } - if (session_fast_forward == SESSION_FORWARD_TYPE_NONE && qpo->create_new_conn == false) { + if (!force_fresh_iam_connection && + session_fast_forward == SESSION_FORWARD_TYPE_NONE && + qpo->create_new_conn == false) { if (qpo->min_gtid) { gtid_uuid = qpo->min_gtid; with_gtid = true; @@ -8850,11 +9396,11 @@ void MySQL_Session::handler___client_DSS_QUERY_SENT___server_DSS_NOT_INITIALIZED } uuid[n]='\0'; #ifndef STRESSTEST_POOL - mc=thread->get_MyConn_local(mybe->hostgroup_id, this, uuid, trxid, -1); + mc=thread->get_MyConn_local(mybe->hostgroup_id, this, uuid, trxid, -1, backend_auth_policy.type); #endif // STRESSTEST_POOL } else { #ifndef STRESSTEST_POOL - mc=thread->get_MyConn_local(mybe->hostgroup_id, this, NULL, 0, (int)qpo->max_lag_ms); + mc=thread->get_MyConn_local(mybe->hostgroup_id, this, NULL, 0, (int)qpo->max_lag_ms, backend_auth_policy.type); #endif // STRESSTEST_POOL } } @@ -8875,9 +9421,13 @@ void MySQL_Session::handler___client_DSS_QUERY_SENT___server_DSS_NOT_INITIALIZED if (mc==NULL) { if (trxid) { - mc=MyHGM->get_MyConn_from_pool(mybe->hostgroup_id, this, (session_fast_forward || qpo->create_new_conn), uuid, trxid, -1); + mc=MyHGM->get_MyConn_from_pool(mybe->hostgroup_id, this, + (session_fast_forward || qpo->create_new_conn || force_fresh_iam_connection), + uuid, trxid, -1, backend_auth_policy.type); } else { - mc=MyHGM->get_MyConn_from_pool(mybe->hostgroup_id, this, (session_fast_forward || qpo->create_new_conn), NULL, 0, (int)qpo->max_lag_ms); + mc=MyHGM->get_MyConn_from_pool(mybe->hostgroup_id, this, + (session_fast_forward || qpo->create_new_conn || force_fresh_iam_connection), + NULL, 0, (int)qpo->max_lag_ms, backend_auth_policy.type); } thread->note_pool_attempt(mc == NULL); #ifdef STRESSTEST_POOL @@ -8902,7 +9452,15 @@ void MySQL_Session::handler___client_DSS_QUERY_SENT___server_DSS_NOT_INITIALIZED #endif // STRESSTESTPOOL_MEASURE } #endif // STRESSTEST_POOL + if (mc != nullptr && mc->fd != -1 && + mc->backend_auth_type() != backend_auth_policy.type) { + // Authentication mode is part of pool compatibility. Never let a + // password connection satisfy an IAM request (or vice versa). + MyHGM->destroy_MyConn_from_pool(mc); + mc = nullptr; + } if (mc) { + mc->set_rowless_passthrough_authorized(rowless_passthrough); mybe->server_myds->attach_connection(mc); thread->status_variables.stvar[st_var_ConnPool_get_conn_success]++; } else { @@ -8949,6 +9507,69 @@ void MySQL_Session::handler___client_DSS_QUERY_SENT___server_DSS_NOT_INITIALIZED proxy_debug(PROXY_DEBUG_MYSQL_CONNECTION, 5, "Sess=%p -- MySQL Connection has no FD\n", this); MySQL_Connection *myconn=mybe->server_myds->myconn; myconn->userinfo->set(client_myds->myconn->userinfo); + myconn->set_backend_auth_type(backend_auth_policy.type); + + if (backend_auth_policy.type == MySQLBackendAuthType::AWS_IAM) { + MySrvC *server = myconn->parent; + aws_iam_connection = myconn; + std::unique_ptr ssl_params( + MyHGM->get_Server_SSL_Params( + server->address, server->port, myconn->userinfo->username)); + AwsIamConnectionConfigInput input; + input.database_user = backend_auth_policy.database_user; + input.configured_endpoint = server->address != nullptr ? server->address : ""; + input.port = server->port; + input.region = server->myhgc != nullptr && + server->myhgc->attributes.aws_iam_region != nullptr + ? server->myhgc->attributes.aws_iam_region : ""; + input.use_ssl = server->use_ssl != 0; + input.ssl_ca = ssl_params != nullptr + ? ssl_params->ssl_ca + : (mysql_thread___ssl_p2s_ca != nullptr ? mysql_thread___ssl_p2s_ca : ""); + input.ssl_capath = ssl_params != nullptr + ? ssl_params->ssl_capath + : (mysql_thread___ssl_p2s_capath != nullptr ? mysql_thread___ssl_p2s_capath : ""); + AwsIamTokenSourceLease lease = acquire_global_aws_iam_token_source(); + input.support_compiled = lease && lease->support_compiled(); + + AwsIamConnectionConfigResult config = + validate_mysql_aws_iam_connection(input); + aws_iam_token_key = config.key; + if (config.status != AwsIamConnectionConfigStatus::OK || !lease) { + if (aws_iam_token_key.database_user.empty()) { + aws_iam_token_key = { + input.configured_endpoint, input.port, input.region, + input.database_user }; + } + fail_aws_iam_backend( + !lease + ? "token_source_unavailable" : config.failure_code.c_str()); + return; + } + + aws_iam_deadline_us = thread->curtime + 5000000ULL; + aws_iam_waiter_id = thread->register_aws_iam_waiter(this); + if (aws_iam_waiter_id == 0) { + fail_aws_iam_backend("worker_inbox_unavailable"); + return; + } + + aws_iam_token_source_lease = std::move(lease); + previous_status.push(CONNECTING_SERVER); + set_status(WAITING_AWS_IAM_TOKEN); + aws_iam_token_source_lease->record_waiting_session(true); + aws_iam_waiting_session_counted = true; + unsigned long long wake_deadline = aws_iam_deadline_us; + if (mybe->server_myds->max_connect_time != 0 && + mybe->server_myds->max_connect_time < wake_deadline) { + wake_deadline = mybe->server_myds->max_connect_time; + } + pause_until = wake_deadline; + aws_iam_request_handle = aws_iam_token_source_lease->request( + aws_iam_token_key, aws_iam_waiter_id, + thread->aws_iam_completion_sink()); + return; + } myconn->handler(0); mybe->server_myds->fd=myconn->fd; @@ -9400,6 +10021,10 @@ void MySQL_Session::Memory_Stats() { void MySQL_Session::create_new_session_and_reset_connection(MySQL_Data_Stream *_myds) { MySQL_Data_Stream *new_myds = NULL; MySQL_Connection * mc = _myds->myconn; + if (mc->backend_auth_type() == MySQLBackendAuthType::AWS_IAM) { + _myds->destroy_MySQL_Connection_From_Pool(false); + return; + } // we remove the connection from the original data stream _myds->detach_connection(); _myds->unplug_backend(); diff --git a/lib/MySQL_Thread.cpp b/lib/MySQL_Thread.cpp index d5ed704b57..fcc8aca8d8 100644 --- a/lib/MySQL_Thread.cpp +++ b/lib/MySQL_Thread.cpp @@ -511,6 +511,9 @@ static char * mysql_thread_variables_names[]= { (char *)"passthrough_auth_empty_password", (char *)"passthrough_auth_unknown_users", (char *)"passthrough_auth_require_tls", +#ifdef PROXYSQL40 + (char *)"aws_locality_awareness", +#endif (char *)"passthrough_default_hg", (char *)"passthrough_default_schema", (char *)"passthrough_auth_cache_ttl_s", @@ -1541,6 +1544,9 @@ MySQL_Threads_Handler::MySQL_Threads_Handler() { variables.passthrough_auth_empty_password = true; variables.passthrough_auth_unknown_users = false; variables.passthrough_auth_require_tls = true; +#ifdef PROXYSQL40 + variables.aws_locality_awareness = false; +#endif variables.passthrough_default_hg = 0; variables.passthrough_default_schema = strdup((char *)""); variables.passthrough_auth_cache_ttl_s = 0; @@ -2892,6 +2898,9 @@ char ** MySQL_Threads_Handler::get_variables_list() { VariablesPointers_bool["passthrough_auth_empty_password"] = make_tuple(&variables.passthrough_auth_empty_password, false); VariablesPointers_bool["passthrough_auth_unknown_users"] = make_tuple(&variables.passthrough_auth_unknown_users, false); VariablesPointers_bool["passthrough_auth_require_tls"] = make_tuple(&variables.passthrough_auth_require_tls, false); +#ifdef PROXYSQL40 + VariablesPointers_bool["aws_locality_awareness"] = make_tuple(&variables.aws_locality_awareness, false); +#endif #ifdef PROXYSQL31 VariablesPointers_bool["caching_sha2_password_auto_generate_rsa_keys"] = make_tuple(&variables.caching_sha2_password_auto_generate_rsa_keys, false); @@ -3521,6 +3530,22 @@ MySQL_Threads_Handler::~MySQL_Threads_Handler() { } MySQL_Thread::~MySQL_Thread() { + // First sever every session-to-request association while the sessions and + // token source are still alive. Only then close the independently-held + // inbox duplicate so late provider publications are harmless drops. + while (!aws_iam_waiters.empty()) { + auto waiter = aws_iam_waiters.begin(); + MySQL_Session *session = waiter->second; + if (session != nullptr) { + session->cancel_aws_iam_wait(); + } else { + aws_iam_waiters.erase(waiter); + } + } + if (aws_iam_inbox) { + aws_iam_inbox->close(); + aws_iam_inbox.reset(); + } if (mysql_sessions) { while(mysql_sessions->len) { @@ -3696,10 +3721,11 @@ bool MySQL_Thread::init() { GloMyQPro->init_thread(); refresh_variables(); i=pipe(pipefd); + assert(i==0); ioctl_FIONBIO(pipefd[0],1); ioctl_FIONBIO(pipefd[1],1); mypolls.add(POLLIN, pipefd[0], NULL, 0); - assert(i==0); + aws_iam_inbox = std::make_shared(pipefd[1]); thr_SetParser = new MySQL_Set_Stmt_Parser(""); match_regexes=(Session_Regex **)malloc(sizeof(Session_Regex *)*4); @@ -3716,6 +3742,34 @@ bool MySQL_Thread::init() { return true; } +uint64_t MySQL_Thread::register_aws_iam_waiter(MySQL_Session *session) { + if (session == nullptr || !aws_iam_inbox || !aws_iam_inbox->available()) return 0; + for (;;) { + uint64_t opaque_id = next_aws_iam_waiter_id++; + if (opaque_id == 0) continue; + if (aws_iam_waiters.emplace(opaque_id, session).second) return opaque_id; + } +} + +void MySQL_Thread::cancel_aws_iam_waiter(uint64_t opaque_id) { + if (opaque_id != 0) aws_iam_waiters.erase(opaque_id); +} + +void MySQL_Thread::drain_aws_iam_completions() { + if (!aws_iam_inbox) return; + auto completions = aws_iam_inbox->drain(); + for (auto& completion : completions) { + auto waiter = aws_iam_waiters.find(completion.opaque_id); + if (waiter == aws_iam_waiters.end()) continue; + MySQL_Session *session = waiter->second; + aws_iam_waiters.erase(waiter); + if (session != nullptr) { + session->accept_aws_iam_completion( + completion.opaque_id, std::move(completion.result)); + } + } +} + struct pollfd * MySQL_Thread::get_pollfd(unsigned int i) { return &mypolls.fds[i]; } @@ -4155,6 +4209,9 @@ void MySQL_Thread::run() { } else { #endif // IDLE_THREADS ProcessAllMyDS_AfterPoll(); + // IAM providers only enqueue opaque completions. Resolve them to + // live sessions here, on the owning worker, before session dispatch. + drain_aws_iam_completions(); // iterate through all sessions and process the session logic process_all_sessions(); return_local_connections(); @@ -5098,6 +5155,9 @@ void MySQL_Thread::refresh_variables() { REFRESH_VARIABLE_BOOL(passthrough_auth_empty_password); REFRESH_VARIABLE_BOOL(passthrough_auth_unknown_users); REFRESH_VARIABLE_BOOL(passthrough_auth_require_tls); +#ifdef PROXYSQL40 + REFRESH_VARIABLE_BOOL(aws_locality_awareness); +#endif REFRESH_VARIABLE_INT(passthrough_default_hg); REFRESH_VARIABLE_INT(passthrough_auth_cache_ttl_s); REFRESH_VARIABLE_INT(passthrough_auth_max_inflight_probes); @@ -6212,6 +6272,9 @@ SQLite3_result * MySQL_Threads_Handler::SQL3_Processlist(processlist_config_t ar case CONNECTING_SERVER: pta[11]=strdup("Connect"); break; + case WAITING_AWS_IAM_TOKEN: + pta[11]=strdup("Waiting AWS IAM token"); + break; case PROCESSING_QUERY: if (sess->pause_until > sess->thread->curtime) { pta[11]=strdup("Delay"); @@ -6755,7 +6818,11 @@ void MySQL_Thread::Get_Memory_Stats() { * @param max_lag_ms The maximum lag time allowed for the connection in milliseconds. * @return A pointer to the retrieved MySQL connection if found; otherwise, NULL. */ -MySQL_Connection * MySQL_Thread::get_MyConn_local(unsigned int _hid, MySQL_Session *sess, char *gtid_uuid, uint64_t gtid_trxid, int max_lag_ms) { +MySQL_Connection * MySQL_Thread::get_MyConn_local( + unsigned int _hid, MySQL_Session *sess, char *gtid_uuid, + uint64_t gtid_trxid, int max_lag_ms, + MySQLBackendAuthType requested_type) +{ // some sanity check if (sess == NULL) return NULL; if (sess->client_myds == NULL) return NULL; @@ -6771,11 +6838,47 @@ MySQL_Connection * MySQL_Thread::get_MyConn_local(unsigned int _hid, MySQL_Sessi (mysql_thread___session_track_variables == session_track_variables::ENFORCED); std::vector parents; // this is a vector of srvers that needs to be excluded in case gtid_uuid is used MySQL_Connection *c=NULL; - for (i=0; ilen; i++) { + MySQL_Connection *client_conn = sess->client_myds->myconn; + bool use_aws_locality = false; +#ifdef PROXYSQL40 + std::shared_ptr aws_locality_snapshot; + if (mysql_thread___aws_locality_awareness && MyHGM != nullptr && + MyHGM->aws_locality_manager() != nullptr) { + aws_locality_snapshot = MyHGM->aws_locality_manager()->snapshot(); + use_aws_locality = aws_locality_snapshot != nullptr && + aws_locality_snapshot->enabled && + aws_locality_snapshot->has_hostgroup(_hid); + } +#endif + if (!use_aws_locality) { + for (i=0; ilen;) { c = (MySQL_Connection *) cached_connections->index(i); + const char *candidate_username = + c->userinfo != nullptr ? c->userinfo->username : nullptr; + const char *requested_username = client_conn->userinfo->username; + const bool same_username = candidate_username != nullptr && + requested_username != nullptr && + strcmp(candidate_username, requested_username) == 0; + if (c->parent->myhgc->hid == _hid && + same_username && + (c->backend_auth_type() != requested_type || + (requested_type == MySQLBackendAuthType::AWS_IAM && + c->requires_CHANGE_USER(client_conn, requested_type)))) { + cached_connections->remove_index_fast(i); + c->send_quit = false; + MyHGM->destroy_MyConn_from_pool(c); + continue; + } + if (c->backend_auth_type() != requested_type || + (requested_type == MySQLBackendAuthType::AWS_IAM && + c->requires_CHANGE_USER(client_conn, requested_type))) { + ++i; + continue; + } // Skip unhealthy or non-reusable connections if (!c->healthy || !c->reusable) { + ++i; continue; } @@ -6786,6 +6889,7 @@ MySQL_Connection * MySQL_Thread::get_MyConn_local(unsigned int _hid, MySQL_Sessi if (check_session_track_backoff) { session_track_backoff_until = c->parent->session_track_backoff_until.load(std::memory_order_relaxed); if (session_track_backoff_until > curtime) { + ++i; continue; } } @@ -6795,8 +6899,7 @@ MySQL_Connection * MySQL_Thread::get_MyConn_local(unsigned int _hid, MySQL_Sessi (gtid_uuid == NULL) || // gtid_uuid is not used (gtid_uuid && find(parents.begin(), parents.end(), c->parent) == parents.end()) // the server is currently not excluded ) { - MySQL_Connection *client_conn = sess->client_myds->myconn; - if (c->requires_CHANGE_USER(client_conn)==false) { // CHANGE_USER is not required + if (c->requires_CHANGE_USER(client_conn, requested_type)==false) { // CHANGE_USER is not required char *schema = client_conn->userinfo->schemaname; if (strcmp(c->userinfo->schemaname,schema)==0) { // same schema unsigned int not_match = 0; // number of not matching session variables @@ -6822,6 +6925,7 @@ MySQL_Connection * MySQL_Thread::get_MyConn_local(unsigned int _hid, MySQL_Sessi if (max_lag_ms >= 0) { if ((unsigned int)max_lag_ms < (c->parent->aws_aurora_current_lag_us / 1000)) { status_variables.stvar[st_var_aws_aurora_replicas_skipped_during_query]++; + ++i; continue; } } @@ -6834,8 +6938,156 @@ MySQL_Connection * MySQL_Thread::get_MyConn_local(unsigned int _hid, MySQL_Sessi } } } + ++i; } return NULL; + } + +#ifdef PROXYSQL40 + // Remove mode-incompatible connections before the allocation-free scoring + // passes below, so their indices remain stable throughout the lottery. + for (i = 0; i < cached_connections->len;) { + c = static_cast(cached_connections->index(i)); + const char* candidate_username = + c->userinfo != nullptr ? c->userinfo->username : nullptr; + const char* requested_username = client_conn->userinfo->username; + const bool same_username = candidate_username != nullptr && + requested_username != nullptr && + strcmp(candidate_username, requested_username) == 0; + if (c->parent->myhgc->hid == _hid && same_username && + (c->backend_auth_type() != requested_type || + (requested_type == MySQLBackendAuthType::AWS_IAM && + c->requires_CHANGE_USER(client_conn, requested_type)))) { + cached_connections->remove_index_fast(i); + c->send_quit = false; + MyHGM->destroy_MyConn_from_pool(c); + continue; + } + ++i; + } + + auto connection_is_eligible = [&](MySQL_Connection* candidate, + bool record_lag_skip) -> bool { + if (candidate->backend_auth_type() != requested_type || + (requested_type == MySQLBackendAuthType::AWS_IAM && + candidate->requires_CHANGE_USER(client_conn, requested_type)) || + !candidate->healthy || !candidate->reusable) { + return false; + } + if (check_session_track_backoff && + candidate->parent->session_track_backoff_until.load( + std::memory_order_relaxed) > curtime) { + return false; + } + if (candidate->parent->myhgc->hid != _hid || + !client_conn->match_tracked_options(candidate) || + candidate->requires_CHANGE_USER(client_conn, requested_type)) { + return false; + } + char* schema = client_conn->userinfo->schemaname; + if (strcmp(candidate->userinfo->schemaname, schema) != 0) { + return false; + } + unsigned int not_match = 0; + candidate->number_of_matching_session_variables(client_conn, not_match); + if (not_match != 0) { + return false; + } + if (gtid_uuid == nullptr && max_lag_ms >= 0 && + static_cast(max_lag_ms) < + (candidate->parent->aws_aurora_current_lag_us / 1000)) { + if (record_lag_skip) { + status_variables.stvar[ + st_var_aws_aurora_replicas_skipped_during_query]++; + } + return false; + } + return true; + }; + + aws_locality_candidates.clear(); + // push_MyConn_local() grows this reusable storage before inserting into the + // cache. A direct cache mutation would violate that invariant; fail neutral + // instead of allocating in the selection path. + if (aws_locality_candidates.capacity() < cached_connections->len) { + return NULL; + } + for (i = 0; i < cached_connections->len; ++i) { + auto* candidate = static_cast(cached_connections->index(i)); + if (connection_is_eligible(candidate, true)) { + aws_locality_candidates.push_back({candidate->parent, i}); + } + } + std::sort(aws_locality_candidates.begin(), aws_locality_candidates.end(), + [](const AwsLocalityCachedCandidate& lhs, + const AwsLocalityCachedCandidate& rhs) { + if (lhs.parent != rhs.parent) { + return std::less()(lhs.parent, rhs.parent); + } + return lhs.cached_index < rhs.cached_index; + }); + + uint64_t total_weight = 0; + unsigned int num_candidates = 0; + for (i = 0; i < aws_locality_candidates.size();) { + MySrvC* parent = aws_locality_candidates[i].parent; + unsigned int next = i + 1; + while (next < aws_locality_candidates.size() && + aws_locality_candidates[next].parent == parent) { + ++next; + } + if (gtid_uuid != nullptr && !MyHGM->gtid_exists(parent, gtid_uuid, gtid_trxid)) { + i = next; + continue; + } + total_weight = aws_locality_saturating_add(total_weight, + aws_locality_snapshot->effective_weight( + _hid, parent->address, parent->port, parent->weight)); + ++num_candidates; + i = next; + } + const uint64_t random_value = + (static_cast(rand_fast()) << 32) | + static_cast(rand_fast()); + if (num_candidates == 0) { + return NULL; + } + const bool uniform_fallback = total_weight == 0; + const uint64_t target = uniform_fallback + ? random_value % num_candidates : random_value % total_weight; + uint64_t cumulative = 0; + unsigned int ordinal = 0; + for (i = 0; i < aws_locality_candidates.size();) { + const auto& candidate = aws_locality_candidates[i]; + MySrvC* parent = candidate.parent; + unsigned int next = i + 1; + while (next < aws_locality_candidates.size() && + aws_locality_candidates[next].parent == parent) { + ++next; + } + if (gtid_uuid != nullptr && !MyHGM->gtid_exists(parent, gtid_uuid, gtid_trxid)) { + i = next; + continue; + } + if (uniform_fallback) { + if (ordinal++ == target) { + return static_cast( + cached_connections->remove_index_fast(candidate.cached_index)); + } + i = next; + continue; + } + cumulative = aws_locality_saturating_add(cumulative, + aws_locality_snapshot->effective_weight( + _hid, parent->address, parent->port, parent->weight)); + if (target < cumulative) { + return static_cast( + cached_connections->remove_index_fast(candidate.cached_index)); + } + i = next; + } +#endif + return NULL; } @@ -6870,8 +7122,17 @@ void MySQL_Thread::push_MyConn_local(MySQL_Connection *c) { if (mysrvc->get_status() == MYSQL_SERVER_STATUS_ONLINE) { if (c->async_state_machine==ASYNC_IDLE) { unsigned int n = (GloMTH && GloMTH->num_threads > 0) ? GloMTH->num_threads : 1; - if ((push_local_counter++ % n) == 0) { - cached_connections->add(c); + if ((push_local_counter++ % n) == 0) { +#ifdef PROXYSQL40 + const size_t required_capacity = cached_connections->len + 1; + if (aws_locality_candidates.capacity() < required_capacity) { + const size_t grown_capacity = std::max( + 32, std::max(required_capacity, + aws_locality_candidates.capacity() * 2)); + aws_locality_candidates.reserve(grown_capacity); + } +#endif + cached_connections->add(c); return; } } diff --git a/lib/MySrvConnList.cpp b/lib/MySrvConnList.cpp index 779fd804c1..2fcb8d8be7 100644 --- a/lib/MySrvConnList.cpp +++ b/lib/MySrvConnList.cpp @@ -115,19 +115,25 @@ ConnectionPoolDecision evaluate_pool_state( return decision; } -void MySrvConnList::get_random_MyConn_inner_search(unsigned int start, unsigned int end, unsigned int& conn_found_idx, unsigned int& connection_quality_level, unsigned int& number_of_matching_session_variables, const MySQL_Connection * client_conn) { +void MySrvConnList::get_random_MyConn_inner_search(unsigned int start, unsigned int end, unsigned int& conn_found_idx, unsigned int& connection_quality_level, unsigned int& number_of_matching_session_variables, const MySQL_Connection * client_conn, MySQLBackendAuthType requested_type) { char *schema = client_conn->userinfo->schemaname; MySQL_Connection * conn=NULL; unsigned int k; for (k = start; k < end; k++) { conn = (MySQL_Connection *)conns->index(k); + // A candidate from another auth mode is not usable, but may belong to + // another backend user whose policy did not change. Leave it idle for + // that identity instead of turning this checkout into cross-user churn. + if (conn->backend_auth_type() != requested_type) continue; + if (requested_type == MySQLBackendAuthType::AWS_IAM && + conn->requires_CHANGE_USER(client_conn, requested_type)) continue; if (conn->match_tracked_options(client_conn)) { if (connection_quality_level == 0) { // this is our best candidate so far connection_quality_level = 1; conn_found_idx = k; } - if (conn->requires_CHANGE_USER(client_conn)==false) { + if (conn->requires_CHANGE_USER(client_conn, requested_type)==false) { if (connection_quality_level == 1) { // this is our best candidate so far connection_quality_level = 2; @@ -180,7 +186,9 @@ void MySrvConnList::get_random_MyConn_inner_search(unsigned int start, unsigned -MySQL_Connection * MySrvConnList::get_random_MyConn(MySQL_Session *sess, bool ff) { +MySQL_Connection * MySrvConnList::get_random_MyConn( + MySQL_Session *sess, bool ff, MySQLBackendAuthType requested_type) +{ MySQL_Connection * conn=NULL; unsigned int i; unsigned int conn_found_idx = 0; @@ -199,6 +207,33 @@ MySQL_Connection * MySrvConnList::get_random_MyConn(MySQL_Session *sess, bool ff connection_warming = mysrvc->myhgc->attributes.connection_warming; free_connections_pct = mysrvc->myhgc->attributes.free_connections_pct; } + if (l && ff == false && sess && sess->client_myds && + sess->client_myds->myconn && sess->client_myds->myconn->userinfo) { + const MySQL_Connection *client_conn = sess->client_myds->myconn; + for (unsigned int candidate_idx = 0; candidate_idx < l;) { + MySQL_Connection *candidate = + (MySQL_Connection *)conns->index(candidate_idx); + const char *candidate_username = + candidate->userinfo != nullptr ? candidate->userinfo->username : nullptr; + const char *requested_username = client_conn->userinfo->username; + const bool same_username = candidate_username != nullptr && + requested_username != nullptr && + strcmp(candidate_username, requested_username) == 0; + const bool wrong_mode = + same_username && candidate->backend_auth_type() != requested_type; + const bool iam_requires_reset = + same_username && requested_type == MySQLBackendAuthType::AWS_IAM && + candidate->requires_CHANGE_USER(client_conn, requested_type); + if (wrong_mode || iam_requires_reset) { + conns->remove_index_fast(candidate_idx); + candidate->send_quit = false; + delete candidate; + --l; + continue; + } + ++candidate_idx; + } + } unsigned int conns_free = mysrvc->ConnectionsFree->conns_length(); unsigned int conns_used = mysrvc->ConnectionsUsed->conns_length(); bool needs_warming = false; @@ -213,9 +248,9 @@ MySQL_Connection * MySrvConnList::get_random_MyConn(MySQL_Session *sess, bool ff i=rand_fast()%l; if (sess && sess->client_myds && sess->client_myds->myconn && sess->client_myds->myconn->userinfo) { MySQL_Connection * client_conn = sess->client_myds->myconn; - get_random_MyConn_inner_search(i, l, conn_found_idx, connection_quality_level, number_of_matching_session_variables, client_conn); + get_random_MyConn_inner_search(i, l, conn_found_idx, connection_quality_level, number_of_matching_session_variables, client_conn, requested_type); if (connection_quality_level !=3 ) { // we didn't find the perfect connection - get_random_MyConn_inner_search(0, i, conn_found_idx, connection_quality_level, number_of_matching_session_variables, client_conn); + get_random_MyConn_inner_search(0, i, conn_found_idx, connection_quality_level, number_of_matching_session_variables, client_conn, requested_type); } // Evaluate pool state to determine create-vs-reuse and eviction (warming already handled above) ConnectionPoolDecision decision = evaluate_pool_state( diff --git a/lib/ProxySQL_Admin.cpp b/lib/ProxySQL_Admin.cpp index b7a9ce5562..8d519075a7 100644 --- a/lib/ProxySQL_Admin.cpp +++ b/lib/ProxySQL_Admin.cpp @@ -1613,10 +1613,11 @@ bool ProxySQL_Admin::GenericRefreshStatistics(const char *query_no_space, unsign } #ifdef PROXYSQL40 // Plugin-registered runtime views: if the query references any chassis- - // registered runtime view (e.g. runtime_mysqlx_users), refresh it on - // the admin path BEFORE the SELECT runs against admindb. We always - // invoke the dispatcher when the session is on the admin port; the - // chassis itself decides whether to fire any plugin's refresh + // registered runtime view (e.g. runtime_mysqlx_users or an on-demand stats + // table), refresh it BEFORE the SELECT runs. Admin sessions can project all + // three DB kinds; stats sessions receive only the stats handle, so they cannot + // trigger an admin/config projection. The chassis decides whether to fire any + // plugin's refresh // callback by per-view substring match against query_no_space, so a // query that touches no registered view is a cheap no-op (one shared // lock + N substring scans, N == registered-view count). @@ -1628,9 +1629,8 @@ bool ProxySQL_Admin::GenericRefreshStatistics(const char *query_no_space, unsign // that touches only a plugin view (e.g. SELECT * FROM runtime_mysqlx_ // users with no other runtime_* mention) still gets its projection // fired. - if (admin) { - proxysql_refresh_configured_plugin_runtime_views(query_no_space, admindb, configdb, statsdb); - } + proxysql_refresh_configured_plugin_runtime_views(query_no_space, + admin ? admindb : nullptr, admin ? configdb : nullptr, statsdb); #endif /* PROXYSQL40 */ // if (stats_mysql_processlist || stats_mysql_connection_pool || stats_mysql_query_digest || stats_mysql_query_digest_reset) { if (refresh==true) { diff --git a/lib/ProxySQL_Admin_Stats.cpp b/lib/ProxySQL_Admin_Stats.cpp index d2738c6608..293137fe90 100644 --- a/lib/ProxySQL_Admin_Stats.cpp +++ b/lib/ProxySQL_Admin_Stats.cpp @@ -8,6 +8,7 @@ #include "cpp.h" #include "MySQL_Authentication.hpp" +#include "Aws_Iam_Provider.h" #include "MySQL_Passthrough_Auth_Cache.h" #include "PgSQL_Authentication.h" #include "MySQL_LDAP_Authentication.hpp" @@ -62,7 +63,13 @@ extern ProxySQL_Statistics *GloProxyStats; extern MySQL_Logger *GloMyLogger; extern PgSQL_Logger *GloPgSQL_Logger; +static AwsIamStatsSnapshot global_aws_iam_stats_snapshot() { + AwsIamTokenSourceLease source = acquire_global_aws_iam_token_source(); + return source ? source->snapshot() : AwsIamStatsSnapshot {}; +} + void ProxySQL_Admin::p_update_metrics() { + update_aws_iam_prometheus_metrics(global_aws_iam_stats_snapshot()); // Update proxysql_uptime auto t1 = monotonic_time(); auto new_uptime = (t1 - GloVars.global.start_time)/1000/1000; @@ -659,6 +666,11 @@ void ProxySQL_Admin::stats___mysql_global() { } } + for (const AwsIamNamedStat& row : + aws_iam_stats_mysql_global_rows(global_aws_iam_stats_snapshot())) { + sqlite3_global_stats_row_step(statsdb, row_stmt, row.name, row.value); + } + statsdb->execute("COMMIT"); } diff --git a/lib/ProxySQL_PluginManager.cpp b/lib/ProxySQL_PluginManager.cpp index 1d22a4264a..b581ffb217 100644 --- a/lib/ProxySQL_PluginManager.cpp +++ b/lib/ProxySQL_PluginManager.cpp @@ -5,6 +5,10 @@ #ifdef PROXYSQL40 #include "ProxySQL_PluginManager.h" +#include "Aws_Iam_Provider.h" +#include "Aws_Locality_Manager.h" +#include "MySQL_HostGroups_Manager.h" +#include "MySQL_Thread.h" #include #include @@ -20,6 +24,7 @@ #include "prometheus/registry.h" extern ProxySQL_GlobalVariables GloVars; +extern MySQL_Threads_Handler *GloMTH; SQLite3DB* proxysql_plugin_get_admindb(); SQLite3DB* proxysql_plugin_get_configdb(); @@ -176,6 +181,51 @@ bool register_runtime_view_service(const ProxySQL_PluginRuntimeView& view) { } return true; } + +bool install_aws_iam_token_source_service( + AwsIamTokenSource *source, void (*destroy)(AwsIamTokenSource *), void *module_handle) { + if (g_registry_target == nullptr) { + proxy_warning("AWS IAM token source installation attempted outside plugin init phase\n"); + return false; + } + return install_global_aws_iam_token_source(source, destroy, module_handle); +} + +bool uninstall_aws_iam_token_source_service(AwsIamTokenSource *expected_source) { + if (g_registry_target == nullptr) { + proxy_warning("AWS IAM token source removal attempted outside plugin init phase\n"); + return false; + } + return uninstall_global_aws_iam_token_source(expected_source); +} + +void get_aws_iam_limits_service(size_t *max_total_waiters, size_t *max_waiters_per_key) { + const size_t maximum = GloMTH != nullptr && GloMTH->variables.max_connections > 0 + ? static_cast(GloMTH->variables.max_connections) + : 1; + if (max_total_waiters != nullptr) *max_total_waiters = maximum; + if (max_waiters_per_key != nullptr) *max_waiters_per_key = maximum; +} + +bool install_aws_metadata_provider_service( + AwsMetadataProvider *provider, + void (*destroy)(AwsMetadataProvider *), + void *module_handle) { + if (g_registry_target == nullptr) { + proxy_warning("AWS metadata provider installation attempted outside plugin init phase\n"); + return false; + } + return install_global_aws_metadata_provider(provider, destroy, module_handle); +} + +void refresh_mysql_aws_locality_stats_service(SQLite3DB* statsdb) { + if (statsdb == nullptr) return; + if (MyHGM != nullptr) { + MyHGM->refresh_aws_locality_stats(statsdb); + return; + } + MySQL_HostGroups_Manager::project_aws_locality_stats(statsdb, {}); +} #endif /* PROXYSQL40 */ SQLite3DB* get_admindb_service() { @@ -302,6 +352,11 @@ ProxySQL_PluginManager::ProxySQL_PluginManager() { services_.get_prometheus_registry = &get_prometheus_registry_service; services_.register_command_alias = ®ister_command_alias_service; services_.register_runtime_view = ®ister_runtime_view_service; + services_.install_aws_iam_token_source = &install_aws_iam_token_source_service; + services_.get_aws_iam_limits = &get_aws_iam_limits_service; + services_.install_aws_metadata_provider = &install_aws_metadata_provider_service; + services_.refresh_mysql_aws_locality_stats = &refresh_mysql_aws_locality_stats_service; + services_.uninstall_aws_iam_token_source = &uninstall_aws_iam_token_source_service; // Phase-B (register_schemas) services: same layout as init(), but DB // handle getters and the query-hook registrar are stubbed -- see the @@ -327,6 +382,8 @@ ProxySQL_PluginManager::ProxySQL_PluginManager() { // refresh callback won't fire until Admin handles a SELECT, by which // point admin module bootstrap has long since completed. services_phase_b_.register_runtime_view = ®ister_runtime_view_service; + services_phase_b_.refresh_mysql_aws_locality_stats = + &refresh_mysql_aws_locality_stats_service; #endif /* PROXYSQL40 */ } diff --git a/lib/mysql_connection.cpp b/lib/mysql_connection.cpp index 55d4ee4ba0..95543d6a8f 100644 --- a/lib/mysql_connection.cpp +++ b/lib/mysql_connection.cpp @@ -9,6 +9,7 @@ using json = nlohmann::json; #include #include #include +#include #include "MySQL_PreparedStatement.h" #include "MySQL_Data_Stream.h" @@ -516,6 +517,7 @@ MySQL_Connection::MySQL_Connection() { MySQL_Connection::~MySQL_Connection() { proxy_debug(PROXY_DEBUG_MYSQL_CONNPOOL, 4, "Destroying MySQL_Connection %p\n", this); + clear_aws_iam_handshake_secret(); if (options.server_version) free(options.server_version); if (options.init_connect) free(options.init_connect); if (options.ldap_user_variable) free(options.ldap_user_variable); @@ -583,6 +585,79 @@ MySQL_Connection::~MySQL_Connection() { } }; +void MySQL_Connection::set_backend_auth_type(MySQLBackendAuthType type) { + if (type != MySQLBackendAuthType::AWS_IAM) { + clear_aws_iam_handshake_secret(); + aws_iam_identity_.reset(); + } + if (type != MySQLBackendAuthType::PASSWORD) { + rowless_passthrough_authorized_ = false; + } + backend_auth_type_ = type; +} + +MySQLBackendAuthType MySQL_Connection::backend_auth_type() const { + return backend_auth_type_; +} + +void MySQL_Connection::set_rowless_passthrough_authorized(bool authorized) { + rowless_passthrough_authorized_ = + authorized && backend_auth_type_ == MySQLBackendAuthType::PASSWORD; +} + +bool MySQL_Connection::can_reset_for_backend_auth_policy( + const MySQLBackendAuthPolicy& policy) const +{ + if (backend_auth_type_ != MySQLBackendAuthType::PASSWORD) return false; + if (policy.type == MySQLBackendAuthType::PASSWORD) return true; + return rowless_passthrough_authorized_ && + policy.type == MySQLBackendAuthType::INVALID && + policy.failure_code == "backend_user_not_found"; +} + +void MySQL_Connection::attach_aws_iam_token( + const AwsIamTokenKey& key, AwsIamTokenResult&& result) +{ + clear_aws_iam_handshake_secret(); + auto identity = std::make_unique(); + identity->key = key; + identity->token_generation = result.generation; + identity->handshake_token = std::move(result.token); + aws_iam_identity_ = std::move(identity); +} + +void MySQL_Connection::clear_aws_iam_handshake_secret() { + if (mysql != nullptr && mysql->passwd != nullptr && + aws_iam_connector_secret_active_) { + OPENSSL_cleanse(mysql->passwd, strlen(mysql->passwd)); + free(mysql->passwd); + mysql->passwd = nullptr; + } + + // Connector/C's nonblocking coroutine retains the original passwd pointer + // across yields. Destroy the suspended operation before cleansing that + // caller-owned buffer so no later mysql_real_connect_cont() can dereference + // it. mysql_close_no_command() also destroys the async context without + // attempting a blocking COM_QUIT. + if (mysql != nullptr && aws_iam_connector_secret_active_ && + aws_iam_async_connect_pending_) { + mysql_close_no_command(mysql); + mysql = nullptr; + ret_mysql = nullptr; + fd = -1; + } + + if (aws_iam_identity_) { + aws_iam_identity_->handshake_token.clear(); + } + aws_iam_connector_secret_active_ = false; + aws_iam_async_connect_pending_ = false; +} + +bool MySQL_Connection::has_aws_iam_handshake_secret() const { + return aws_iam_identity_ && !aws_iam_identity_->handshake_token.empty(); +} + bool MySQL_Connection::set_autocommit(bool _ac) { proxy_debug(PROXY_DEBUG_MYSQL_CONNPOOL, 4, "Setting autocommit %d\n", _ac); options.autocommit=_ac; @@ -675,11 +750,25 @@ bool MySQL_Connection::get_status_sql_log_bin0() { return status_flags & STATUS_MYSQL_CONNECTION_SQL_LOG_BIN0; } -bool MySQL_Connection::requires_CHANGE_USER(const MySQL_Connection *client_conn) { - char *username = client_conn->userinfo->username; - if (strcmp(userinfo->username,username)) { +bool MySQL_Connection::backend_auth_compatible( + const char *requested_username, MySQLBackendAuthType requested_type) const +{ + return requested_username != nullptr && userinfo != nullptr && + userinfo->username != nullptr && + backend_auth_type_ == requested_type && + strcmp(userinfo->username, requested_username) == 0; +} + +bool MySQL_Connection::requires_CHANGE_USER( + const MySQL_Connection *client_conn, + MySQLBackendAuthType requested_type) const +{ + const char *username = client_conn != nullptr && client_conn->userinfo != nullptr + ? client_conn->userinfo->username : nullptr; + if (!backend_auth_compatible(username, requested_type)) { // the two connections use different usernames - // The connection need to be reset with CHANGE_USER + // or authentication modes. The caller decides whether CHANGE_USER is + // permitted for that mode. return true; } for (auto i = 0; i < SQL_NAME_LAST_LOW_WM; i++) { @@ -1084,6 +1173,32 @@ void MySQL_Connection::connect_start() { } } #endif + if (backend_auth_type_ == MySQLBackendAuthType::AWS_IAM) { + const bool valid_iam_handshake = parent->port != 0 && aws_iam_identity_ && + !aws_iam_identity_->handshake_token.empty(); + if (!valid_iam_handshake) { + mysql->net.last_errno = CR_CONNECTION_ERROR; + std::snprintf(mysql->net.last_error, sizeof(mysql->net.last_error), + "AWS IAM backend authentication requires a TCP endpoint and handshake token"); + std::strncpy(mysql->net.sqlstate, "HY000", sizeof(mysql->net.sqlstate)); + mysql->net.sqlstate[sizeof(mysql->net.sqlstate) - 1] = '\0'; + ret_mysql = nullptr; + async_exit_status = 0; + fd = mysql_get_socket(mysql); + return; + } + + auth_password = const_cast(aws_iam_identity_->handshake_token.c_str()); + my_bool enabled = 1; + my_bool reconnect = 0; + mysql_options(mysql, MYSQL_OPT_SSL_ENFORCE, &enabled); + mysql_options(mysql, MYSQL_OPT_SSL_VERIFY_SERVER_CERT, &enabled); + mysql_options(mysql, MYSQL_ENABLE_CLEARTEXT_PLUGIN, &enabled); + mysql_options(mysql, MYSQL_OPT_RECONNECT, &reconnect); + mysql_options(mysql, MARIADB_OPT_TLS_SERVER_NAME, + aws_iam_identity_->key.endpoint.c_str()); + aws_iam_connector_secret_active_ = true; + } if (parent->port) { char* host_ip = connect_start_DNS_lookup(); async_exit_status=mysql_real_connect_start(&ret_mysql, mysql, host_ip, userinfo->username, auth_password, userinfo->schemaname, parent->port, NULL, client_flags); @@ -1094,6 +1209,9 @@ void MySQL_Connection::connect_start() { } async_exit_status=mysql_real_connect_start(&ret_mysql, mysql, "localhost", userinfo->username, auth_password, userinfo->schemaname, parent->port, parent->address, client_flags); } + if (aws_iam_connector_secret_active_) { + aws_iam_async_connect_pending_ = async_exit_status != 0; + } fd=mysql_get_socket(mysql); // { // // FIXME: THIS IS FOR TESTING PURPOSE ONLY @@ -1112,10 +1230,18 @@ void MySQL_Connection::connect_start() { void MySQL_Connection::connect_cont(short event) { proxy_debug(PROXY_DEBUG_MYSQL_PROTOCOL, 6,"event=%d\n", event); async_exit_status = mysql_real_connect_cont(&ret_mysql, mysql, mysql_status(event, true)); + if (aws_iam_connector_secret_active_) { + aws_iam_async_connect_pending_ = async_exit_status != 0; + } } void MySQL_Connection::change_user_start() { PROXY_TRACE(); + // IAM credentials exist only for the initial TLS handshake. Reaching this + // path with an IAM identity would turn an ephemeral token into a reusable + // password; all session/reset callers must replace the connection instead. + assert(backend_auth_type_ != MySQLBackendAuthType::AWS_IAM); + assert(!has_aws_iam_handshake_secret()); //fprintf(stderr,"change_user_start FD %d\n", fd); MySQL_Connection_userinfo *_ui = NULL; if (myds->sess->client_myds == NULL) { @@ -1332,6 +1458,8 @@ MDB_ASYNC_ST MySQL_Connection::handler(short event) { break; break; case ASYNC_CONNECT_END: + aws_iam_async_connect_pending_ = false; + clear_aws_iam_handshake_secret(); if (myds) { if (myds->sess) { if (myds->sess->thread) { @@ -1342,7 +1470,11 @@ MDB_ASYNC_ST MySQL_Connection::handler(short event) { } if (!ret_mysql) { int myerr = mysql_errno(mysql); - if (ssl_params != NULL && myerr == 2026) { + if (backend_auth_type_ == MySQLBackendAuthType::AWS_IAM) { + proxy_error("Failed to connect IAM backend on %u:%s:%d , FD (Conn:%d , MyDS:%d) , %d: authentication details redacted.\n", + parent->myhgc->hid, parent->address, parent->port, + mysql->net.fd, myds->fd, myerr); + } else if (ssl_params != NULL && myerr == 2026) { proxy_error("Failed to mysql_real_connect() on %u:%s:%d , FD (Conn:%d , MyDS:%d) , %d: %s. SSL Params: %s , %s , %s , %s , %s , %s , %s , %s\n", parent->myhgc->hid, parent->address, parent->port, mysql->net.fd , myds->fd, mysql_errno(mysql), mysql_error(mysql), ssl_params->ssl_ca.c_str() , ssl_params->ssl_cert.c_str() , ssl_params->ssl_key.c_str() , ssl_params->ssl_capath.c_str() , @@ -1432,10 +1564,14 @@ MDB_ASYNC_ST MySQL_Connection::handler(short event) { parent->connect_error(mysql_errno(mysql)); break; case ASYNC_CONNECT_TIMEOUT: + { + const int myerr = mysql != nullptr ? mysql_errno(mysql) : CR_CONNECTION_ERROR; + clear_aws_iam_handshake_secret(); //proxy_error("Connect timeout on %s:%d : %llu - %llu = %llu\n", parent->address, parent->port, myds->sess->thread->curtime , myds->wait_until, myds->sess->thread->curtime - myds->wait_until); proxy_error("Connect timeout on %s:%d : exceeded by %lluus\n", parent->address, parent->port, myds->sess->thread->curtime - myds->wait_until); - MyHGM->p_update_mysql_error_counter(p_mysql_error_type::mysql, parent->myhgc->hid, parent->address, parent->port, mysql_errno(mysql)); - parent->connect_error(mysql_errno(mysql)); + MyHGM->p_update_mysql_error_counter(p_mysql_error_type::mysql, parent->myhgc->hid, parent->address, parent->port, myerr); + parent->connect_error(myerr); + } break; case ASYNC_CHANGE_USER_START: change_user_start(); diff --git a/src/Makefile b/src/Makefile index 5f03691f96..4b6c9442aa 100644 --- a/src/Makefile +++ b/src/Makefile @@ -1,5 +1,6 @@ #!/bin/make -f +.DEFAULT_GOAL := default PROXYSQL_PATH := $(shell while [ ! -f ./src/proxysql_global.cpp ]; do cd ..; done; pwd) @@ -160,7 +161,6 @@ endif ifeq ($(CENTOSVER),6) MYLIBS += -lgcrypt endif - SQLITE_VEC_OBJ := $(DEPS_PATH)/sqlite3/sqlite3/vec.o LIBPROXYSQLAR := $(PROXYSQL_LDIR)/libproxysql.a ifeq ($(UNAME_S),Darwin) diff --git a/src/main.cpp b/src/main.cpp index d126ce662a..e80bad85fd 100644 --- a/src/main.cpp +++ b/src/main.cpp @@ -34,6 +34,7 @@ using json = nlohmann::json; #include "PgSQL_Query_Processor.h" #include "MySQL_Authentication.hpp" #include "MySQL_Passthrough_Auth_Cache.h" +#include "Aws_Iam_Provider.h" #include "PgSQL_Authentication.h" #include "MySQL_LDAP_Authentication.hpp" #include "MySQL_Query_Cache.h" @@ -42,6 +43,7 @@ using json = nlohmann::json; #include "Web_Interface.hpp" #ifdef PROXYSQL40 #include "ProxySQL_PluginManager.h" +#include "Aws_Locality_Manager.h" #endif /* PROXYSQL40 */ #include "proxysql_utils.h" #include "PgSQL_Monitor.hpp" @@ -501,6 +503,7 @@ MySQL_Query_Processor* GloMyQPro; PgSQL_Query_Processor* GloPgQPro; ProxySQL_Admin *GloAdmin; MySQL_Threads_Handler *GloMTH = NULL; +AwsIamTokenSource* GloAwsIamTokenSource = NULL; PgSQL_Threads_Handler* GloPTH = NULL; // GloMCPH removed in Step 4.C; GloGATH/GloAI removed in Step 5 — the @@ -1594,7 +1597,6 @@ void ProxySQL_Main_init_phase2___not_started(const bootstrap_info_t& boostrap_in } ProxySQL_Main_init_Auth_module(); - if (GloVars.global.nostart) { pthread_mutex_lock(&GloVars.global.start_mutex); } @@ -1807,6 +1809,15 @@ bool ProxySQL_Main_init_phase3___start_all() { void ProxySQL_Main_init_phase4___shutdown() { cpu_timer t; ProxySQL_Main_join_all_threads(); +#ifdef PROXYSQL40 + // The locality manager can retain a provider lease between refreshes. + // Drain it before shutting down the plugin-owned provider registry. + if (MyHGM != nullptr && MyHGM->aws_locality_manager() != nullptr) { + MyHGM->aws_locality_manager()->shutdown(); + } + shutdown_global_aws_metadata_provider(); +#endif + shutdown_global_aws_iam_token_source(); //write(GloAdmin->pipefd[1], &GloAdmin->pipefd[1], 1); // write a random byte if (GloVars.global.nostart) { @@ -3219,7 +3230,7 @@ int main(int argc, const char * argv[]) { { cpu_timer t; if (ProxySQL_Main_init_phase3___start_all() == false) { - goto finish; + goto __shutdown; } #ifdef DEBUG std::cerr << "Main init phase3 completed in "; diff --git a/test/deps/aws_iam_mysql_server/Makefile b/test/deps/aws_iam_mysql_server/Makefile new file mode 100644 index 0000000000..54ac3abf6f --- /dev/null +++ b/test/deps/aws_iam_mysql_server/Makefile @@ -0,0 +1,22 @@ +#!/bin/make -f + +PROXYSQL_PATH := $(shell while [ ! -f ./src/proxysql_global.cpp ]; do cd ..; done; pwd) + +include $(PROXYSQL_PATH)/include/makefiles_vars.mk +include $(PROXYSQL_PATH)/include/makefiles_paths.mk + +.DEFAULT_GOAL := all + +CPPFLAGS := -I$(SSL_IDIR) +CXXFLAGS := $(STDCPP) -O0 -ggdb -Wall -Wextra -Werror $(WASAN) +LDFLAGS := -L$(SSL_LDIR) $(WASAN) +LDLIBS := -lssl -lcrypto -lpthread + +.PHONY: all clean +all: aws_iam_mysql_server-t + +aws_iam_mysql_server-t: aws_iam_mysql_server.cpp + $(CXX) $(CPPFLAGS) $(CXXFLAGS) $< $(LDFLAGS) $(LDLIBS) -o $@ + +clean: + rm -f aws_iam_mysql_server-t diff --git a/test/deps/aws_iam_mysql_server/aws_iam_mysql_server.cpp b/test/deps/aws_iam_mysql_server/aws_iam_mysql_server.cpp new file mode 100644 index 0000000000..c0f7ec5b17 --- /dev/null +++ b/test/deps/aws_iam_mysql_server/aws_iam_mysql_server.cpp @@ -0,0 +1,449 @@ +/** + * @file aws_iam_mysql_server.cpp + * @brief One-shot TLS MySQL server for AWS IAM backend protocol tests. + * + * This is intentionally not a SQL server. It implements only the initial + * protocol-10 exchange, SSLRequest, TLS upgrade, mysql_clear_password auth + * switch, and a terminal OK or ERR 1045 packet. + */ + +#include +#include +#include + +#include +#include +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace { + +constexpr uint32_t CLIENT_LONG_PASSWORD = 0x00000001U; +constexpr uint32_t CLIENT_PROTOCOL_41 = 0x00000200U; +constexpr uint32_t CLIENT_SSL = 0x00000800U; +constexpr uint32_t CLIENT_SECURE_CONNECTION = 0x00008000U; +constexpr uint32_t CLIENT_PLUGIN_AUTH = 0x00080000U; + +struct Options { + std::string mode; + std::string certificate; + std::string private_key; + unsigned int delay_ms { 1500 }; + int stage_fd { -1 }; +}; + +struct Observation { + bool tls { false }; + bool pre_tls_payload { false }; + int minimum_tls_version { 0 }; + std::string username; + std::string sni; + std::string peer; + size_t token_length { 0 }; + std::string token_sha256; + std::string error { "none" }; +}; + +uint32_t read_le32(const unsigned char *p) { + return static_cast(p[0]) | + (static_cast(p[1]) << 8U) | + (static_cast(p[2]) << 16U) | + (static_cast(p[3]) << 24U); +} + +void append_le16(std::vector& out, uint16_t value) { + out.push_back(static_cast(value & 0xffU)); + out.push_back(static_cast((value >> 8U) & 0xffU)); +} + +void append_le32(std::vector& out, uint32_t value) { + for (unsigned int i = 0; i != 4; ++i) { + out.push_back(static_cast((value >> (8U * i)) & 0xffU)); + } +} + +bool wait_fd(int fd, short events, int timeout_ms = 5000) { + pollfd descriptor { fd, events, 0 }; + for (;;) { + const int result = poll(&descriptor, 1, timeout_ms); + if (result > 0) return (descriptor.revents & events) != 0; + if (result == 0) return false; + if (errno != EINTR) return false; + } +} + +bool read_exact_fd(int fd, unsigned char *data, size_t size) { + while (size != 0) { + if (!wait_fd(fd, POLLIN)) return false; + const ssize_t count = recv(fd, data, size, 0); + if (count <= 0) return false; + data += count; + size -= static_cast(count); + } + return true; +} + +bool write_exact_fd(int fd, const unsigned char *data, size_t size) { + while (size != 0) { + if (!wait_fd(fd, POLLOUT)) return false; + const ssize_t count = send(fd, data, size, 0); + if (count <= 0) return false; + data += count; + size -= static_cast(count); + } + return true; +} + +bool read_exact_ssl(SSL *ssl, unsigned char *data, size_t size) { + while (size != 0) { + const int count = SSL_read(ssl, data, static_cast(size)); + if (count <= 0) return false; + data += count; + size -= static_cast(count); + } + return true; +} + +bool write_exact_ssl(SSL *ssl, const unsigned char *data, size_t size) { + while (size != 0) { + const int count = SSL_write(ssl, data, static_cast(size)); + if (count <= 0) return false; + data += count; + size -= static_cast(count); + } + return true; +} + +template +bool read_packet(Reader&& reader, std::vector& payload, unsigned char& sequence) { + unsigned char header[4] {}; + if (!reader(header, sizeof(header))) return false; + const size_t length = static_cast(header[0]) | + (static_cast(header[1]) << 8U) | + (static_cast(header[2]) << 16U); + if (length > 16U * 1024U * 1024U) return false; + sequence = header[3]; + payload.assign(length, 0); + return length == 0 || reader(payload.data(), payload.size()); +} + +template +bool write_packet(Writer&& writer, const std::vector& payload, + unsigned char sequence) { + if (payload.size() > 0xffffffU) return false; + unsigned char header[4] { + static_cast(payload.size() & 0xffU), + static_cast((payload.size() >> 8U) & 0xffU), + static_cast((payload.size() >> 16U) & 0xffU), + sequence + }; + return writer(header, sizeof(header)) && + (payload.empty() || writer(payload.data(), payload.size())); +} + +std::vector greeting() { + const uint32_t capabilities = CLIENT_LONG_PASSWORD | CLIENT_PROTOCOL_41 | + CLIENT_SSL | CLIENT_SECURE_CONNECTION | CLIENT_PLUGIN_AUTH; + const char server_version[] = "8.0.36-aws-iam-test"; + const unsigned char seed1[] = "12345678"; + const unsigned char seed2[] = "abcdefghijkl"; + const char plugin[] = "mysql_native_password"; + + std::vector packet; + packet.push_back(10); + packet.insert(packet.end(), server_version, server_version + sizeof(server_version)); + append_le32(packet, 4242); + packet.insert(packet.end(), seed1, seed1 + 8); + packet.push_back(0); + append_le16(packet, static_cast(capabilities & 0xffffU)); + packet.push_back(45); + append_le16(packet, 2); + append_le16(packet, static_cast((capabilities >> 16U) & 0xffffU)); + packet.push_back(21); + packet.insert(packet.end(), 10, 0); + packet.insert(packet.end(), seed2, seed2 + 12); + packet.push_back(0); + packet.insert(packet.end(), plugin, plugin + sizeof(plugin)); + return packet; +} + +std::vector auth_switch() { + const char plugin[] = "mysql_clear_password"; + std::vector packet { 0xfe }; + packet.insert(packet.end(), plugin, plugin + sizeof(plugin)); + packet.push_back(0); + return packet; +} + +std::vector ok_packet() { + return { 0x00, 0x00, 0x00, 0x02, 0x00, 0x00, 0x00 }; +} + +std::vector access_denied_packet() { + std::vector packet { 0xff, 0x15, 0x04, '#' }; + const char state[] = "28000"; + const char message[] = "Access denied"; + packet.insert(packet.end(), state, state + 5); + packet.insert(packet.end(), message, message + sizeof(message) - 1); + return packet; +} + +std::string hex(const std::string& value) { + static const char digits[] = "0123456789abcdef"; + std::string out; + out.reserve(value.size() * 2); + for (unsigned char byte : value) { + out.push_back(digits[byte >> 4U]); + out.push_back(digits[byte & 0x0fU]); + } + return out; +} + +std::string sha256(const unsigned char *data, size_t size) { + unsigned char digest[SHA256_DIGEST_LENGTH] {}; + SHA256(data, size, digest); + return hex(std::string(reinterpret_cast(digest), sizeof(digest))); +} + +std::string parse_username(const std::vector& response) { + constexpr size_t username_offset = 4 + 4 + 1 + 23; + if (response.size() <= username_offset) return {}; + const auto begin = response.begin() + static_cast(username_offset); + const auto end = std::find(begin, response.end(), 0); + if (end == response.end()) return {}; + return std::string(begin, end); +} + +void report(const Observation& observation) { + std::printf( + "RESULT tls=%d pre_tls=%d min_tls=%d username_hex=%s token_len=%zu token_sha256=%s " + "sni_hex=%s peer=%s error=%s\n", + observation.tls ? 1 : 0, + observation.pre_tls_payload ? 1 : 0, + observation.minimum_tls_version, + hex(observation.username).c_str(), observation.token_length, + observation.token_sha256.empty() ? "-" : observation.token_sha256.c_str(), + hex(observation.sni).c_str(), + observation.peer.empty() ? "-" : observation.peer.c_str(), + observation.error.c_str()); + std::fflush(stdout); +} + +bool parse_options(int argc, char **argv, Options& options) { + for (int i = 1; i < argc; ++i) { + const std::string argument { argv[i] }; + if ((argument == "--mode" || argument == "--cert" || + argument == "--key" || argument == "--delay-ms" || + argument == "--stage-fd") && i + 1 < argc) { + const std::string value { argv[++i] }; + if (argument == "--mode") options.mode = value; + else if (argument == "--cert") options.certificate = value; + else if (argument == "--key") options.private_key = value; + else if (argument == "--delay-ms") { + options.delay_ms = static_cast(std::strtoul(value.c_str(), nullptr, 10)); + } else { + options.stage_fd = static_cast(std::strtol(value.c_str(), nullptr, 10)); + } + } else { + return false; + } + } + return !options.mode.empty() && !options.certificate.empty() && + !options.private_key.empty(); +} + +int create_listener(uint16_t& port) { + const int listener = socket(AF_INET, SOCK_STREAM, 0); + if (listener < 0) return -1; + int enabled = 1; + setsockopt(listener, SOL_SOCKET, SO_REUSEADDR, &enabled, sizeof(enabled)); + sockaddr_in address {}; + address.sin_family = AF_INET; + address.sin_addr.s_addr = htonl(INADDR_LOOPBACK); + address.sin_port = 0; + if (bind(listener, reinterpret_cast(&address), sizeof(address)) != 0 || + listen(listener, 1) != 0) { + close(listener); + return -1; + } + socklen_t length = sizeof(address); + if (getsockname(listener, reinterpret_cast(&address), &length) != 0) { + close(listener); + return -1; + } + port = ntohs(address.sin_port); + return listener; +} + +} // namespace + +int main(int argc, char **argv) { + std::signal(SIGPIPE, SIG_IGN); + Options options; + if (!parse_options(argc, argv, options)) { + std::fprintf(stderr, + "usage: %s --mode MODE --cert FILE --key FILE [--delay-ms N] [--stage-fd FD]\n", + argv[0]); + return 2; + } + + SSL_CTX *context = SSL_CTX_new(TLS_server_method()); + if (context == nullptr || + SSL_CTX_set_min_proto_version(context, TLS1_2_VERSION) != 1 || + SSL_CTX_use_certificate_chain_file(context, options.certificate.c_str()) != 1 || + SSL_CTX_use_PrivateKey_file(context, options.private_key.c_str(), SSL_FILETYPE_PEM) != 1 || + SSL_CTX_check_private_key(context) != 1) { + ERR_print_errors_fp(stderr); + if (context != nullptr) SSL_CTX_free(context); + return 2; + } + + uint16_t port = 0; + const int listener = create_listener(port); + if (listener < 0) { + SSL_CTX_free(context); + return 2; + } + std::printf("READY port=%u\n", port); + std::fflush(stdout); + + Observation observation; + observation.minimum_tls_version = SSL_CTX_get_min_proto_version(context); + if (!wait_fd(listener, POLLIN, 10000)) { + observation.error = "accept_timeout"; + report(observation); + close(listener); + SSL_CTX_free(context); + return 1; + } + sockaddr_in peer {}; + socklen_t peer_length = sizeof(peer); + const int client = accept(listener, reinterpret_cast(&peer), &peer_length); + close(listener); + if (client < 0) { + observation.error = "accept_failed"; + report(observation); + SSL_CTX_free(context); + return 1; + } + const timeval io_timeout { 6, 0 }; + (void)setsockopt(client, SOL_SOCKET, SO_RCVTIMEO, &io_timeout, sizeof(io_timeout)); + (void)setsockopt(client, SOL_SOCKET, SO_SNDTIMEO, &io_timeout, sizeof(io_timeout)); + char peer_text[INET_ADDRSTRLEN] {}; + if (inet_ntop(AF_INET, &peer.sin_addr, peer_text, sizeof(peer_text)) != nullptr) { + observation.peer = peer_text; + } + + auto fd_reader = [client](unsigned char *data, size_t size) { + return read_exact_fd(client, data, size); + }; + auto fd_writer = [client](const unsigned char *data, size_t size) { + return write_exact_fd(client, data, size); + }; + if (options.mode == "delay_handshake") { + if (options.stage_fd >= 0) { + const char stage = 'S'; + (void)write(options.stage_fd, &stage, 1); + } + std::this_thread::sleep_for(std::chrono::milliseconds(options.delay_ms)); + } + if (!write_packet(fd_writer, greeting(), 0)) { + observation.error = "greeting_write_failed"; + report(observation); + close(client); + SSL_CTX_free(context); + return 1; + } + std::vector request; + unsigned char sequence = 0; + if (!read_packet(fd_reader, request, sequence)) { + observation.error = "initial_read_failed"; + report(observation); + close(client); + SSL_CTX_free(context); + return 1; + } + const uint32_t client_capabilities = request.size() >= 4 ? read_le32(request.data()) : 0; + if (request.size() != 32 || (client_capabilities & CLIENT_SSL) == 0) { + observation.pre_tls_payload = true; + observation.error = "non_tls_handshake_refused"; + report(observation); + close(client); + SSL_CTX_free(context); + return 1; + } + if (options.mode == "close_transport") { + observation.error = "transport_closed"; + report(observation); + close(client); + SSL_CTX_free(context); + return 0; + } + SSL *ssl = SSL_new(context); + SSL_set_fd(ssl, client); + if (SSL_accept(ssl) != 1) { + observation.error = "tls_accept_failed"; + report(observation); + SSL_free(ssl); + close(client); + SSL_CTX_free(context); + return options.mode == "wrong_hostname" || options.mode == "untrusted_ca" || + options.mode == "delay_handshake" ? 0 : 1; + } + observation.tls = true; + const char *sni = SSL_get_servername(ssl, TLSEXT_NAMETYPE_host_name); + if (sni != nullptr) observation.sni = sni; + + auto ssl_reader = [ssl](unsigned char *data, size_t size) { + return read_exact_ssl(ssl, data, size); + }; + auto ssl_writer = [ssl](const unsigned char *data, size_t size) { + return write_exact_ssl(ssl, data, size); + }; + if (!read_packet(ssl_reader, request, sequence)) { + observation.error = "tls_handshake_response_failed"; + report(observation); + SSL_shutdown(ssl); + SSL_free(ssl); + close(client); + SSL_CTX_free(context); + return 1; + } + observation.username = parse_username(request); + if (!write_packet(ssl_writer, auth_switch(), static_cast(sequence + 1))) { + observation.error = "auth_switch_write_failed"; + } else if (!read_packet(ssl_reader, request, sequence)) { + observation.error = "auth_response_read_failed"; + } else if (request.empty() || request.back() != 0) { + observation.error = "auth_response_not_nul_terminated"; + } else { + observation.token_length = request.size() - 1; + observation.token_sha256 = sha256(request.data(), observation.token_length); + const std::vector terminal = options.mode == "access_denied" + ? access_denied_packet() : ok_packet(); + if (!write_packet(ssl_writer, terminal, static_cast(sequence + 1))) { + observation.error = "terminal_write_failed"; + } + } + + report(observation); + SSL_shutdown(ssl); + SSL_free(ssl); + close(client); + SSL_CTX_free(context); + return observation.error == "none" ? 0 : 1; +} diff --git a/test/infra/control/run-unit-tests-asan-coverage.bash b/test/infra/control/run-unit-tests-asan-coverage.bash index bd2697d38e..0b73ee5dbd 100755 --- a/test/infra/control/run-unit-tests-asan-coverage.bash +++ b/test/infra/control/run-unit-tests-asan-coverage.bash @@ -145,6 +145,14 @@ if [ -s coverage/lcov.info ]; then --output-file coverage/lcov.info \ --ignore-errors unused || true + # gcov records the container checkout prefix in SF: entries even though + # capture is rooted at lib/. Normalize those paths to the git-tree shape + # used by the TAP coverage producer so Codecov merges both reports. + sed -i \ + -e 's|^SF:/opt/proxysql/|SF:|' \ + -e "s|^SF:${REPO_ROOT}/|SF:|" \ + coverage/lcov.info + genhtml --quiet \ --output-directory coverage/html \ --title "ProxySQL unit tests coverage" \ diff --git a/test/tap/groups/groups.json b/test/tap/groups/groups.json index f7a4223ff0..8a0a4d7833 100644 --- a/test/tap/groups/groups.json +++ b/test/tap/groups/groups.json @@ -12,6 +12,20 @@ "ai_llm_retry_scenarios-t" : [ "ai-g1","@proxysql_min_version:4.0" ], "ai_validation-t" : [ "ai-g1","@proxysql_min_version:4.0" ], "auth_unit-t" : [ "unit-tests-g1" ], + "aws_iam_completion_queue_unit-t" : [ "unit-tests-g1","mysqlx-tsan-g1","@proxysql_min_version:4.0" ], + "aws_iam_connection_config_unit-t" : [ "unit-tests-g1","@proxysql_min_version:4.0" ], + "aws_iam_connection_secret_unit-t" : [ "unit-tests-g1","@proxysql_min_version:4.0" ], + "aws_iam_failure_unit-t" : [ "unit-tests-g1","@proxysql_min_version:4.0" ], + "aws_iam_kill_helper_unit-t" : [ "unit-tests-g1","@proxysql_min_version:4.0" ], + "aws_iam_policy_unit-t" : [ "unit-tests-g1","@proxysql_min_version:4.0" ], + "aws_iam_pool_unit-t" : [ "unit-tests-g1","@proxysql_min_version:4.0" ], + "aws_iam_provider_boundary_unit-t" : [ "unit-tests-g1","mysqlx-tsan-g1","@proxysql_min_version:4.0" ], + "aws_iam_session_state_unit-t" : [ "unit-tests-g1","mysqlx-tsan-g1","@proxysql_min_version:4.0" ], + "aws_locality_config_unit-t" : [ "unit-tests-g1","@proxysql_min_version:4.0" ], + "aws_locality_manager_unit-t" : [ "unit-tests-g1","@proxysql_min_version:4.0" ], + "aws_locality_policy_unit-t" : [ "unit-tests-g1","@proxysql_min_version:4.0" ], + "aws_locality_selection_unit-t" : [ "unit-tests-g1","@proxysql_min_version:4.0" ], + "aws_locality_stats_unit-t" : [ "unit-tests-g1","@proxysql_min_version:4.0" ], "backend_sync_unit-t" : [ "unit-tests-g1" ], "basic-t" : [ "legacy-g1","mariadb10-galera-g1","mysql-auto_increment_delay_multiplex=0-g1","mysql-multiplexing=false-g1","mysql-query_digests=0-g1","mysql-query_digests_keep_comment=1-g1","mysql84-g1","mysql84-gr-g1","mysql90-g1","mysql90-gr-g1","mysql93-g1","mysql93-gr-g1","mysql95-g1","mysql95-gr-g1" ], "c_tokenizer_unit-t" : [ "unit-tests-g1" ], @@ -69,6 +83,7 @@ "listener_conflicts_validation-t" : [ "no-infra-g1" ], "llm_bridge_accuracy-t" : [ "ai-g1","@proxysql_min_version:4.0" ], "log_utils_unit-t" : [ "unit-tests-g1" ], + "mariadb_tls_server_name_unit-t" : [ "unit-tests-g1" ], "max_connections_ff-t" : [ "legacy-g1","mysql-auto_increment_delay_multiplex=0-g1","mysql-multiplexing=false-g1","mysql-query_digests=0-g1","mysql-query_digests_keep_comment=1-g1","mysql84-g1","mysql90-g1","mysql95-g1" ], "mcp_mixed_mysql_pgsql_concurrency_stress-t" : [ "ai-g1","@proxysql_min_version:4.0" ], "mcp_mixed_stats_cap_churn-t" : [ "ai-g1","@proxysql_min_version:4.0" ], @@ -351,6 +366,8 @@ "test_aurora_query_routing-t" : [ "cluster_sim_aurora-g1" ], "test_auth_methods-t" : [ "mysql-auto_increment_delay_multiplex=0-g2","mysql-multiplexing=false-g2","mysql-query_digests=0-g2","mysql-query_digests_keep_comment=1-g2","mysql84-g7","mysql90-g2","mysql95-g2" ], "test_auto_increment_delay_multiplex-t" : [ "legacy-g7","mysql-auto_increment_delay_multiplex=0-g2","mysql-multiplexing=false-g2","mysql-query_digests=0-g2","mysql-query_digests_keep_comment=1-g2","mysql84-g7","mysql90-g2","mysql95-g2" ], + "test_aws_iam_backend_auth-t" : [ "unit-tests-g1","@proxysql_min_version:4.0" ], + "test_aws_iam_metrics-t" : [ "unit-tests-g1","@proxysql_min_version:4.0" ], "test_backend_conn_ping-t" : [ "legacy-g7","mysql-auto_increment_delay_multiplex=0-g2","mysql-multiplexing=false-g2","mysql-query_digests=0-g2","mysql-query_digests_keep_comment=1-g2","mysql84-g7","mysql90-g2","mysql95-g2" ], "test_binlog_dump_multi_backend_crash-t" : [ "legacy-binlog-g1" ], "test_binlog_fast_forward-t" : [ "legacy-binlog-g1","mysql84-binlog-g2","mysql90-binlog-g2","mysql95-binlog-g2" ], diff --git a/test/tap/test_helpers/test_globals.cpp b/test/tap/test_helpers/test_globals.cpp index b7b92cb7f8..f67a88cfac 100644 --- a/test/tap/test_helpers/test_globals.cpp +++ b/test/tap/test_helpers/test_globals.cpp @@ -62,6 +62,7 @@ using json = nlohmann::json; #include "PgSQL_Authentication.h" #include "MySQL_LDAP_Authentication.hpp" #include "MySQL_Passthrough_Auth_Cache.h" +#include "Aws_Iam_Provider.h" #include "MySQL_Query_Cache.h" #include "PgSQL_Query_Cache.h" #include "proxysql_restapi.h" @@ -93,6 +94,7 @@ MySQL_Query_Processor *GloMyQPro = nullptr; PgSQL_Query_Processor *GloPgQPro = nullptr; ProxySQL_Admin *GloAdmin = nullptr; MySQL_Threads_Handler *GloMTH = nullptr; +AwsIamTokenSource *GloAwsIamTokenSource = nullptr; PgSQL_Threads_Handler *GloPTH = nullptr; // MCP_Threads_Handler *GloMCPH stub removed in Step 4.C; GloGATH and diff --git a/test/tap/tests/Makefile b/test/tap/tests/Makefile index 849f2abaaf..52a5fe6db4 100644 --- a/test/tap/tests/Makefile +++ b/test/tap/tests/Makefile @@ -317,6 +317,22 @@ test_ignore_min_gtid-t: test_ignore_min_gtid-t.cpp $(TAP_LDIR)/libtap$(SHLIB_EXT test_admin_prometheus_metrics_dump-t: test_admin_prometheus_metrics_dump-t.cpp $(TAP_LDIR)/libtap$(SHLIB_EXT) $(CXX) $< $(IDIRS) $(LDIRS) $(OPT) $(MYLIBS) -o $@ +# Keep the IAM observability contract deterministic and independent of an +# external provider while exercising the production admin/process-global +# integrations through the platform-aware unit-test linker configuration. +.PHONY: test_aws_iam_metrics-t +test_aws_iam_metrics-t: + $(MAKE) -C unit PROXYSQL40=$(PROXYSQL40) test_aws_iam_metrics-t + ln -fs unit/test_aws_iam_metrics-t $@ + +# Build the controlled server first, then use the component-test linker so the +# production MySQL_Connection path can run without a daemon or external provider. +.PHONY: test_aws_iam_backend_auth-t +test_aws_iam_backend_auth-t: + $(MAKE) -C $(PROXYSQL_PATH)/test/deps/aws_iam_mysql_server PROXYSQL40=$(PROXYSQL40) + $(MAKE) -C unit PROXYSQL40=$(PROXYSQL40) test_aws_iam_backend_auth-t + ln -fs unit/test_aws_iam_backend_auth-t $@ + create_connection_annotation: test_connection_annotation-t.cpp $(TAP_LDIR)/libtap$(SHLIB_EXT) $(CXX) -DTEST_AURORA $< $(IDIRS) $(LDIRS) $(OPT) $(OBJ) $(MYLIBS) $(STATIC_LIBS) -o $@ diff --git a/test/tap/tests/mysql_hostgroup_attributes_config_file-t.cpp b/test/tap/tests/mysql_hostgroup_attributes_config_file-t.cpp index 358f205ce8..d0e2fd425a 100644 --- a/test/tap/tests/mysql_hostgroup_attributes_config_file-t.cpp +++ b/test/tap/tests/mysql_hostgroup_attributes_config_file-t.cpp @@ -24,7 +24,7 @@ using std::fstream; int validate_mysql_hostgroup_attributes_from_config(MYSQL* admin) { string hostgroup_attributes_values[5][12] = { - {"900000", "900000", "-1", "11", "ic1", "0", "1", "9001", "{\"isv\":100}", "{\"hs\":200}", "{\"weight\":100,\"max_connections\":500}", "attributes test hostgroup 900000"}, + {"900000", "900000", "-1", "11", "ic1", "0", "1", "9001", "{\"isv\":100}", "{\"aws_iam_region\":\"us-east-1\"}", "{\"weight\":100,\"max_connections\":500}", "attributes test hostgroup 900000"}, {"900001", "900001", "0", "12", "ic2", "1", "0", "9002", "{\"isv\":101}", "{\"hs\":201}", "{\"weight\":101,\"max_connections\":501}", "attributes test hostgroup 900001"}, {"900002", "900002", "1", "13", "ic3", "0", "1", "9003", "{\"isv\":102}", "{\"hs\":202}", "{\"weight\":102,\"max_connections\":502}", "attributes test hostgroup 900002"}, {"900003", "900003", "-1", "14", "ic4", "1", "0", "9004", "{\"isv\":103}", "{\"hs\":203}", "{\"weight\":103,\"max_connections\":503}", "attributes test hostgroup 900003"}, @@ -97,7 +97,7 @@ void make_hostgroup_attributes_config_lines(std::vector& config_lin " connection_warming=1", " throttle_connections_per_sec=9001", " ignore_session_variables=\"{\"isv\":100}\"", - " hostgroup_settings=\"{\"hs\":200}\"", + " hostgroup_settings=\"{\"aws_iam_region\":\"us-east-1\"}\"", " servers_defaults=\"{\"weight\":100,\"max_connections\":500}\"", " comment=\"attributes test hostgroup 900000\"", " },", @@ -287,4 +287,4 @@ int main(int, char**) { write_mysql_hostgroup_attributes_to_config(admin); return run_cleanup_and_exit(); -} \ No newline at end of file +} diff --git a/test/tap/tests/test_aws_iam_backend_auth-t.cpp b/test/tap/tests/test_aws_iam_backend_auth-t.cpp new file mode 100644 index 0000000000..70851009ba --- /dev/null +++ b/test/tap/tests/test_aws_iam_backend_auth-t.cpp @@ -0,0 +1,621 @@ +/** + * @file test_aws_iam_backend_auth-t.cpp + * @brief Controlled TLS/MySQL-protocol coverage for AWS IAM backend auth. + * + * The production SDK is intentionally not required here. A deterministic fake + * source supplies a recognizable 2 KiB value to the real MySQL_Connection and + * bundled Connector/C. A one-shot loopback server performs the real MySQL TLS + * and mysql_clear_password exchange and reports only length/digest metadata. + */ + +#include "tap.h" +#include "test_globals.h" +#include "test_init.h" +#include "frontend_x509_test_utils.h" + +#include "proxysql.h" +#include "cpp.h" +#include "Aws_Iam_Provider.h" +#include "MySQL_Data_Stream.h" +#include "MySQL_Logger.hpp" +#include "MySQL_Monitor.hpp" + +#include + +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +extern MySQL_HostGroups_Manager *MyHGM; +extern MySQL_Logger *GloMyLogger; +extern MySQL_Monitor *GloMyMon; + +namespace { + +using Clock = std::chrono::steady_clock; + +constexpr const char *kEndpoint = "db.cluster-test.us-east-1.rds.amazonaws.com"; +constexpr const char *kWrongEndpoint = "wrong.cluster-test.us-east-1.rds.amazonaws.com"; +constexpr const char *kUsername = "iam_protocol_user"; + +struct ChildServer { + pid_t pid { -1 }; + FILE *output { nullptr }; + int stage_fd { -1 }; + unsigned int port { 0 }; + std::string result; +}; + +struct CertificateFixture { + std::string directory; + std::string ca; + std::string untrusted_ca; + std::string valid_cert; + std::string valid_key; + std::string wrong_cert; + std::string wrong_key; + std::string untrusted_cert; + std::string untrusted_key; + bool ready { false }; + + CertificateFixture() = default; + CertificateFixture(const CertificateFixture&) = delete; + CertificateFixture& operator=(const CertificateFixture&) = delete; + CertificateFixture& operator=(CertificateFixture&&) = delete; + CertificateFixture(CertificateFixture&& other) noexcept : + directory(std::move(other.directory)), ca(std::move(other.ca)), + untrusted_ca(std::move(other.untrusted_ca)), + valid_cert(std::move(other.valid_cert)), valid_key(std::move(other.valid_key)), + wrong_cert(std::move(other.wrong_cert)), wrong_key(std::move(other.wrong_key)), + untrusted_cert(std::move(other.untrusted_cert)), + untrusted_key(std::move(other.untrusted_key)), ready(other.ready) { + other.directory.clear(); + other.ready = false; + } + + ~CertificateFixture() { + if (directory.empty()) return; + const char *files[] { + "ca.key", "ca.pem", "ca.srl", "untrusted-ca.key", "untrusted-ca.pem", + "untrusted-ca.srl", "valid.key", "valid.csr", "valid.pem", "valid.ext", + "wrong.key", "wrong.csr", "wrong.pem", "wrong.ext", + "untrusted.key", "untrusted.csr", "untrusted.pem", "untrusted.ext" + }; + for (const char *file : files) unlink((directory + "/" + file).c_str()); + rmdir(directory.c_str()); + } +}; + +class FakeSource final : public AwsIamTokenSource { +public: + AwsIamRequestHandle request(const AwsIamTokenKey&, uint64_t, + std::weak_ptr) override { return {}; } + AwsIamTokenResult request_blocking(const AwsIamTokenKey&, Clock::time_point) override { + ++blocking_requests; + AwsIamTokenResult result; + result.status = AwsIamStatus::OK; + result.generation = 1; + result.token = SecureString(token); + return result; + } + void cancel(AwsIamRequestHandle) override {} + void invalidate(const AwsIamTokenKey&, uint64_t) override {} + void record_backend_connection(bool) override {} + void record_waiting_session(bool) override {} + AwsIamStatsSnapshot snapshot() const override { return {}; } + + std::string token; + unsigned int blocking_requests { 0 }; +}; + +class ScopedPublishedTokenSource { +public: + explicit ScopedPublishedTokenSource(AwsIamTokenSource *source) { + publish_global_aws_iam_token_source(source); + } + ~ScopedPublishedTokenSource() { publish_global_aws_iam_token_source(nullptr); } + + ScopedPublishedTokenSource(const ScopedPublishedTokenSource&) = delete; + ScopedPublishedTokenSource& operator=(const ScopedPublishedTokenSource&) = delete; +}; + +bool write_text_file(const std::string& path, const std::string& contents) { + std::ofstream output(path); + output << contents; + return output.good(); +} + +bool make_leaf(const CertificateFixture& fixture, const std::string& name, + const std::string& common_name, const std::string& ca, const std::string& ca_key, + unsigned int serial) { + const std::string prefix = fixture.directory + "/" + name; + const std::string ext = prefix + ".ext"; + return write_text_file(ext, "subjectAltName=DNS:" + common_name + "\n") && + run_openssl({ "req", "-new", "-newkey", "rsa:2048", "-nodes", + "-subj", "/CN=" + common_name, "-keyout", prefix + ".key", "-out", prefix + ".csr" }) && + run_openssl({ "x509", "-req", "-days", "1", "-set_serial", + std::to_string(serial), "-in", prefix + ".csr", "-CA", ca, "-CAkey", ca_key, + "-extfile", ext, "-out", prefix + ".pem" }); +} + +CertificateFixture create_certificates() { + CertificateFixture fixture; + const std::string path_template = "./proxysql-aws-iam-protocol-XXXXXX"; + std::vector path(path_template.begin(), path_template.end()); + path.push_back('\0'); + char *directory = mkdtemp(path.data()); + if (directory == nullptr) return fixture; + fixture.directory = directory; + fixture.ca = fixture.directory + "/ca.pem"; + fixture.untrusted_ca = fixture.directory + "/untrusted-ca.pem"; + const std::string ca_key = fixture.directory + "/ca.key"; + const std::string untrusted_ca_key = fixture.directory + "/untrusted-ca.key"; + + const bool roots = + run_openssl({ "req", "-x509", "-newkey", "rsa:2048", "-nodes", + "-days", "1", "-subj", "/CN=ProxySQL AWS IAM Test CA", + "-keyout", ca_key, "-out", fixture.ca }) && + run_openssl({ "req", "-x509", "-newkey", "rsa:2048", "-nodes", + "-days", "1", "-subj", "/CN=ProxySQL AWS IAM Untrusted CA", + "-keyout", untrusted_ca_key, "-out", fixture.untrusted_ca }); + fixture.valid_cert = fixture.directory + "/valid.pem"; + fixture.valid_key = fixture.directory + "/valid.key"; + fixture.wrong_cert = fixture.directory + "/wrong.pem"; + fixture.wrong_key = fixture.directory + "/wrong.key"; + fixture.untrusted_cert = fixture.directory + "/untrusted.pem"; + fixture.untrusted_key = fixture.directory + "/untrusted.key"; + fixture.ready = roots && + make_leaf(fixture, "valid", kEndpoint, fixture.ca, ca_key, 820001) && + make_leaf(fixture, "wrong", kWrongEndpoint, fixture.ca, ca_key, 820002) && + make_leaf(fixture, "untrusted", kEndpoint, fixture.untrusted_ca, untrusted_ca_key, 820003); + return fixture; +} + +std::string sha256(const std::string& input) { + unsigned char digest[SHA256_DIGEST_LENGTH] {}; + SHA256(reinterpret_cast(input.data()), input.size(), digest); + static const char digits[] = "0123456789abcdef"; + std::string output; + output.reserve(sizeof(digest) * 2); + for (unsigned char byte : digest) { + output.push_back(digits[byte >> 4U]); + output.push_back(digits[byte & 0x0fU]); + } + return output; +} + +std::string hex(const std::string& input) { + static const char digits[] = "0123456789abcdef"; + std::string output; + output.reserve(input.size() * 2); + for (unsigned char byte : input) { + output.push_back(digits[byte >> 4U]); + output.push_back(digits[byte & 0x0fU]); + } + return output; +} + +std::string field(const std::string& record, const std::string& name) { + const std::string prefix = name + "="; + const size_t begin = record.find(prefix); + if (begin == std::string::npos) return {}; + const size_t value_begin = begin + prefix.size(); + const size_t end = record.find(' ', value_begin); + return record.substr(value_begin, end == std::string::npos ? std::string::npos : end - value_begin); +} + +bool read_control_line(int fd, std::string& line, unsigned int timeout_ms) { + const auto deadline = Clock::now() + std::chrono::milliseconds(timeout_ms); + while (Clock::now() < deadline && line.size() < 2048) { + pollfd descriptor { fd, POLLIN, 0 }; + const auto remaining = std::chrono::duration_cast( + deadline - Clock::now()).count(); + const int ready = poll(&descriptor, 1, static_cast(remaining > 0 ? remaining : 1)); + if (ready < 0 && errno == EINTR) continue; + if (ready <= 0 || (descriptor.revents & POLLIN) == 0) return false; + char byte = 0; + if (read(fd, &byte, 1) != 1) return false; + line.push_back(byte); + if (byte == '\n') return true; + } + return false; +} + +ChildServer start_server(const std::string& mode, const std::string& certificate, + const std::string& private_key, unsigned int delay_ms = 1500) { + ChildServer child; + int pipefd[2] {}; + int stage_pipe[2] {}; + if (pipe(pipefd) != 0) return child; + if (pipe(stage_pipe) != 0) { + close(pipefd[0]); + close(pipefd[1]); + return child; + } + const std::string server { AWS_IAM_MYSQL_SERVER_PATH }; + child.pid = fork(); + if (child.pid == 0) { + dup2(pipefd[1], STDOUT_FILENO); + close(pipefd[0]); + close(pipefd[1]); + close(stage_pipe[0]); + execl(server.c_str(), server.c_str(), "--mode", mode.c_str(), + "--cert", certificate.c_str(), "--key", private_key.c_str(), + "--delay-ms", std::to_string(delay_ms).c_str(), "--stage-fd", + std::to_string(stage_pipe[1]).c_str(), static_cast(nullptr)); + _exit(127); + } + close(pipefd[1]); + close(stage_pipe[1]); + if (child.pid < 0) { + close(pipefd[0]); + close(stage_pipe[0]); + return child; + } + child.stage_fd = stage_pipe[0]; + std::string ready_line; + if (read_control_line(pipefd[0], ready_line, 5000)) { + unsigned int port = 0; + if (sscanf(ready_line.c_str(), "READY port=%u", &port) == 1) child.port = port; + } + child.output = fdopen(pipefd[0], "r"); + if (child.output != nullptr) setvbuf(child.output, nullptr, _IONBF, 0); + return child; +} + +void finish_server(ChildServer& child) { + if (child.pid > 0) { + const auto deadline = Clock::now() + std::chrono::seconds(8); + int status = 0; + while (Clock::now() < deadline) { + const pid_t result = waitpid(child.pid, &status, WNOHANG); + if (result == child.pid || (result < 0 && errno == ECHILD)) { + child.pid = -1; + break; + } + if (result < 0 && errno != EINTR) break; + pollfd descriptor { child.output == nullptr ? -1 : fileno(child.output), POLLIN, 0 }; + (void)poll(&descriptor, 1, 25); + } + if (child.pid > 0) { + (void)kill(child.pid, SIGTERM); + const auto terminate_deadline = Clock::now() + std::chrono::milliseconds(250); + while (Clock::now() < terminate_deadline) { + const pid_t result = waitpid(child.pid, &status, WNOHANG); + if (result == child.pid || (result < 0 && errno == ECHILD)) { + child.pid = -1; + break; + } + if (result < 0 && errno != EINTR) break; + poll(nullptr, 0, 10); + } + } + if (child.pid > 0) { + (void)kill(child.pid, SIGKILL); + while (waitpid(child.pid, &status, 0) < 0 && errno == EINTR) {} + child.pid = -1; + } + } + if (child.output != nullptr) { + char line[2048] {}; + while (fgets(line, sizeof(line), child.output) != nullptr) child.result = line; + fclose(child.output); + child.output = nullptr; + } + if (child.stage_fd >= 0) { + close(child.stage_fd); + child.stage_fd = -1; + } + while (!child.result.empty() && + (child.result.back() == '\n' || child.result.back() == '\r')) { + child.result.pop_back(); + } +} + +MySrvC *create_server(unsigned int hostgroup, const char *endpoint, unsigned int port, + bool use_ssl) { + srv_info_t info; + info.addr = const_cast(endpoint); + info.port = port; + info.kind = "aws-iam-protocol-test"; + srv_opts_t opts; + opts.weigth = 1; + opts.max_conns = 100; + opts.use_ssl = use_ssl ? 1 : 0; + MyHGM->wrlock(); + const int result = MyHGM->create_new_server_in_hg(hostgroup, info, opts); + MyHGC *group = MyHGM->MyHGC_find(hostgroup); + MyHGM->wrunlock(); + if (result != 0 || group == nullptr) return nullptr; + MySrvC *server = nullptr; + for (unsigned int i = 0; i < group->mysrvs->cnt(); ++i) { + MySrvC *candidate = group->mysrvs->idx(i); + if (candidate != nullptr && candidate->port == port && + strcmp(candidate->address, endpoint) == 0) { + server = candidate; + break; + } + } + if (server == nullptr) return nullptr; + server->use_ssl = use_ssl ? 1 : 0; + return server; +} + +class ConnectionFixture { +public: + ConnectionFixture(MySQL_Thread& worker, MySrvC *server) { + session = new MySQL_Session(); + session->thread = &worker; + session->connections_handler = true; + frontend_stream = new MySQL_Data_Stream(); + frontend_stream->init(MYDS_FRONTEND, session, -1); + frontend = new MySQL_Connection(); + frontend_stream->attach_connection(frontend); + session->client_myds = frontend_stream; + frontend->userinfo->set(const_cast("frontend_user"), + const_cast("frontend-password-must-not-change"), + const_cast("frontend_schema"), nullptr); + stream = new MySQL_Data_Stream(); + stream->init(MYDS_BACKEND_NOT_CONNECTED, session, -1); + connection = new MySQL_Connection(); + connection->send_quit = false; + connection->parent = server; + stream->attach_connection(connection); + session->mybe = session->create_backend(server->myhgc->hid, stream); + connection->userinfo->set(const_cast(kUsername), + const_cast("ordinary-password-must-not-leak"), + const_cast(""), nullptr); + } + ~ConnectionFixture() { + if (connection != nullptr) { + stream->myconn = nullptr; + connection->myds = nullptr; + delete connection; + } + delete session; + } + + MySQL_Session *session { nullptr }; + MySQL_Data_Stream *frontend_stream { nullptr }; + MySQL_Connection *frontend { nullptr }; + MySQL_Data_Stream *stream { nullptr }; + MySQL_Connection *connection { nullptr }; +}; + +MDB_ASYNC_ST drive(MySQL_Thread& worker, MySQL_Connection *connection, + unsigned int timeout_ms = 5000) { + worker.curtime = monotonic_time(); + MDB_ASYNC_ST state = connection->handler(0); + const auto deadline = Clock::now() + std::chrono::milliseconds(timeout_ms); + while (state == ASYNC_CONNECT_CONT && Clock::now() < deadline) { + pollfd descriptor { connection->fd, + connection->wait_events != 0 ? connection->wait_events : static_cast(POLLIN | POLLOUT), 0 }; + const int ready = poll(&descriptor, 1, 50); + short events = 0; + if (ready > 0) { + if ((descriptor.revents & POLLIN) != 0) events |= POLLIN; + if ((descriptor.revents & POLLOUT) != 0) events |= POLLOUT; + } + worker.curtime = monotonic_time(); + state = connection->handler(events); + } + return state; +} + +struct StagedConnect { + MDB_ASYNC_ST state { ASYNC_CONNECT_FAILED }; + bool server_reached_stage { false }; +}; + +StagedConnect drive_until_server_stage(MySQL_Thread& worker, MySQL_Connection *connection, + ChildServer& server) { + worker.curtime = monotonic_time(); + MDB_ASYNC_ST state = connection->handler(0); + for (unsigned int attempt = 0; state == ASYNC_CONNECT_CONT && attempt != 100; ++attempt) { + pollfd descriptors[2] { + { connection->fd, connection->wait_events != 0 ? connection->wait_events : + static_cast(POLLIN | POLLOUT), 0 }, + { server.stage_fd, POLLIN, 0 } + }; + const int ready = poll(descriptors, 2, 20); + if (ready <= 0) continue; + if ((descriptors[1].revents & POLLIN) != 0) { + char stage = 0; + if (read(server.stage_fd, &stage, 1) == 1 && stage == 'S') return { state, true }; + } + short events = 0; + if ((descriptors[0].revents & POLLIN) != 0) events |= POLLIN; + if ((descriptors[0].revents & POLLOUT) != 0) events |= POLLOUT; + if (events == 0) continue; + worker.curtime = monotonic_time(); + state = connection->handler(events); + } + return { state, false }; +} + +AwsIamTokenKey key(unsigned int port) { + return { kEndpoint, static_cast(port), "us-east-1", kUsername }; +} + +AwsIamTokenResult token_result(const std::string& token) { + AwsIamTokenResult result; + result.status = AwsIamStatus::OK; + result.generation = 7; + result.token = SecureString(token); + return result; +} + +void run_success_case(MySQL_Thread& worker, const CertificateFixture& certificates, + const std::string& token) { + ChildServer server = start_server("success", certificates.valid_cert, certificates.valid_key); + if (server.port == 0) BAIL_OUT("controlled server did not become ready"); + MySrvC *backend = create_server(1101, kEndpoint, server.port, true); + ConnectionFixture fixture(worker, backend); + fixture.connection->set_backend_auth_type(MySQLBackendAuthType::AWS_IAM); + FakeSource source; + source.token = token; + fixture.connection->attach_aws_iam_token(key(server.port), + source.request_blocking(key(server.port), Clock::now() + std::chrono::seconds(1))); + const MDB_ASYNC_ST status = drive(worker, fixture.connection); + finish_server(server); + + ok(status == ASYNC_CONNECT_SUCCESSFUL && source.blocking_requests == 1, + "one fake-source result authenticates through the real Connector/C TLS path"); + ok(field(server.result, "tls") == "1" && field(server.result, "pre_tls") == "0", + "mysql_clear_password payload is requested only after SSL_accept"); + ok(field(server.result, "min_tls") == std::to_string(TLS1_2_VERSION), + "the controlled IAM backend refuses TLS protocol versions older than TLS 1.2"); + ok(field(server.result, "username_hex") == hex(kUsername) && + field(server.result, "token_len") == std::to_string(token.size()) && + field(server.result, "token_sha256") == sha256(token), + "the expected backend user and 2 KiB-plus token arrive byte-for-byte without being logged"); + ok(field(server.result, "sni_hex") == "64622e636c75737465722d746573742e75732d656173742d312e7264732e616d617a6f6e6177732e636f6d" && + field(server.result, "peer") == "127.0.0.1", + "socket uses loopback while SNI and certificate verification use the configured RDS endpoint"); + ok(!fixture.connection->has_aws_iam_handshake_secret() && + (fixture.connection->mysql == nullptr || fixture.connection->mysql->passwd == nullptr) && + fixture.connection->userinfo->password != nullptr && + strcmp(fixture.connection->userinfo->password, "ordinary-password-must-not-leak") == 0, + "successful authentication cleanses transient token copies without replacing the stored password"); +} + +void run_certificate_failure(MySQL_Thread& worker, const CertificateFixture& certificates, + const std::string& mode, const std::string& certificate, const std::string& private_key, + unsigned int hostgroup, const char *label) { + ChildServer server = start_server(mode, certificate, private_key); + if (server.port == 0) BAIL_OUT("controlled certificate-failure server did not become ready"); + MySrvC *backend = create_server(hostgroup, kEndpoint, server.port, true); + ConnectionFixture fixture(worker, backend); + fixture.connection->set_backend_auth_type(MySQLBackendAuthType::AWS_IAM); + fixture.connection->attach_aws_iam_token(key(server.port), token_result("TOKEN_MUST_NOT_CROSS_FAILED_TLS")); + const MDB_ASYNC_ST status = drive(worker, fixture.connection); + finish_server(server); + ok(status == ASYNC_CONNECT_FAILED && field(server.result, "token_len") == "0", + "%s fails before token transmission", label); +} + +void run_pre_auth_abort_cases(MySQL_Thread& worker, const CertificateFixture& certificates, + const std::string& token) { + ChildServer closed = start_server("close_transport", certificates.valid_cert, certificates.valid_key); + if (closed.port == 0) BAIL_OUT("controlled transport-close server did not become ready"); + MySrvC *closed_backend = create_server(1106, kEndpoint, closed.port, true); + { + ConnectionFixture fixture(worker, closed_backend); + fixture.connection->set_backend_auth_type(MySQLBackendAuthType::AWS_IAM); + fixture.connection->attach_aws_iam_token(key(closed.port), token_result(token)); + const MDB_ASYNC_ST status = drive(worker, fixture.connection); + finish_server(closed); + ok(status == ASYNC_CONNECT_FAILED && field(closed.result, "token_len") == "0" && + !fixture.connection->has_aws_iam_handshake_secret(), + "transport loss before TLS authentication sends no token and clears the handshake secret"); + } + + ChildServer delayed = start_server("delay_handshake", certificates.valid_cert, + certificates.valid_key, 2000); + if (delayed.port == 0) BAIL_OUT("controlled delayed-handshake server did not become ready"); + MySrvC *delayed_backend = create_server(1107, kEndpoint, delayed.port, true); + { + ConnectionFixture fixture(worker, delayed_backend); + fixture.connection->set_backend_auth_type(MySQLBackendAuthType::AWS_IAM); + fixture.connection->attach_aws_iam_token(key(delayed.port), token_result(token)); + const StagedConnect initial = drive_until_server_stage(worker, fixture.connection, delayed); + fixture.connection->clear_aws_iam_handshake_secret(); + finish_server(delayed); + ok(initial.state == ASYNC_CONNECT_CONT && initial.server_reached_stage && + field(delayed.result, "token_len") == "0" && + !fixture.connection->has_aws_iam_handshake_secret() && fixture.connection->mysql == nullptr, + "connection cancellation during a delayed server handshake leaves no live Connector/C token owner"); + } +} + +} // namespace + +int main() { + plan(17); + diag("SDK-independent controlled TLS/protocol coverage; no SDK-on provider/signing verification is claimed"); + if (test_init_minimal() != 0 || test_init_query_processor() != 0 || + test_init_hostgroups() != 0) BAIL_OUT("failed to initialize component globals"); + GloMyLogger = new MySQL_Logger(); + GloMyMon = new MySQL_Monitor(); + GloMyMon->dns_cache->pin(kEndpoint, "127.0.0.1"); + + CertificateFixture certificates = create_certificates(); + ok(certificates.ready, "temporary CA and SAN server certificates were generated"); + if (!certificates.ready) BAIL_OUT("could not generate controlled TLS certificates"); + + { + MySQL_Thread worker; + if (!worker.init()) BAIL_OUT("MySQL_Thread::init failed"); + worker.curtime = monotonic_time(); + mysql_thread___ssl_p2s_ca = strdup(certificates.ca.c_str()); + const std::string token = std::string(kEndpoint) + ":3306/?Action=connect&" + std::string(2048, 'x'); + run_success_case(worker, certificates, token); + run_certificate_failure(worker, certificates, "wrong_hostname", + certificates.wrong_cert, certificates.wrong_key, 1102, "wrong hostname"); + run_certificate_failure(worker, certificates, "untrusted_ca", + certificates.untrusted_cert, certificates.untrusted_key, 1103, "untrusted CA"); + run_pre_auth_abort_cases(worker, certificates, token); + + ChildServer denied = start_server("access_denied", certificates.valid_cert, certificates.valid_key); + if (denied.port == 0) BAIL_OUT("controlled access-denied server did not become ready"); + MySrvC *denied_backend = create_server(1104, kEndpoint, denied.port, true); + { + ConnectionFixture fixture(worker, denied_backend); + fixture.connection->set_backend_auth_type(MySQLBackendAuthType::AWS_IAM); + fixture.connection->attach_aws_iam_token(key(denied.port), token_result(token)); + const MDB_ASYNC_ST status = drive(worker, fixture.connection); + finish_server(denied); + ok(status == ASYNC_CONNECT_FAILED && field(denied.result, "token_len") == std::to_string(token.size()), + "access denied occurs after one TLS-protected clear-password transmission"); + ok(!fixture.connection->has_aws_iam_handshake_secret() && + (fixture.connection->mysql == nullptr || fixture.connection->mysql->passwd == nullptr), + "access denied cleanses both token owners"); + } + + ChildServer ordinary = start_server("success", certificates.valid_cert, certificates.valid_key); + if (ordinary.port == 0) BAIL_OUT("controlled ordinary-password server did not become ready"); + MySrvC *ordinary_backend = create_server(1101, kEndpoint, ordinary.port, false); + FakeSource fake; + fake.token = token; + ScopedPublishedTokenSource published(&fake); + { + ConnectionFixture fixture(worker, ordinary_backend); + const MDB_ASYNC_ST status = drive(worker, fixture.connection); + finish_server(ordinary); + ok(status == ASYNC_CONNECT_FAILED && field(ordinary.result, "pre_tls") == "1" && + field(ordinary.result, "token_len") == "0", + "ordinary password mode coexists in the IAM hostgroup, keeps cleartext auth disabled, and sends no IAM token"); + ok(fake.blocking_requests == 0, "ordinary password mode never calls the published fake IAM source"); + } + + AwsIamConnectionConfigInput invalid { kUsername, kEndpoint, 3306, "us-east-1", false, + certificates.ca, "", true }; + ok(validate_mysql_aws_iam_connection(invalid).status == AwsIamConnectionConfigStatus::TLS_REQUIRED, + "use_ssl=0 is rejected before token acquisition"); + invalid.use_ssl = true; + invalid.ssl_ca.clear(); + ok(validate_mysql_aws_iam_connection(invalid).status == AwsIamConnectionConfigStatus::CA_TRUST_REQUIRED, + "missing CA trust is rejected before token acquisition"); + } + + delete GloMyMon; + GloMyMon = nullptr; + delete GloMyLogger; + GloMyLogger = nullptr; + test_cleanup_hostgroups(); + test_cleanup_query_processor(); + test_cleanup_minimal(); + return exit_status(); +} diff --git a/test/tap/tests/test_aws_iam_backend_auth-t.env b/test/tap/tests/test_aws_iam_backend_auth-t.env new file mode 100644 index 0000000000..464001b410 --- /dev/null +++ b/test/tap/tests/test_aws_iam_backend_auth-t.env @@ -0,0 +1,7 @@ +# Reserved for an optional external-provider daemon variant. +# The always-runnable controlled component test does not read these values. +AWS_ACCESS_KEY_ID=TAPONLYACCESSKEY +AWS_SECRET_ACCESS_KEY=tap-only-secret-never-use +AWS_SESSION_TOKEN=tap-only-session-never-use +AWS_EC2_METADATA_DISABLED=true +TAP_QUIET_ENVLOAD=1 diff --git a/test/tap/tests/test_aws_iam_metrics-t.cpp b/test/tap/tests/test_aws_iam_metrics-t.cpp new file mode 100644 index 0000000000..4974caa5a8 --- /dev/null +++ b/test/tap/tests/test_aws_iam_metrics-t.cpp @@ -0,0 +1,209 @@ +#include "tap.h" + +#include "Aws_Iam_Provider.h" +#include "ProxySQL_Statistics.hpp" +#include "cpp.h" +#include "proxysql.h" +#include "test_globals.h" +#include "test_init.h" + +#include "prometheus/registry.h" +#include "prometheus/text_serializer.h" + +#include +#include +#include +#include + +extern ProxySQL_Admin *GloAdmin; +extern ProxySQL_Statistics *GloProxyStats; + +namespace { + +constexpr const char *kSensitiveEndpoint = + "metrics-secret.cluster-abcdefghijkl.us-east-1.rds.amazonaws.com"; +constexpr const char *kSensitiveRegion = "metrics-secret-region"; +constexpr const char *kSensitiveUser = "metrics-secret-user"; +constexpr const char *kSensitiveToken = "FAKE_AWS_SESSION_TOKEN_METRICS"; + +class ScriptedSource final : public AwsIamTokenSource { +public: + explicit ScriptedSource(AwsIamStatsSnapshot stats) : stats_(stats) {} + + AwsIamRequestHandle request(const AwsIamTokenKey&, uint64_t, + std::weak_ptr) override { + return {}; + } + AwsIamTokenResult request_blocking(const AwsIamTokenKey&, + std::chrono::steady_clock::time_point) override { + return {}; + } + void cancel(AwsIamRequestHandle) override {} + void invalidate(const AwsIamTokenKey&, uint64_t) override {} + void record_backend_connection(bool) override {} + void record_waiting_session(bool) override {} + AwsIamStatsSnapshot snapshot() const override { return stats_; } + +private: + AwsIamStatsSnapshot stats_; +}; + +std::map query_aws_iam_stats(SQLite3DB *db) { + std::map values; + char *error = nullptr; + int columns = 0; + int affected_rows = 0; + SQLite3_result *result = nullptr; + db->execute_statement( + "SELECT Variable_Name, Variable_Value FROM stats_mysql_global " + "WHERE Variable_Name LIKE 'AwsIam_%' ORDER BY Variable_Name", + &error, &columns, &affected_rows, &result); + if (error != nullptr) { + free(error); + delete result; + return values; + } + if (result != nullptr) { + for (SQLite3_row *row : result->rows) { + if (row->fields[0] != nullptr && row->fields[1] != nullptr) { + values.emplace(row->fields[0], std::stoull(row->fields[1])); + } + } + } + delete result; + return values; +} + +bool has_unlabelled_sample(const std::string& text, const std::string& name, + uint64_t value) { + const std::string sample = name + " " + std::to_string(value) + "\n"; + return text.find(sample) != std::string::npos && + text.find(name + "{") == std::string::npos; +} + +} // namespace + +int main() { + plan(7); + + const bool initialized = test_init_minimal() == 0 && + test_init_query_processor() == 0 && test_init_hostgroups() == 0; + if (initialized) { + GloVars.statsdb_disk = strdup(":memory:"); + GloProxyStats = new ProxySQL_Statistics(); + GloProxyStats->init(); + GloAdmin = new ProxySQL_Admin(); // NOSONAR: process-scoped partial fixture + GloAdmin->statsdb = new SQLite3DB(); + char memory_db[] = ":memory:"; + GloAdmin->statsdb->open( + memory_db, SQLITE_OPEN_READWRITE | SQLITE_OPEN_CREATE | SQLITE_OPEN_FULLMUTEX); + GloAdmin->statsdb->execute( + "CREATE TABLE stats_mysql_global (Variable_Name VARCHAR NOT NULL PRIMARY KEY, " + "Variable_Value VARCHAR NOT NULL)"); + } + ok(initialized && GloAdmin != nullptr && GloAdmin->statsdb != nullptr, + "production admin and process-registry fixture initializes"); + + AwsIamStatsSnapshot chosen; + chosen.token_requests = 17; + chosen.token_cache_hits = 11; + chosen.token_refresh_successes = 7; + chosen.token_refresh_failures = 5; + chosen.credential_provider_failures = 3; + chosen.queue_rejections = 2; + chosen.backend_connection_successes = 13; + chosen.backend_connection_failures = 4; + chosen.token_cache_entries = 6; + chosen.in_flight_generations = 1; + chosen.queued_generations = 8; + chosen.waiting_sessions = 9; + ScriptedSource source(chosen); + publish_global_aws_iam_token_source(&source); + auto lease = acquire_global_aws_iam_token_source(); + ok(lease && lease->snapshot().token_requests == 17 && + lease->snapshot().waiting_sessions == 9, + "public stats consumers acquire the scripted provider snapshot"); + lease = AwsIamTokenSourceLease {}; + + const std::map expected { + { "AwsIam_Token_requests", 17 }, + { "AwsIam_Token_cache_hits", 11 }, + { "AwsIam_Token_refresh_successes", 7 }, + { "AwsIam_Token_refresh_failures", 5 }, + { "AwsIam_Credential_provider_failures", 3 }, + { "AwsIam_Queue_rejections", 2 }, + { "AwsIam_Backend_connection_successes", 13 }, + { "AwsIam_Backend_connection_failures", 4 }, + { "AwsIam_Token_cache_entries", 6 }, + { "AwsIam_In_flight_generations", 1 }, + { "AwsIam_Queued_generations", 8 }, + { "AwsIam_Waiting_sessions", 9 }, + }; + GloAdmin->stats___mysql_global(); + const std::map stats = query_aws_iam_stats(GloAdmin->statsdb); + ok(stats == expected, "stats_mysql_global projects all twelve provider values"); + + MySQL_Threads_Handler *saved_threads = GloMTH; + GloMTH = nullptr; + GloAdmin->p_update_metrics(); + GloMTH = saved_threads; + prometheus::TextSerializer serializer; + const std::string metrics = serializer.Serialize(GloVars.prometheus_registry->Collect()); + const std::map prometheus_expected { + { "proxysql_mysql_aws_iam_token_requests_total", 17 }, + { "proxysql_mysql_aws_iam_token_cache_hits_total", 11 }, + { "proxysql_mysql_aws_iam_token_refresh_successes_total", 7 }, + { "proxysql_mysql_aws_iam_token_refresh_failures_total", 5 }, + { "proxysql_mysql_aws_iam_credential_provider_failures_total", 3 }, + { "proxysql_mysql_aws_iam_queue_rejections_total", 2 }, + { "proxysql_mysql_aws_iam_backend_connection_successes_total", 13 }, + { "proxysql_mysql_aws_iam_backend_connection_failures_total", 4 }, + { "proxysql_mysql_aws_iam_token_cache_entries", 6 }, + { "proxysql_mysql_aws_iam_in_flight_generations", 1 }, + { "proxysql_mysql_aws_iam_queued_generations", 8 }, + { "proxysql_mysql_aws_iam_waiting_sessions", 9 }, + }; + bool prometheus_exact = true; + for (const auto& metric : prometheus_expected) { + prometheus_exact = prometheus_exact && + has_unlabelled_sample(metrics, metric.first, metric.second); + } + ok(prometheus_exact, + "Prometheus projects all twelve exact label-free provider values"); + ok(metrics.find(kSensitiveEndpoint) == std::string::npos && + metrics.find(kSensitiveRegion) == std::string::npos && + metrics.find(kSensitiveUser) == std::string::npos && + metrics.find(kSensitiveToken) == std::string::npos, + "provider metrics contain no endpoint, region, user, or token text"); + + publish_global_aws_iam_token_source(nullptr); + auto active_registry = GloVars.prometheus_registry; + GloVars.prometheus_registry = std::make_shared(); + auto unavailable = create_aws_iam_token_source({ 16, 16 }); + const AwsIamTokenResult unavailable_result = unavailable->request_blocking( + {}, std::chrono::steady_clock::now()); + ok(!unavailable->support_compiled() && + unavailable_result.status == AwsIamStatus::SUPPORT_NOT_COMPILED && + unavailable_result.failure.category == "support_not_compiled", + "the provider-neutral fallback remains fail closed"); + publish_global_aws_iam_token_source(unavailable.get()); + GloAdmin->stats___mysql_global(); + saved_threads = GloMTH; + GloMTH = nullptr; + GloAdmin->p_update_metrics(); + GloMTH = saved_threads; + const auto unavailable_rows = query_aws_iam_stats(GloAdmin->statsdb); + const std::string unavailable_metrics = + serializer.Serialize(GloVars.prometheus_registry->Collect()); + bool unavailable_zero = unavailable_rows.size() == 12; + for (const auto& row : unavailable_rows) unavailable_zero = unavailable_zero && row.second == 0; + for (const auto& metric : prometheus_expected) { + unavailable_zero = unavailable_zero && + has_unlabelled_sample(unavailable_metrics, metric.first, 0); + } + ok(unavailable_zero, + "the unavailable provider projects fixed zero admin and Prometheus values"); + shutdown_global_aws_iam_token_source(); + + return exit_status(); +} diff --git a/test/tap/tests/test_cluster_sync-t.cpp b/test/tap/tests/test_cluster_sync-t.cpp index be3fdcf6c6..f2c2e7fef2 100644 --- a/test/tap/tests/test_cluster_sync-t.cpp +++ b/test/tap/tests/test_cluster_sync-t.cpp @@ -111,6 +111,25 @@ const uint32_t R_PORT = 16062; // Use 127.0.0.1 to connect to it, not cl.host (which may point to a different container). const char* R_HOST = "127.0.0.1"; +void restore_task11_mysql_users(MYSQL* admin, int& backup_stage) { + if (admin == nullptr || backup_stage == 0) return; + const auto best_effort = [admin](const char* query) { + if (mysql_query(admin, query) != 0) { + diag("Task 11 mysql_users restore query failed: %s", mysql_error(admin)); + } + }; + best_effort("DELETE FROM mysql_users"); + best_effort("INSERT INTO mysql_users SELECT * FROM mysql_users_sync_test_task11"); + if (backup_stage >= 2) { + best_effort("DELETE FROM disk.mysql_users"); + best_effort("INSERT INTO disk.mysql_users SELECT * FROM mysql_users_disk_sync_test_task11"); + } + best_effort("DROP TABLE IF EXISTS mysql_users_sync_test_task11"); + best_effort("DROP TABLE IF EXISTS mysql_users_disk_sync_test_task11"); + best_effort("LOAD MYSQL USERS TO RUNTIME"); + backup_stage = 0; +} + // Hostname visible to other containers on the Docker network. // Used when registering the replica in proxysql_servers on the primary so the // primary's cluster monitor can reach the replica. Falls back to R_HOST when @@ -1181,6 +1200,7 @@ int main(int, char**) { int res = 0; CommandLine cl; std::atomic save_proxy_stderr(false); + int task11_mysql_users_backup_stage = 0; if (cl.getEnv()) { diag("Failed to get the required environmental variables."); @@ -1213,7 +1233,7 @@ int main(int, char**) { plan( // Sync tests by values - 16 + + 18 + // Module checkums tests; enabled and disabled checksums check_modules_checksums_sync__tests + (cl.use_noise ? 3 : 0) @@ -1397,7 +1417,7 @@ int main(int, char**) { std::make_tuple(18, 2, -1, 20, "SET sql_mode = \"\"", 0, 0, 100, "", "", ""), std::make_tuple(19, 2, -1, 20, "SET sql_mode = \"\"", 0, 0, 100, "{}", "{}", "{}"), std::make_tuple(20, 0, 0, 30, "SET long_query_time = 0", 1, 0, 123, "{\"session_variables\":[\"tmp_table_size\",\"join_buffer_size\"]}", "", ""), - std::make_tuple(21, 2, -1, 50, "SET sql_mode = \"\"", 1, 0, 125, "{\"session_variables\":[\"tmp_table_size\",\"join_buffer_size\"]}", "{\"handle_warnings\":1}", ""), + std::make_tuple(21, 2, -1, 50, "SET sql_mode = \"\"", 1, 0, 125, "{\"session_variables\":[\"tmp_table_size\",\"join_buffer_size\"]}", "{\"handle_warnings\":1,\"aws_iam_region\":\"us-east-1\"}", ""), std::make_tuple(22, 3, -1, 40, "SET sql_mode = \"\"", 1, 0, 124, "{\"session_variables\":[\"tmp_table_size\",\"join_buffer_size\"]}", "", "{\"weight\": 100, \"max_connections\": 1000}") }; std::vector insert_mysql_hostgroup_attributes_queries{}; @@ -1498,6 +1518,53 @@ int main(int, char**) { system(print_replica_hostgroup_attributes.c_str()); ok(not_synced_query == false, "'mysql_hostgroup_attributes' should be synced."); + // IAM policy is configuration state. Keep the hostgroup-region row live + // while exercising a user save/load and cluster round-trip so the replica + // must observe both halves of the effective IAM policy exactly. + const char* iam_username = "tap_aws_iam_cluster_sync"; + const char* iam_attributes = "{\"backend_auth\":{\"type\":\"aws_iam\"}}"; + MYSQL_QUERY__(proxy_admin, "DROP TABLE IF EXISTS mysql_users_sync_test_task11"); + MYSQL_QUERY__(proxy_admin, "DROP TABLE IF EXISTS mysql_users_disk_sync_test_task11"); + MYSQL_QUERY__(proxy_admin, "CREATE TABLE mysql_users_sync_test_task11 AS SELECT * FROM mysql_users"); + task11_mysql_users_backup_stage = 1; + MYSQL_QUERY__(proxy_admin, + "CREATE TABLE mysql_users_disk_sync_test_task11 AS SELECT * FROM disk.mysql_users"); + task11_mysql_users_backup_stage = 2; + MYSQL_QUERY__(proxy_admin, "DELETE FROM mysql_users WHERE username='tap_aws_iam_cluster_sync'"); + MYSQL_QUERY__(proxy_admin, + "INSERT INTO mysql_users (username,password,active,default_hostgroup,backend,frontend,attributes) " + "VALUES ('tap_aws_iam_cluster_sync','',1,21,1,0,'{\"backend_auth\":{\"type\":\"aws_iam\"}}')"); + MYSQL_QUERY__(proxy_admin, "SAVE MYSQL USERS TO DISK"); + MYSQL_QUERY__(proxy_admin, "DELETE FROM mysql_users WHERE username='tap_aws_iam_cluster_sync'"); + MYSQL_QUERY__(proxy_admin, "LOAD MYSQL USERS FROM DISK"); + + const std::string primary_iam_policy_query { + "SELECT COUNT(*) FROM mysql_users WHERE username='" + std::string(iam_username) + + "' AND password='' AND active=1 AND default_hostgroup=21 AND backend=1 AND frontend=0 " + "AND attributes='" + std::string(iam_attributes) + "'" + }; + MYSQL_QUERY__(proxy_admin, primary_iam_policy_query.c_str()); + MYSQL_RES* primary_iam_res = mysql_store_result(proxy_admin); + MYSQL_ROW primary_iam_row = mysql_fetch_row(primary_iam_res); + const bool primary_iam_policy_ok = primary_iam_row && primary_iam_row[0] && + std::atoi(primary_iam_row[0]) == 1; + mysql_free_result(primary_iam_res); + ok(primary_iam_policy_ok, "IAM user policy survives save/load exactly"); + + MYSQL_QUERY__(proxy_admin, "LOAD MYSQL USERS TO RUNTIME"); + const std::string replica_iam_policy_query { + "SELECT ((SELECT COUNT(*) FROM mysql_users WHERE username='" + std::string(iam_username) + + "' AND password='' AND active=1 AND default_hostgroup=21 AND backend=1 AND frontend=0 " + "AND attributes='" + std::string(iam_attributes) + "')=1 " + "AND (SELECT COUNT(*) FROM mysql_hostgroup_attributes WHERE hostgroup_id=21 " + "AND hostgroup_settings='{\"handle_warnings\":1,\"aws_iam_region\":\"us-east-1\"}')=1)" + }; + const bool iam_cluster_not_synced = + wait_for_node_sync(r_proxy_admin, { replica_iam_policy_query }, SYNC_TIMEOUT) != 0; + ok(!iam_cluster_not_synced, "cluster sync carries exact IAM user policy and hostgroup region"); + + restore_task11_mysql_users(proxy_admin, task11_mysql_users_backup_stage); + // TEARDOWN CONFIG MYSQL_QUERY__(proxy_admin, "DELETE FROM mysql_hostgroup_attributes"); MYSQL_QUERY__(proxy_admin, "INSERT INTO mysql_hostgroup_attributes SELECT * FROM mysql_hostgroup_attributes_sync_test_2687"); @@ -2775,6 +2842,7 @@ int main(int, char**) { cleanup: // Teardown config + restore_task11_mysql_users(proxy_admin, task11_mysql_users_backup_stage); // In case of test failing, save the stderr output from the spawned proxysql instance if (tests_failed() != 0) { diff --git a/test/tap/tests/test_mysql_hostgroup_attributes-1-t.cpp b/test/tap/tests/test_mysql_hostgroup_attributes-1-t.cpp index 8825126eb8..582aa163e3 100644 --- a/test/tap/tests/test_mysql_hostgroup_attributes-1-t.cpp +++ b/test/tap/tests/test_mysql_hostgroup_attributes-1-t.cpp @@ -51,6 +51,36 @@ int run_one_test(MYSQL *mysqladmin, const char *expected_checksum, const char *q return 0; } +int test_aws_iam_region_roundtrip(MYSQL *mysqladmin) { + const char* const settings = "{\"aws_iam_region\":\"us-east-1\"}"; + const char* const insert = + "INSERT INTO mysql_hostgroup_attributes (hostgroup_id, hostgroup_settings) " + "VALUES (999901, '{\"aws_iam_region\":\"us-east-1\"}')"; + MYSQL_QUERY(mysqladmin, "DELETE FROM mysql_hostgroup_attributes"); + MYSQL_QUERY(mysqladmin, insert); + MYSQL_QUERY(mysqladmin, "LOAD MYSQL SERVERS TO RUNTIME"); + MYSQL_QUERY(mysqladmin, + "SELECT hostgroup_settings FROM runtime_mysql_hostgroup_attributes WHERE hostgroup_id=999901"); + MYSQL_RES* result = mysql_store_result(mysqladmin); + MYSQL_ROW row = mysql_fetch_row(result); + ok(row != nullptr && strcmp(row[0], settings) == 0, + "LOAD MYSQL SERVERS TO RUNTIME preserves aws_iam_region hostgroup settings"); + mysql_free_result(result); + + MYSQL_QUERY(mysqladmin, "SAVE MYSQL SERVERS FROM RUNTIME"); + MYSQL_QUERY(mysqladmin, "DELETE FROM mysql_hostgroup_attributes"); + MYSQL_QUERY(mysqladmin, "SAVE MYSQL SERVERS FROM RUNTIME"); + MYSQL_QUERY(mysqladmin, "LOAD MYSQL SERVERS TO RUNTIME"); + MYSQL_QUERY(mysqladmin, + "SELECT hostgroup_settings FROM runtime_mysql_hostgroup_attributes WHERE hostgroup_id=999901"); + result = mysql_store_result(mysqladmin); + row = mysql_fetch_row(result); + ok(row != nullptr && strcmp(row[0], settings) == 0, + "SAVE and reload preserve aws_iam_region hostgroup settings"); + mysql_free_result(result); + return 0; +} + int main(int argc, char** argv) { CommandLine cl; @@ -88,7 +118,7 @@ int main(int argc, char** argv) { } }; - plan(queries_and_checksums.size()*4); + plan(queries_and_checksums.size()*4 + 2); diag("Testing the loading of mysql_hostgroup_attributes"); MYSQL* mysqladmin = mysql_init(NULL); @@ -111,8 +141,8 @@ int main(int argc, char** argv) { sleep(10); } } + test_aws_iam_region_roundtrip(mysqladmin); mysql_close(mysqladmin); return exit_status(); } - diff --git a/test/tap/tests/unit/Makefile b/test/tap/tests/unit/Makefile index 1125125d1b..e313f3f2b5 100644 --- a/test/tap/tests/unit/Makefile +++ b/test/tap/tests/unit/Makefile @@ -8,6 +8,28 @@ # # See: GitHub issue #5473 (Phase 2.1: Test Infrastructure Foundation) +.DEFAULT_GOAL := all + +# The public build owns only the provider-neutral IAM core. Keep an explicit +# artifact gate because it exercises the actual build output rather than a +# source-code convention. Capture nm completely before matching so grep +# cannot terminate the producer early under pipefail. +.PHONY: aws_public_provider_boundary-t +aws_public_provider_boundary-t: + @test ! -d $(PROXYSQL_PATH)/plugins/aws || { \ + echo "FAIL: public tree owns the concrete AWS provider" >&2; exit 1; \ + } + @test ! -e $(PROXYSQL_PATH)/test/tap/tests/unit/aws_locality_plugin_unit-t.cpp || { \ + echo "FAIL: public tree owns the real AWS locality provider test" >&2; exit 1; \ + } + @test -x $(PROXYSQL_PATH)/src/proxysql + @nm -C $(PROXYSQL_PATH)/src/proxysql > /tmp/proxysql-public.nm + @if grep -Fq 'Aws::' /tmp/proxysql-public.nm; then \ + echo "FAIL: src/proxysql contains AWS C++ SDK symbols" >&2; exit 1; \ + fi + @test ! -e $(PROXYSQL_PATH)/plugins/aws/ProxySQL_Aws_Plugin.so || { \ + echo "FAIL: public build produced the AWS plugin" >&2; exit 1; \ + } PROXYSQL_PATH := $(shell while [ ! -f ./src/proxysql_global.cpp ]; do cd ..; done; pwd) @@ -140,7 +162,6 @@ LZ4_LDIR := $(DEPS_PATH)/lz4/lz4/lib IDIRS += -I$(CLICKHOUSE_CPP_IDIR) - # =========================================================================== # libproxysql.a — the core library under test # =========================================================================== @@ -246,7 +267,6 @@ ifeq ($(UNAME_S),Linux) endif endif - # =========================================================================== # Compiler flags # =========================================================================== @@ -322,12 +342,12 @@ endif # misused (most importantly identity-forgery setters), while tests that include # the plugin headers and link against the plugin sources still see the helpers # they need. -OPT := $(STDCPP) -O0 -ggdb $(PSQLCH) $(PSQLGA) $(PSQL40) $(PSQL31) $(PSQLFFTO) $(PSQLTSDB) $(PSQLED25519) $(PSQLDEBUG) \ +OPT := $(STDCPP) -O0 -ggdb $(PSQLCH) $(PSQLGA) $(PSQL40) $(PSQL31) $(PSQLFFTO) $(PSQLTSDB) $(PSQLED25519) $(PSQLAWSIAM) $(PSQLDEBUG) \ -DGITVERSION=\"$(GIT_VERSION)\" -DMYSQLX_TEST_BUILD $(NOJEM) $(WGCOV) $(WASAN) \ -Wl,--no-as-needed -Wl,-rpath,$(TAP_LDIR) ifeq ($(UNAME_S),Darwin) - OPT := $(STDCPP) -O0 -ggdb $(PSQLCH) $(PSQLGA) $(PSQL40) $(PSQL31) $(PSQLFFTO) $(PSQLTSDB) $(PSQLED25519) $(PSQLDEBUG) \ + OPT := $(STDCPP) -O0 -ggdb $(PSQLCH) $(PSQLGA) $(PSQL40) $(PSQL31) $(PSQLFFTO) $(PSQLTSDB) $(PSQLED25519) $(PSQLAWSIAM) $(PSQLDEBUG) \ -DGITVERSION=\"$(GIT_VERSION)\" -DMYSQLX_TEST_BUILD $(NOJEM) $(WGCOV) $(WASAN) endif @@ -402,7 +422,7 @@ $(LIBPROXYSQLAR): FORCE # =========================================================================== UNIT_TESTS := smoke_test-t query_cache_unit-t query_processor_unit-t \ - protocol_unit-t auth_unit-t connection_pool_unit-t \ + protocol_unit-t auth_unit-t aws_iam_policy_unit-t aws_iam_connection_config_unit-t aws_iam_provider_boundary_unit-t aws_iam_completion_queue_unit-t aws_iam_session_state_unit-t aws_iam_connection_secret_unit-t aws_iam_pool_unit-t aws_iam_failure_unit-t aws_iam_kill_helper_unit-t aws_locality_policy_unit-t aws_locality_manager_unit-t connection_pool_unit-t \ rule_matching_unit-t hostgroups_unit-t monitor_health_unit-t \ pgsql_command_complete_unit-t \ ffto_protocol_unit-t \ @@ -418,6 +438,7 @@ UNIT_TESTS := smoke_test-t query_cache_unit-t query_processor_unit-t \ mysql_user_variables_unit-t \ mysql_decompress_payload_unit-t \ mysql_resolution_unit-t \ + mariadb_tls_server_name_unit-t \ pgsql_variables_validator_unit-t \ proxysql_utils_unit-t \ gen_utils_unit-t \ @@ -477,6 +498,9 @@ UNIT_TESTS += \ plugin_prometheus_unit-t \ plugin_lifecycle_unit-t \ plugin_runtime_views_unit-t \ + aws_locality_config_unit-t \ + aws_locality_selection_unit-t \ + aws_locality_stats_unit-t \ test_mysqlx_plugin_load-t \ test_mysqlx_admin_tables-t \ mysqlx_config_store_unit-t \ @@ -750,6 +774,11 @@ mysql_resolution_unit-t: mysql_resolution_unit-t.cpp $(ODIR)/tap.o $(ODIR)/tap_n -I$(TAP_IDIR) -I$(PROXYSQL_PATH)/include \ $(STDCPP) -O0 -ggdb $(WGCOV) $(LWGCOV) $(WASAN) -lpthread -o $@ +mariadb_tls_server_name_unit-t: mariadb_tls_server_name_unit-t.cpp $(ODIR)/tap.o $(ODIR)/tap_noise_stubs.o + $(CXX) $< $(ODIR)/tap.o $(ODIR)/tap_noise_stubs.o $(IDIRS) $(LDIRS) $(OPT) \ + -Wl,-Bstatic -lmariadbclient -Wl,-Bdynamic -lssl -lcrypto -lz -lpthread -ldl \ + $(LWGCOV) -o $@ + plugin_query_hook_unit-t: plugin_query_hook_unit-t.cpp $(FAKE_PLUGIN_SO) $(ODIR)/tap.o $(ODIR)/test_globals.o $(ODIR)/test_init.o $(LIBPROXYSQLAR) $(CXX) $< $(ODIR)/tap.o $(ODIR)/test_globals.o $(ODIR)/test_init.o \ -DPROXYSQL_FAKE_PLUGIN_PATH=\"$(FAKE_PLUGIN_SO)\" \ @@ -889,6 +918,71 @@ genai_mysql_catalog_unit-t: genai_mysql_catalog_unit-t.cpp $(GENAI_ALL_SRCS) $(T ifeq ($(UNAME_S),Linux) # Let the RSA unit test observe the exact flock attempt without a production hook. caching_sha2_rsa_unit-t: ALLOW_MULTI_DEF += -Wl,--wrap=flock + +# Observe the Connector/C handshake boundary and credential cleansing without a +# production-only test hook. +aws_iam_connection_secret_unit-t: ALLOW_MULTI_DEF += \ + -Wl,--wrap=mysql_options \ + -Wl,--wrap=mysql_real_connect_start \ + -Wl,--wrap=mysql_close_no_command \ + -Wl,--wrap=OPENSSL_cleanse +endif + +aws_iam_provider_boundary_unit-t: aws_iam_provider_boundary_unit-t.cpp \ + $(TEST_HELPERS_OBJ) $(LIBPROXYSQLAR) + $(CXX) $< $(TEST_HELPERS_OBJ) $(IDIRS) $(LDIRS) $(OPT) \ + $(WHOLE_LIBPROXYSQL) $(STATIC_LIBS) $(MYLIBS) -ldl \ + $(ALLOW_MULTI_DEF) -o $@ + +# The completion inbox is intentionally independent of sessions and the core +# archive so its producer/lifetime contract can run under TSan in isolation. +aws_iam_completion_queue_unit-t: aws_iam_completion_queue_unit-t.cpp \ + $(ODIR)/tap.o $(ODIR)/tap_noise_stubs.o + $(CXX) $< $(ODIR)/tap.o $(ODIR)/tap_noise_stubs.o $(IDIRS) $(OPT) \ + -lpthread -lcrypto -o $@ + +aws_locality_policy_unit-t: aws_locality_policy_unit-t.cpp \ + $(PROXYSQL_PATH)/lib/Aws_Locality_Manager.cpp $(ODIR)/tap.o $(ODIR)/tap_noise_stubs.o + $(CXX) $< $(PROXYSQL_PATH)/lib/Aws_Locality_Manager.cpp \ + $(ODIR)/tap.o $(ODIR)/tap_noise_stubs.o \ + $(IDIRS) $(OPT) -lpthread -ldl -o $@ + +aws_locality_manager_unit-t: aws_locality_manager_unit-t.cpp \ + $(PROXYSQL_PATH)/lib/Aws_Locality_Manager.cpp $(ODIR)/tap.o $(ODIR)/tap_noise_stubs.o + $(CXX) $< $(PROXYSQL_PATH)/lib/Aws_Locality_Manager.cpp \ + $(ODIR)/tap.o $(ODIR)/tap_noise_stubs.o \ + $(IDIRS) $(OPT) -lpthread -ldl -o $@ + +test_aws_iam_metrics-t: ../test_aws_iam_metrics-t.cpp $(TEST_HELPERS_OBJ) $(LIBPROXYSQLAR) + $(CXX) $< $(TEST_HELPERS_OBJ) -I$(TEST_HELPERS_DIR) \ + $(IDIRS) $(LDIRS) $(OPT) $(WHOLE_LIBPROXYSQL) $(STATIC_LIBS) \ + $(MYLIBS) $(ALLOW_MULTI_DEF) -o $@ + +test_aws_iam_backend_auth-t: ../test_aws_iam_backend_auth-t.cpp $(TEST_HELPERS_OBJ) $(LIBPROXYSQLAR) + $(CXX) -DAWS_IAM_MYSQL_SERVER_PATH=\"$(PROXYSQL_PATH)/test/deps/aws_iam_mysql_server/aws_iam_mysql_server-t\" \ + $< $(TEST_HELPERS_OBJ) -I$(TEST_HELPERS_DIR) \ + $(IDIRS) $(LDIRS) $(OPT) $(WHOLE_LIBPROXYSQL) $(STATIC_LIBS) \ + $(MYLIBS) $(ALLOW_MULTI_DEF) -o $@ + +aws_locality_stats_unit-t: aws_locality_stats_unit-t.cpp $(TEST_HELPERS_OBJ) $(LIBPROXYSQLAR) + $(CXX) $< $(TEST_HELPERS_OBJ) $(IDIRS) $(LDIRS) $(OPT) \ + $(WHOLE_LIBPROXYSQL) $(STATIC_LIBS) $(MYLIBS) $(ALLOW_MULTI_DEF) -o $@ + +ifeq ($(UNAME_S),Linux) +aws_iam_session_state_unit-t: ALLOW_MULTI_DEF += -Wl,--wrap=mysql_real_connect_start -pthread +aws_iam_pool_unit-t: ALLOW_MULTI_DEF += \ + -Wl,--wrap=mysql_change_user_start \ + -Wl,--wrap=mysql_real_connect_start \ + -Wl,--wrap=pthread_create +aws_iam_failure_unit-t: ALLOW_MULTI_DEF += \ + -Wl,--wrap=mysql_real_connect_start \ + -Wl,--wrap=_ZN14MySQL_Protocol16generate_pkt_ERREbPPvPjhtPKcS4_b +aws_iam_kill_helper_unit-t: ALLOW_MULTI_DEF += \ + -Wl,--wrap=mysql_options \ + -Wl,--wrap=mysql_real_connect \ + -Wl,--wrap=mysql_query \ + -Wl,--wrap=mysql_close \ + -Wl,--wrap=OPENSSL_cleanse endif # Pattern rule: all unit tests use the same compile + link flags. diff --git a/test/tap/tests/unit/aws_iam_completion_queue_unit-t.cpp b/test/tap/tests/unit/aws_iam_completion_queue_unit-t.cpp new file mode 100644 index 0000000000..40049472ae --- /dev/null +++ b/test/tap/tests/unit/aws_iam_completion_queue_unit-t.cpp @@ -0,0 +1,173 @@ +/** + * @file aws_iam_completion_queue_unit-t.cpp + * @brief Concurrency and lifetime tests for the per-worker IAM completion inbox. + */ + +#include "tap.h" + +#include "MySQL_Thread.h" + +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace { + +std::atomic cleanse_calls { 0 }; + +void tracked_cleanse(void *ptr, size_t size) { + OPENSSL_cleanse(ptr, size); + cleanse_calls.fetch_add(1, std::memory_order_relaxed); +} + +AwsIamCompletion completion(uint64_t opaque_id, bool with_token = false) { + AwsIamCompletion value; + value.opaque_id = opaque_id; + value.result.status = AwsIamStatus::OK; + if (with_token) { + value.result.token = SecureString("sensitive-queue-token", tracked_cleanse); + } + return value; +} + +struct PipePair { + int fds[2] { -1, -1 }; + PipePair() { + if (pipe(fds) != 0) { + BAIL_OUT("pipe() failed"); + } + const int flags = fcntl(fds[1], F_GETFL, 0); + if (flags < 0 || fcntl(fds[1], F_SETFL, flags | O_NONBLOCK) != 0) { + BAIL_OUT("failed to make completion pipe nonblocking"); + } + } + ~PipePair() { + if (fds[0] >= 0) close(fds[0]); + if (fds[1] >= 0) close(fds[1]); + } + unsigned int wake_count() { + const int flags = fcntl(fds[0], F_GETFL, 0); + fcntl(fds[0], F_SETFL, flags | O_NONBLOCK); + unsigned int count = 0; + unsigned char byte = 0; + while (read(fds[0], &byte, 1) == 1) ++count; + return count; + } +}; + +bool run_multi_producer_once() { + PipePair pipe_pair; + auto inbox = std::make_shared(pipe_pair.fds[1], 512); + constexpr unsigned int producers = 8; + constexpr unsigned int per_producer = 32; + std::vector threads; + threads.reserve(producers); + for (unsigned int producer = 0; producer < producers; ++producer) { + threads.emplace_back([inbox, producer] { + for (unsigned int item = 0; item < per_producer; ++item) { + const uint64_t id = producer * per_producer + item + 1; + inbox->post(completion(id)); + } + }); + } + for (auto& thread : threads) thread.join(); + + auto values = inbox->drain(); + std::unordered_set ids; + for (const auto& value : values) ids.insert(value.opaque_id); + return values.size() == producers * per_producer && + ids.size() == producers * per_producer && pipe_pair.wake_count() == 1; +} + +void test_multi_producer_and_wake_coalescing() { + bool passed = true; + for (unsigned int iteration = 0; iteration < 100; ++iteration) { + passed = passed && run_multi_producer_once(); + } + ok(passed, + "100 multi-producer runs deliver every completion with one empty-to-nonempty wake each"); +} + +void test_fifo_drain() { + PipePair pipe_pair; + AwsIamWorkerInbox inbox(pipe_pair.fds[1], 4); + inbox.post(completion(11)); + inbox.post(completion(12)); + inbox.post(completion(13)); + auto values = inbox.drain(); + ok(values.size() == 3 && values[0].opaque_id == 11 && + values[1].opaque_id == 12 && values[2].opaque_id == 13, + "drain preserves completion FIFO order"); + ok(pipe_pair.wake_count() == 1, + "multiple posts before a drain write only one pipe wake"); +} + +void test_bounded_overflow_cleanses() { + cleanse_calls.store(0, std::memory_order_relaxed); + PipePair pipe_pair; + AwsIamWorkerInbox inbox(pipe_pair.fds[1], 2); + inbox.post(completion(1, true)); + inbox.post(completion(2, true)); + inbox.post(completion(3, true)); + ok(cleanse_calls.load(std::memory_order_relaxed) == 1, + "a completion rejected by the queue bound is cleansed immediately"); + auto values = inbox.drain(); + ok(values.size() == 2 && values[0].opaque_id == 1 && values[1].opaque_id == 2, + "bounded overflow does not displace accepted FIFO completions"); + values.clear(); + ok(cleanse_calls.load(std::memory_order_relaxed) == 3, + "drained completion tokens remain independently owned and are cleansed on release"); +} + +void test_close_and_late_post() { + cleanse_calls.store(0, std::memory_order_relaxed); + PipePair pipe_pair; + AwsIamWorkerInbox inbox(pipe_pair.fds[1], 2); + inbox.close(); + inbox.close(); + inbox.post(completion(9, true)); + ok(cleanse_calls.load(std::memory_order_relaxed) == 1 && inbox.drain().empty(), + "close is idempotent and a late post is cleansed instead of retained"); + ok(pipe_pair.wake_count() == 0, + "a post after close never writes to the worker pipe"); +} + +void test_expired_weak_producers() { + PipePair pipe_pair; + auto inbox = std::make_shared(pipe_pair.fds[1], 64); + std::weak_ptr weak = inbox; + std::atomic release { false }; + std::vector producers; + for (unsigned int i = 0; i < 16; ++i) { + producers.emplace_back([weak, &release, i] { + while (!release.load(std::memory_order_acquire)) std::this_thread::yield(); + if (auto live = weak.lock()) live->post(completion(i + 1)); + }); + } + inbox.reset(); + release.store(true, std::memory_order_release); + for (auto& producer : producers) producer.join(); + ok(weak.expired(), + "inbox destruction is safe while producer threads hold only expired weak pointers"); +} + +} // namespace + +int main() { + plan(9); + test_multi_producer_and_wake_coalescing(); + test_fifo_drain(); + test_bounded_overflow_cleanses(); + test_close_and_late_post(); + test_expired_weak_producers(); + return exit_status(); +} diff --git a/test/tap/tests/unit/aws_iam_connection_config_unit-t.cpp b/test/tap/tests/unit/aws_iam_connection_config_unit-t.cpp new file mode 100644 index 0000000000..253b4d9c07 --- /dev/null +++ b/test/tap/tests/unit/aws_iam_connection_config_unit-t.cpp @@ -0,0 +1,183 @@ +#include "tap.h" + +#include "MySQL_Backend_Auth.h" +#include "MySQL_HostGroups_Manager.h" + +#include +#include +#include +#include + +void init_myhgc_hostgroup_settings(const char* hostgroup_settings, MyHGC* myhgc); + +namespace { + +AwsIamConnectionConfigInput valid_input(const char* endpoint, const char* region) { + AwsIamConnectionConfigInput input; + input.database_user = "iam_database_user"; + input.configured_endpoint = endpoint; + input.port = 3306; + input.region = region; + input.use_ssl = true; + input.ssl_ca = "/etc/ssl/rds-ca.pem"; + input.support_compiled = true; + return input; +} + +void ok_status(const AwsIamConnectionConfigInput& input, const char* name) { + const AwsIamConnectionConfigResult result = validate_mysql_aws_iam_connection(input); + ok(result.status == AwsIamConnectionConfigStatus::OK, "%s", name); + if (result.status == AwsIamConnectionConfigStatus::OK) { + ok(result.key.endpoint == input.configured_endpoint && + result.key.port == input.port && + result.key.region == input.region && + result.key.database_user == input.database_user, + "%s preserves the exact IAM token key", name); + } +} + +void failure_status(const AwsIamConnectionConfigInput& input, + AwsIamConnectionConfigStatus expected_status, const char* expected_code, const char* name) { + const AwsIamConnectionConfigResult result = validate_mysql_aws_iam_connection(input); + ok(result.status == expected_status && result.failure_code == expected_code, "%s", name); +} + +void test_valid_rds_endpoint_shapes() { + ok_status(valid_input("instance.abcdef.us-east-1.rds.amazonaws.com", "us-east-1"), + "RDS instance endpoint is accepted"); + ok_status(valid_input("cluster-main.abcdef.us-east-1.rds.amazonaws.com", "us-east-1"), + "Aurora cluster endpoint is accepted"); + ok_status(valid_input("cluster-ro-main.abcdef.us-east-1.rds.amazonaws.com", "us-east-1"), + "Aurora reader endpoint is accepted"); + ok_status(valid_input("cluster-custom-main.abcdef.us-east-1.rds.amazonaws.com", "us-east-1"), + "Aurora custom-cluster endpoint is accepted"); + ok_status(valid_input("instance.abcdef.cn-north-1.rds.amazonaws.com.cn", "cn-north-1"), + "China RDS endpoint is accepted"); + ok_status(valid_input("instance.abcdef.us-gov-west-1.rds.amazonaws.com", "us-gov-west-1"), + "GovCloud RDS endpoint is accepted"); + ok_status(valid_input("instance.abcdef.us-iso-east-1.rds.c2s.ic.gov", "us-iso-east-1"), + "ISO RDS endpoint is accepted"); + ok_status(valid_input("instance.abcdef.us-isob-east-1.rds.sc2s.sgov.gov", "us-isob-east-1"), + "ISOB RDS endpoint is accepted"); +} + +void test_validation_failures() { + AwsIamConnectionConfigInput input = valid_input("instance.abcdef.us-east-1.rds.amazonaws.com", "us-east-1"); + input.support_compiled = false; + failure_status(input, AwsIamConnectionConfigStatus::SUPPORT_NOT_COMPILED, "support_not_compiled", + "uncompiled AWS IAM support is rejected before tuple validation"); + + input = valid_input("instance.abcdef.us-east-1.rds.amazonaws.com", ""); + failure_status(input, AwsIamConnectionConfigStatus::MISSING_REGION, "missing_region", + "missing IAM region is rejected"); + + input = valid_input("instance.abcdef.us-west-2.rds.amazonaws.com", "us-east-1"); + failure_status(input, AwsIamConnectionConfigStatus::REGION_ENDPOINT_MISMATCH, "region_endpoint_mismatch", + "endpoint region must equal IAM region"); + + input = valid_input("", "us-east-1"); + failure_status(input, AwsIamConnectionConfigStatus::INVALID_ENDPOINT, "invalid_endpoint", + "empty endpoint is rejected"); + input = valid_input("instance.abcdef.us-east-1.rds.amazonaws.com.", "us-east-1"); + failure_status(input, AwsIamConnectionConfigStatus::INVALID_ENDPOINT, "invalid_endpoint", + "trailing-dot endpoint is rejected"); + input = valid_input("192.0.2.10", "us-east-1"); + failure_status(input, AwsIamConnectionConfigStatus::INVALID_ENDPOINT, "invalid_endpoint", + "IPv4 endpoint is rejected"); + input = valid_input("2001:db8::1", "us-east-1"); + failure_status(input, AwsIamConnectionConfigStatus::INVALID_ENDPOINT, "invalid_endpoint", + "IPv6 endpoint is rejected"); + input = valid_input("database.example.com", "us-east-1"); + failure_status(input, AwsIamConnectionConfigStatus::INVALID_ENDPOINT, "invalid_endpoint", + "custom CNAME endpoint is rejected"); + const std::string label(63, 'a'); + const std::string endpoint_over_dns_limit = label + "." + label + "." + label + "." + label + + ".us-east-1.rds.amazonaws.com"; + input = valid_input(endpoint_over_dns_limit.c_str(), "us-east-1"); + failure_status(input, AwsIamConnectionConfigStatus::INVALID_ENDPOINT, "invalid_endpoint", + "endpoint exceeding the DNS presentation limit is rejected"); + + input = valid_input("instance.abcdef.us-east-1.rds.amazonaws.com", "us-east-1"); + input.port = 0; + failure_status(input, AwsIamConnectionConfigStatus::UNIX_SOCKET_NOT_ALLOWED, "unix_socket_not_allowed", + "Unix-socket IAM connection is rejected"); + + input = valid_input("instance.abcdef.us-east-1.rds.amazonaws.com", "us-east-1"); + input.use_ssl = false; + failure_status(input, AwsIamConnectionConfigStatus::TLS_REQUIRED, "tls_required", + "IAM connection requires TLS"); + + input = valid_input("instance.abcdef.us-east-1.rds.amazonaws.com", "us-east-1"); + input.ssl_ca.clear(); + input.ssl_capath.clear(); + failure_status(input, AwsIamConnectionConfigStatus::CA_TRUST_REQUIRED, "ca_trust_required", + "IAM connection requires a CA file or CA path"); +} + +void test_malformed_hostgroup_settings_diagnostics_are_redacted() { + MyHGC hostgroup(43); + FILE* captured = tmpfile(); + if (captured == nullptr) { + ok(false, "temporary stderr capture file is available"); + return; + } + fflush(stderr); + const int saved_stderr = dup(STDERR_FILENO); + if (saved_stderr < 0 || dup2(fileno(captured), STDERR_FILENO) < 0) { + if (saved_stderr >= 0) { + close(saved_stderr); + } + fclose(captured); + ok(false, "stderr is redirected for malformed-settings diagnostics"); + return; + } + + init_myhgc_hostgroup_settings("{\"aws_iam_region\":FAKE_AWS_SECRET}", &hostgroup); + fflush(stderr); + dup2(saved_stderr, STDERR_FILENO); + close(saved_stderr); + + std::string diagnostics; + char buffer[256]; + rewind(captured); + while (fgets(buffer, sizeof(buffer), captured) != nullptr) { + diagnostics += buffer; + } + fclose(captured); + ok(diagnostics.find("hostgroup_settings_parse_failed") != std::string::npos && + diagnostics.find("FAKE_AWS_SECRET") == std::string::npos, + "malformed hostgroup settings diagnostics use a redacted parse failure category without sensitive tokens"); + ok(diagnostics.find("hostgroup 43") != std::string::npos, + "malformed hostgroup settings diagnostics identify the hostgroup without payload details"); +} + +void test_hostgroup_region_parser_clears_rejected_values() { + MyHGC hostgroup(42); + init_myhgc_hostgroup_settings("{\"aws_iam_region\":\"us-east-1\"}", &hostgroup); + ok(hostgroup.attributes.aws_iam_region != nullptr && + strcmp(hostgroup.attributes.aws_iam_region, "us-east-1") == 0, + "valid hostgroup IAM region is owned by the hostgroup"); + + const char* invalid_settings[] = { + "{\"aws_iam_region\":1}", + "{\"aws_iam_region\":\"\"}", + "{\"aws_iam_region\":\"us east 1\"}", + }; + for (const char* settings : invalid_settings) { + init_myhgc_hostgroup_settings(settings, &hostgroup); + ok(hostgroup.attributes.aws_iam_region == nullptr, + "rejected hostgroup IAM region does not retain a previous value"); + init_myhgc_hostgroup_settings("{\"aws_iam_region\":\"us-east-1\"}", &hostgroup); + } +} + +} // namespace + +int main() { + plan(0); + test_valid_rds_endpoint_shapes(); + test_validation_failures(); + test_hostgroup_region_parser_clears_rejected_values(); + test_malformed_hostgroup_settings_diagnostics_are_redacted(); + return exit_status(); +} diff --git a/test/tap/tests/unit/aws_iam_connection_secret_unit-t.cpp b/test/tap/tests/unit/aws_iam_connection_secret_unit-t.cpp new file mode 100644 index 0000000000..2d243f6077 --- /dev/null +++ b/test/tap/tests/unit/aws_iam_connection_secret_unit-t.cpp @@ -0,0 +1,526 @@ +/** + * @file aws_iam_connection_secret_unit-t.cpp + * @brief Verify AWS IAM tokens exist only for the backend handshake. + */ + +#include "tap.h" +#include "test_globals.h" +#include "test_init.h" + +#include "proxysql.h" +#include "cpp.h" +#include "Aws_Iam_Provider.h" +#include "MySQL_Backend_Auth.h" +#include "MySQL_Data_Stream.h" +#include "MySQL_Logger.hpp" + +#include + +#include +#include +#include +#include +#include +#include + +#ifndef __linux__ + +int main() { + plan(1); + skip(1, "requires GNU ld --wrap support"); + return exit_status(); +} + +#else + +extern MySQL_HostGroups_Manager *MyHGM; +extern MySQL_Logger *GloMyLogger; + +namespace { + +struct ConnectorObservation { + unsigned int connect_calls { 0 }; + std::string host; + std::string password; + std::string tls_server_name; + bool ssl_enforce_seen { false }; + bool ssl_enforce { false }; + bool ssl_verify_seen { false }; + bool ssl_verify { false }; + bool cleartext_seen { false }; + bool cleartext { false }; + bool reconnect_seen { false }; + bool reconnect { false }; + bool defer_password_copy { false }; +}; + +ConnectorObservation connector; +void *connector_password_copy = nullptr; +size_t connector_password_size = 0; +unsigned int connector_password_cleanse_calls = 0; +bool connector_password_was_zeroed = false; +unsigned int token_cleanse_calls = 0; +size_t token_cleanse_size = 0; +bool token_was_zeroed = false; +const char *connector_retained_password = nullptr; +bool connector_abort_seen = false; +bool token_cleansed_after_abort = false; + +void reset_observations() { + connector = ConnectorObservation {}; + connector_password_copy = nullptr; + connector_password_size = 0; + connector_password_cleanse_calls = 0; + connector_password_was_zeroed = false; + token_cleanse_calls = 0; + token_cleanse_size = 0; + token_was_zeroed = false; + connector_retained_password = nullptr; + connector_abort_seen = false; + token_cleansed_after_abort = false; +} + +bool all_zero(const void *ptr, size_t size) { + const unsigned char *bytes = static_cast(ptr); + for (size_t i = 0; i < size; ++i) { + if (bytes[i] != 0) { + return false; + } + } + return true; +} + +void tracked_token_cleanse(void *ptr, size_t size) { + ++token_cleanse_calls; + token_cleanse_size = size; + OPENSSL_cleanse(ptr, size); + token_was_zeroed = all_zero(ptr, size); + token_cleansed_after_abort = connector_abort_seen; +} + +std::string capture_stderr(const std::function& action) { + FILE *captured = tmpfile(); + if (captured == nullptr) { + BAIL_OUT("failed to create stderr capture file"); + } + fflush(stderr); + const int saved_stderr = dup(STDERR_FILENO); + if (saved_stderr < 0 || dup2(fileno(captured), STDERR_FILENO) < 0) { + BAIL_OUT("failed to redirect stderr"); + } + action(); + fflush(stderr); + dup2(saved_stderr, STDERR_FILENO); + close(saved_stderr); + + std::string output; + char buffer[256]; + rewind(captured); + while (fgets(buffer, sizeof(buffer), captured) != nullptr) { + output += buffer; + } + fclose(captured); + return output; +} + +AwsIamTokenKey token_key() { + return { "orders.cluster-abcdefghijkl.us-east-1.rds.amazonaws.com", + 3306, "us-east-1", "iam_backend" }; +} + +AwsIamTokenResult token_result(const std::string& token, uint64_t generation = 73) { + AwsIamTokenResult result; + result.status = AwsIamStatus::OK; + result.generation = generation; + result.token = SecureString(token, tracked_token_cleanse); + return result; +} + +MySrvC *create_server() { + srv_info_t info; + info.addr = "198.51.100.23"; + info.port = 3306; + info.kind = "aws-iam-connection-secret-unit"; + + srv_opts_t opts; + opts.weigth = 1; + opts.max_conns = 100; + opts.use_ssl = 0; + + MyHGM->wrlock(); + int rc = MyHGM->create_new_server_in_hg(606, info, opts); + MyHGC *hostgroup = MyHGM->MyHGC_find(606); + MyHGM->wrunlock(); + if (rc != 0 || hostgroup == nullptr || hostgroup->mysrvs->cnt() != 1) { + BAIL_OUT("failed to create the IAM unit-test backend"); + } + return hostgroup->mysrvs->idx(0); +} + +class ConnectionFixture { +public: + ConnectionFixture(MySQL_Thread& worker, MySrvC *server) { + session = new MySQL_Session(); + session->thread = &worker; + // Avoid changing process-wide frontend counters in this component test. + session->connections_handler = true; + + frontend_stream = new MySQL_Data_Stream(); + frontend_stream->init(MYDS_FRONTEND, session, -1); + frontend = new MySQL_Connection(); + frontend_stream->myconn = frontend; + frontend->myds = frontend_stream; + session->client_myds = frontend_stream; + frontend->userinfo->set( + const_cast("frontend_user"), + const_cast("frontend-password-must-not-change"), + const_cast("frontend_schema"), nullptr); + + backend_stream = new MySQL_Data_Stream(); + backend_stream->init(MYDS_BACKEND_NOT_CONNECTED, session, -1); + connection = new MySQL_Connection(); + connection->send_quit = false; + connection->parent = server; + connection->myds = backend_stream; + backend_stream->myconn = connection; + MySQL_Backend *backend = session->create_backend(606, backend_stream); + session->mybe = backend; + connection->userinfo->set( + const_cast("iam_backend"), + const_cast("ordinary-backend-password"), + const_cast("orders"), nullptr); + } + + ~ConnectionFixture() { + destroy_connection(); + delete session; + } + + void destroy_connection() { + if (connection != nullptr) { + backend_stream->myconn = nullptr; + connection->myds = nullptr; + delete connection; + connection = nullptr; + } + connector_password_copy = nullptr; + connector_password_size = 0; + } + + const char *frontend_password() const { + return frontend->userinfo->password; + } + + MySQL_Session *session { nullptr }; + MySQL_Data_Stream *frontend_stream { nullptr }; + MySQL_Connection *frontend { nullptr }; + MySQL_Data_Stream *backend_stream { nullptr }; + MySQL_Connection *connection { nullptr }; +}; + +void attach_iam(ConnectionFixture& fixture, const std::string& token) { + fixture.connection->set_backend_auth_type(MySQLBackendAuthType::AWS_IAM); + fixture.connection->attach_aws_iam_token(token_key(), token_result(token)); +} + +void check_cleanup(ConnectionFixture& fixture, const std::string& label, size_t token_size) { + ok(!fixture.connection->has_aws_iam_handshake_secret(), + "%s clears the ProxySQL handshake token", label.c_str()); + ok(fixture.connection->mysql == nullptr || fixture.connection->mysql->passwd == nullptr, + "%s clears MYSQL::passwd", label.c_str()); + ok(token_cleanse_calls == 1 && token_cleanse_size == token_size && token_was_zeroed, + "%s cleanses the complete ProxySQL token buffer", label.c_str()); + ok(connector_password_cleanse_calls == 1 && connector_password_was_zeroed, + "%s cleanses the Connector/C password copy", label.c_str()); +} + +void test_password_mode(MySQL_Thread& worker, MySrvC *server) { + reset_observations(); + ConnectionFixture fixture(worker, server); + fixture.connection->connect_start(); + + ok(fixture.connection->backend_auth_type() == MySQLBackendAuthType::PASSWORD, + "password authentication remains the default backend mode"); + ok(connector.password == "ordinary-backend-password", + "password mode still passes userinfo->password to Connector/C"); + ok(std::strcmp(fixture.frontend_password(), "frontend-password-must-not-change") == 0, + "password-mode backend startup does not change the frontend password"); + connector_password_copy = nullptr; +} + +void test_sha1_password_mode(MySQL_Thread& worker, MySrvC *server) { + reset_observations(); + ConnectionFixture fixture(worker, server); + fixture.connection->userinfo->set( + const_cast("iam_backend"), + const_cast("*0123456789012345678901234567890123456789"), + const_cast("orders"), + const_cast("sha1-stage-one-secret")); + fixture.connection->connect_start(); + + ok(connector.password == "sha1-stage-one-secret", + "password mode preserves the existing SHA1 credential path"); + connector_password_copy = nullptr; +} + +void test_iam_handshake_and_explicit_clear(MySQL_Thread& worker, MySrvC *server) { + reset_observations(); + ConnectionFixture fixture(worker, server); + const std::string token = "orders.cluster:3306/?Action=connect&" + std::string(2048, 'x'); + attach_iam(fixture, token); + connector.defer_password_copy = true; + fixture.connection->connect_start(); + + ok(fixture.connection->backend_auth_type() == MySQLBackendAuthType::AWS_IAM, + "IAM authentication is stored independently from userinfo"); + ok(connector.password == token, + "IAM startup passes the full token instead of userinfo->password"); + ok(connector.host == "198.51.100.23", + "Connector/C keeps using the DNS-cache transport IP"); + ok(connector.tls_server_name == token_key().endpoint, + "TLS verification uses the configured RDS endpoint identity"); + ok(connector.ssl_enforce_seen && connector.ssl_enforce, + "IAM startup enforces TLS"); + ok(connector.ssl_verify_seen && connector.ssl_verify, + "IAM startup enables server-certificate verification"); + ok(connector.cleartext_seen && connector.cleartext, + "IAM startup enables the cleartext authentication plugin"); + ok(connector.reconnect_seen && !connector.reconnect, + "IAM startup explicitly disables Connector/C auto-reconnect"); + ok(std::strcmp(fixture.frontend_password(), "frontend-password-must-not-change") == 0, + "IAM backend startup never changes the frontend password"); + + ok(connector_retained_password != nullptr, + "connector fake retains the exact IAM input pointer across the async yield"); + fixture.connection->clear_aws_iam_handshake_secret(); + ok(connector_abort_seen, + "explicit clear aborts the connector coroutine before invalidating its password pointer"); + ok(token_cleanse_calls == 1 && token_cleanse_size == token.size() && token_was_zeroed, + "explicit clear eventually cleanses the complete ProxySQL token buffer"); + ok(token_cleansed_after_abort, + "explicit clear cleanses the token only after aborting the connector coroutine"); + ok(fixture.connection->mysql == nullptr, + "explicit clear releases the aborted Connector/C handle"); +} + +void test_iam_to_password_transition(MySQL_Thread& worker, MySrvC *server) { + reset_observations(); + ConnectionFixture fixture(worker, server); + const std::string token = "transition-token"; + attach_iam(fixture, token); + fixture.connection->set_backend_auth_type(MySQLBackendAuthType::PASSWORD); + + ok(!fixture.connection->has_aws_iam_handshake_secret() && token_cleanse_calls == 1, + "IAM to password transition cleanses and drops the IAM handshake secret"); + fixture.connection->connect_start(); + ok(connector.password == "ordinary-backend-password", + "IAM to password transition restores ordinary password authentication"); + fixture.connection->async_state_machine = ASYNC_CONNECT_END; + fixture.connection->ret_mysql = fixture.connection->mysql; + fixture.connection->handler(0); + ok(fixture.connection->mysql->passwd != nullptr && + std::strcmp(fixture.connection->mysql->passwd, "ordinary-backend-password") == 0, + "password-mode terminal cleanup preserves Connector/C's ordinary password state"); + ok(connector_password_cleanse_calls == 0, + "password-mode terminal cleanup never treats the ordinary password as an IAM token"); + connector_password_copy = nullptr; +} + +void test_unix_socket_rejected(MySQL_Thread& worker, MySrvC *tcp_server) { + reset_observations(); + MySrvC unix_server( + const_cast("/var/lib/proxysql/proxysql-iam-unit.sock"), 0, 0, 1, + MYSQL_SERVER_STATUS_ONLINE, 0, 100, 0, 0, 0, + const_cast("aws-iam-unix-unit")); + unix_server.myhgc = tcp_server->myhgc; + ConnectionFixture fixture(worker, &unix_server); + const std::string token = "unix-socket-token"; + attach_iam(fixture, token); + fixture.connection->handler(0); + ok(connector.connect_calls == 0, + "IAM mode rejects a Unix socket before Connector/C startup"); + ok(fixture.connection->ret_mysql == nullptr && fixture.connection->async_exit_status == 0, + "Unix-socket rejection is a terminal connection error"); + ok(std::strcmp(fixture.frontend_password(), "frontend-password-must-not-change") == 0, + "Unix-socket rejection does not change the frontend password"); + ok(token_cleanse_calls == 1 && token_was_zeroed, + "Unix-socket rejection automatically cleanses the handshake token"); +} + +void test_missing_token_rejected(MySQL_Thread& worker, MySrvC *server) { + reset_observations(); + ConnectionFixture fixture(worker, server); + fixture.connection->set_backend_auth_type(MySQLBackendAuthType::AWS_IAM); + fixture.connection->handler(0); + ok(connector.connect_calls == 0, + "IAM mode rejects a missing token before Connector/C startup"); + ok(fixture.connection->async_state_machine == ASYNC_CONNECT_FAILED, + "missing IAM token reaches the terminal failed state"); +} + +void test_terminal_cleanup(MySQL_Thread& worker, MySrvC *server, + MDB_ASYNC_ST terminal_state, bool success, const char *label) +{ + reset_observations(); + ConnectionFixture fixture(worker, server); + const std::string token = std::string(label) + "-terminal-token"; + attach_iam(fixture, token); + fixture.connection->connect_start(); + fixture.connection->async_state_machine = terminal_state; + fixture.connection->ret_mysql = success ? fixture.connection->mysql : nullptr; + fixture.backend_stream->wait_until = 1; + worker.curtime = 2; + fixture.connection->handler(0); + + check_cleanup(fixture, label, token.size()); + const MDB_ASYNC_ST expected = terminal_state == ASYNC_CONNECT_END + ? (success ? ASYNC_CONNECT_SUCCESSFUL : ASYNC_CONNECT_FAILED) + : ASYNC_CONNECT_TIMEOUT; + ok(fixture.connection->async_state_machine == expected, + "%s preserves the expected terminal state", label); +} + +void test_terminal_error_is_redacted(MySQL_Thread& worker, MySrvC *server) { + reset_observations(); + ConnectionFixture fixture(worker, server); + const std::string token = "sensitive-token-not-for-logs"; + attach_iam(fixture, token); + fixture.connection->connect_start(); + std::snprintf(fixture.connection->mysql->net.last_error, + sizeof(fixture.connection->mysql->net.last_error), + "backend reflected credential: %s", token.c_str()); + fixture.connection->mysql->net.last_errno = 1045; + fixture.connection->async_state_machine = ASYNC_CONNECT_END; + fixture.connection->ret_mysql = nullptr; + + const std::string log = capture_stderr([&fixture]() { fixture.connection->handler(0); }); + ok(log.find(token) == std::string::npos && + log.find("backend reflected credential") == std::string::npos, + "IAM terminal failure never logs raw backend or credential text"); + ok(log.find("details redacted") != std::string::npos, + "IAM terminal failure emits a fixed redacted diagnostic"); +} + +void test_destructor_cleanup(MySQL_Thread& worker, MySrvC *server) { + reset_observations(); + ConnectionFixture fixture(worker, server); + const std::string token = "destructor-terminal-token"; + attach_iam(fixture, token); + fixture.connection->connect_start(); + fixture.destroy_connection(); + + ok(token_cleanse_calls == 1 && token_cleanse_size == token.size() && token_was_zeroed, + "destructor cleanses the complete ProxySQL token buffer"); + ok(connector_password_cleanse_calls == 1 && connector_password_was_zeroed, + "destructor cleanses the Connector/C password copy"); +} + +} // namespace + +extern "C" { + +int __real_mysql_options(MYSQL *, enum mysql_option, const void *); +void __real_OPENSSL_cleanse(void *, size_t); +void __real_mysql_close_no_command(MYSQL *); + +int __wrap_mysql_options(MYSQL *mysql, enum mysql_option option, const void *arg) { + switch (option) { + case MARIADB_OPT_TLS_SERVER_NAME: + connector.tls_server_name = arg != nullptr ? static_cast(arg) : ""; + break; + case MYSQL_OPT_SSL_ENFORCE: + connector.ssl_enforce_seen = true; + connector.ssl_enforce = arg != nullptr && *static_cast(arg) != 0; + break; + case MYSQL_OPT_SSL_VERIFY_SERVER_CERT: + connector.ssl_verify_seen = true; + connector.ssl_verify = arg != nullptr && *static_cast(arg) != 0; + break; + case MYSQL_ENABLE_CLEARTEXT_PLUGIN: + connector.cleartext_seen = true; + connector.cleartext = arg != nullptr && *static_cast(arg) != 0; + break; + case MYSQL_OPT_RECONNECT: + connector.reconnect_seen = true; + connector.reconnect = arg != nullptr && *static_cast(arg) != 0; + break; + default: + break; + } + return __real_mysql_options(mysql, option, arg); +} + +int __wrap_mysql_real_connect_start(MYSQL **ret, MYSQL *mysql, const char *host, + const char *, const char *password, const char *, unsigned int port, + const char *, unsigned long) +{ + ++connector.connect_calls; + connector.host = host != nullptr ? host : ""; + connector.password = password != nullptr ? password : ""; + connector_retained_password = password; + *ret = nullptr; + + if (!connector.defer_password_copy) { + mysql->passwd = strdup(password != nullptr ? password : ""); + } + mysql->host = strdup(host != nullptr ? host : ""); + mysql->port = port; + connector_password_copy = mysql->passwd; + connector_password_size = mysql->passwd != nullptr ? std::strlen(mysql->passwd) : 0; + return MYSQL_WAIT_READ; +} + +void __wrap_mysql_close_no_command(MYSQL *mysql) { + if (connector_retained_password != nullptr) { + connector_abort_seen = true; + connector_retained_password = nullptr; + } + __real_mysql_close_no_command(mysql); +} + +void __wrap_OPENSSL_cleanse(void *ptr, size_t size) { + __real_OPENSSL_cleanse(ptr, size); + if (ptr == connector_password_copy && size == connector_password_size) { + ++connector_password_cleanse_calls; + connector_password_was_zeroed = all_zero(ptr, size); + } +} + +} // extern "C" + +int main() { + plan(47); + if (test_init_minimal() != 0 || test_init_query_processor() != 0 || + test_init_hostgroups() != 0) { + BAIL_OUT("failed to initialize the unit-test component globals"); + } + GloMyLogger = new MySQL_Logger(); + + MySrvC *server = create_server(); + { + MySQL_Thread worker; + if (!worker.init()) { + BAIL_OUT("MySQL_Thread::init() failed"); + } + test_password_mode(worker, server); + test_sha1_password_mode(worker, server); + test_iam_handshake_and_explicit_clear(worker, server); + test_iam_to_password_transition(worker, server); + test_unix_socket_rejected(worker, server); + test_missing_token_rejected(worker, server); + test_terminal_cleanup(worker, server, ASYNC_CONNECT_END, true, "terminal success"); + test_terminal_cleanup(worker, server, ASYNC_CONNECT_END, false, "terminal error"); + test_terminal_error_is_redacted(worker, server); + test_terminal_cleanup(worker, server, ASYNC_CONNECT_TIMEOUT, false, "timeout"); + test_destructor_cleanup(worker, server); + } + + delete GloMyLogger; + GloMyLogger = nullptr; + test_cleanup_hostgroups(); + test_cleanup_query_processor(); + test_cleanup_minimal(); + return exit_status(); +} + +#endif // __linux__ diff --git a/test/tap/tests/unit/aws_iam_failure_unit-t.cpp b/test/tap/tests/unit/aws_iam_failure_unit-t.cpp new file mode 100644 index 0000000000..5a060c100e --- /dev/null +++ b/test/tap/tests/unit/aws_iam_failure_unit-t.cpp @@ -0,0 +1,618 @@ +/** + * @file aws_iam_failure_unit-t.cpp + * @brief IAM-only backend authentication retry and redaction regressions. + */ + +#include "tap.h" +#include "test_globals.h" +#include "test_init.h" + +#include "proxysql.h" +#include "cpp.h" +#include "Aws_Iam_Provider.h" +#include "MySQL_Authentication.hpp" +#include "MySQL_Data_Stream.h" +#include "MySQL_HostGroups_Manager.h" +#include "MySQL_Logger.hpp" +#include "mysqld_error.h" +#include "errmsg.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +using namespace std::chrono_literals; + +extern MySQL_Authentication *GloMyAuth; +extern MySQL_HostGroups_Manager *MyHGM; +extern MySQL_Logger *GloMyLogger; + +namespace { + +constexpr int kHostgroup = 909; +constexpr const char *kEndpoint = + "failure.cluster-abcdefghijkl.us-east-1.rds.amazonaws.com"; +constexpr const char *kRegion = "us-east-1"; +constexpr const char *kIamUser = "iam_failure_backend"; +constexpr const char *kPasswordUser = "password_failure_backend"; +constexpr const char *kTokenOne = "FAKE_IAM_TOKEN_GENERATION_ONE"; +constexpr const char *kTokenTwo = "FAKE_IAM_TOKEN_GENERATION_TWO"; +constexpr const char *kBackendText = + "backend reflected AKIAFAKEACCESSKEY and FAKE_SESSION_TOKEN"; +MySrvC *failure_server = nullptr; +std::atomic blocking_source_destroyed { false }; + +struct ConnectOutcome { + unsigned int error; + const char *message; + bool pending; +}; + +std::vector connect_outcomes; +size_t connect_outcome_index = 0; +unsigned int connector_calls = 0; +uint16_t client_error_code = 0; +std::string client_error_state; +std::string client_error_message; + +AwsIamTokenResult token_result(uint64_t generation, const char *token) { + AwsIamTokenResult result; + result.status = AwsIamStatus::OK; + result.generation = generation; + result.token = SecureString(token); + return result; +} + +class FakeTokenSource final : public AwsIamTokenSource { +public: + AwsIamRequestHandle request(const AwsIamTokenKey& key, uint64_t opaque_id, + std::weak_ptr sink) override { + keys.push_back(key); + const size_t index = keys.size() - 1; + AwsIamCompletion completion; + completion.opaque_id = opaque_id; + completion.result = token_result( + index < generations.size() ? generations[index] : generations.back(), + index == 0 ? kTokenOne : kTokenTwo); + if (auto target = sink.lock()) target->post(std::move(completion)); + return { next_handle++ }; + } + + AwsIamTokenResult request_blocking(const AwsIamTokenKey&, + std::chrono::steady_clock::time_point) override { + return {}; + } + + void cancel(AwsIamRequestHandle) override {} + + void invalidate(const AwsIamTokenKey& key, uint64_t generation) override { + if (block_invalidation) { + std::unique_lock lock(invalidation_mutex); + invalidation_entered = true; + invalidation_cv.notify_all(); + invalidation_cv.wait(lock, [this] { return shutdown_started; }); + const bool shutdown_returned = invalidation_cv.wait_for(lock, 20ms, + [this] { return shutdown_completed; }); + retirement_waited_for_invalidation = !shutdown_returned && + !blocking_source_destroyed.load(std::memory_order_acquire); + } + invalidated_keys.push_back(key); + invalidated_generations.push_back(generation); + if (cached_key == key && cached_generation == generation) { + cached_generation = 0; + } + } + + void record_backend_connection(bool success) override { + if (success) ++backend_successes; + else { + ++backend_failures; + if (corrupt_port_on_failure != nullptr) { + corrupt_port_on_failure->aws_iam_connect_token_key.port = 0; + } + } + } + void record_waiting_session(bool waiting) override { + if (waiting) ++waiting_sessions; + else if (waiting_sessions != 0) --waiting_sessions; + } + + AwsIamStatsSnapshot snapshot() const override { + AwsIamStatsSnapshot result; + result.waiting_sessions = waiting_sessions; + return result; + } + + bool wait_for_invalidation() { + std::unique_lock lock(invalidation_mutex); + return invalidation_cv.wait_for(lock, 1s, + [this] { return invalidation_entered; }); + } + + void mark_shutdown_completed() { + std::lock_guard lock(invalidation_mutex); + shutdown_completed = true; + invalidation_cv.notify_all(); + } + + void mark_shutdown_started() { + std::lock_guard lock(invalidation_mutex); + shutdown_started = true; + invalidation_cv.notify_all(); + } + + std::vector generations { 1, 2 }; + uint64_t next_handle { 1 }; + std::vector keys; + std::vector invalidated_keys; + std::vector invalidated_generations; + AwsIamTokenKey cached_key; + uint64_t cached_generation { 0 }; + unsigned int backend_successes { 0 }; + unsigned int backend_failures { 0 }; + uint64_t waiting_sessions { 0 }; + MySQL_Session *corrupt_port_on_failure { nullptr }; + bool block_invalidation { false }; + std::mutex invalidation_mutex; + std::condition_variable invalidation_cv; + bool invalidation_entered { false }; + bool shutdown_started { false }; + bool shutdown_completed { false }; + bool retirement_waited_for_invalidation { false }; +}; + +void mark_blocking_source_destroyed(AwsIamTokenSource *) { + blocking_source_destroyed.store(true, std::memory_order_release); +} + +class ScopedPublishedTokenSource { +public: + explicit ScopedPublishedTokenSource(AwsIamTokenSource *source) { + publish_global_aws_iam_token_source(source); + } + ~ScopedPublishedTokenSource() { publish_global_aws_iam_token_source(nullptr); } + + ScopedPublishedTokenSource(const ScopedPublishedTokenSource&) = delete; + ScopedPublishedTokenSource& operator=(const ScopedPublishedTokenSource&) = delete; +}; + +bool add_backend_user(const char *username, const char *password, + const char *attributes) { + return GloMyAuth->add( + (char *)username, (char *)password, USERNAME_BACKEND, + false, 0, (char *)"", false, false, false, 100, + (char *)attributes, (char *)""); +} + +void add_server() { + srv_info_t info; + info.addr = kEndpoint; + info.port = 3306; + info.kind = "aws-iam-failure-unit"; + srv_opts_t opts; + opts.weigth = 1; + opts.max_conns = 100; + opts.use_ssl = 1; + MyHGM->wrlock(); + const int rc = MyHGM->create_new_server_in_hg(kHostgroup, info, opts); + MyHGC *hostgroup = MyHGM->MyHGC_find(kHostgroup); + MyHGM->wrunlock(); + if (rc != 0 || hostgroup == nullptr) BAIL_OUT("failed to create failure fixture"); + free(hostgroup->attributes.aws_iam_region); + hostgroup->attributes.aws_iam_region = strdup(kRegion); + failure_server = hostgroup->mysrvs->idx(0); +} + +MySQL_Connection *established_iam_connection(int fd) { + MySQL_Connection *connection = new MySQL_Connection(); + connection->mysql = mysql_init(nullptr); + if (connection->mysql == nullptr) BAIL_OUT("mysql_init() failed for retry pool fixture"); + connection->ret_mysql = connection->mysql; + connection->mysql->charset = mariadb_get_charset_by_name("utf8mb4"); + if (connection->mysql->charset == nullptr) BAIL_OUT("charset fixture failed"); + connection->parent = failure_server; + connection->userinfo->set( + const_cast(kIamUser), const_cast(""), + const_cast("orders"), nullptr); + connection->set_backend_auth_type(MySQLBackendAuthType::AWS_IAM); + connection->healthy = true; + connection->reusable = true; + connection->send_quit = false; + connection->fd = fd; + connection->async_state_machine = ASYNC_IDLE; + return connection; +} + +void destroy_used(MySQL_Connection *connection) { + if (connection == nullptr) return; + connection->send_quit = false; + MyHGM->destroy_MyConn_from_pool(connection); +} + +std::string capture_stderr(const std::function& action) { + FILE *captured = tmpfile(); + if (captured == nullptr) BAIL_OUT("tmpfile() failed"); + fflush(stderr); + const int saved = dup(STDERR_FILENO); + if (saved < 0 || dup2(fileno(captured), STDERR_FILENO) < 0) { + BAIL_OUT("stderr redirect failed"); + } + action(); + fflush(stderr); + dup2(saved, STDERR_FILENO); + close(saved); + std::string output; + char buffer[256]; + rewind(captured); + while (fgets(buffer, sizeof(buffer), captured) != nullptr) output += buffer; + fclose(captured); + return output; +} + +class SessionFixture { +public: + SessionFixture(MySQL_Thread& worker, const char *username) : worker_(worker) { + session = new MySQL_Session(); + session->thread = &worker_; + session->connections_handler = true; + frontend_stream = new MySQL_Data_Stream(); + frontend_stream->init(MYDS_FRONTEND, session, -1); + frontend = new MySQL_Connection(); + frontend_stream->attach_connection(frontend); + frontend_stream->myprot.init(&frontend_stream, frontend->userinfo, session); + session->client_myds = frontend_stream; + frontend->userinfo->set( + const_cast(username), const_cast("ordinary-password"), + const_cast("orders"), nullptr); + + session->mybe = session->create_backend(kHostgroup); + session->current_hostgroup = kHostgroup; + session->default_hostgroup = kHostgroup; + session->CurrentQuery.start_time = worker_.curtime; + session->previous_status.push(PROCESSING_QUERY); + session->set_status(CONNECTING_SERVER); + } + + ~SessionFixture() { delete session; } + + int run() { + session->to_process = 1; + return session->handler(); + } + + MySQL_Data_Stream *backend() const { + return session != nullptr && session->mybe != nullptr + ? session->mybe->server_myds : nullptr; + } + + MySQL_Thread& worker_; + MySQL_Session *session { nullptr }; + MySQL_Data_Stream *frontend_stream { nullptr }; + MySQL_Connection *frontend { nullptr }; +}; + +void reset_connector(std::initializer_list outcomes) { + connect_outcomes.assign(outcomes); + connect_outcome_index = 0; + connector_calls = 0; +} + +void start_iam_attempt(SessionFixture& fixture, MySQL_Thread& worker) { + fixture.run(); + worker.drain_aws_iam_completions(); + fixture.run(); +} + +std::string process_terminal_connect(SessionFixture& fixture) { + return capture_stderr([&fixture] { fixture.run(); }); +} + +void test_first_1045_retries_once(MySQL_Thread& worker) { + FakeTokenSource source; + ScopedPublishedTokenSource published_source(&source); + reset_connector({ { ER_ACCESS_DENIED_ERROR, kBackendText, false }, + { ER_ACCESS_DENIED_ERROR, kBackendText, false } }); + SessionFixture fixture(worker, kIamUser); + fixture.backend()->connect_retries_on_failure = 3; + start_iam_attempt(fixture, worker); + process_terminal_connect(fixture); + + ok(source.invalidated_generations.size() == 1 && + source.invalidated_generations[0] == 1 && + source.invalidated_keys[0] == source.keys[0], + "the first IAM 1045 conditionally invalidates its exact key and generation"); + ok(source.keys.size() == 2 && fixture.session->status == WAITING_AWS_IAM_TOKEN, + "the first IAM 1045 acquires exactly one fresh token"); + ok(fixture.backend() != nullptr && + fixture.backend()->connect_retries_on_failure == 0, + "the IAM fresh-token retry disables the ordinary multi-server retry budget"); + + if (source.keys.size() == 2 && fixture.session->status == WAITING_AWS_IAM_TOKEN) { + worker.drain_aws_iam_completions(); + client_error_code = 0; + client_error_state.clear(); + client_error_message.clear(); + const std::string log = capture_stderr([&fixture] { fixture.run(); }); + ok(source.keys.size() == 2 && source.invalidated_generations.size() == 1 && + fixture.session->status == WAITING_CLIENT_DATA, + "a repeated fresh-token 1045 is terminal with no third attempt"); + ok(client_error_code == 9002 && client_error_state == "HY000" && + client_error_message == "Unable to connect to backend" && + client_error_message.find(kEndpoint) == std::string::npos && + client_error_message.find(kRegion) == std::string::npos && + client_error_message.find(kBackendText) == std::string::npos && + client_error_message.find(kTokenOne) == std::string::npos && + client_error_message.find(kTokenTwo) == std::string::npos && + client_error_message.find("AKIAFAKEACCESSKEY") == std::string::npos && + client_error_message.find("FAKE_SESSION_TOKEN") == std::string::npos, + "the IAM client error is fixed and contains no backend, AWS, or token detail"); + ok(log.find("user='iam_failure_backend'") != std::string::npos && + log.find("hostgroup=909") != std::string::npos && + log.find(kEndpoint) != std::string::npos && + log.find("region='us-east-1'") != std::string::npos && + log.find("category='backend_auth_rejected'") != std::string::npos && + log.find("code=''") != std::string::npos && + log.find("request_id=''") != std::string::npos && + log.find("clock") != std::string::npos, + "the repeated 1045 diagnostic is redacted and includes a clock-skew hint"); + ok(log.find(kBackendText) == std::string::npos && + log.find(kTokenOne) == std::string::npos && + log.find(kTokenTwo) == std::string::npos && + log.find("AKIAFAKEACCESSKEY") == std::string::npos && + log.find("FAKE_SESSION_TOKEN") == std::string::npos, + "the operator diagnostic contains no credential or backend error text"); + } else { + ok(false, "a repeated fresh-token 1045 is terminal with no third attempt"); + ok(false, "the IAM client error is fixed and contains no backend, AWS, or token detail"); + ok(false, "the repeated 1045 diagnostic is redacted and includes a clock-skew hint"); + ok(false, "the operator diagnostic contains no credential or backend error text"); + } +} + +void test_stale_generation_cannot_evict_newer(MySQL_Thread& worker) { + FakeTokenSource source; + source.generations = { 41, 42 }; + ScopedPublishedTokenSource published_source(&source); + reset_connector({ { ER_ACCESS_DENIED_ERROR, kBackendText, false } }); + SessionFixture fixture(worker, kIamUser); + fixture.backend()->connect_retries_on_failure = 3; + start_iam_attempt(fixture, worker); + source.cached_key = source.keys[0]; + source.cached_generation = 42; + process_terminal_connect(fixture); + ok(source.invalidated_generations.size() == 1 && + source.invalidated_generations[0] == 41 && source.cached_generation == 42, + "a delayed generation-N 1045 cannot evict the cached generation N+1"); +} + +void test_transport_failure_does_not_invalidate(MySQL_Thread& worker) { + FakeTokenSource source; + ScopedPublishedTokenSource published_source(&source); + reset_connector({ { CR_SSL_CONNECTION_ERROR, "TLS transport failed", false } }); + SessionFixture fixture(worker, kIamUser); + fixture.backend()->connect_retries_on_failure = 3; + start_iam_attempt(fixture, worker); + process_terminal_connect(fixture); + ok(source.invalidated_generations.empty() && source.keys.size() == 1 && + fixture.session->status == WAITING_CLIENT_DATA, + "an IAM TLS or transport failure is terminal without token invalidation or retry"); +} + +void test_password_1045_keeps_normal_retry(MySQL_Thread& worker) { + FakeTokenSource source; + ScopedPublishedTokenSource published_source(&source); + reset_connector({ { ER_ACCESS_DENIED_ERROR, "ordinary password rejected", false }, + { 0, "", true } }); + SessionFixture fixture(worker, kPasswordUser); + fixture.backend()->connect_retries_on_failure = 1; + fixture.run(); + if (fixture.backend() == nullptr) BAIL_OUT("password fixture did not acquire a backend"); + fixture.run(); + ok(source.keys.empty() && source.invalidated_generations.empty() && + connector_calls == 2 && fixture.backend()->connect_retries_on_failure == 0 && + fixture.session->status == CONNECTING_SERVER, + "password-mode 1045 retains the existing ordinary connection retry behavior"); +} + +void test_fresh_retry_bypasses_local_and_global_idle_iam(MySQL_Thread& worker) { + FakeTokenSource source; + ScopedPublishedTokenSource published_source(&source); + SessionFixture fixture(worker, kIamUser); + MySQL_Connection *local = established_iam_connection(601); + MySQL_Connection *global = established_iam_connection(602); + failure_server->ConnectionsUsed->add(local); + worker.push_MyConn_local(local); + failure_server->ConnectionsFree->add(global); + fixture.session->aws_iam_fresh_token_retry_attempted = true; + fixture.session->previous_status.pop(); + fixture.session->previous_status.push(WAITING_CLIENT_DATA); + + fixture.frontend_stream->active = 0; + fixture.backend()->active = 0; + fixture.run(); + fixture.frontend_stream->active = 1; + fixture.backend()->active = 1; + MySQL_Connection *selected = fixture.backend()->myconn; + ok(selected != nullptr && selected->fd == -1 && + local->myds == nullptr && + failure_server->ConnectionsFree->conns_length() == 1 && + source.keys.size() == 1 && fixture.session->status == WAITING_AWS_IAM_TOKEN, + "an IAM fresh-token retry bypasses compatible local and global idle connections and starts a new handshake"); + + if (fixture.backend()->myconn != nullptr) { + fixture.backend()->destroy_MySQL_Connection_From_Pool(false); + } + local = worker.get_MyConn_local( + kHostgroup, fixture.session, nullptr, 0, -1, + MySQLBackendAuthType::AWS_IAM); + destroy_used(local); + global = MyHGM->get_MyConn_from_pool( + kHostgroup, fixture.session, false, nullptr, 0, -1, + MySQLBackendAuthType::AWS_IAM); + destroy_used(global); +} + +void test_pooled_success_clears_latch_for_later_1045(MySQL_Thread& worker) { + FakeTokenSource source; + ScopedPublishedTokenSource published_source(&source); + reset_connector({ { ER_ACCESS_DENIED_ERROR, kBackendText, false } }); + SessionFixture fixture(worker, kIamUser); + MySQL_Connection *pooled = established_iam_connection(603); + failure_server->ConnectionsUsed->add(pooled); + fixture.backend()->attach_connection(pooled); + fixture.session->aws_iam_fresh_token_retry_attempted = true; + fixture.session->previous_status.pop(); + fixture.session->previous_status.push(WAITING_CLIENT_DATA); + fixture.frontend_stream->active = 0; + fixture.backend()->active = 0; + fixture.run(); + fixture.frontend_stream->active = 1; + fixture.backend()->active = 1; + fixture.backend()->destroy_MySQL_Connection_From_Pool(false); + fixture.session->previous_status.push(PROCESSING_QUERY); + fixture.session->set_status(CONNECTING_SERVER); + start_iam_attempt(fixture, worker); + process_terminal_connect(fixture); + + ok(source.invalidated_generations.size() == 1 && + source.invalidated_generations[0] == 1 && source.keys.size() == 2 && + fixture.session->aws_iam_fresh_token_retry_attempted && + fixture.session->status == WAITING_AWS_IAM_TOKEN, + "a pooled IAM acquisition success clears the retry latch so a later independent 1045 gets one fresh attempt"); +} + +void test_missing_port_cannot_retry_or_invalidate(MySQL_Thread& worker) { + FakeTokenSource source; + ScopedPublishedTokenSource published_source(&source); + reset_connector({ { ER_ACCESS_DENIED_ERROR, kBackendText, false } }); + SessionFixture fixture(worker, kIamUser); + source.corrupt_port_on_failure = fixture.session; + start_iam_attempt(fixture, worker); + ok(source.invalidated_generations.empty() && source.keys.size() == 1 && + fixture.session->status == WAITING_CLIENT_DATA, + "an IAM 1045 with a missing key port is terminal without invalidation or retry"); +} + +void test_provider_retirement_waits_for_1045_invalidation(MySQL_Thread& worker) { + blocking_source_destroyed.store(false, std::memory_order_release); + auto *source = new FakeTokenSource(); + source->block_invalidation = true; + void *module_handle = dlopen(nullptr, RTLD_NOW | RTLD_LOCAL); + if (module_handle == nullptr || !install_global_aws_iam_token_source( + source, mark_blocking_source_destroyed, module_handle)) { + BAIL_OUT("failed to install blocking IAM source"); + } + + reset_connector({ { ER_ACCESS_DENIED_ERROR, kBackendText, false } }); + SessionFixture fixture(worker, kIamUser); + fixture.run(); + worker.drain_aws_iam_completions(); + auto shutdown = std::async(std::launch::async, [source] { + if (!source->wait_for_invalidation()) return false; + source->mark_shutdown_started(); + shutdown_global_aws_iam_token_source(); + source->mark_shutdown_completed(); + return true; + }); + fixture.run(); + const bool invalidation_triggered_shutdown = shutdown.get(); + ok(invalidation_triggered_shutdown && + source->retirement_waited_for_invalidation, + "provider retirement cannot destroy or unload a source during 1045 invalidation"); + ok(blocking_source_destroyed.load(std::memory_order_acquire) && + source->invalidated_generations.size() == 1, + "provider retirement completes after the retained invalidation lease drains"); + delete source; +} + +} // namespace + +extern "C" { + +bool __real__ZN14MySQL_Protocol16generate_pkt_ERREbPPvPjhtPKcS4_b( + MySQL_Protocol *, bool, void **, unsigned int *, uint8_t, uint16_t, + const char *, const char *, bool); + +bool __wrap__ZN14MySQL_Protocol16generate_pkt_ERREbPPvPjhtPKcS4_b( + MySQL_Protocol *protocol, bool send, void **ptr, unsigned int *len, + uint8_t sequence_id, uint16_t error_code, const char *sql_state, + const char *sql_message, bool track) { + client_error_code = error_code; + client_error_state = sql_state != nullptr ? sql_state : ""; + client_error_message = sql_message != nullptr ? sql_message : ""; + return __real__ZN14MySQL_Protocol16generate_pkt_ERREbPPvPjhtPKcS4_b( + protocol, send, ptr, len, sequence_id, error_code, sql_state, + sql_message, track); +} + +int __wrap_mysql_real_connect_start(MYSQL **ret, MYSQL *mysql, const char *host, + const char *, const char *password, const char *, unsigned int port, + const char *, unsigned long) { + ++connector_calls; + const ConnectOutcome outcome = connect_outcome_index < connect_outcomes.size() + ? connect_outcomes[connect_outcome_index++] + : ConnectOutcome { CR_CONNECTION_ERROR, "unexpected connector call", false }; + mysql->host = strdup(host != nullptr ? host : ""); + mysql->passwd = strdup(password != nullptr ? password : ""); + mysql->port = port; + mysql->net.last_errno = outcome.error; + std::snprintf(mysql->net.last_error, sizeof(mysql->net.last_error), "%s", + outcome.message != nullptr ? outcome.message : ""); + std::snprintf(mysql->net.sqlstate, sizeof(mysql->net.sqlstate), "%s", + outcome.error == ER_ACCESS_DENIED_ERROR ? "28000" : "HY000"); + *ret = nullptr; + return outcome.pending ? MYSQL_WAIT_READ : 0; +} + +} // extern "C" + +int main() { + plan(15); + if (test_init_minimal() != 0 || test_init_auth() != 0 || + test_init_query_processor() != 0 || test_init_hostgroups() != 0) { + BAIL_OUT("failed to initialize unit-test globals"); + } + GloMyLogger = new MySQL_Logger(); + if (!add_backend_user(kIamUser, "", "{\"backend_auth\":{\"type\":\"aws_iam\"}}") || + !add_backend_user(kPasswordUser, "ordinary-password", "")) { + BAIL_OUT("failed to load backend user fixtures"); + } + add_server(); + + { + MySQL_Thread worker; + if (!worker.init()) BAIL_OUT("MySQL_Thread::init() failed"); + free(mysql_thread___ssl_p2s_ca); + mysql_thread___ssl_p2s_ca = strdup("/unit/fake-ca.pem"); + worker.curtime = 10000000; + test_first_1045_retries_once(worker); + test_stale_generation_cannot_evict_newer(worker); + test_transport_failure_does_not_invalidate(worker); + test_password_1045_keeps_normal_retry(worker); + test_fresh_retry_bypasses_local_and_global_idle_iam(worker); + test_pooled_success_clears_latch_for_later_1045(worker); + test_missing_port_cannot_retry_or_invalidate(worker); + test_provider_retirement_waits_for_1045_invalidation(worker); + } + GloAwsIamTokenSource = nullptr; + + delete GloMyLogger; + GloMyLogger = nullptr; + test_cleanup_hostgroups(); + test_cleanup_query_processor(); + test_cleanup_auth(); + test_cleanup_minimal(); + return exit_status(); +} diff --git a/test/tap/tests/unit/aws_iam_kill_helper_unit-t.cpp b/test/tap/tests/unit/aws_iam_kill_helper_unit-t.cpp new file mode 100644 index 0000000000..220e99aa57 --- /dev/null +++ b/test/tap/tests/unit/aws_iam_kill_helper_unit-t.cpp @@ -0,0 +1,404 @@ +/** + * @file aws_iam_kill_helper_unit-t.cpp + * @brief Detached IAM query/connection kill credential regressions. + */ + +#include "tap.h" +#include "test_globals.h" +#include "test_init.h" + +#include "proxysql.h" +#include "cpp.h" +#include "Aws_Iam_Provider.h" +#include "MySQL_HostGroups_Manager.h" +#include "MySQL_Logger.hpp" + +#include + +#include +#include +#include +#include +#include +#include +#include +#include + +extern MySQL_HostGroups_Manager *MyHGM; +extern MySQL_Logger *GloMyLogger; + +namespace { + +using Clock = std::chrono::steady_clock; + +constexpr unsigned int kHostgroup = 919; +constexpr const char *kEndpoint = + "killer.cluster-abcdefghijkl.us-east-1.rds.amazonaws.com"; +constexpr const char *kTransportIp = "198.51.100.44"; +constexpr const char *kRegion = "us-east-1"; +constexpr const char *kDatabaseUser = "iam_kill_backend"; +constexpr const char *kOriginalHandshakeToken = "ORIGINAL_HANDSHAKE_TOKEN_MUST_STAY_OWNED"; +constexpr const char *kQueryToken = "CURRENT_QUERY_KILL_TOKEN"; +constexpr const char *kConnectionToken = "CURRENT_CONNECTION_KILL_TOKEN"; +constexpr const char *kPassword = "ordinary-password"; + +unsigned int secure_token_cleanse_calls = 0; +size_t secure_token_cleanse_size = 0; +bool secure_token_was_zeroed = false; + +bool all_zero(const void *ptr, size_t size) { + const auto *bytes = static_cast(ptr); + for (size_t i = 0; i < size; ++i) { + if (bytes[i] != 0) return false; + } + return true; +} + +void tracked_token_cleanse(void *ptr, size_t size) { + OPENSSL_cleanse(ptr, size); + ++secure_token_cleanse_calls; + secure_token_cleanse_size = size; + secure_token_was_zeroed = all_zero(ptr, size); +} + +class FakeBlockingTokenSource final : public AwsIamTokenSource { +public: + AwsIamRequestHandle request(const AwsIamTokenKey&, uint64_t, + std::weak_ptr) override { return {}; } + + AwsIamTokenResult request_blocking(const AwsIamTokenKey& key, + Clock::time_point deadline) override { + keys.push_back(key); + deadlines.push_back(deadline); + if (block_request) { + std::unique_lock lock(block_mutex); + request_entered = true; + block_cv.notify_all(); + block_cv.wait(lock, [this] { return request_released; }); + } + if (wait_until_deadline) std::this_thread::sleep_until(deadline); + AwsIamTokenResult result; + result.status = status; + if (status == AwsIamStatus::OK) { + result.generation = keys.size(); + result.token = SecureString(next_token, tracked_token_cleanse); + } else { + result.failure.category = "credential_provider"; + result.failure.aws_error_code = "FakeProviderFailure"; + result.failure.request_id = "fake-request-id"; + } + return result; + } + + void cancel(AwsIamRequestHandle) override {} + void invalidate(const AwsIamTokenKey&, uint64_t) override {} + void record_backend_connection(bool success) override { + if (success) ++successes; + else ++failures; + } + void record_waiting_session(bool) override {} + AwsIamStatsSnapshot snapshot() const override { return {}; } + + AwsIamStatus status { AwsIamStatus::OK }; + std::string next_token { kQueryToken }; + std::vector keys; + std::vector deadlines; + unsigned int successes { 0 }; + unsigned int failures { 0 }; + bool wait_until_deadline { false }; + bool block_request { false }; + bool request_entered { false }; + bool request_released { false }; + std::mutex block_mutex; + std::condition_variable block_cv; +}; + +struct ConnectorObservation { + unsigned int connect_calls { 0 }; + unsigned int query_calls { 0 }; + std::string expected_password; + bool password_matches { false }; + bool original_password_reused { false }; + std::string host; + std::string username; + unsigned int port { 0 }; + std::string query; + bool ssl_enforce { false }; + bool ssl_verify { false }; + bool cleartext { false }; + bool reconnect_seen { false }; + bool reconnect { true }; + bool connect_timeout_seen { false }; + unsigned int connect_timeout { 0 }; + std::string tls_server_name; + char *connector_password { nullptr }; + size_t connector_password_size { 0 }; + unsigned int connector_cleanse_calls { 0 }; + bool connector_password_was_zeroed { false }; + bool passwd_null_at_close { false }; +}; + +ConnectorObservation connector; + +void reset_observations(const char *expected_password) { + connector = ConnectorObservation {}; + connector.expected_password = expected_password != nullptr ? expected_password : ""; + secure_token_cleanse_calls = 0; + secure_token_cleanse_size = 0; + secure_token_was_zeroed = false; +} + +KillArgs *iam_args(int kill_type, unsigned long id, Clock::time_point deadline) { + return new KillArgs( + const_cast(kDatabaseUser), nullptr, + const_cast(kEndpoint), 3306, kHostgroup, id, kill_type, 1, + nullptr, const_cast(kTransportIp), MySQLBackendAuthType::AWS_IAM, + kEndpoint, kRegion, kDatabaseUser, deadline); +} + +void test_iam_kill(FakeBlockingTokenSource& source, int kill_type, + unsigned long id, const char *token, const char *expected_query, + const char *label) { + reset_observations(token); + source.status = AwsIamStatus::OK; + source.wait_until_deadline = false; + source.next_token = token; + const Clock::time_point deadline = Clock::now() + std::chrono::seconds(2); + KillArgs *args = iam_args(kill_type, id, deadline); + ok(args->password == nullptr, + "%s carries IAM mode metadata without an original password or token", label); + kill_query_thread(args); + + const AwsIamTokenKey expected_key { kEndpoint, 3306, kRegion, kDatabaseUser }; + ok(!source.keys.empty() && source.keys.back() == expected_key && + source.deadlines.back() == deadline && connector.connect_calls == 1 && + connector.host == kTransportIp && connector.username == kDatabaseUser && + connector.port == 3306 && connector.password_matches && + !connector.original_password_reused, + "%s obtains its own current token with the exact key and helper deadline", label); + ok(connector.ssl_enforce && connector.ssl_verify && connector.cleartext && + connector.reconnect_seen && !connector.reconnect && + connector.tls_server_name == kEndpoint && + connector.connect_timeout_seen && connector.connect_timeout >= 1 && + connector.connect_timeout <= 2, + "%s enforces TLS, hostname verification, cleartext auth, no reconnect, and a deadline-bounded connect timeout", label); + ok(connector.query_calls == 1 && connector.query == expected_query, + "%s connects and issues only the requested KILL command", label); + ok(secure_token_cleanse_calls == 1 && + secure_token_cleanse_size == std::strlen(token) && secure_token_was_zeroed && + connector.connector_cleanse_calls == 1 && + connector.connector_password_was_zeroed && connector.passwd_null_at_close, + "%s cleanses the secure token and Connector/C password copies", label); +} + +void test_helper_deadline(FakeBlockingTokenSource& source) { + reset_observations(kQueryToken); + source.status = AwsIamStatus::TIMEOUT; + const Clock::time_point deadline = Clock::now() + std::chrono::milliseconds(5); + KillArgs *args = iam_args(KILL_QUERY, 333, deadline); + const size_t requests_before = source.keys.size(); + kill_query_thread(args); + ok(source.keys.size() == requests_before + 1 && + source.deadlines.back() == deadline && connector.connect_calls == 0 && + connector.query_calls == 0, + "an IAM kill helper passes through its deadline and never connects after timeout"); +} + +void test_deadline_expiring_during_token_request(FakeBlockingTokenSource& source) { + reset_observations(kQueryToken); + source.status = AwsIamStatus::OK; + source.next_token = kQueryToken; + source.wait_until_deadline = true; + KillArgs *args = iam_args( + KILL_QUERY, 334, Clock::now() + std::chrono::milliseconds(5)); + kill_query_thread(args); + source.wait_until_deadline = false; + ok(connector.connect_calls == 0 && connector.query_calls == 0, + "an IAM helper does not start TCP connect after its deadline expires during token acquisition"); +} + +void test_password_mode_unchanged(FakeBlockingTokenSource& source) { + reset_observations(kPassword); + const size_t requests_before = source.keys.size(); + KillArgs *args = new KillArgs( + const_cast("password_backend"), const_cast(kPassword), + const_cast(kEndpoint), 3306, kHostgroup, 444, KILL_QUERY, 0, + nullptr, const_cast(kTransportIp)); + kill_query_thread(args); + ok(source.keys.size() == requests_before && connector.connect_calls == 1 && + connector.password_matches && !connector.ssl_enforce && !connector.ssl_verify && + !connector.cleartext && connector.tls_server_name.empty() && + connector.query == "KILL QUERY 444", + "password-mode kill helpers retain their existing password connector behavior"); +} + +void test_helper_shutdown_lifetime(FakeBlockingTokenSource& source) { + reset_observations(kQueryToken); + source.status = AwsIamStatus::OK; + source.next_token = kQueryToken; + source.block_request = true; + source.request_entered = false; + source.request_released = false; + std::thread helper([&] { + kill_query_thread(iam_args( + KILL_QUERY, 555, Clock::now() + std::chrono::seconds(2))); + }); + { + std::unique_lock lock(source.block_mutex); + if (!source.block_cv.wait_for(lock, std::chrono::seconds(1), + [&source] { return source.request_entered; })) { + BAIL_OUT("IAM helper did not enter the blocking token source"); + } + } + std::atomic shutdown_started { false }; + std::atomic shutdown_returned { false }; + std::thread shutdown([&] { + shutdown_started.store(true, std::memory_order_release); + shutdown_global_aws_iam_token_source(); + shutdown_returned.store(true, std::memory_order_release); + }); + while (!shutdown_started.load(std::memory_order_acquire)) { + std::this_thread::yield(); + } + for (;;) { + AwsIamTokenSourceLease probe = acquire_global_aws_iam_token_source(); + if (!probe) break; + std::this_thread::yield(); + } + const bool returned_while_helper_active = + shutdown_returned.load(std::memory_order_acquire); + { + std::lock_guard lock(source.block_mutex); + source.request_released = true; + } + source.block_cv.notify_all(); + helper.join(); + shutdown.join(); + source.block_request = false; + ok(!returned_while_helper_active && + shutdown_returned.load(std::memory_order_acquire) && + GloAwsIamTokenSource == nullptr && connector.connect_calls == 1 && + connector.query_calls == 1, + "token-source shutdown waits until an already-running detached IAM helper finishes safely"); + + const size_t requests_before = source.keys.size(); + reset_observations(kQueryToken); + kill_query_thread(iam_args( + KILL_QUERY, 556, Clock::now() + std::chrono::seconds(2))); + ok(source.keys.size() == requests_before && connector.connect_calls == 0 && + connector.query_calls == 0, + "an IAM helper starting after token-source shutdown is rejected safely"); +} + +} // namespace + +extern "C" { + +int __real_mysql_options(MYSQL *, enum mysql_option, const void *); +void __real_OPENSSL_cleanse(void *, size_t); +void __real_mysql_close(MYSQL *); + +int __wrap_mysql_options(MYSQL *mysql, enum mysql_option option, const void *arg) { + switch (option) { + case MARIADB_OPT_TLS_SERVER_NAME: + connector.tls_server_name = arg != nullptr + ? static_cast(arg) : ""; + break; + case MYSQL_OPT_SSL_ENFORCE: + connector.ssl_enforce = arg != nullptr && + *static_cast(arg) != 0; + break; + case MYSQL_OPT_SSL_VERIFY_SERVER_CERT: + connector.ssl_verify = arg != nullptr && + *static_cast(arg) != 0; + break; + case MYSQL_ENABLE_CLEARTEXT_PLUGIN: + connector.cleartext = arg != nullptr && + *static_cast(arg) != 0; + break; + case MYSQL_OPT_RECONNECT: + connector.reconnect_seen = true; + connector.reconnect = arg != nullptr && + *static_cast(arg) != 0; + break; + case MYSQL_OPT_CONNECT_TIMEOUT: + connector.connect_timeout_seen = true; + connector.connect_timeout = arg != nullptr + ? *static_cast(arg) : 0; + break; + default: + break; + } + return __real_mysql_options(mysql, option, arg); +} + +MYSQL *__wrap_mysql_real_connect(MYSQL *mysql, const char *host, + const char *user, const char *password, const char *, unsigned int port, + const char *, unsigned long) { + ++connector.connect_calls; + connector.host = host != nullptr ? host : ""; + connector.username = user != nullptr ? user : ""; + connector.port = port; + connector.password_matches = password != nullptr && + connector.expected_password == password; + connector.original_password_reused = password != nullptr && + std::strcmp(password, kOriginalHandshakeToken) == 0; + mysql->host = strdup(host != nullptr ? host : ""); + mysql->passwd = strdup(password != nullptr ? password : ""); + mysql->port = port; + connector.connector_password = mysql->passwd; + connector.connector_password_size = std::strlen(mysql->passwd); + return mysql; +} + +int __wrap_mysql_query(MYSQL *, const char *query) { + ++connector.query_calls; + connector.query = query != nullptr ? query : ""; + return 0; +} + +void __wrap_OPENSSL_cleanse(void *ptr, size_t size) { + __real_OPENSSL_cleanse(ptr, size); + if (ptr == connector.connector_password && + size == connector.connector_password_size) { + ++connector.connector_cleanse_calls; + connector.connector_password_was_zeroed = all_zero(ptr, size); + } +} + +void __wrap_mysql_close(MYSQL *mysql) { + connector.passwd_null_at_close = mysql == nullptr || mysql->passwd == nullptr; + __real_mysql_close(mysql); +} + +} // extern "C" + +int main() { + plan(15); + if (test_init_minimal() != 0 || test_init_query_processor() != 0 || + test_init_hostgroups() != 0) { + BAIL_OUT("failed to initialize unit-test globals"); + } + GloMyLogger = new MySQL_Logger(); + if (!GloMTH->set_variable("ssl_p2s_ca", "/unit/fake-ca.pem")) { + BAIL_OUT("failed to configure helper CA fixture"); + } + + FakeBlockingTokenSource source; + publish_global_aws_iam_token_source(&source); + test_iam_kill(source, KILL_QUERY, 111, kQueryToken, + "KILL QUERY 111", "IAM query kill"); + test_iam_kill(source, KILL_CONNECTION, 222, kConnectionToken, + "KILL CONNECTION 222", "IAM connection kill"); + test_helper_deadline(source); + test_deadline_expiring_during_token_request(source); + test_password_mode_unchanged(source); + test_helper_shutdown_lifetime(source); + + delete GloMyLogger; + GloMyLogger = nullptr; + test_cleanup_hostgroups(); + test_cleanup_query_processor(); + test_cleanup_minimal(); + return exit_status(); +} diff --git a/test/tap/tests/unit/aws_iam_policy_unit-t.cpp b/test/tap/tests/unit/aws_iam_policy_unit-t.cpp new file mode 100644 index 0000000000..8b3bd8ff85 --- /dev/null +++ b/test/tap/tests/unit/aws_iam_policy_unit-t.cpp @@ -0,0 +1,197 @@ +#include "tap.h" +#include "test_globals.h" +#include "test_init.h" + +#include "MySQL_Authentication.hpp" +#include "MySQL_Backend_Auth.h" + +#include +#include +#include +#include + +extern MySQL_Authentication *GloMyAuth; + +static bool add_backend_user(const char *username, const char *password, + const char *attributes) { + return GloMyAuth->add( + (char *)username, (char *)password, USERNAME_BACKEND, + false, 0, (char *)"", false, false, false, 100, + (char *)attributes, (char *)""); +} + +static bool add_backend_user_capturing_stderr(const char *username, const char *password, + const char *attributes, std::string& log) { + FILE *captured = tmpfile(); + if (captured == nullptr) { + return false; + } + + fflush(stderr); + const int saved_stderr = dup(STDERR_FILENO); + if (saved_stderr < 0 || dup2(fileno(captured), STDERR_FILENO) < 0) { + if (saved_stderr >= 0) { + close(saved_stderr); + } + fclose(captured); + return false; + } + + const bool added = add_backend_user(username, password, attributes); + fflush(stderr); + dup2(saved_stderr, STDERR_FILENO); + close(saved_stderr); + + char buffer[256]; + rewind(captured); + while (fgets(buffer, sizeof(buffer), captured) != nullptr) { + log += buffer; + } + fclose(captured); + return added; +} + +static void test_password_defaults() { + ok(parse_mysql_backend_auth_policy("db_user", "", false).type == MySQLBackendAuthType::PASSWORD, + "empty attributes use password authentication"); + ok(parse_mysql_backend_auth_policy("db_user", "{}", false).type == MySQLBackendAuthType::PASSWORD, + "empty attributes object uses password authentication"); +} + +static void test_aws_iam_policy() { + const MySQLBackendAuthPolicy policy = parse_mysql_backend_auth_policy( + "iam_user", "{\"backend_auth\":{\"type\":\"aws_iam\"}}", false); + ok(policy.type == MySQLBackendAuthType::AWS_IAM, + "aws_iam backend_auth type selects IAM authentication"); + ok(policy.database_user == "iam_user", + "policy preserves the mapped backend username"); + ok(!policy.ignored_password, + "IAM policy with an empty password does not mark a password ignored"); + + const MySQLBackendAuthPolicy with_password = parse_mysql_backend_auth_policy( + "iam_user", "{\"backend_auth\":{\"type\":\"aws_iam\"}}", true); + ok(with_password.type == MySQLBackendAuthType::AWS_IAM && with_password.ignored_password, + "IAM policy marks a configured backend password as ignored"); +} + +static void test_invalid_attribute_shapes() { + const MySQLBackendAuthPolicy scalar = parse_mysql_backend_auth_policy("db_user", "\"value\"", false); + ok(scalar.type == MySQLBackendAuthType::INVALID && scalar.failure_code == "attributes_not_object", + "scalar attributes are rejected without parsing credentials"); + const MySQLBackendAuthPolicy array = parse_mysql_backend_auth_policy("db_user", "[]", false); + ok(array.type == MySQLBackendAuthType::INVALID && array.failure_code == "attributes_not_object", + "array attributes are rejected without parsing credentials"); + const MySQLBackendAuthPolicy malformed = parse_mysql_backend_auth_policy( + "db_user", "{\"backend_auth\":FAKE_AWS_SECRET}", false); + ok(malformed.type == MySQLBackendAuthType::INVALID && malformed.failure_code == "attributes_not_object", + "malformed attributes are rejected without exposing their contents"); +} + +static void test_invalid_backend_auth_shapes() { + const char *invalid_values[] = { "\"aws_iam\"", "[]", "null" }; + for (const char *value : invalid_values) { + const std::string attributes = std::string("{\"backend_auth\":") + value + "}"; + const MySQLBackendAuthPolicy policy = parse_mysql_backend_auth_policy("db_user", attributes, false); + ok(policy.type == MySQLBackendAuthType::INVALID && policy.failure_code == "backend_auth_not_object", + "non-object backend_auth is rejected"); + } +} + +static void test_invalid_type_values() { + const MySQLBackendAuthPolicy missing = parse_mysql_backend_auth_policy("db_user", "{\"backend_auth\":{}}", false); + ok(missing.type == MySQLBackendAuthType::INVALID && missing.failure_code == "type_missing", + "backend_auth without type is rejected"); + const MySQLBackendAuthPolicy non_string = parse_mysql_backend_auth_policy( + "db_user", "{\"backend_auth\":{\"type\":1}}", false); + ok(non_string.type == MySQLBackendAuthType::INVALID && non_string.failure_code == "type_not_string", + "non-string backend_auth type is rejected"); + const MySQLBackendAuthPolicy unknown = parse_mysql_backend_auth_policy( + "db_user", "{\"backend_auth\":{\"type\":\"kerberos\"}}", false); + ok(unknown.type == MySQLBackendAuthType::INVALID && unknown.failure_code == "type_unsupported", + "unknown backend_auth type is rejected"); + const MySQLBackendAuthPolicy differently_cased = parse_mysql_backend_auth_policy( + "db_user", "{\"backend_auth\":{\"type\":\"AWS_IAM\"}}", false); + ok(differently_cased.type == MySQLBackendAuthType::INVALID && differently_cased.failure_code == "type_unsupported", + "differently-cased backend_auth type is rejected"); +} + +static void test_diagnostics_do_not_leak_attributes() { + const MySQLBackendAuthPolicy policy = parse_mysql_backend_auth_policy( + "db_user", "{\"backend_auth\":FAKE_AWS_SECRET,\"token\":\"FAKE_SESSION_TOKEN\"}", false); + ok(policy.failure_code.find("FAKE_AWS_SECRET") == std::string::npos, + "diagnostics do not expose fake AWS secrets"); + ok(policy.failure_code.find("FAKE_SESSION_TOKEN") == std::string::npos, + "diagnostics do not expose fake session tokens"); + ok(policy.failure_code.find("backend_auth") == std::string::npos, + "diagnostics do not expose raw malformed JSON"); +} + +static void test_resolver_uses_backend_account_only() { + ok(add_backend_user("iam_backend", "unused-password", + "{\"backend_auth\":{\"type\":\"aws_iam\"}}"), + "backend IAM user is added to the real authentication store"); + const MySQLBackendAuthPolicy policy = resolve_mysql_backend_auth_policy(*GloMyAuth, "iam_backend"); + ok(policy.type == MySQLBackendAuthType::AWS_IAM, + "resolver reads IAM policy from the backend account"); + ok(policy.database_user == "iam_backend" && policy.ignored_password, + "resolver preserves username and ignores the backend password for IAM"); +} + +static void test_resolver_rejects_missing_or_inactive_backend_account() { + const MySQLBackendAuthPolicy missing = resolve_mysql_backend_auth_policy(*GloMyAuth, "missing_backend"); + ok(missing.type == MySQLBackendAuthType::INVALID && missing.failure_code == "backend_user_not_found", + "resolver rejects a missing backend account"); + + ok(add_backend_user("inactive_backend", "password", ""), + "inactive backend fixture is added"); + GloMyAuth->set_all_inactive(USERNAME_BACKEND); + GloMyAuth->remove_inactives(USERNAME_BACKEND); + const MySQLBackendAuthPolicy inactive = resolve_mysql_backend_auth_policy(*GloMyAuth, "inactive_backend"); + ok(inactive.type == MySQLBackendAuthType::INVALID && inactive.failure_code == "backend_user_not_found", + "resolver rejects an inactive backend account after runtime removal"); +} + +static void test_resolver_rejects_malformed_loaded_backend_attributes() { + ok(add_backend_user("malformed_backend", "configured-password", + "{\"backend_auth\":FAKE_AWS_SECRET,\"token\":\"FAKE_SESSION_TOKEN\"}"), + "malformed backend user is loaded into the real authentication store"); + const MySQLBackendAuthPolicy policy = resolve_mysql_backend_auth_policy(*GloMyAuth, "malformed_backend"); + ok(policy.type == MySQLBackendAuthType::INVALID && policy.failure_code == "attributes_not_object", + "resolver fails closed when loaded backend attributes are malformed"); +} + +static void test_rejected_backend_policy_does_not_emit_iam_password_warning() { + ok(add_backend_user("normalized_backend", "configured-password", ""), + "backend user exists before its runtime attributes are updated"); + + std::string log; + ok(add_backend_user_capturing_stderr("normalized_backend", "configured-password", + "{\"backend_auth\":{\"type\":\"aws_iam\"},\"default-transaction_isolation\":1}", log), + "backend user reload with invalid default transaction isolation completes"); + const MySQLBackendAuthPolicy policy = resolve_mysql_backend_auth_policy(*GloMyAuth, "normalized_backend"); + ok(policy.type == MySQLBackendAuthType::INVALID && policy.failure_code == "attributes_not_object", + "post-validation backend policy is invalid"); + ok(log.find("clear the unused backend password") == std::string::npos, + "invalid post-validation backend policy does not emit an IAM password warning"); +} + +int main() { + plan(31); + test_init_minimal(); + test_init_auth(); + + test_password_defaults(); + test_aws_iam_policy(); + test_invalid_attribute_shapes(); + test_invalid_backend_auth_shapes(); + test_invalid_type_values(); + test_diagnostics_do_not_leak_attributes(); + test_resolver_uses_backend_account_only(); + test_resolver_rejects_missing_or_inactive_backend_account(); + test_resolver_rejects_malformed_loaded_backend_attributes(); + test_rejected_backend_policy_does_not_emit_iam_password_warning(); + + test_cleanup_auth(); + test_cleanup_minimal(); + return exit_status(); +} diff --git a/test/tap/tests/unit/aws_iam_pool_unit-t.cpp b/test/tap/tests/unit/aws_iam_pool_unit-t.cpp new file mode 100644 index 0000000000..30ea4fa5e7 --- /dev/null +++ b/test/tap/tests/unit/aws_iam_pool_unit-t.cpp @@ -0,0 +1,794 @@ +/** + * @file aws_iam_pool_unit-t.cpp + * @brief Pool-identity and reset-safety tests for MySQL AWS IAM backends. + */ + +#include "tap.h" +#include "test_globals.h" +#include "test_init.h" + +#include "proxysql.h" +#include "cpp.h" +#include "Aws_Iam_Provider.h" +#include "MySQL_Authentication.hpp" +#include "MySQL_Data_Stream.h" +#include "MySQL_Logger.hpp" + +#include +#include +#include +#include +#include + +extern MySQL_HostGroups_Manager *MyHGM; +extern MySQL_Threads_Handler *GloMTH; +extern MySQL_Logger *GloMyLogger; +extern MySQL_Authentication *GloMyAuth; +extern void *HGCU_thread_run(); + +namespace { + +constexpr const char *kUser = "pool_user"; +constexpr const char *kOtherUser = "other_pool_user"; +constexpr const char *kInvalidReloadUser = "invalid_reload_pool_user"; +constexpr const char *kRowlessPassthroughUser = "rowless_passthrough_pool_user"; +constexpr const char *kSchema = "pool_schema"; + +std::atomic change_user_calls { 0 }; +std::atomic change_user_immediate_success { false }; +MySQL_Connection *kill_source_connection = nullptr; +unsigned int kill_helper_dispatches = 0; +bool kill_source_token_present_at_dispatch = false; +bool kill_args_password_absent = false; +unsigned long kill_args_id = 0; +int kill_args_type = 0; +MySQLBackendAuthType kill_args_auth_type = MySQLBackendAuthType::PASSWORD; +std::string kill_args_endpoint; +std::string kill_args_region; +std::string kill_args_database_user; +std::string kill_args_transport; +unsigned int kill_args_port = 0; +int kill_args_use_ssl = 0; +std::chrono::steady_clock::time_point kill_args_deadline; + +MySrvC *create_server(unsigned int hostgroup_id, const char *address) { + srv_info_t info; + info.addr = address; + info.port = 3306; + info.kind = "aws-iam-pool-unit"; + srv_opts_t opts; + opts.weigth = 1; + opts.max_conns = 100; + opts.use_ssl = 0; + + MyHGM->wrlock(); + const int rc = MyHGM->create_new_server_in_hg(hostgroup_id, info, opts); + MyHGC *hostgroup = MyHGM->MyHGC_find(hostgroup_id); + MyHGM->wrunlock(); + if (rc != 0 || hostgroup == nullptr || hostgroup->mysrvs->cnt() != 1) { + BAIL_OUT("failed to create pool fixture for hostgroup %u", hostgroup_id); + } + return hostgroup->mysrvs->idx(0); +} + +MySQL_Connection *create_established_connection( + MySrvC *server, const char *username, MySQLBackendAuthType type) +{ + MySQL_Connection *connection = new MySQL_Connection(); + connection->mysql = mysql_init(nullptr); + if (connection->mysql == nullptr) { + delete connection; + BAIL_OUT("mysql_init() failed for pool fixture"); + } + connection->ret_mysql = connection->mysql; + connection->mysql->charset = mariadb_get_charset_by_name("utf8mb4"); + if (connection->mysql->charset == nullptr) { + delete connection; + BAIL_OUT("failed to initialize connector charset for pool fixture"); + } + connection->parent = server; + connection->userinfo->set( + const_cast(username), const_cast("password"), + const_cast(kSchema), nullptr); + connection->set_backend_auth_type(type); + connection->healthy = true; + connection->reusable = true; + connection->send_quit = false; + connection->fd = 123; + connection->async_state_machine = ASYNC_IDLE; + return connection; +} + +class SessionFixture { +public: + SessionFixture(MySQL_Thread& worker, unsigned int hostgroup_id, + MySQLBackendAuthType requested_type, const char *username = kUser) + : worker(worker) + { + session = new MySQL_Session(); + session->thread = &worker; + session->connections_handler = true; + + frontend_stream = new MySQL_Data_Stream(); + frontend_stream->init(MYDS_FRONTEND, session, -1); + frontend = new MySQL_Connection(); + frontend_stream->attach_connection(frontend); + frontend_stream->myprot.init(&frontend_stream, frontend->userinfo, session); + session->client_myds = frontend_stream; + frontend->userinfo->set( + const_cast(username), const_cast("password"), + const_cast(kSchema), nullptr); + frontend->set_backend_auth_type(requested_type); + + session->mybe = session->create_backend(hostgroup_id); + session->current_hostgroup = hostgroup_id; + session->default_hostgroup = hostgroup_id; + } + + ~SessionFixture() { + delete session; + } + + void attach_backend(MySQL_Connection *connection) { + MySQL_Data_Stream *stream = session->mybe->server_myds; + stream->attach_connection(connection); + stream->assign_fd_from_mysql_conn(); + stream->myds_type = MYDS_BACKEND; + stream->DSS = STATE_MARIADB_QUERY; + } + + MySQL_Connection *selected() const { + return session->mybe != nullptr && session->mybe->server_myds != nullptr + ? session->mybe->server_myds->myconn : nullptr; + } + + MySQL_Thread& worker; + MySQL_Session *session { nullptr }; + MySQL_Data_Stream *frontend_stream { nullptr }; + MySQL_Connection *frontend { nullptr }; +}; + +void destroy_used(MySQL_Connection *connection) { + if (connection == nullptr) return; + connection->send_quit = false; + MyHGM->destroy_MyConn_from_pool(connection); +} + +void test_identity_matrix(MySQL_Thread& worker) { + MySrvC *server = create_server(801, "identity-matrix"); + SessionFixture same_password(worker, 801, MySQLBackendAuthType::PASSWORD, kUser); + SessionFixture same_iam(worker, 801, MySQLBackendAuthType::AWS_IAM, kUser); + SessionFixture other_password(worker, 801, MySQLBackendAuthType::PASSWORD, kOtherUser); + + MySQL_Connection *password = create_established_connection( + server, kUser, MySQLBackendAuthType::PASSWORD); + MySQL_Connection *iam = create_established_connection( + server, kUser, MySQLBackendAuthType::AWS_IAM); + + ok(!password->requires_CHANGE_USER( + same_password.frontend, MySQLBackendAuthType::PASSWORD), + "same username and password mode are compatible without CHANGE_USER"); + ok(password->requires_CHANGE_USER( + other_password.frontend, MySQLBackendAuthType::PASSWORD), + "different username in password mode still requires ordinary CHANGE_USER"); + ok(password->requires_CHANGE_USER( + same_iam.frontend, MySQLBackendAuthType::AWS_IAM), + "password connection cannot satisfy the same username in IAM mode"); + ok(iam->requires_CHANGE_USER( + same_password.frontend, MySQLBackendAuthType::PASSWORD), + "IAM connection cannot satisfy the same username in password mode"); + ok(!iam->requires_CHANGE_USER( + same_iam.frontend, MySQLBackendAuthType::AWS_IAM), + "same username and IAM mode are compatible without token-age checks"); + ok(iam->requires_CHANGE_USER( + other_password.frontend, MySQLBackendAuthType::AWS_IAM), + "different username in IAM mode cannot reuse the established connection"); + + worker.curtime += 16ULL * 60ULL * 1000000ULL; + ok(!iam->requires_CHANGE_USER( + same_iam.frontend, MySQLBackendAuthType::AWS_IAM), + "an established IAM connection remains reusable beyond token lifetime"); + + delete password; + delete iam; +} + +void test_runtime_mode_changes_global(MySQL_Thread& worker) { + MySrvC *password_server = create_server(802, "password-to-iam"); + SessionFixture iam_request(worker, 802, MySQLBackendAuthType::AWS_IAM); + MySQL_Connection *old_password = create_established_connection( + password_server, kUser, MySQLBackendAuthType::PASSWORD); + password_server->ConnectionsFree->add(old_password); + MySQL_Connection *selected = MyHGM->get_MyConn_from_pool( + 802, iam_request.session, false, nullptr, 0, -1, + MySQLBackendAuthType::AWS_IAM); + ok(selected != nullptr && selected->fd == -1 && + password_server->ConnectionsFree->conns_length() == 0, + "PASSWORD to IAM policy change lazily destroys the old global entry and creates fresh"); + destroy_used(selected); + + MySrvC *iam_server = create_server(803, "iam-to-password"); + SessionFixture password_request(worker, 803, MySQLBackendAuthType::PASSWORD); + MySQL_Connection *old_iam = create_established_connection( + iam_server, kUser, MySQLBackendAuthType::AWS_IAM); + iam_server->ConnectionsFree->add(old_iam); + selected = MyHGM->get_MyConn_from_pool( + 803, password_request.session, false, nullptr, 0, -1, + MySQLBackendAuthType::PASSWORD); + ok(selected != nullptr && selected->fd == -1 && + iam_server->ConnectionsFree->conns_length() == 0, + "IAM to PASSWORD policy change lazily destroys the old global entry and creates fresh"); + destroy_used(selected); +} + +void test_runtime_mode_change_local(MySQL_Thread& worker) { + MySrvC *server = create_server(804, "local-password-to-iam"); + SessionFixture request(worker, 804, MySQLBackendAuthType::AWS_IAM); + MySQL_Connection *old_password = create_established_connection( + server, kUser, MySQLBackendAuthType::PASSWORD); + server->ConnectionsUsed->add(old_password); + worker.push_MyConn_local(old_password); + + MySQL_Connection *selected = worker.get_MyConn_local( + 804, request.session, nullptr, 0, -1, + MySQLBackendAuthType::AWS_IAM); + ok(selected == nullptr && server->ConnectionsUsed->conns_length() == 0, + "local checkout lazily destroys a connection from the previous auth mode"); + if (selected != nullptr) destroy_used(selected); +} + +void test_mixed_user_mode_global(MySQL_Thread& worker) { + MySrvC *password_server = create_server(809, "mixed-global-password"); + SessionFixture password_request( + worker, 809, MySQLBackendAuthType::PASSWORD, kUser); + MySQL_Connection *unrelated_iam = create_established_connection( + password_server, kOtherUser, MySQLBackendAuthType::AWS_IAM); + unrelated_iam->fd = 201; + MySQL_Connection *exact_password = create_established_connection( + password_server, kUser, MySQLBackendAuthType::PASSWORD); + exact_password->fd = 202; + password_server->ConnectionsFree->add(unrelated_iam); + password_server->ConnectionsFree->add(exact_password); + MySQL_Connection *selected = MyHGM->get_MyConn_from_pool( + 809, password_request.session, false, nullptr, 0, -1, + MySQLBackendAuthType::PASSWORD); + ok(selected != nullptr && selected->fd == 202 && + password_server->ConnectionsFree->conns_length() == 1, + "PASSWORD checkout preserves another user's idle IAM connection globally"); + destroy_used(selected); + + SessionFixture iam_request( + worker, 809, MySQLBackendAuthType::AWS_IAM, kOtherUser); + selected = MyHGM->get_MyConn_from_pool( + 809, iam_request.session, false, nullptr, 0, -1, + MySQLBackendAuthType::AWS_IAM); + ok(selected != nullptr && selected->fd == 201, + "unrelated global IAM entry remains reusable by its exact identity"); + destroy_used(selected); + + MySrvC *iam_server = create_server(810, "mixed-global-iam"); + SessionFixture exact_iam_request( + worker, 810, MySQLBackendAuthType::AWS_IAM, kUser); + MySQL_Connection *unrelated_password = create_established_connection( + iam_server, kOtherUser, MySQLBackendAuthType::PASSWORD); + unrelated_password->fd = 211; + MySQL_Connection *exact_iam = create_established_connection( + iam_server, kUser, MySQLBackendAuthType::AWS_IAM); + exact_iam->fd = 212; + iam_server->ConnectionsFree->add(unrelated_password); + iam_server->ConnectionsFree->add(exact_iam); + selected = MyHGM->get_MyConn_from_pool( + 810, exact_iam_request.session, false, nullptr, 0, -1, + MySQLBackendAuthType::AWS_IAM); + ok(selected != nullptr && selected->fd == 212 && + iam_server->ConnectionsFree->conns_length() == 1, + "IAM checkout preserves another user's idle PASSWORD connection globally"); + destroy_used(selected); + + SessionFixture unrelated_password_request( + worker, 810, MySQLBackendAuthType::PASSWORD, kOtherUser); + selected = MyHGM->get_MyConn_from_pool( + 810, unrelated_password_request.session, false, nullptr, 0, -1, + MySQLBackendAuthType::PASSWORD); + ok(selected != nullptr && selected->fd == 211, + "unrelated global PASSWORD entry remains reusable by its exact identity"); + destroy_used(selected); +} + +void test_mixed_user_mode_local(MySQL_Thread& worker) { + MySrvC *password_server = create_server(811, "mixed-local-password"); + SessionFixture password_request( + worker, 811, MySQLBackendAuthType::PASSWORD, kUser); + MySQL_Connection *unrelated_iam = create_established_connection( + password_server, kOtherUser, MySQLBackendAuthType::AWS_IAM); + unrelated_iam->fd = 301; + MySQL_Connection *exact_password = create_established_connection( + password_server, kUser, MySQLBackendAuthType::PASSWORD); + exact_password->fd = 302; + password_server->ConnectionsUsed->add(unrelated_iam); + password_server->ConnectionsUsed->add(exact_password); + worker.push_MyConn_local(unrelated_iam); + worker.push_MyConn_local(exact_password); + MySQL_Connection *selected = worker.get_MyConn_local( + 811, password_request.session, nullptr, 0, -1, + MySQLBackendAuthType::PASSWORD); + ok(selected != nullptr && selected->fd == 302, + "local PASSWORD checkout selects its exact identity in a mixed-mode cache"); + destroy_used(selected); + + SessionFixture iam_request( + worker, 811, MySQLBackendAuthType::AWS_IAM, kOtherUser); + selected = worker.get_MyConn_local( + 811, iam_request.session, nullptr, 0, -1, + MySQLBackendAuthType::AWS_IAM); + ok(selected != nullptr && selected->fd == 301, + "local PASSWORD checkout preserves another user's reusable IAM entry"); + destroy_used(selected); + + MySrvC *iam_server = create_server(812, "mixed-local-iam"); + SessionFixture exact_iam_request( + worker, 812, MySQLBackendAuthType::AWS_IAM, kUser); + MySQL_Connection *unrelated_password = create_established_connection( + iam_server, kOtherUser, MySQLBackendAuthType::PASSWORD); + unrelated_password->fd = 311; + MySQL_Connection *exact_iam = create_established_connection( + iam_server, kUser, MySQLBackendAuthType::AWS_IAM); + exact_iam->fd = 312; + iam_server->ConnectionsUsed->add(unrelated_password); + iam_server->ConnectionsUsed->add(exact_iam); + worker.push_MyConn_local(unrelated_password); + worker.push_MyConn_local(exact_iam); + selected = worker.get_MyConn_local( + 812, exact_iam_request.session, nullptr, 0, -1, + MySQLBackendAuthType::AWS_IAM); + ok(selected != nullptr && selected->fd == 312, + "local IAM checkout selects its exact identity in a mixed-mode cache"); + destroy_used(selected); + + SessionFixture unrelated_password_request( + worker, 812, MySQLBackendAuthType::PASSWORD, kOtherUser); + selected = worker.get_MyConn_local( + 812, unrelated_password_request.session, nullptr, 0, -1, + MySQLBackendAuthType::PASSWORD); + ok(selected != nullptr && selected->fd == 311, + "local IAM checkout preserves another user's reusable PASSWORD entry"); + destroy_used(selected); +} + +void load_invalid_reload_policy() { + if (!GloMyAuth->add( + const_cast(kInvalidReloadUser), const_cast("password"), + USERNAME_BACKEND, false, 0, const_cast(""), false, false, + false, 100, const_cast("{\"backend_auth\":{\"type\":17}}"), + const_cast(""))) { + BAIL_OUT("failed to load malformed backend policy fixture"); + } +} + +void test_attached_password_reload_to_invalid_fails_closed(MySQL_Thread& worker) { + MySrvC *server = create_server(813, "attached-invalid-policy"); + SessionFixture fixture( + worker, 813, MySQLBackendAuthType::PASSWORD, kInvalidReloadUser); + MySQL_Connection *password = create_established_connection( + server, kInvalidReloadUser, MySQLBackendAuthType::PASSWORD); + server->ConnectionsUsed->add(password); + fixture.attach_backend(password); + fixture.session->set_status(PROCESSING_QUERY); + load_invalid_reload_policy(); + change_user_calls.store(0, std::memory_order_relaxed); + fixture.session->to_process = 1; + fixture.session->handler(); + ok(fixture.session->status == WAITING_CLIENT_DATA && + fixture.selected() == nullptr && fixture.session->previous_status.empty() && + server->ConnectionsUsed->conns_length() == 0 && + change_user_calls.load(std::memory_order_relaxed) == 0, + "attached PASSWORD connection fails terminally when reload makes policy INVALID"); +} + +void test_invalid_policy_cannot_continue_change_user(MySQL_Thread& worker) { + MySrvC *server = create_server(814, "changing-user-invalid-policy"); + SessionFixture fixture( + worker, 814, MySQLBackendAuthType::PASSWORD, kInvalidReloadUser); + MySQL_Connection *password = create_established_connection( + server, kInvalidReloadUser, MySQLBackendAuthType::PASSWORD); + server->ConnectionsUsed->add(password); + fixture.attach_backend(password); + fixture.session->previous_status.push(PROCESSING_QUERY); + fixture.session->set_status(CHANGING_USER_SERVER); + change_user_calls.store(0, std::memory_order_relaxed); + fixture.session->to_process = 1; + fixture.session->handler(); + ok(fixture.session->status == WAITING_CLIENT_DATA && + fixture.selected() == nullptr && fixture.session->previous_status.empty() && + server->ConnectionsUsed->conns_length() == 0 && + change_user_calls.load(std::memory_order_relaxed) == 0, + "INVALID policy discovered in CHANGING_USER_SERVER cannot send a password"); +} + +void test_invalid_policy_cannot_enter_reset(MySQL_Thread& worker) { + MySrvC *server = create_server(815, "reset-invalid-policy"); + SessionFixture fixture( + worker, 815, MySQLBackendAuthType::PASSWORD, kInvalidReloadUser); + MySQL_Connection *password = create_established_connection( + server, kInvalidReloadUser, MySQLBackendAuthType::PASSWORD); + server->ConnectionsUsed->add(password); + fixture.attach_backend(password); + fixture.session->set_status(RESETTING_CONNECTION); + change_user_calls.store(0, std::memory_order_relaxed); + fixture.session->to_process = 1; + const int rc = fixture.session->handler(); + ok(rc == -1 && fixture.session->status == session_status___NONE && + fixture.selected() == nullptr && + server->ConnectionsUsed->conns_length() == 0 && + change_user_calls.load(std::memory_order_relaxed) == 0, + "INVALID backend policy destroys reset work without COM_CHANGE_USER"); +} + +void test_authorized_rowless_passthrough_can_enter_detached_reset( + MySQL_Thread& worker) +{ + MySrvC *server = create_server(817, "reset-rowless-passthrough"); + SessionFixture fixture( + worker, 817, MySQLBackendAuthType::PASSWORD, + kRowlessPassthroughUser); + MySQL_Connection *password = create_established_connection( + server, kRowlessPassthroughUser, MySQLBackendAuthType::PASSWORD); + password->set_rowless_passthrough_authorized(true); + server->ConnectionsUsed->add(password); + fixture.attach_backend(password); + fixture.session->set_status(RESETTING_CONNECTION); + change_user_calls.store(0, std::memory_order_relaxed); + fixture.session->to_process = 1; + const int rc = fixture.session->handler(); + ok(rc == 0 && fixture.session->status == RESETTING_CONNECTION && + fixture.selected() == password && + server->ConnectionsUsed->conns_length() == 1 && + change_user_calls.load(std::memory_order_relaxed) == 1, + "authorized rowless pass-through PASSWORD can enter detached reset"); +} + +void test_destroy_path_never_queues_iam() { + MySrvC *server = create_server(805, "destroy-iam"); + GloMTH->variables.connpoll_reset_queue_length = 50; + MySQL_Connection *iam = create_established_connection( + server, kUser, MySQLBackendAuthType::AWS_IAM); + iam->send_quit = true; + server->ConnectionsUsed->add(iam); + MyHGM->destroy_MyConn_from_pool(iam); + ok(MyHGM->queue.size() == 0 && server->ConnectionsUsed->conns_length() == 0, + "destroy_MyConn_from_pool deletes IAM instead of queueing COM_CHANGE_USER reset"); + + MySQL_Connection *password = create_established_connection( + server, kUser, MySQLBackendAuthType::PASSWORD); + password->send_quit = true; + server->ConnectionsUsed->add(password); + MyHGM->destroy_MyConn_from_pool(password); + ok(MyHGM->queue.size() == 1 && server->ConnectionsUsed->conns_length() == 1, + "ordinary password connection keeps the existing reset-queue behavior"); + if (MyHGM->queue.size() != 0) { + MySQL_Connection *queued = MyHGM->queue.remove(); + queued->send_quit = false; + MyHGM->destroy_MyConn_from_pool(queued); + } +} + +void test_destroy_path_dispatches_detached_iam_connection_kill() { + MySrvC *server = create_server(820, "kill-dispatch-iam"); + server->use_ssl = 1; + free(server->myhgc->attributes.aws_iam_region); + server->myhgc->attributes.aws_iam_region = strdup("us-east-1"); + GloMTH->variables.connpoll_reset_queue_length = 50; + mysql_thread___kill_backend_connection_when_disconnect = true; + + MySQL_Connection *iam = create_established_connection( + server, kUser, MySQLBackendAuthType::AWS_IAM); + AwsIamTokenResult token; + token.status = AwsIamStatus::OK; + token.generation = 77; + token.token = SecureString("ORIGINAL_IAM_HANDSHAKE_TOKEN"); + iam->attach_aws_iam_token( + { server->address, server->port, "us-east-1", kUser }, + std::move(token)); + iam->send_quit = true; + iam->async_state_machine = ASYNC_QUERY_CONT; + iam->mysql->thread_id = 741; + iam->connected_host_details.ip = strdup("198.51.100.28"); + server->ConnectionsUsed->add(iam); + + kill_source_connection = iam; + kill_helper_dispatches = 0; + kill_source_token_present_at_dispatch = false; + kill_args_password_absent = false; + kill_args_id = 0; + kill_args_type = 0; + kill_args_auth_type = MySQLBackendAuthType::PASSWORD; + kill_args_endpoint.clear(); + kill_args_region.clear(); + kill_args_database_user.clear(); + kill_args_transport.clear(); + kill_args_port = 0; + kill_args_use_ssl = 0; + kill_args_deadline = {}; + const auto before = std::chrono::steady_clock::now(); + MyHGM->destroy_MyConn_from_pool(iam); + const auto after = std::chrono::steady_clock::now(); + kill_source_connection = nullptr; + mysql_thread___kill_backend_connection_when_disconnect = false; + + ok(kill_helper_dispatches == 1 && kill_source_token_present_at_dispatch && + kill_args_password_absent && kill_args_id == 741 && + kill_args_type == KILL_CONNECTION && + kill_args_auth_type == MySQLBackendAuthType::AWS_IAM && + kill_args_endpoint == "kill-dispatch-iam" && + kill_args_region == "us-east-1" && kill_args_database_user == kUser && + kill_args_transport == "198.51.100.28" && kill_args_port == 3306 && + kill_args_use_ssl == 1 && + kill_args_deadline > before && kill_args_deadline > after && + MyHGM->queue.size() == 0 && + server->ConnectionsUsed->conns_length() == 0, + "production destroy dispatches an IAM KILL_CONNECTION helper without changing the source token or entering reset"); +} + +void test_reset_queue_worker_never_changes_iam() { + MySrvC *server = create_server(808, "reset-worker-iam"); + MySQL_Connection *iam = create_established_connection( + server, kUser, MySQLBackendAuthType::AWS_IAM); + iam->set_rowless_passthrough_authorized(true); + server->ConnectionsUsed->add(iam); + MyHGM->queue.add(iam); + change_user_calls.store(0, std::memory_order_relaxed); + std::thread reset_worker([]() { HGCU_thread_run(); }); + bool removed = false; + for (unsigned int attempt = 0; attempt < 1000 && !removed; ++attempt) { + MyHGM->wrlock(); + removed = server->ConnectionsUsed->conns_length() == 0; + MyHGM->wrunlock(); + if (!removed) usleep(1000); + } + MyHGM->queue.add(nullptr); + reset_worker.join(); + ok(MyHGM->queue.size() == 0 && + removed && server->ConnectionsUsed->conns_length() == 0 && + change_user_calls.load(std::memory_order_relaxed) == 0, + "reset queue worker destroys IAM without invoking COM_CHANGE_USER"); +} + +void test_reset_queue_worker_never_resets_invalid_policy() { + MySrvC *server = create_server(816, "reset-worker-invalid-policy"); + MySQL_Connection *password = create_established_connection( + server, kInvalidReloadUser, MySQLBackendAuthType::PASSWORD); + password->set_rowless_passthrough_authorized(true); + server->ConnectionsUsed->add(password); + MyHGM->queue.add(password); + const unsigned long resets_before = MyHGM->status.myconnpoll_reset; + change_user_calls.store(0, std::memory_order_relaxed); + std::thread reset_worker([]() { HGCU_thread_run(); }); + bool removed = false; + for (unsigned int attempt = 0; attempt < 1000 && !removed; ++attempt) { + MyHGM->wrlock(); + removed = server->ConnectionsUsed->conns_length() == 0; + MyHGM->wrunlock(); + if (!removed) usleep(1000); + } + MyHGM->queue.add(nullptr); + reset_worker.join(); + ok(removed && MyHGM->status.myconnpoll_reset == resets_before && + change_user_calls.load(std::memory_order_relaxed) == 0, + "reset queue worker discards INVALID policy before reset processing"); +} + +void test_reset_queue_worker_preserves_authorized_rowless_passthrough() { + MySrvC *server = create_server(818, "reset-worker-rowless-passthrough"); + MySQL_Connection *password = create_established_connection( + server, kRowlessPassthroughUser, MySQLBackendAuthType::PASSWORD); + password->set_rowless_passthrough_authorized(true); + password->mysql->net.pvio = + reinterpret_castmysql->net.pvio)>(1); + password->mysql->net.fd = 123; + password->mysql->net.buff = + reinterpret_castmysql->net.buff)>(1); + server->ConnectionsUsed->add(password); + MyHGM->queue.add(password); + const unsigned long resets_before = MyHGM->status.myconnpoll_reset; + change_user_calls.store(0, std::memory_order_relaxed); + change_user_immediate_success.store(true, std::memory_order_relaxed); + std::thread reset_worker([]() { HGCU_thread_run(); }); + for (unsigned int attempt = 0; + attempt < 1000 && change_user_calls.load(std::memory_order_relaxed) == 0; + ++attempt) { + usleep(1000); + } + MyHGM->queue.add(nullptr); + reset_worker.join(); + change_user_immediate_success.store(false, std::memory_order_relaxed); + const bool returned_to_pool = + server->ConnectionsFree->conns_length() == 1 && + server->ConnectionsUsed->conns_length() == 0; + ok(MyHGM->status.myconnpoll_reset == resets_before + 1 && + change_user_calls.load(std::memory_order_relaxed) == 1 && returned_to_pool, + "reset worker preserves authorized rowless pass-through PASSWORD semantics"); + if (returned_to_pool) { + password->mysql->net.pvio = nullptr; + password->mysql->net.fd = 0; + password->mysql->net.buff = nullptr; + } +} + +void test_reset_queue_worker_discards_unmarked_rowless_password() { + MySrvC *server = create_server(819, "reset-worker-unmarked-rowless"); + MySQL_Connection *password = create_established_connection( + server, kRowlessPassthroughUser, MySQLBackendAuthType::PASSWORD); + server->ConnectionsUsed->add(password); + MyHGM->queue.add(password); + const unsigned long resets_before = MyHGM->status.myconnpoll_reset; + change_user_calls.store(0, std::memory_order_relaxed); + std::thread reset_worker([]() { HGCU_thread_run(); }); + bool removed = false; + for (unsigned int attempt = 0; attempt < 1000 && !removed; ++attempt) { + MyHGM->wrlock(); + removed = server->ConnectionsUsed->conns_length() == 0; + MyHGM->wrunlock(); + if (!removed) usleep(1000); + } + MyHGM->queue.add(nullptr); + reset_worker.join(); + ok(removed && MyHGM->status.myconnpoll_reset == resets_before && + change_user_calls.load(std::memory_order_relaxed) == 0, + "unmarked rowless PASSWORD cannot bypass reset policy validation"); +} + +void test_change_user_state_replaces_iam(MySQL_Thread& worker) { + MySrvC *server = create_server(806, "change-user-iam"); + SessionFixture fixture(worker, 806, MySQLBackendAuthType::PASSWORD, kOtherUser); + MySQL_Connection *iam = create_established_connection( + server, kUser, MySQLBackendAuthType::AWS_IAM); + server->ConnectionsUsed->add(iam); + fixture.attach_backend(iam); + fixture.session->previous_status.push(WAITING_CLIENT_DATA); + fixture.session->set_status(CHANGING_USER_SERVER); + change_user_calls.store(0, std::memory_order_relaxed); + fixture.session->to_process = 1; + fixture.session->handler(); + ok(fixture.selected() != nullptr && fixture.selected()->fd == -1 && + fixture.selected()->backend_auth_type() == MySQLBackendAuthType::PASSWORD && + server->ConnectionsUsed->conns_length() == 1 && + change_user_calls.load(std::memory_order_relaxed) == 0, + "CHANGING_USER_SERVER destroys IAM and returns through fresh acquisition"); +} + +void test_resetting_state_destroys_iam(MySQL_Thread& worker) { + MySrvC *server = create_server(807, "resetting-iam"); + SessionFixture fixture(worker, 807, MySQLBackendAuthType::AWS_IAM); + MySQL_Connection *iam = create_established_connection( + server, kUser, MySQLBackendAuthType::AWS_IAM); + server->ConnectionsUsed->add(iam); + fixture.attach_backend(iam); + fixture.session->set_status(RESETTING_CONNECTION); + change_user_calls.store(0, std::memory_order_relaxed); + fixture.session->to_process = 1; + const int rc = fixture.session->handler(); + ok(rc == -1 && fixture.selected() == nullptr && + change_user_calls.load(std::memory_order_relaxed) == 0 && + server->ConnectionsUsed->conns_length() == 0 && + server->ConnectionsFree->conns_length() == 0, + "RESETTING_CONNECTION destroys IAM without invoking COM_CHANGE_USER"); +} + +} // namespace + +#ifdef __linux__ +extern "C" { + +int __real_mysql_change_user_start( + my_bool *, MYSQL *, const char *, const char *, const char *); +int __real_mysql_real_connect_start(MYSQL **, MYSQL *, const char *, const char *, + const char *, const char *, unsigned int, const char *, unsigned long); +int __real_pthread_create(pthread_t *, const pthread_attr_t *, + void *(*)(void *), void *); + +int __wrap_mysql_change_user_start( + my_bool *ret, MYSQL *, const char *, const char *, const char *) +{ + change_user_calls.fetch_add(1, std::memory_order_relaxed); + *ret = 0; + return change_user_immediate_success.load(std::memory_order_relaxed) + ? 0 : MYSQL_WAIT_READ; +} + +int __wrap_mysql_real_connect_start(MYSQL **ret, MYSQL *, const char *, + const char *, const char *, const char *, unsigned int, const char *, + unsigned long) +{ + *ret = nullptr; + return MYSQL_WAIT_READ; +} + +int __wrap_pthread_create(pthread_t *thread, const pthread_attr_t *attr, + void *(*start_routine)(void *), void *arg) +{ + if (start_routine != &kill_query_thread) { + return __real_pthread_create(thread, attr, start_routine, arg); + } + ++kill_helper_dispatches; + KillArgs *kill_args = static_cast(arg); + kill_source_token_present_at_dispatch = + kill_source_connection != nullptr && + kill_source_connection->has_aws_iam_handshake_secret(); + kill_args_password_absent = kill_args->password == nullptr; + kill_args_id = kill_args->id; + kill_args_type = kill_args->kill_type; + kill_args_auth_type = kill_args->backend_auth_type; + kill_args_endpoint = kill_args->configured_endpoint; + kill_args_region = kill_args->region; + kill_args_database_user = kill_args->database_user; + kill_args_transport = kill_args->get_host_address(); + kill_args_port = kill_args->port; + kill_args_use_ssl = kill_args->use_ssl; + kill_args_deadline = kill_args->token_deadline; + delete kill_args; + return 0; +} + +} // extern "C" +#endif + +int main() { +#ifndef __linux__ + plan(1); + skip(1, "requires GNU ld --wrap support"); + return exit_status(); +#else + plan(31); + if (test_init_minimal() != 0 || test_init_auth() != 0 || + test_init_query_processor() != 0 || + test_init_hostgroups() != 0) { + BAIL_OUT("failed to initialize unit-test globals"); + } + GloMyLogger = new MySQL_Logger(); + if (!GloMyAuth->add( + const_cast(kUser), const_cast("password"), USERNAME_BACKEND, + false, 0, const_cast(""), false, false, false, 100, + const_cast(""), const_cast("")) || + !GloMyAuth->add( + const_cast(kOtherUser), const_cast("password"), USERNAME_BACKEND, + false, 0, const_cast(""), false, false, false, 100, + const_cast(""), const_cast("")) || + !GloMyAuth->add( + const_cast(kInvalidReloadUser), const_cast("password"), USERNAME_BACKEND, + false, 0, const_cast(""), false, false, false, 100, + const_cast(""), const_cast(""))) { + BAIL_OUT("failed to load backend user fixtures"); + } + GloMTH->num_threads = 1; + { + MySQL_Thread worker; + if (!worker.init()) BAIL_OUT("MySQL_Thread::init() failed"); + worker.curtime = 10000000; + test_identity_matrix(worker); // 7 + test_runtime_mode_changes_global(worker); // 2 + test_runtime_mode_change_local(worker); // 1 + test_mixed_user_mode_global(worker); // 4 + test_mixed_user_mode_local(worker); // 4 + test_attached_password_reload_to_invalid_fails_closed(worker); // 1 + test_invalid_policy_cannot_continue_change_user(worker); // 1 + test_invalid_policy_cannot_enter_reset(worker); // 1 + test_authorized_rowless_passthrough_can_enter_detached_reset(worker); // 1 + test_destroy_path_never_queues_iam(); // 2 + test_destroy_path_dispatches_detached_iam_connection_kill(); // 1 + test_reset_queue_worker_never_changes_iam(); // 1 + test_reset_queue_worker_never_resets_invalid_policy(); // 1 + test_reset_queue_worker_preserves_authorized_rowless_passthrough(); // 1 + test_reset_queue_worker_discards_unmarked_rowless_password(); // 1 + test_change_user_state_replaces_iam(worker); // 1 + test_resetting_state_destroys_iam(worker); // 1 + } + + delete GloMyLogger; + GloMyLogger = nullptr; + test_cleanup_hostgroups(); + test_cleanup_query_processor(); + test_cleanup_auth(); + test_cleanup_minimal(); + return exit_status(); +#endif +} diff --git a/test/tap/tests/unit/aws_iam_provider_boundary_unit-t.cpp b/test/tap/tests/unit/aws_iam_provider_boundary_unit-t.cpp new file mode 100644 index 0000000000..7859cabdc2 --- /dev/null +++ b/test/tap/tests/unit/aws_iam_provider_boundary_unit-t.cpp @@ -0,0 +1,227 @@ +#include "tap.h" + +#include "Aws_Iam_Provider.h" + +#include + +#include +#include +#include +#include +#include +#include +#include +#include + +using namespace std::chrono_literals; + +namespace { + +constexpr std::string_view kToken { "FAKE_PROVIDER_BOUNDARY_TOKEN" }; +std::atomic cleanse_calls { 0 }; +std::atomic source_destroyed { false }; +std::atomic replacement_install_attempted { false }; +std::atomic replacement_install_succeeded { false }; +std::atomic replacement_destroyed { false }; + +void tracked_cleanse(void *memory, size_t size) { + if (size == kToken.size() && std::memcmp(memory, kToken.data(), size) == 0) { + cleanse_calls.fetch_add(1, std::memory_order_relaxed); + } + OPENSSL_cleanse(memory, size); +} + +class CapturingSink final : public AwsIamCompletionSink { +public: + void post(AwsIamCompletion&& completion) override { + called = true; + value = std::move(completion); + } + + bool called { false }; + AwsIamCompletion value; +}; + +class FakeSource final : public AwsIamTokenSource { +public: + AwsIamRequestHandle request(const AwsIamTokenKey&, uint64_t opaque_id, + std::weak_ptr sink) override { + AwsIamCompletion completion; + completion.opaque_id = opaque_id; + completion.result.status = AwsIamStatus::OK; + completion.result.token = SecureString(kToken, tracked_cleanse); + if (auto live = sink.lock()) live->post(std::move(completion)); + return { 1 }; + } + + AwsIamTokenResult request_blocking(const AwsIamTokenKey&, + std::chrono::steady_clock::time_point) override { + AwsIamTokenResult result; + result.status = AwsIamStatus::OK; + result.token = SecureString(kToken, tracked_cleanse); + return result; + } + + void cancel(AwsIamRequestHandle) override {} + void invalidate(const AwsIamTokenKey&, uint64_t) override {} + void record_backend_connection(bool) override {} + void record_waiting_session(bool) override {} + AwsIamStatsSnapshot snapshot() const override { return {}; } +}; + +void destroy_source(AwsIamTokenSource *source) { + delete source; + source_destroyed.store(true, std::memory_order_release); +} + +void destroy_replacement(AwsIamTokenSource *source) { + delete source; + replacement_destroyed.store(true, std::memory_order_release); +} + +void destroy_source_and_attempt_replacement(AwsIamTokenSource *source) { + delete source; + void *module_handle = dlopen(nullptr, RTLD_NOW | RTLD_LOCAL); + auto *replacement = new FakeSource(); + replacement_install_attempted.store(true, std::memory_order_release); + const bool installed = module_handle != nullptr && + install_global_aws_iam_token_source( + replacement, destroy_replacement, module_handle); + replacement_install_succeeded.store(installed, std::memory_order_release); + if (!installed) { + delete replacement; + if (module_handle != nullptr) dlclose(module_handle); + } +} + +bool all_stats_zero(const AwsIamStatsSnapshot& snapshot) { + const AwsIamNamedStats rows = aws_iam_stats_mysql_global_rows(snapshot); + for (const AwsIamNamedStat& row : rows) { + if (row.value != 0) return false; + } + return true; +} + +} // namespace + +int main() { + plan(14); + shutdown_global_aws_iam_token_source(); + ok(!acquire_global_aws_iam_token_source(), + "provider registry starts without an installed source"); + + void *module_handle = dlopen(nullptr, RTLD_NOW | RTLD_LOCAL); + auto *source = new FakeSource(); + ok(module_handle != nullptr && install_global_aws_iam_token_source( + source, destroy_source, module_handle), + "provider registry accepts a source with a retained module handle"); + + AwsIamTokenSourceLease lease = acquire_global_aws_iam_token_source(); + ok(lease && lease.get() == source, + "acquiring the registry returns a retained source lease"); + + auto sink = std::make_shared(); + const AwsIamTokenKey key { + "boundary.example", 3306, "boundary-region", "boundary-user" + }; + const AwsIamRequestHandle request = lease->request(key, 42, sink); + ok(request.value == 1 && sink->called && sink->value.opaque_id == 42 && + sink->value.result.status == AwsIamStatus::OK && + std::string_view(sink->value.result.token.c_str(), + sink->value.result.token.size()) == kToken, + "a provider token moves through the public completion sink"); + sink->value.result.token.clear(); + ok(cleanse_calls.load(std::memory_order_relaxed) == 1, + "the moved token is cleansed by the public secure-string contract"); + + auto uninstall = std::async(std::launch::async, [source] { + return uninstall_global_aws_iam_token_source(source); + }); + const auto stop_deadline = std::chrono::steady_clock::now() + 1s; + while (acquire_global_aws_iam_token_source() && + std::chrono::steady_clock::now() < stop_deadline) { + std::this_thread::yield(); + } + ok(!acquire_global_aws_iam_token_source(), + "uninstall rejects new leases before draining the retained source"); + ok(uninstall.wait_for(20ms) == std::future_status::timeout && + !source_destroyed.load(std::memory_order_acquire), + "uninstall waits for the outstanding source lease before destruction"); + lease = AwsIamTokenSourceLease {}; + ok(uninstall.get() && source_destroyed.load(std::memory_order_acquire), + "uninstall destroys the source only after its final lease drains"); + + std::unique_ptr unavailable = + create_aws_iam_token_source({ 16, 8 }); + const AwsIamTokenResult unavailable_result = unavailable->request_blocking( + key, std::chrono::steady_clock::now() + 1s); + ok(!unavailable->support_compiled() && + unavailable_result.status == AwsIamStatus::SUPPORT_NOT_COMPILED && + unavailable_result.failure.category == "support_not_compiled", + "the provider-neutral fallback fails closed with a fixed diagnostic"); + + publish_global_aws_iam_token_source(unavailable.get()); + shutdown_global_aws_iam_token_source(); + ok(!acquire_global_aws_iam_token_source() && + all_stats_zero(unavailable->snapshot()), + "shutdown removes the fallback source and leaves all twelve stats zero"); + + replacement_install_attempted.store(false, std::memory_order_release); + replacement_install_succeeded.store(false, std::memory_order_release); + replacement_destroyed.store(false, std::memory_order_release); + void *retiring_handle = dlopen(nullptr, RTLD_NOW | RTLD_LOCAL); + auto *retiring_source = new FakeSource(); + if (retiring_handle == nullptr || !install_global_aws_iam_token_source( + retiring_source, destroy_source_and_attempt_replacement, retiring_handle)) { + BAIL_OUT("failed to install overlapping-retirement source"); + } + AwsIamTokenSourceLease retirement_lease = + acquire_global_aws_iam_token_source(); + auto overlapping_uninstall = std::async(std::launch::async, [retiring_source] { + return uninstall_global_aws_iam_token_source(retiring_source); + }); + auto overlapping_shutdown = std::async(std::launch::async, [] { + shutdown_global_aws_iam_token_source(); + }); + const auto retirement_deadline = std::chrono::steady_clock::now() + 1s; + while (acquire_global_aws_iam_token_source() && + std::chrono::steady_clock::now() < retirement_deadline) { + std::this_thread::yield(); + } + ok(overlapping_uninstall.wait_for(20ms) == std::future_status::timeout && + overlapping_shutdown.wait_for(20ms) == std::future_status::timeout, + "overlapping uninstall and shutdown both wait for the claimed source lease"); + retirement_lease = AwsIamTokenSourceLease {}; + overlapping_uninstall.get(); + overlapping_shutdown.get(); + ok(replacement_install_attempted.load(std::memory_order_acquire) && + !replacement_install_succeeded.load(std::memory_order_acquire) && + !replacement_destroyed.load(std::memory_order_acquire), + "replacement publication is rejected until overlapping retirement finishes"); + + // Clean up the replacement that an unfixed registry may have accepted before + // checking that publication reopens after both retirement callers return. + shutdown_global_aws_iam_token_source(); + replacement_destroyed.store(false, std::memory_order_release); + void *survivor_handle = dlopen(nullptr, RTLD_NOW | RTLD_LOCAL); + auto survivor = std::make_unique(); + FakeSource *survivor_observer = survivor.get(); + const bool survivor_installed = survivor_handle != nullptr && + install_global_aws_iam_token_source( + survivor_observer, destroy_replacement, survivor_handle); + if (survivor_installed) { + (void)survivor.release(); + } else if (survivor_handle != nullptr) { + dlclose(survivor_handle); + } + ok(survivor_installed, + "replacement publication reopens after every retirement caller finishes"); + AwsIamTokenSourceLease survivor_lease = acquire_global_aws_iam_token_source(); + ok(survivor_lease && survivor_lease.get() == survivor_observer && + !replacement_destroyed.load(std::memory_order_acquire), + "the post-retirement replacement remains published and alive"); + survivor_lease = AwsIamTokenSourceLease {}; + shutdown_global_aws_iam_token_source(); + + return exit_status(); +} diff --git a/test/tap/tests/unit/aws_iam_session_state_unit-t.cpp b/test/tap/tests/unit/aws_iam_session_state_unit-t.cpp new file mode 100644 index 0000000000..8beccf76bb --- /dev/null +++ b/test/tap/tests/unit/aws_iam_session_state_unit-t.cpp @@ -0,0 +1,764 @@ +/** + * @file aws_iam_session_state_unit-t.cpp + * @brief State-machine tests for asynchronous backend IAM token acquisition. + */ + +#include "tap.h" +#include "test_globals.h" +#include "test_init.h" + +#include "proxysql.h" +#include "cpp.h" +#include "Aws_Iam_Provider.h" +#include "MySQL_Authentication.hpp" +#include "MySQL_Data_Stream.h" +#include "MySQL_HostGroups_Manager.h" +#include "MySQL_Logger.hpp" + +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +extern MySQL_Authentication *GloMyAuth; +extern MySQL_HostGroups_Manager *MyHGM; +extern MySQL_Logger *GloMyLogger; + +namespace { + +constexpr int kHostgroup = 707; +constexpr const char *kEndpointA = "orders.cluster-abcdefghijkl.us-east-1.rds.amazonaws.com"; +constexpr const char *kEndpointB = "orders-ro.cluster-ro-abcdefghijkl.us-east-1.rds.amazonaws.com"; +constexpr const char *kRegion = "us-east-1"; +constexpr const char *kIamUser = "iam_backend"; +constexpr const char *kPasswordUser = "password_backend"; +constexpr const char *kUnknownPassthroughUser = "passthrough_unknown_backend"; +constexpr const char *kMalformedPassthroughUser = "passthrough_malformed_backend"; +constexpr const char *kSensitiveToken = "FAKE_AWS_SESSION_TOKEN_MUST_NOT_ESCAPE"; + +std::atomic token_cleanse_calls { 0 }; +std::atomic connector_calls { 0 }; + +void tracked_cleanse(void *ptr, size_t size) { + OPENSSL_cleanse(ptr, size); + token_cleanse_calls.fetch_add(1, std::memory_order_relaxed); +} + +AwsIamTokenResult result(AwsIamStatus status, bool include_token = false) { + AwsIamTokenResult value; + value.status = status; + if (include_token) { + value.token = SecureString(kSensitiveToken, tracked_cleanse); + } + if (status != AwsIamStatus::OK) { + value.failure.category = "credential_provider"; + value.failure.aws_error_code = "NoCredentials"; + value.failure.request_id = "request-123"; + } + return value; +} + +class FakeTokenSource final : public AwsIamTokenSource { +public: + enum class Mode { DELAYED, IMMEDIATE_OK, IMMEDIATE_QUEUE_FULL }; + + explicit FakeTokenSource(Mode mode = Mode::DELAYED) : mode_(mode) {} + + AwsIamRequestHandle request(const AwsIamTokenKey& key, uint64_t opaque_id, + std::weak_ptr sink) override { + keys.push_back(key); + opaque_ids.push_back(opaque_id); + sinks.push_back(sink); + AwsIamRequestHandle handle { next_handle++ }; + handles.push_back(handle); + if (mode_ == Mode::IMMEDIATE_OK) { + post(sinks.size() - 1, result(AwsIamStatus::OK, true)); + } else if (mode_ == Mode::IMMEDIATE_QUEUE_FULL) { + post(sinks.size() - 1, result(AwsIamStatus::QUEUE_FULL)); + } + return handle; + } + + AwsIamTokenResult request_blocking(const AwsIamTokenKey&, + std::chrono::steady_clock::time_point) override { + return result(AwsIamStatus::PROVIDER_ERROR); + } + + void cancel(AwsIamRequestHandle handle) override { + if (handle.value != 0) { + if (cancel_observed_session != nullptr && cancel_observed_worker != nullptr) { + const uint64_t waiter_id = cancel_observed_session->aws_iam_waiter_id; + cancel_saw_live_waiter = waiter_id != 0 && + cancel_observed_worker->aws_iam_waiters.count(waiter_id) == 1; + cancel_saw_waiting_metric = waiting_sessions == 1; + } + canceled.push_back(handle.value); + } + } + void invalidate(const AwsIamTokenKey&, uint64_t) override {} + void record_backend_connection(bool success) override { + if (success) ++backend_successes; + else ++backend_failures; + } + void record_waiting_session(bool waiting) override { + if (waiting) ++waiting_sessions; + else if (waiting_sessions != 0) --waiting_sessions; + } + AwsIamStatsSnapshot snapshot() const override { + AwsIamStatsSnapshot result; + result.waiting_sessions = waiting_sessions; + return result; + } + + void post(size_t request_index, AwsIamTokenResult value) { + if (auto sink = sinks.at(request_index).lock()) { + AwsIamCompletion completion; + completion.opaque_id = opaque_ids.at(request_index); + completion.result = std::move(value); + sink->post(std::move(completion)); + } + } + + Mode mode_; + uint64_t next_handle { 1 }; + std::vector keys; + std::vector opaque_ids; + std::vector> sinks; + std::vector handles; + std::vector canceled; + MySQL_Session *cancel_observed_session { nullptr }; + MySQL_Thread *cancel_observed_worker { nullptr }; + bool cancel_saw_live_waiter { false }; + bool cancel_saw_waiting_metric { false }; + unsigned int backend_successes { 0 }; + unsigned int backend_failures { 0 }; + uint64_t waiting_sessions { 0 }; +}; + +class BlockingFakeTokenSource final : public AwsIamTokenSource { +public: + ~BlockingFakeTokenSource() override { release_request(); } + + AwsIamRequestHandle request(const AwsIamTokenKey& key, uint64_t opaque_id, + std::weak_ptr sink) override { + { + std::lock_guard lock(request_mutex); + request_entered = true; + requests.push_back(key); + opaque_ids.push_back(opaque_id); + sinks.push_back(std::move(sink)); + } + request_cv.notify_all(); + request_worker = std::thread([this] { + std::unique_lock lock(request_mutex); + request_worker_blocked = true; + request_cv.notify_all(); + request_cv.wait(lock, [this] { return request_released; }); + }); + return AwsIamRequestHandle { next_handle++ }; + } + + AwsIamTokenResult request_blocking(const AwsIamTokenKey&, + std::chrono::steady_clock::time_point) override { + return result(AwsIamStatus::PROVIDER_ERROR); + } + + void cancel(AwsIamRequestHandle handle) override { + if (handle.value != 0) { + { + std::lock_guard lock(request_mutex); + canceled.push_back(handle.value); + } + release_request(); + } + } + void invalidate(const AwsIamTokenKey&, uint64_t) override {} + void record_backend_connection(bool) override {} + void record_waiting_session(bool waiting) override { + if (waiting) ++waiting_sessions; + else if (waiting_sessions != 0) --waiting_sessions; + } + AwsIamStatsSnapshot snapshot() const override { + AwsIamStatsSnapshot result; + result.waiting_sessions = waiting_sessions; + return result; + } + + void wait_for_request() { + std::unique_lock lock(request_mutex); + if (!request_cv.wait_for(lock, std::chrono::seconds(1), + [this] { return request_entered; })) { + BAIL_OUT("session did not enter the blocking IAM token source"); + } + } + + void wait_for_request_worker_to_block() { + std::unique_lock lock(request_mutex); + if (!request_cv.wait_for(lock, std::chrono::seconds(1), + [this] { return request_worker_blocked; })) { + BAIL_OUT("IAM request worker did not reach its blocking predicate"); + } + } + + void release_request() { + { + std::lock_guard lock(request_mutex); + request_released = true; + } + request_cv.notify_all(); + if (request_worker.joinable()) request_worker.join(); + } + + std::mutex request_mutex; + std::condition_variable request_cv; + bool request_entered { false }; + bool request_worker_blocked { false }; + bool request_released { false }; + std::thread request_worker; + uint64_t next_handle { 1 }; + std::vector requests; + std::vector opaque_ids; + std::vector> sinks; + std::vector canceled; + uint64_t waiting_sessions { 0 }; +}; + +bool add_backend_user(const char *username, const char *password, const char *attributes) { + return GloMyAuth->add( + (char *)username, (char *)password, USERNAME_BACKEND, + false, 0, (char *)"", false, false, false, 100, + (char *)attributes, (char *)""); +} + +MySrvC *add_server(const char *endpoint) { + srv_info_t info; + info.addr = endpoint; + info.port = 3306; + info.kind = "aws-iam-session-state-unit"; + srv_opts_t opts; + opts.weigth = 1; + opts.max_conns = 100; + opts.use_ssl = 1; + MyHGM->wrlock(); + const int rc = MyHGM->create_new_server_in_hg(kHostgroup, info, opts); + MyHGC *hostgroup = MyHGM->MyHGC_find(kHostgroup); + MyHGM->wrunlock(); + if (rc != 0 || hostgroup == nullptr) BAIL_OUT("failed to create IAM backend fixture"); + return MyHGM->find_server_in_hg(kHostgroup, endpoint, 3306); +} + +std::string capture_stderr(const std::function& action) { + FILE *captured = tmpfile(); + if (captured == nullptr) BAIL_OUT("tmpfile() failed"); + fflush(stderr); + const int saved = dup(STDERR_FILENO); + if (saved < 0 || dup2(fileno(captured), STDERR_FILENO) < 0) BAIL_OUT("stderr redirect failed"); + action(); + fflush(stderr); + dup2(saved, STDERR_FILENO); + close(saved); + std::string output; + char buffer[256]; + rewind(captured); + while (fgets(buffer, sizeof(buffer), captured) != nullptr) output += buffer; + fclose(captured); + return output; +} + +std::string client_output(MySQL_Data_Stream *stream) { + std::string bytes; + if (stream == nullptr || stream->PSarrayOUT == nullptr) return bytes; + for (unsigned int i = 0; i < stream->PSarrayOUT->len; ++i) { + const PtrSize_t& packet = stream->PSarrayOUT->pdata[i]; + bytes.append(static_cast(packet.ptr), packet.size); + } + return bytes; +} + +class SessionFixture { +public: + SessionFixture(MySQL_Thread& worker, const char *username = kIamUser) : worker_(worker) { + session = new MySQL_Session(); + session->thread = &worker_; + session->connections_handler = true; + frontend_stream = new MySQL_Data_Stream(); + frontend_stream->init(MYDS_FRONTEND, session, -1); + frontend = new MySQL_Connection(); + frontend_stream->attach_connection(frontend); + frontend_stream->myprot.init(&frontend_stream, frontend->userinfo, session); + session->client_myds = frontend_stream; + frontend->userinfo->set( + const_cast(username), const_cast("ordinary-password"), + const_cast("orders"), nullptr); + + session->mybe = session->find_or_create_backend(kHostgroup); + session->current_hostgroup = kHostgroup; + session->default_hostgroup = kHostgroup; + session->CurrentQuery.start_time = worker_.curtime; + session->previous_status.push(PROCESSING_QUERY); + session->set_status(CONNECTING_SERVER); + } + + ~SessionFixture() { + if (session != nullptr) delete session; + } + + void start() { + run(); + } + + int run() { + session->to_process = 1; + return session->handler(); + } + + MySQL_Connection *selected_connection() const { + return session && session->mybe && session->mybe->server_myds + ? session->mybe->server_myds->myconn : nullptr; + } + + MySrvC *selected_server() const { + return selected_connection() ? selected_connection()->parent : nullptr; + } + + MySQL_Thread& worker_; + MySQL_Session *session { nullptr }; + MySQL_Data_Stream *frontend_stream { nullptr }; + MySQL_Connection *frontend { nullptr }; +}; + +void complete_and_drain(MySQL_Thread& worker, FakeTokenSource& source, + AwsIamTokenResult value, size_t request_index = 0) { + source.post(request_index, std::move(value)); + worker.drain_aws_iam_completions(); +} + +void make_fast_forward(SessionFixture& fixture) { + fixture.session->session_fast_forward = SESSION_FORWARD_TYPE_PERMANENT; + while (!fixture.session->previous_status.empty()) { + fixture.session->previous_status.pop(); + } + fixture.session->previous_status.push(FAST_FORWARD); +} + +void test_immediate_cache_hit(MySQL_Thread& worker) { + FakeTokenSource source(FakeTokenSource::Mode::IMMEDIATE_OK); + publish_global_aws_iam_token_source(&source); + SessionFixture fixture(worker); + fixture.start(); + MySQL_Connection *selected = fixture.selected_connection(); + MySrvC *server = fixture.selected_server(); + ok(fixture.session->status == WAITING_AWS_IAM_TOKEN && selected != nullptr && source.keys.size() == 1, + "an immediate cache completion still enters the owner-thread waiting state"); + worker.drain_aws_iam_completions(); + const int rc = fixture.run(); + ok(rc == 0 && fixture.session->status == CONNECTING_SERVER && + fixture.selected_connection() == selected && fixture.selected_server() == server && + selected->has_aws_iam_handshake_secret(), + "cache-hit completion resumes the selected fresh connection without reselection"); +} + +void test_delayed_completion(MySQL_Thread& worker) { + FakeTokenSource source; + publish_global_aws_iam_token_source(&source); + SessionFixture fixture(worker); + fixture.start(); + MySQL_Connection *selected = fixture.selected_connection(); + ok(fixture.run() == 0 && + fixture.session->status == WAITING_AWS_IAM_TOKEN && source.waiting_sessions == 1, + "a session remains parked while its delayed token is unfinished"); + source.post(0, result(AwsIamStatus::OK, true)); + ok(fixture.session->status == WAITING_AWS_IAM_TOKEN && source.waiting_sessions == 1, + "a queued completion leaves the live-session gauge set until owner-thread exit"); + worker.drain_aws_iam_completions(); + fixture.run(); + ok(fixture.session->status == CONNECTING_SERVER && fixture.selected_connection() == selected && + source.waiting_sessions == 0, + "a delayed completion resumes the originally selected connection"); +} + +void test_provider_error_is_generic(MySQL_Thread& worker) { + FakeTokenSource source; + publish_global_aws_iam_token_source(&source); + SessionFixture fixture(worker); + fixture.start(); + complete_and_drain(worker, source, result(AwsIamStatus::PROVIDER_ERROR, true)); + const std::string log = capture_stderr([&fixture] { + fixture.run(); + }); + const std::string output = client_output(fixture.frontend_stream); + ok(fixture.session->status == WAITING_CLIENT_DATA && fixture.selected_connection() == nullptr, + "provider error destroys the held fresh connection and returns to the client state"); + ok(output.find(kSensitiveToken) == std::string::npos && + output.find("NoCredentials") == std::string::npos && + log.find(kSensitiveToken) == std::string::npos, + "provider failure never exposes a token or provider details to the client or log"); +} + +void test_queue_rejection(MySQL_Thread& worker) { + FakeTokenSource source(FakeTokenSource::Mode::IMMEDIATE_QUEUE_FULL); + publish_global_aws_iam_token_source(&source); + SessionFixture fixture(worker); + fixture.start(); + worker.drain_aws_iam_completions(); + fixture.run(); + ok(fixture.session->status == WAITING_CLIENT_DATA && fixture.selected_connection() == nullptr, + "queue rejection follows the generic failure path and releases the fresh connection"); +} + +void test_five_second_deadline(MySQL_Thread& worker) { + FakeTokenSource source; + publish_global_aws_iam_token_source(&source); + SessionFixture fixture(worker); + fixture.start(); + worker.curtime += 5000000; + fixture.run(); + ok(fixture.session->status == WAITING_CLIENT_DATA && source.canceled.size() == 1 && + fixture.selected_connection() == nullptr, + "the five-second IAM deadline cancels the request and destroys the retained connection"); +} + +void test_existing_backend_deadline_wins(MySQL_Thread& worker) { + FakeTokenSource source; + publish_global_aws_iam_token_source(&source); + SessionFixture fixture(worker); + fixture.session->mybe->server_myds->max_connect_time = worker.curtime + 1000000; + fixture.start(); + worker.curtime += 1000000; + fixture.run(); + ok(fixture.session->status == WAITING_CLIENT_DATA && source.canceled.size() == 1 && + fixture.selected_connection() == nullptr, + "an earlier backend-acquisition deadline wins over the five-second token deadline"); +} + +void test_frontend_disconnect(MySQL_Thread& worker) { + FakeTokenSource source; + publish_global_aws_iam_token_source(&source); + auto fixture = std::make_unique(worker); + fixture->start(); + worker.register_session(&worker, fixture->session, false); + fixture->session->healthy = 0; + fixture->session = nullptr; + worker.process_all_sessions(); + ok(source.canceled.size() == 1, + "frontend teardown cancels and unregisters an in-flight IAM waiter"); +} + +void test_late_completion_is_dropped(MySQL_Thread& worker) { + token_cleanse_calls.store(0, std::memory_order_relaxed); + FakeTokenSource source; + publish_global_aws_iam_token_source(&source); + SessionFixture fixture(worker); + fixture.start(); + worker.curtime += 5000000; + fixture.run(); + complete_and_drain(worker, source, result(AwsIamStatus::OK, true)); + ok(token_cleanse_calls.load(std::memory_order_relaxed) == 1 && + fixture.selected_connection() == nullptr, + "a late success after timeout is cleansed and cannot reattach a connection"); +} + +void test_shutdown_completion(MySQL_Thread& worker) { + FakeTokenSource source; + publish_global_aws_iam_token_source(&source); + SessionFixture fixture(worker); + fixture.start(); + complete_and_drain(worker, source, result(AwsIamStatus::SHUTDOWN)); + fixture.run(); + ok(fixture.session->status == WAITING_CLIENT_DATA && fixture.selected_connection() == nullptr, + "token-source shutdown resumes the owner thread only to perform generic cleanup"); +} + +void test_cancel_keeps_wait_state_until_provider_cancel(MySQL_Thread& worker) { + FakeTokenSource source; + publish_global_aws_iam_token_source(&source); + SessionFixture fixture(worker); + fixture.start(); + source.cancel_observed_session = fixture.session; + source.cancel_observed_worker = &worker; + fixture.session->cancel_aws_iam_wait(); + ok(source.canceled.size() == 1 && source.cancel_saw_live_waiter && + source.cancel_saw_waiting_metric, + "provider cancellation observes the live worker waiter and waiting metric"); +} + +void test_session_wait_keeps_original_source_leased(MySQL_Thread& worker) { + BlockingFakeTokenSource original_source; + publish_global_aws_iam_token_source(&original_source); + SessionFixture fixture(worker); + fixture.start(); + original_source.wait_for_request(); + original_source.wait_for_request_worker_to_block(); + + std::promise shutdown_started; + std::future shutdown_started_future = shutdown_started.get_future(); + auto shutdown = std::async(std::launch::async, [&shutdown_started] { + shutdown_started.set_value(); + shutdown_global_aws_iam_token_source(); + }); + shutdown_started_future.wait(); + const auto shutdown_entry_deadline = + std::chrono::steady_clock::now() + std::chrono::seconds(1); + for (;;) { + AwsIamTokenSourceLease probe = acquire_global_aws_iam_token_source(); + if (!probe) break; + if (std::chrono::steady_clock::now() >= shutdown_entry_deadline) { + BAIL_OUT("global IAM shutdown did not disable new leases"); + } + std::this_thread::yield(); + } + const bool shutdown_waits_for_session_lease = + shutdown.wait_for(std::chrono::milliseconds(20)) == std::future_status::timeout; + + ok(shutdown_waits_for_session_lease, + "global shutdown waits for the session-owned IAM lease"); + fixture.session->cancel_aws_iam_wait(); + shutdown.get(); + + FakeTokenSource republished_source; + publish_global_aws_iam_token_source(&republished_source); + ok(original_source.requests.size() == 1 && original_source.canceled.size() == 1 && + republished_source.keys.empty() && republished_source.canceled.empty(), + "the old IAM wait never requests or cancels a republished token source"); + publish_global_aws_iam_token_source(nullptr); +} + +void test_fast_forward_provider_failure_is_terminal(MySQL_Thread& worker) { + FakeTokenSource source; + publish_global_aws_iam_token_source(&source); + SessionFixture fixture(worker); + make_fast_forward(fixture); + fixture.start(); + complete_and_drain(worker, source, result(AwsIamStatus::PROVIDER_ERROR)); + fixture.run(); + const size_t output_after_failure = client_output(fixture.frontend_stream).size(); + fixture.run(); + ok(fixture.session->status == WAITING_CLIENT_DATA && + fixture.selected_connection() == nullptr && worker.aws_iam_waiters.empty() && + client_output(fixture.frontend_stream).size() == output_after_failure, + "fast-forward provider failure is terminal and cannot emit a second error"); +} + +void test_fast_forward_timeout_is_terminal(MySQL_Thread& worker) { + FakeTokenSource source; + publish_global_aws_iam_token_source(&source); + SessionFixture fixture(worker); + make_fast_forward(fixture); + fixture.start(); + worker.curtime += 5000000; + fixture.run(); + ok(fixture.session->status == WAITING_CLIENT_DATA && source.canceled.size() == 1 && + fixture.selected_connection() == nullptr && worker.aws_iam_waiters.empty(), + "fast-forward IAM timeout cancels once and reaches a terminal client state"); +} + +void test_fast_forward_config_failure_is_terminal(MySQL_Thread& worker) { + FakeTokenSource source; + publish_global_aws_iam_token_source(&source); + char *saved_ca = mysql_thread___ssl_p2s_ca; + mysql_thread___ssl_p2s_ca = strdup(""); + SessionFixture fixture(worker); + make_fast_forward(fixture); + fixture.start(); + free(mysql_thread___ssl_p2s_ca); + mysql_thread___ssl_p2s_ca = saved_ca; + ok(fixture.session->status == WAITING_CLIENT_DATA && source.keys.empty() && + fixture.selected_connection() == nullptr && worker.aws_iam_waiters.empty(), + "fast-forward IAM configuration failure releases the connection and is terminal"); +} + +void test_worker_shutdown_closes_delivery_boundary() { + token_cleanse_calls.store(0, std::memory_order_relaxed); + FakeTokenSource source; + publish_global_aws_iam_token_source(&source); + std::weak_ptr delivery; + { + auto worker = std::make_unique(); + if (!worker->init()) BAIL_OUT("shutdown worker init failed"); + free(mysql_thread___ssl_p2s_ca); + mysql_thread___ssl_p2s_ca = strdup("/unit/fake-ca.pem"); + worker->curtime = 10000000; + auto fixture = std::make_unique(*worker); + fixture->start(); + delivery = source.sinks.at(0); + worker->register_session(worker.get(), fixture->session, false); + fixture->session = nullptr; + worker.reset(); + } + source.post(0, result(AwsIamStatus::OK, true)); + ok(source.canceled.size() == 1 && delivery.expired() && + token_cleanse_calls.load(std::memory_order_relaxed) == 1, + "worker shutdown cancels waiters before closing the inbox and late results cleanse/drop"); +} + +void test_selected_server_retention(MySQL_Thread& worker) { + FakeTokenSource source; + publish_global_aws_iam_token_source(&source); + SessionFixture fixture(worker); + fixture.start(); + MySQL_Connection *selected = fixture.selected_connection(); + MySrvC *server = fixture.selected_server(); + complete_and_drain(worker, source, result(AwsIamStatus::OK, true)); + fixture.run(); + ok(source.keys.size() == 1 && fixture.selected_connection() == selected && + fixture.selected_server() == server && connector_calls.load(std::memory_order_relaxed) > 0, + "request-through-connect retains one fresh connection and never chooses an alternate server"); +} + +void test_password_mode_unchanged(MySQL_Thread& worker) { + FakeTokenSource source; + publish_global_aws_iam_token_source(&source); + const unsigned int calls_before = connector_calls.load(std::memory_order_relaxed); + SessionFixture fixture(worker, kPasswordUser); + fixture.start(); + ok(source.keys.empty() && fixture.session->status == CONNECTING_SERVER && + fixture.selected_connection() != nullptr && + fixture.selected_connection()->backend_auth_type() == MySQLBackendAuthType::PASSWORD && + connector_calls.load(std::memory_order_relaxed) == calls_before + 1, + "ordinary password-mode fresh connection acquisition remains synchronous and unchanged"); +} + +void test_unknown_user_passthrough_uses_password(MySQL_Thread& worker) { + FakeTokenSource source; + publish_global_aws_iam_token_source(&source); + const MySQLBackendAuthPolicy missing_policy = + resolve_mysql_backend_auth_policy(*GloMyAuth, kUnknownPassthroughUser); + const unsigned int calls_before = connector_calls.load(std::memory_order_relaxed); + SessionFixture fixture(worker, kUnknownPassthroughUser); + fixture.session->passthrough_credential = true; + fixture.start(); + ok(missing_policy.type == MySQLBackendAuthType::INVALID && + missing_policy.failure_code == "backend_user_not_found" && + source.keys.empty() && fixture.session->status == CONNECTING_SERVER && + fixture.selected_connection() != nullptr && + fixture.selected_connection()->backend_auth_type() == MySQLBackendAuthType::PASSWORD && + connector_calls.load(std::memory_order_relaxed) == calls_before + 1, + "authorized unknown-user pass-through keeps password backend semantics without a backend row"); +} + +void test_malformed_policy_stays_fail_closed_for_passthrough(MySQL_Thread& worker) { + FakeTokenSource source; + publish_global_aws_iam_token_source(&source); + const unsigned int calls_before = connector_calls.load(std::memory_order_relaxed); + SessionFixture fixture(worker, kMalformedPassthroughUser); + fixture.session->passthrough_credential = true; + fixture.start(); + ok(source.keys.empty() && fixture.session->status == WAITING_CLIENT_DATA && + fixture.selected_connection() == nullptr && + connector_calls.load(std::memory_order_relaxed) == calls_before, + "pass-through authorization never overrides a malformed backend IAM policy"); +} + +void test_sdk_off_source_reports_support_not_compiled(MySQL_Thread& worker) { + AwsIamRuntimeConfig config; + config.max_total_waiters = 128; + config.max_waiters_per_key = 8; + std::unique_ptr source = create_aws_iam_token_source(config); + publish_global_aws_iam_token_source(source.get()); + SessionFixture fixture(worker); + const std::string log = capture_stderr([&] { + fixture.start(); + if (fixture.session->status == WAITING_AWS_IAM_TOKEN) { + worker.drain_aws_iam_completions(); + fixture.run(); + } + }); + publish_global_aws_iam_token_source(nullptr); + ok(fixture.session->status == WAITING_CLIENT_DATA && + fixture.selected_connection() == nullptr && + log.find("category='support_not_compiled'") != std::string::npos, + "the SDK-off source fails closed with the documented support_not_compiled operator reason"); +} + +} // namespace + +extern "C" { + +int __real_mysql_real_connect_start(MYSQL **, MYSQL *, const char *, const char *, + const char *, const char *, unsigned int, const char *, unsigned long); + +int __wrap_mysql_real_connect_start(MYSQL **ret, MYSQL *mysql, const char *host, + const char *, const char *password, const char *, unsigned int port, + const char *, unsigned long) { + connector_calls.fetch_add(1, std::memory_order_relaxed); + *ret = nullptr; + mysql->host = strdup(host != nullptr ? host : ""); + mysql->passwd = strdup(password != nullptr ? password : ""); + mysql->port = port; + return MYSQL_WAIT_READ; +} + +} // extern "C" + +int main() { + plan(27); + if (test_init_minimal() != 0 || test_init_auth() != 0 || + test_init_query_processor() != 0 || test_init_hostgroups() != 0) { + BAIL_OUT("failed to initialize unit-test globals"); + } + GloMyLogger = new MySQL_Logger(); + ok(add_backend_user(kIamUser, "", "{\"backend_auth\":{\"type\":\"aws_iam\"}}"), + "IAM backend account fixture is loaded"); + ok(add_backend_user(kPasswordUser, "ordinary-password", ""), + "password backend account fixture is loaded"); + if (!add_backend_user(kMalformedPassthroughUser, "ordinary-password", + "{\"backend_auth\":{\"type\":17}}")) { + BAIL_OUT("malformed backend account fixture failed to load"); + } + add_server(kEndpointA); + add_server(kEndpointB); + MyHGC *hostgroup = MyHGM->MyHGC_find(kHostgroup); + free(hostgroup->attributes.aws_iam_region); + hostgroup->attributes.aws_iam_region = strdup(kRegion); + + { + MySQL_Thread worker; + if (!worker.init()) BAIL_OUT("MySQL_Thread::init() failed"); + free(mysql_thread___ssl_p2s_ca); + mysql_thread___ssl_p2s_ca = strdup("/unit/fake-ca.pem"); + worker.curtime = 10000000; + test_session_wait_keeps_original_source_leased(worker); + test_cancel_keeps_wait_state_until_provider_cancel(worker); + test_immediate_cache_hit(worker); + test_delayed_completion(worker); + test_provider_error_is_generic(worker); + test_queue_rejection(worker); + test_five_second_deadline(worker); + test_existing_backend_deadline_wins(worker); + test_frontend_disconnect(worker); + test_late_completion_is_dropped(worker); + test_shutdown_completion(worker); + test_fast_forward_provider_failure_is_terminal(worker); + test_fast_forward_timeout_is_terminal(worker); + test_fast_forward_config_failure_is_terminal(worker); + test_selected_server_retention(worker); + test_password_mode_unchanged(worker); + test_unknown_user_passthrough_uses_password(worker); + test_malformed_policy_stays_fail_closed_for_passthrough(worker); + test_sdk_off_source_reports_support_not_compiled(worker); + publish_global_aws_iam_token_source(nullptr); + } + test_worker_shutdown_closes_delivery_boundary(); + publish_global_aws_iam_token_source(nullptr); + + delete GloMyLogger; + GloMyLogger = nullptr; + test_cleanup_hostgroups(); + test_cleanup_query_processor(); + test_cleanup_auth(); + test_cleanup_minimal(); + return exit_status(); +} diff --git a/test/tap/tests/unit/aws_locality_config_unit-t.cpp b/test/tap/tests/unit/aws_locality_config_unit-t.cpp new file mode 100644 index 0000000000..c31ac9f254 --- /dev/null +++ b/test/tap/tests/unit/aws_locality_config_unit-t.cpp @@ -0,0 +1,195 @@ +#include "tap.h" +#include "test_globals.h" + +#include "Aws_Locality_Manager.h" +#include "MySQL_HostGroups_Manager.h" +#include "MySQL_Thread.h" +#include "proxysql_utils.h" + +#include +#include +#include +#include + +void init_myhgc_hostgroup_settings(const char* hostgroup_settings, MyHGC* myhgc); + +namespace { + +bool contains_variable(char** variables, const char* expected) { + for (char** item = variables; item != nullptr && *item != nullptr; ++item) { + if (strcmp(*item, expected) == 0) { + return true; + } + } + return false; +} + +void free_variables(char** variables) { + if (variables == nullptr) { + return; + } + for (char** item = variables; *item != nullptr; ++item) { + free(*item); + } + free(variables); +} + +std::string capture_invalid_policy_log(MyHGC& hostgroup) { + FILE* captured = tmpfile(); + if (captured == nullptr) { + return {}; + } + fflush(stderr); + const int saved_stderr = dup(STDERR_FILENO); + if (saved_stderr < 0 || dup2(fileno(captured), STDERR_FILENO) < 0) { + if (saved_stderr >= 0) { + close(saved_stderr); + } + fclose(captured); + return {}; + } + + init_myhgc_hostgroup_settings( + R"({"aws":{"locality_awareness":{"same_region_multiplier":2.0,"same_az_multiplier":"FAKE_SECRET_MULTIPLIER"}}})", + &hostgroup); + fflush(stderr); + dup2(saved_stderr, STDERR_FILENO); + close(saved_stderr); + + std::string output; + char buffer[256]; + rewind(captured); + while (fgets(buffer, sizeof(buffer), captured) != nullptr) { + output += buffer; + } + fclose(captured); + return output; +} + +} // namespace + +int main() { + plan(20); + + MyHGC hostgroup(42); + init_myhgc_hostgroup_settings( + R"({"aws_iam_region":"us-east-1","aws":{"locality_awareness":{"same_region_multiplier":2.5,"same_az_multiplier":4.75}}})", + &hostgroup); + ok(hostgroup.attributes.aws_locality_policy.valid, + "valid nested AWS locality policy is owned by the hostgroup"); + ok(hostgroup.attributes.aws_locality_policy.same_region_multiplier == 2.5 && + hostgroup.attributes.aws_locality_policy.same_az_multiplier == 4.75, + "hostgroup retains both floating-point multipliers"); + ok(hostgroup.attributes.aws_locality_policy.refresh_interval_seconds == 300 && + hostgroup.attributes.aws_locality_policy.stale_ttl_seconds == 1800, + "hostgroup policy receives the documented timing defaults"); + ok(hostgroup.attributes.aws_iam_region != nullptr && + strcmp(hostgroup.attributes.aws_iam_region, "us-east-1") == 0, + "locality policy does not alter the independent IAM authentication Region"); + + init_myhgc_hostgroup_settings( + R"({"aws":{"locality_awareness":{"same_region_multiplier":1.0,"same_az_multiplier":10.0,"refresh_interval_seconds":30,"stale_ttl_seconds":604800}}})", + &hostgroup); + ok(hostgroup.attributes.aws_locality_policy.valid && + hostgroup.attributes.aws_locality_policy.refresh_interval_seconds == 30 && + hostgroup.attributes.aws_locality_policy.stale_ttl_seconds == 604800, + "inclusive policy bounds survive hostgroup parsing"); + + init_myhgc_hostgroup_settings( + R"({"aws":{"locality_awareness":{"same_region_multiplier":5.0,"same_az_multiplier":4.0}}})", + &hostgroup); + ok(!hostgroup.attributes.aws_locality_policy.valid, + "invalid reload clears the previously accepted policy"); + init_myhgc_hostgroup_settings(R"({"aws":[]})", &hostgroup); + ok(!hostgroup.attributes.aws_locality_policy.valid, + "non-object AWS settings remain disabled"); + init_myhgc_hostgroup_settings("{}", &hostgroup); + ok(!hostgroup.attributes.aws_locality_policy.valid, + "removing locality settings removes the runtime policy"); + + const std::string diagnostics = capture_invalid_policy_log(hostgroup); + ok(diagnostics.find("hostgroup 42") != std::string::npos && + diagnostics.find("same_az_multiplier") != std::string::npos, + "invalid locality diagnostic identifies only hostgroup and field"); + ok(diagnostics.find("FAKE_SECRET_MULTIPLIER") == std::string::npos, + "invalid locality diagnostic never prints the hostgroup JSON payload"); + ok(!hostgroup.attributes.aws_locality_policy.valid, + "rejected diagnostic case leaves no active locality policy"); + + test_globals_init(); + { + MySQL_Threads_Handler handler; + char** variables = handler.get_variables_list(); +#ifdef PROXYSQL31 + handler.set_variable("caching_sha2_password_auto_generate_rsa_keys", "false"); + handler.set_variable("caching_sha2_password_private_key_path", ""); + handler.set_variable("caching_sha2_password_public_key_path", ""); +#endif + mf_unique_ptr default_value { handler.get_variable("aws_locality_awareness") }; + ok(contains_variable(variables, "aws_locality_awareness"), + "v4 MySQL variable list exposes aws_locality_awareness"); + ok(default_value != nullptr && strcmp(default_value.get(), "false") == 0, + "aws_locality_awareness defaults to false"); + ok(handler.set_variable("aws_locality_awareness", "true") && + handler.get_variable_int("aws_locality_awareness") == 1, + "master switch accepts true"); + ok(handler.set_variable("AWS_LOCALITY_AWARENESS", "0") && + handler.get_variable_int("aws_locality_awareness") == 0, + "master switch is case-insensitive and accepts zero"); + ok(!handler.set_variable("aws_locality_awareness", "yes") && + handler.get_variable_int("aws_locality_awareness") == 0, + "master switch rejects non-boolean spelling"); + handler.set_variable("aws_locality_awareness", "1"); + ok(handler.commit().rejected_variables.empty() && + handler.get_variable_int("aws_locality_awareness") == 1, + "MySQL variable commit preserves the accepted master switch"); + free_variables(variables); + } + test_globals_cleanup(); + + GloVars.prometheus_registry = std::make_shared(); + { + MySQL_HostGroups_Manager manager; + MySQL_HostGroups_Manager *previous_hgm = MyHGM; + MyHGM = &manager; + srv_info_t info; + info.addr = "db.abcdef.us-east-1.rds.amazonaws.com"; + info.port = 3306; + info.kind = "AWS locality test"; + srv_opts_t options; + options.weigth = 7; + options.max_conns = 10; + options.use_ssl = 1; + + manager.wrlock(); + manager.create_new_server_in_hg(420, info, options); + MyHGC* configured_hostgroup = manager.MyHGC_find(420); + init_myhgc_hostgroup_settings( + R"({"aws":{"locality_awareness":{"same_region_multiplier":2.0,"same_az_multiplier":4.0}}})", + configured_hostgroup); + MySrvC* configured_server = static_cast( + configured_hostgroup->mysrvs->servers->index(0)); + manager.wrunlock(); + + manager.refresh_aws_locality_configuration(); + auto snapshot = manager.aws_locality_manager()->snapshot(); + const auto* entry = snapshot->find(420, info.addr, info.port); + ok(entry != nullptr && entry->configured_weight == 7, + "post-commit refresh copies hostgroup policy and backend identity into the manager"); + + manager.set_aws_locality_awareness_enabled(true); + manager.set_aws_locality_awareness_enabled(false); + ok(!manager.aws_locality_manager()->snapshot()->enabled && + configured_server->weight == 7, + "master-switch transitions never mutate the configured runtime server weight"); + + init_myhgc_hostgroup_settings("{}", configured_hostgroup); + manager.refresh_aws_locality_configuration(); + ok(manager.aws_locality_manager()->snapshot()->entries.empty(), + "removing the hostgroup policy removes its manager configuration"); + MyHGM = previous_hgm; + } + GloVars.prometheus_registry.reset(); + + return exit_status(); +} diff --git a/test/tap/tests/unit/aws_locality_manager_unit-t.cpp b/test/tap/tests/unit/aws_locality_manager_unit-t.cpp new file mode 100644 index 0000000000..f5fef8a6d8 --- /dev/null +++ b/test/tap/tests/unit/aws_locality_manager_unit-t.cpp @@ -0,0 +1,677 @@ +#include "tap.h" + +#include "Aws_Locality_Manager.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +using namespace std::chrono_literals; + +namespace { + +struct FakeProviderState { + struct Pending { + AwsMetadataRequestHandle handle; + AwsMetadataRequest request; + std::weak_ptr sink; + }; + + std::mutex mutex; + std::condition_variable cv; + std::vector requests; + std::vector canceled; + uint64_t next_handle { 1 }; + bool shut_down { false }; + bool destroyed { false }; + bool block_requests { false }; + bool request_entered { false }; + bool release_requests { false }; +}; + +class FakeProvider final : public AwsMetadataProvider { +public: + explicit FakeProvider(std::shared_ptr state) + : state_(std::move(state)) {} + + ~FakeProvider() override { + std::lock_guard lock(state_->mutex); + state_->destroyed = true; + state_->cv.notify_all(); + } + + AwsMetadataRequestHandle request( + const AwsMetadataRequest& request, + std::weak_ptr sink) override { + std::unique_lock lock(state_->mutex); + const AwsMetadataRequestHandle handle { state_->next_handle++ }; + state_->requests.push_back({handle, request, std::move(sink)}); + state_->request_entered = true; + state_->cv.notify_all(); + state_->cv.wait(lock, [&] { + return !state_->block_requests || state_->release_requests; + }); + return handle; + } + + void cancel(AwsMetadataRequestHandle handle) override { + std::lock_guard lock(state_->mutex); + state_->canceled.push_back(handle.value); + state_->cv.notify_all(); + } + + void shutdown() override { + std::lock_guard lock(state_->mutex); + state_->shut_down = true; + state_->cv.notify_all(); + } + +private: + std::shared_ptr state_; +}; + +void destroy_fake_provider(AwsMetadataProvider* provider) { + delete provider; +} + +bool wait_for_request_count( + const std::shared_ptr& state, + size_t count) { + std::unique_lock lock(state->mutex); + return state->cv.wait_for(lock, 2s, [&] { + return state->requests.size() >= count; + }); +} + +std::vector requests_copy( + const std::shared_ptr& state) { + std::lock_guard lock(state->mutex); + return state->requests; +} + +bool complete_request( + const std::shared_ptr& state, + AwsMetadataRequestKind kind, + const std::string& region, + AwsMetadataResult result, + size_t occurrence = 0) { + FakeProviderState::Pending pending; + { + std::lock_guard lock(state->mutex); + size_t found = 0; + for (const auto& item : state->requests) { + if (item.request.kind == kind && item.request.region == region) { + if (found++ == occurrence) { + pending = item; + break; + } + } + } + } + if (pending.handle.value == 0) { + return false; + } + if (auto sink = pending.sink.lock()) { + AwsMetadataCompletion completion; + completion.opaque_id = pending.request.opaque_id; + completion.generation = pending.request.generation; + completion.result = std::move(result); + sink->post(std::move(completion)); + return true; + } + return false; +} + +bool complete_request_after( + const std::shared_ptr& state, + AwsMetadataRequestKind kind, + const std::string& region, + AwsMetadataResult result, + size_t first_request) { + FakeProviderState::Pending pending; + { + std::lock_guard lock(state->mutex); + for (size_t index = first_request; index < state->requests.size(); ++index) { + const auto& item = state->requests[index]; + if (item.request.kind == kind && item.request.region == region) { + pending = item; + break; + } + } + } + if (pending.handle.value == 0) { + return false; + } + if (auto sink = pending.sink.lock()) { + AwsMetadataCompletion completion; + completion.opaque_id = pending.request.opaque_id; + completion.generation = pending.request.generation; + completion.result = std::move(result); + sink->post(std::move(completion)); + return true; + } + return false; +} + +template +bool wait_until(Predicate predicate) { + const auto deadline = std::chrono::steady_clock::now() + 2s; + while (std::chrono::steady_clock::now() < deadline) { + if (predicate()) { + return true; + } + std::this_thread::sleep_for(1ms); + } + return predicate(); +} + +AwsLocalityHostgroupConfig make_hostgroup( + uint32_t id, + double region_multiplier, + double az_multiplier, + uint32_t refresh, + uint32_t stale, + std::initializer_list endpoints) { + AwsLocalityHostgroupConfig config; + config.hostgroup_id = id; + config.policy.valid = true; + config.policy.same_region_multiplier = region_multiplier; + config.policy.same_az_multiplier = az_multiplier; + config.policy.refresh_interval_seconds = refresh; + config.policy.stale_ttl_seconds = stale; + for (const char* endpoint : endpoints) { + config.backends.push_back(recognize_rds_endpoint(id, endpoint, 3306)); + } + return config; +} + +const AwsLocalitySnapshotEntry* lookup( + const std::shared_ptr& snapshot, + uint32_t hostgroup_id, + const char* hostname) { + return snapshot->find(hostgroup_id, hostname, 3306); +} + +} // namespace + +int main() { + plan(49); + + auto provider_state = std::make_shared(); + ok(install_global_aws_metadata_provider( + new FakeProvider(provider_state), destroy_fake_provider, nullptr), + "metadata provider can be installed without an SDK dependency in core"); + { + auto lease = acquire_global_aws_metadata_provider(); + ok(lease && lease.get() != nullptr, + "installed provider is available through a retained lease"); + } + + std::atomic steady_seconds { 0 }; + std::atomic wall_seconds { 1700000000 }; + AwsLocalityManagerConfig manager_config; + manager_config.steady_clock = [&] { + return std::chrono::steady_clock::time_point( + std::chrono::seconds(steady_seconds.load())); + }; + manager_config.wall_clock = [&] { + return std::chrono::system_clock::time_point( + std::chrono::seconds(wall_seconds.load())); + }; + manager_config.request_timeout = 5s; + + MySQLAwsLocalityManager manager(manager_config); + const auto east_one = "db-1.abcdefghijkl.us-east-1.rds.amazonaws.com"; + const auto east_missing = "db-missing.abcdefghijkl.us-east-1.rds.amazonaws.com"; + const auto west_one = "db-2.abcdefghijkl.eu-west-1.rds.amazonaws.com"; + manager.configure({ + make_hostgroup(10, 2.0, 4.0, 300, 1800, {east_one, east_missing}), + make_hostgroup(11, 1.5, 3.0, 60, 120, {east_one, west_one}), + }); + const auto disabled_snapshot = manager.snapshot(); + const auto* disabled_entry = lookup(disabled_snapshot, 10, east_one); + ok(disabled_entry && disabled_entry->status == AwsLocalityMetadataStatus::disabled && + disabled_entry->multiplier == 1.0, + "configured manager publishes neutral diagnostic rows while disabled"); + + manager.set_enabled(true); + ok(wait_for_request_count(provider_state, 3), + "enable dispatches local identity and one coalesced request per Region"); + const auto initial_requests = requests_copy(provider_state); + const auto east_request = std::find_if(initial_requests.begin(), initial_requests.end(), + [](const auto& pending) { + return pending.request.kind == AwsMetadataRequestKind::rds_region && + pending.request.region == "us-east-1"; + }); + ok(east_request != initial_requests.end() && east_request->request.endpoints.size() == 2, + "duplicate endpoint registrations across hostgroups are coalesced"); + ok(east_request != initial_requests.end() && + east_request->request.deadline == manager_config.steady_clock() + 5s, + "provider request receives the configured monotonic deadline"); + + AwsMetadataResult local_result; + local_result.status = AwsMetadataStatus::ok; + local_result.local = {"us-east-1", "us-east-1a", "111122223333"}; + ok(complete_request(provider_state, AwsMetadataRequestKind::local_location, + "", std::move(local_result)), "local identity completion is accepted"); + + AwsMetadataResult east_result; + east_result.status = AwsMetadataStatus::ok; + east_result.endpoints.push_back({east_one, 3306, AwsEndpointType::instance, + "us-east-1", "us-east-1a", "111122223333"}); + ok(complete_request(provider_state, AwsMetadataRequestKind::rds_region, + "us-east-1", std::move(east_result)), "east Region completion is accepted"); + + AwsMetadataResult west_result; + west_result.status = AwsMetadataStatus::ok; + west_result.endpoints.push_back({west_one, 3306, AwsEndpointType::instance, + "eu-west-1", "eu-west-1a", "111122223333"}); + ok(complete_request(provider_state, AwsMetadataRequestKind::rds_region, + "eu-west-1", std::move(west_result)), "west Region completion is accepted"); + + ok(wait_until([&] { + auto snapshot = manager.snapshot(); + const auto* entry = lookup(snapshot, 10, east_one); + return entry != nullptr && entry->status == AwsLocalityMetadataStatus::fresh; + }), "immutable snapshot is published after normalized completions"); + + auto snapshot = manager.snapshot(); + const auto* hg10_east = lookup(snapshot, 10, east_one); + const auto* hg11_east = lookup(snapshot, 11, east_one); + const auto* hg11_west = lookup(snapshot, 11, west_one); + const auto* missing = lookup(snapshot, 10, east_missing); + ok(hg10_east && hg10_east->locality == AwsLocalityClass::same_az && + hg10_east->multiplier == 4.0, + "same-AZ result uses only the hostgroup's AZ multiplier"); + ok(hg11_east && hg11_east->multiplier == 3.0, + "one backend can use a different policy in another hostgroup"); + ok(hg11_west && hg11_west->locality == AwsLocalityClass::remote && + hg11_west->multiplier == 1.0, + "known remote backend remains neutral"); + ok(missing && missing->status == AwsLocalityMetadataStatus::error && + missing->failure_category == "endpoint_not_found" && missing->multiplier == 1.0, + "authoritative endpoint-not-found is redacted and neutral"); + ok(snapshot->effective_weight(10, east_one, 3306, 10) == 40 && + snapshot->effective_weight(11, west_one, 3306, 30) == 30, + "snapshot computes temporary weights without mutating configuration"); + + steady_seconds.store(61); + wall_seconds.store(1700000061); + manager.request_refresh(); + ok(wait_for_request_count(provider_state, 6), + "shortest configured refresh interval drives the next coalesced cycle"); + snapshot = manager.snapshot(); + hg10_east = lookup(snapshot, 10, east_one); + hg11_east = lookup(snapshot, 11, east_one); + ok(hg10_east && hg10_east->status == AwsLocalityMetadataStatus::fresh && + hg11_east && hg11_east->status == AwsLocalityMetadataStatus::stale, + "freshness is evaluated independently for each hostgroup policy"); + ok(hg11_east && hg11_east->multiplier == 3.0, + "stale but unexpired metadata remains active"); + AwsMetadataResult refresh_failure; + refresh_failure.status = AwsMetadataStatus::timeout; + const bool failed_refreshes_delivered = + complete_request(provider_state, AwsMetadataRequestKind::local_location, + "", refresh_failure, 1) && + complete_request(provider_state, AwsMetadataRequestKind::rds_region, + "us-east-1", refresh_failure, 1) && + complete_request(provider_state, AwsMetadataRequestKind::rds_region, + "eu-west-1", std::move(refresh_failure), 1); + const auto failed_refresh_snapshot = manager.snapshot(); + const auto* failed_refresh_entry = lookup(failed_refresh_snapshot, 11, east_one); + ok(failed_refreshes_delivered && + failed_refresh_entry != nullptr && + failed_refresh_entry->status == AwsLocalityMetadataStatus::stale, + "failed refresh preserves the last successful value within stale TTL"); + + steady_seconds.store(121); + wall_seconds.store(1700000121); + manager.request_refresh(); + ok(wait_for_request_count(provider_state, 9), + "forced scheduler wake dispatches another due refresh cycle"); + snapshot = manager.snapshot(); + hg11_east = lookup(snapshot, 11, east_one); + ok(hg11_east && hg11_east->status == AwsLocalityMetadataStatus::expired && + hg11_east->multiplier == 1.0, + "expired metadata becomes neutral"); + + size_t canceled_before_reload = 0; + { + std::lock_guard lock(provider_state->mutex); + canceled_before_reload = provider_state->canceled.size(); + } + manager.configure({ + make_hostgroup(12, 2.0, 4.0, 300, 1800, + {"db-new.abcdefghijkl.us-east-1.rds.amazonaws.com"}), + }); + ok(wait_until([&] { + std::lock_guard lock(provider_state->mutex); + return provider_state->canceled.size() > canceled_before_reload; + }), "configuration generation change cancels obsolete requests"); + const uint64_t new_generation = manager.snapshot()->generation; + ok(new_generation > snapshot->generation, + "configuration reload advances the immutable snapshot generation"); + + AwsMetadataResult late_result; + late_result.status = AwsMetadataStatus::ok; + late_result.endpoints.push_back({east_one, 3306, AwsEndpointType::instance, + "us-east-1", "us-east-1a", "111122223333"}); + ok(complete_request(provider_state, AwsMetadataRequestKind::rds_region, + "us-east-1", std::move(late_result), 1), + "obsolete provider callback can still reach the completion sink safely"); + ok(manager.snapshot()->find(10, east_one, 3306) == nullptr, + "completion from an obsolete generation cannot repopulate removed policy state"); + + manager.set_enabled(false); + ok(wait_until([&] { + const auto current = manager.snapshot(); + const auto* entry = lookup(current, 12, + "db-new.abcdefghijkl.us-east-1.rds.amazonaws.com"); + return entry && entry->status == AwsLocalityMetadataStatus::disabled && + entry->multiplier == 1.0; + }), "disabling the master switch retains neutral diagnostic rows"); + const size_t requests_while_disabled = requests_copy(provider_state).size(); + manager.request_refresh(); + std::this_thread::yield(); + ok(requests_copy(provider_state).size() == requests_while_disabled, + "disabled manager does not dispatch provider work"); + + manager.set_enabled(true); + ok(wait_for_request_count(provider_state, requests_while_disabled + 2), + "re-enabling resumes local and regional discovery"); + manager.shutdown(); + const auto shutdown_snapshot = manager.snapshot(); + ok(lookup(shutdown_snapshot, 12, + "db-new.abcdefghijkl.us-east-1.rds.amazonaws.com") != nullptr, + "manager shutdown retains its final neutral diagnostic snapshot"); + + std::atomic global_shutdown_done { false }; + auto held_lease = acquire_global_aws_metadata_provider(); + std::thread shutdown_thread([&] { + shutdown_global_aws_metadata_provider(); + global_shutdown_done.store(true); + }); + ok(wait_until([&] { + return !acquire_global_aws_metadata_provider(); + }) && !global_shutdown_done.load(), + "global shutdown rejects new leases while waiting for an active lease"); + held_lease = {}; + shutdown_thread.join(); + ok(global_shutdown_done.load(), + "global shutdown completes after the final provider lease drains"); + ok(provider_state->shut_down && provider_state->destroyed, + "provider shutdown and destruction occur after leases drain"); + + MySQLAwsLocalityManager absent_provider_manager(manager_config); + const auto absent_endpoint = "db-absent.abcdefghijkl.us-east-1.rds.amazonaws.com"; + auto absent_hostgroup = + make_hostgroup(20, 2.0, 4.0, 300, 1800, {absent_endpoint}); + absent_hostgroup.backends[0].configured_weight = 37; + absent_provider_manager.configure({std::move(absent_hostgroup)}); + absent_provider_manager.set_enabled(true); + ok(wait_until([&] { + const auto current = absent_provider_manager.snapshot(); + const auto* entry = lookup(current, 20, absent_endpoint); + return entry && entry->status == AwsLocalityMetadataStatus::error && + entry->failure_category == "provider_unavailable"; + }), "missing plugin provider is reported with a fixed neutral category"); + const auto absent_snapshot = absent_provider_manager.snapshot(); + const auto* absent_entry = lookup(absent_snapshot, 20, absent_endpoint); + ok(absent_entry != nullptr && absent_entry->configured_weight == 37 && + absent_entry->multiplier == 1.0 && + absent_snapshot->effective_weight(20, absent_endpoint, 3306, 37) == 37, + "missing provider preserves configured and effective neutral weights"); + + auto replacement_state = std::make_shared(); + ok(install_global_aws_metadata_provider( + new FakeProvider(replacement_state), destroy_fake_provider, nullptr), + "provider registry accepts a replacement after complete shutdown"); + absent_provider_manager.request_refresh(); + ok(wait_for_request_count(replacement_state, 2), + "manager acquires the replacement provider on its next refresh"); + absent_provider_manager.shutdown(); + + MySQLAwsLocalityManager concurrent_manager(manager_config); + std::vector concurrent_hostgroups; + concurrent_hostgroups.reserve(100); + std::vector concurrent_regions; + std::vector concurrent_endpoints; + for (uint32_t index = 1; index <= 100; ++index) { + const std::string region = "us-test-" + std::to_string(index); + const std::string endpoint = "db-" + std::to_string(index) + + ".abcdefghijkl." + region + ".rds.amazonaws.com"; + AwsLocalityHostgroupConfig hostgroup; + hostgroup.hostgroup_id = 100 + index; + hostgroup.policy = {true, 2.0, 4.0, 300, 1800}; + hostgroup.backends.push_back(recognize_rds_endpoint( + hostgroup.hostgroup_id, endpoint, 3306)); + concurrent_regions.push_back(region); + concurrent_endpoints.push_back(endpoint); + concurrent_hostgroups.push_back(std::move(hostgroup)); + } + concurrent_manager.configure(std::move(concurrent_hostgroups)); + concurrent_manager.set_enabled(true); + ok(wait_for_request_count(replacement_state, 103), + "100-Region fixture publishes one request per distinct Region"); + AwsMetadataResult concurrent_local; + concurrent_local.status = AwsMetadataStatus::ok; + concurrent_local.local = {"us-test-1", "us-test-1a", "111122223333"}; + complete_request(replacement_state, AwsMetadataRequestKind::local_location, + "", std::move(concurrent_local), 1); + std::atomic posted_completions { 0 }; + std::vector producers; + producers.reserve(100); + for (size_t index = 0; index < 100; ++index) { + producers.emplace_back([&, index] { + AwsMetadataResult result; + result.status = AwsMetadataStatus::ok; + result.endpoints.push_back({concurrent_endpoints[index], 3306, + AwsEndpointType::instance, concurrent_regions[index], + concurrent_regions[index] + "a", "111122223333"}); + if (complete_request(replacement_state, + AwsMetadataRequestKind::rds_region, concurrent_regions[index], + std::move(result))) { + posted_completions.fetch_add(1); + } + }); + } + for (auto& producer : producers) { + producer.join(); + } + ok(posted_completions.load() == 100, + "100 concurrent completion producers all reach the shared sink"); + ok(wait_until([&] { + const auto current = concurrent_manager.snapshot(); + return current->entries.size() == 100 && + std::all_of(current->entries.begin(), current->entries.end(), + [](const auto& item) { + return item.second.status == AwsLocalityMetadataStatus::fresh; + }); + }), "concurrent completions publish one consistent immutable snapshot"); + concurrent_manager.shutdown(); + + const size_t startup_request_count = requests_copy(replacement_state).size(); + MySQLAwsLocalityManager startup_manager(manager_config); + startup_manager.set_enabled(true); + startup_manager.configure({ + make_hostgroup(250, 2.0, 4.0, 300, 1800, + {"db-startup.abcdefghijkl.us-east-2.rds.amazonaws.com"}), + }); + ok(wait_for_request_count(replacement_state, startup_request_count + 2), + "configuration starts discovery when the master switch was enabled first"); + startup_manager.shutdown(); + + std::mutex completion_hook_mutex; + std::condition_variable completion_hook_cv; + bool completion_hook_entered = false; + bool release_completion = false; + AwsLocalityManagerConfig blocking_config = manager_config; + blocking_config.before_completion = [&] { + std::unique_lock lock(completion_hook_mutex); + completion_hook_entered = true; + completion_hook_cv.notify_all(); + completion_hook_cv.wait_for(lock, 2s, [&] { return release_completion; }); + }; + MySQLAwsLocalityManager blocking_manager(blocking_config); + const auto blocking_endpoint = "db-block.abcdefghijkl.ap-block-1.rds.amazonaws.com"; + const size_t blocking_request_count = requests_copy(replacement_state).size(); + blocking_manager.configure({ + make_hostgroup(300, 2.0, 4.0, 300, 1800, {blocking_endpoint}), + }); + blocking_manager.set_enabled(true); + ok(wait_for_request_count(replacement_state, blocking_request_count + 2), + "callback-shutdown fixture dispatches local and regional requests"); + std::atomic callback_returned { false }; + std::thread callback_thread([&] { + AwsMetadataResult result; + result.status = AwsMetadataStatus::ok; + result.local = {"ap-block-1", "ap-block-1a", "111122223333"}; + complete_request_after(replacement_state, + AwsMetadataRequestKind::local_location, "", std::move(result), + blocking_request_count); + callback_returned.store(true); + }); + { + std::unique_lock lock(completion_hook_mutex); + if (!completion_hook_cv.wait_for(lock, 2s, [&] { return completion_hook_entered; })) { + BAIL_OUT("completion hook did not enter before shutdown fixture deadline"); + } + } + std::mutex shutdown_mutex; + std::condition_variable shutdown_cv; + bool shutdown_done = false; + std::thread blocking_shutdown([&] { + blocking_manager.shutdown(); + std::lock_guard lock(shutdown_mutex); + shutdown_done = true; + shutdown_cv.notify_all(); + }); + bool shutdown_finished_early = false; + { + std::unique_lock lock(shutdown_mutex); + shutdown_finished_early = shutdown_cv.wait_for(lock, 100ms, [&] { + return shutdown_done; + }); + } + ok(!shutdown_finished_early, + "manager shutdown waits for an active completion callback"); + { + std::lock_guard lock(completion_hook_mutex); + release_completion = true; + completion_hook_cv.notify_all(); + } + callback_thread.join(); + blocking_shutdown.join(); + ok(callback_returned.load() && shutdown_done, + "callback and shutdown finish without accessing detached manager state"); + + AwsLocalityManagerConfig bounded_disable_config = manager_config; + bounded_disable_config.disable_wait_timeout = 50ms; + MySQLAwsLocalityManager bounded_disable_manager(bounded_disable_config); + { + std::lock_guard lock(replacement_state->mutex); + replacement_state->block_requests = true; + replacement_state->request_entered = false; + replacement_state->release_requests = false; + } + bounded_disable_manager.configure({ + make_hostgroup(400, 2.0, 4.0, 300, 1800, + {"db-disable.abcdefghijkl.us-east-2.rds.amazonaws.com"}), + }); + bounded_disable_manager.set_enabled(true); + { + std::unique_lock lock(replacement_state->mutex); + if (!replacement_state->cv.wait_for(lock, 2s, [&] { + return replacement_state->request_entered; + })) { + BAIL_OUT("provider request did not enter bounded-disable fixture"); + } + } + std::mutex disable_mutex; + std::condition_variable disable_cv; + bool disable_returned = false; + std::thread disable_thread([&] { + bounded_disable_manager.set_enabled(false); + std::lock_guard lock(disable_mutex); + disable_returned = true; + disable_cv.notify_all(); + }); + bool disable_returned_within_bound = false; + { + std::unique_lock lock(disable_mutex); + disable_returned_within_bound = disable_cv.wait_for(lock, 250ms, [&] { + return disable_returned; + }); + } + ok(disable_returned_within_bound, + "disabling locality does not block admin on a stalled provider call"); + { + std::lock_guard lock(replacement_state->mutex); + replacement_state->release_requests = true; + replacement_state->cv.notify_all(); + } + disable_thread.join(); + bounded_disable_manager.shutdown(); + + MySQLAwsLocalityManager invalid_identity_manager(manager_config); + AwsLocalityHostgroupConfig invalid_identity_hostgroup; + invalid_identity_hostgroup.hostgroup_id = 500; + invalid_identity_hostgroup.policy = {true, 2.0, 4.0, 300, 1800}; + AwsEndpointCandidate ipv6_identity; + ipv6_identity.hostgroup_id = 500; + ipv6_identity.hostname = "fe80::1"; + ipv6_identity.port = 3306; + AwsEndpointCandidate path_identity = ipv6_identity; + path_identity.hostname = "bad/name"; + invalid_identity_hostgroup.backends.emplace_back(ipv6_identity, 7); + invalid_identity_hostgroup.backends.emplace_back(path_identity, 9); + invalid_identity_manager.configure({std::move(invalid_identity_hostgroup)}); + const auto invalid_identity_snapshot = invalid_identity_manager.snapshot(); + const auto* ipv6_entry = invalid_identity_snapshot->find(500, "fe80::1", 3306); + const auto* path_entry = invalid_identity_snapshot->find(500, "bad/name", 3306); + ok(invalid_identity_snapshot->entries.size() == 2 && ipv6_entry != nullptr && + path_entry != nullptr && ipv6_entry->configured_weight == 7 && + path_entry->configured_weight == 9, + "rejected DNS spellings retain distinct snapshot identities"); + invalid_identity_manager.shutdown(); + + const size_t lease_drain_request_count = requests_copy(replacement_state).size(); + MySQLAwsLocalityManager lease_drain_manager(manager_config); + lease_drain_manager.configure({ + make_hostgroup(600, 2.0, 4.0, 300, 1800, + {"db-drain.abcdefghijkl.us-east-2.rds.amazonaws.com"}), + }); + lease_drain_manager.set_enabled(true); + ok(wait_for_request_count(replacement_state, lease_drain_request_count + 2), + "enabled manager retains the installed provider during discovery"); + std::atomic unload_done { false }; + std::thread unload_thread([&] { + shutdown_global_aws_metadata_provider(); + unload_done.store(true); + }); + ok(wait_until([&] { + return !acquire_global_aws_metadata_provider(); + }) && !unload_done.load(), + "provider unload rejects new leases while the manager lease is active"); + lease_drain_manager.set_enabled(false); + unload_thread.join(); + ok(unload_done.load() && replacement_state->shut_down && + replacement_state->destroyed, + "disabling locality drains the manager lease before provider destruction"); + const size_t requests_after_unload = requests_copy(replacement_state).size(); + lease_drain_manager.shutdown(); + lease_drain_manager.set_enabled(true); + lease_drain_manager.request_refresh(); + std::this_thread::sleep_for(20ms); + ok(requests_copy(replacement_state).size() == requests_after_unload, + "manager shutdown leaves no locality worker able to restart provider work"); + + return exit_status(); +} diff --git a/test/tap/tests/unit/aws_locality_policy_unit-t.cpp b/test/tap/tests/unit/aws_locality_policy_unit-t.cpp new file mode 100644 index 0000000000..152406e07c --- /dev/null +++ b/test/tap/tests/unit/aws_locality_policy_unit-t.cpp @@ -0,0 +1,190 @@ +#include "tap.h" + +#include "json.hpp" +#include "Aws_Locality_Manager.h" + +#include +#include +#include + +using nlohmann::json; + +namespace { + +void test_policy_validation() { + AwsLocalityPolicyError error; + AwsLocalityPolicy policy = parse_aws_locality_policy( + json::parse(R"({"same_region_multiplier":2.5,"same_az_multiplier":4.75})"), + 10, error); + ok(policy.valid && policy.same_region_multiplier == 2.5 && + policy.same_az_multiplier == 4.75 && + policy.refresh_interval_seconds == 300 && + policy.stale_ttl_seconds == 1800 && error.field.empty(), + "valid policy uses explicit multipliers and timing defaults"); + + policy = parse_aws_locality_policy(json::parse( + R"({"same_region_multiplier":1.0,"same_az_multiplier":10.0,"refresh_interval_seconds":30,"stale_ttl_seconds":604800})"), + 11, error); + ok(policy.valid && policy.refresh_interval_seconds == 30 && + policy.stale_ttl_seconds == 604800, + "inclusive multiplier and timing boundaries are accepted"); + + policy = parse_aws_locality_policy(json::parse( + R"({"same_az_multiplier":4.0})"), 12, error); + ok(!policy.valid && error.field == "same_region_multiplier", + "missing same-Region multiplier rejects the complete locality policy"); + + policy = parse_aws_locality_policy(json::parse( + R"({"same_region_multiplier":2.0,"same_az_multiplier":"4"})"), 13, error); + ok(!policy.valid && error.field == "same_az_multiplier", + "a numeric-looking string is not accepted as a multiplier"); + + json non_finite = { + {"same_region_multiplier", std::numeric_limits::quiet_NaN()}, + {"same_az_multiplier", 4.0} + }; + policy = parse_aws_locality_policy(non_finite, 14, error); + ok(!policy.valid && error.field == "same_region_multiplier", + "non-finite multiplier is rejected"); + + policy = parse_aws_locality_policy(json::parse( + R"({"same_region_multiplier":5.0,"same_az_multiplier":4.0})"), 15, error); + ok(!policy.valid && error.field == "same_az_multiplier", + "same-AZ multiplier cannot be lower than same-Region multiplier"); + + policy = parse_aws_locality_policy(json::parse( + R"({"same_region_multiplier":0.99,"same_az_multiplier":4.0})"), 16, error); + ok(!policy.valid && error.field == "same_region_multiplier", + "multiplier below 1.0 is rejected"); + + policy = parse_aws_locality_policy(json::parse( + R"({"same_region_multiplier":2.0,"same_az_multiplier":10.01})"), 17, error); + ok(!policy.valid && error.field == "same_az_multiplier", + "multiplier above 10.0 is rejected"); + + policy = parse_aws_locality_policy(json::parse( + R"({"same_region_multiplier":2.0,"same_az_multiplier":4.0,"refresh_interval_seconds":29})"), + 18, error); + ok(!policy.valid && error.field == "refresh_interval_seconds", + "refresh interval below 30 seconds is rejected"); + + policy = parse_aws_locality_policy(json::parse( + R"({"same_region_multiplier":2.0,"same_az_multiplier":4.0,"refresh_interval_seconds":500,"stale_ttl_seconds":499})"), + 19, error); + ok(!policy.valid && error.field == "stale_ttl_seconds", + "stale TTL shorter than refresh interval is rejected"); + + policy = parse_aws_locality_policy(json::parse( + R"({"same_region_multiplier":2.0,"same_az_multiplier":4.0,"refresh_interval_seconds":3600})"), + 20, error); + ok(!policy.valid && error.field == "stale_ttl_seconds", + "omitted stale TTL cannot undercut an explicit refresh interval"); + + policy = parse_aws_locality_policy(json::array(), 21, error); + ok(!policy.valid && error.field == "locality_awareness", + "non-object locality policy is rejected"); +} + +void test_endpoint_recognition() { + AwsEndpointCandidate candidate = recognize_rds_endpoint( + 10, "DB-1.ABCDEF.us-east-1.rds.amazonaws.com.", 3306); + ok(candidate.recognized && + candidate.hostname == "db-1.abcdef.us-east-1.rds.amazonaws.com" && + candidate.region == "us-east-1" && candidate.partition == "aws" && + candidate.hostgroup_id == 10 && candidate.port == 3306, + "standard RDS endpoint is normalized and routed to its Region"); + + candidate = recognize_rds_endpoint( + 11, "db.cluster-abcdefghijkl.us-gov-west-1.rds.amazonaws.com", 3306); + ok(candidate.recognized && candidate.region == "us-gov-west-1" && + candidate.partition == "aws-us-gov", + "GovCloud RDS endpoint is recognized"); + + candidate = recognize_rds_endpoint( + 12, "db.cluster-ro-abcdefghijkl.cn-north-1.rds.amazonaws.com.cn", 3306); + ok(candidate.recognized && candidate.region == "cn-north-1" && + candidate.partition == "aws-cn", + "China RDS endpoint is recognized"); + + candidate = recognize_rds_endpoint( + 13, "db.cluster-custom-abcdefghijkl.eu-west-1.rds.amazonaws.com", 3306); + ok(candidate.recognized && candidate.region == "eu-west-1", + "Aurora custom endpoint is a valid discovery candidate"); + + ok(!recognize_rds_endpoint(14, "db.proxy-abcdefghijkl.us-east-1.rds.amazonaws.com", 3306).recognized, + "RDS Proxy endpoint is excluded from the first version"); + candidate = recognize_rds_endpoint( + 14, "proxy-app1.abcdefghijkl.us-east-1.rds.amazonaws.com", 3306); + ok(candidate.recognized && candidate.region == "us-east-1", + "an ordinary RDS identifier beginning with proxy- is not misclassified as RDS Proxy"); + ok(!recognize_rds_endpoint(15, "db.internal.example.com", 3306).recognized, + "custom CNAME is not inferred as RDS"); + ok(!recognize_rds_endpoint(16, "db.us-east-1.rds.amazonaws.com.evil.example", 3306).recognized, + "AWS-looking hostname with a false suffix is rejected"); + ok(!recognize_rds_endpoint(17, "db.INVALID_REGION.rds.amazonaws.com", 3306).recognized, + "malformed candidate Region is rejected"); +} + +void test_classification() { + const AwsLocalLocation local {"us-east-1", "us-east-1a", "111122223333"}; + + AwsBackendLocation backend { + AwsEndpointType::instance, "us-east-1", "us-east-1a", "111122223333" + }; + ok(classify_aws_locality(local, backend) == AwsLocalityClass::same_az, + "instance in same account and AZ receives same-AZ classification"); + + backend.account_id = "444455556666"; + ok(classify_aws_locality(local, backend) == AwsLocalityClass::same_region, + "same AZ name in a different account is only same-Region"); + + backend.account_id.clear(); + ok(classify_aws_locality(local, backend) == AwsLocalityClass::same_region, + "unknown backend account cannot receive same-AZ classification"); + + backend = {AwsEndpointType::cluster, "us-east-1", "us-east-1a", "111122223333"}; + ok(classify_aws_locality(local, backend) == AwsLocalityClass::same_region, + "cluster endpoint never receives same-AZ classification"); + + backend = {AwsEndpointType::reader, "us-east-1", "", "111122223333"}; + ok(classify_aws_locality(local, backend) == AwsLocalityClass::same_region, + "reader endpoint can receive same-Region classification"); + + backend = {AwsEndpointType::instance, "eu-west-1", "eu-west-1a", "111122223333"}; + ok(classify_aws_locality(local, backend) == AwsLocalityClass::remote, + "different known Region is remote"); + + backend.region.clear(); + ok(classify_aws_locality(local, backend) == AwsLocalityClass::unknown, + "missing backend Region is neutral"); + + AwsLocalLocation unknown_local {"", "", ""}; + backend.region = "us-east-1"; + ok(classify_aws_locality(unknown_local, backend) == AwsLocalityClass::unknown, + "missing local Region is neutral"); +} + +void test_effective_weight() { + ok(aws_locality_effective_weight(3, 2.5) == 7, + "effective weight truncates toward zero"); + ok(aws_locality_effective_weight(0, 10.0) == 0, + "configured zero weight remains zero"); + ok(aws_locality_effective_weight(-1, 4.0) == 0, + "unexpected negative configured weight is safely neutralized"); + ok(aws_locality_effective_weight(std::numeric_limits::max(), 10.0) == + std::numeric_limits::max(), + "effective weight saturates before the 64-bit accumulator"); + ok(aws_locality_effective_weight(10, 1.0) == 10, + "neutral multiplier preserves configured weight"); +} + +} // namespace + +int main() { + plan(34); + test_policy_validation(); + test_endpoint_recognition(); + test_classification(); + test_effective_weight(); + return exit_status(); +} diff --git a/test/tap/tests/unit/aws_locality_selection_unit-t.cpp b/test/tap/tests/unit/aws_locality_selection_unit-t.cpp new file mode 100644 index 0000000000..4af69f2889 --- /dev/null +++ b/test/tap/tests/unit/aws_locality_selection_unit-t.cpp @@ -0,0 +1,487 @@ +#include "tap.h" +#include "test_globals.h" +#include "test_init.h" + +#include "Aws_Locality_Manager.h" +#include "GTID_Server_Data.h" +#include "MySQL_Data_Stream.h" +#include "MySQL_HostGroups_Manager.h" +#include "MySQL_Logger.hpp" +#include "MySQL_Thread.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +using namespace std::chrono_literals; + +namespace { +thread_local bool track_hot_path_allocations = false; +thread_local size_t hot_path_allocations = 0; +} + +void* operator new(std::size_t size) { + if (track_hot_path_allocations) ++hot_path_allocations; + if (void* memory = std::malloc(size)) return memory; + throw std::bad_alloc(); +} + +void* operator new[](std::size_t size) { + return ::operator new(size); +} + +void operator delete(void* memory) noexcept { std::free(memory); } +void operator delete[](void* memory) noexcept { std::free(memory); } +void operator delete(void* memory, std::size_t) noexcept { std::free(memory); } +void operator delete[](void* memory, std::size_t) noexcept { std::free(memory); } + +extern MySQL_HostGroups_Manager* MyHGM; +extern MySQL_Threads_Handler* GloMTH; +extern MySQL_Logger* GloMyLogger; + +void init_myhgc_hostgroup_settings(const char* hostgroup_settings, MyHGC* myhgc); + +namespace { + +constexpr unsigned int kHostgroup = 740; +constexpr const char* kUser = "locality_user"; +constexpr const char* kSchema = "locality_schema"; +constexpr const char* kLocal = "db-local.abcdefghijkl.us-east-1.rds.amazonaws.com"; +constexpr const char* kRegional = "cluster-regional.abcdefghijkl.us-east-1.rds.amazonaws.com"; +constexpr const char* kRemote = "db-remote.abcdefghijkl.eu-west-1.rds.amazonaws.com"; + +struct ProviderState { + struct Pending { + AwsMetadataRequestHandle handle; + AwsMetadataRequest request; + std::weak_ptr sink; + }; + + std::mutex mutex; + std::condition_variable cv; + std::vector pending; + uint64_t next_handle { 1 }; +}; + +class FakeProvider final : public AwsMetadataProvider { +public: + explicit FakeProvider(std::shared_ptr state) + : state_(std::move(state)) {} + + AwsMetadataRequestHandle request( + const AwsMetadataRequest& request, + std::weak_ptr sink) override { + std::lock_guard lock(state_->mutex); + AwsMetadataRequestHandle handle { state_->next_handle++ }; + state_->pending.push_back({handle, request, std::move(sink)}); + state_->cv.notify_all(); + return handle; + } + + void cancel(AwsMetadataRequestHandle) override {} + void shutdown() override {} + +private: + std::shared_ptr state_; +}; + +void destroy_provider(AwsMetadataProvider* provider) { + delete provider; +} + +bool wait_for_requests(const std::shared_ptr& state, size_t count) { + std::unique_lock lock(state->mutex); + return state->cv.wait_for(lock, 2s, [&] { return state->pending.size() >= count; }); +} + +bool complete( + const std::shared_ptr& state, + AwsMetadataRequestKind kind, + const char* region, + AwsMetadataResult result) { + ProviderState::Pending selected; + { + std::lock_guard lock(state->mutex); + for (const auto& pending : state->pending) { + if (pending.request.kind == kind && pending.request.region == region) { + selected = pending; + break; + } + } + } + if (selected.handle.value == 0) return false; + auto sink = selected.sink.lock(); + if (!sink) return false; + AwsMetadataCompletion completion; + completion.opaque_id = selected.request.opaque_id; + completion.generation = selected.request.generation; + completion.result = std::move(result); + sink->post(std::move(completion)); + return true; +} + +template +bool wait_until(Predicate predicate) { + const auto deadline = std::chrono::steady_clock::now() + 2s; + while (std::chrono::steady_clock::now() < deadline) { + if (predicate()) return true; + std::this_thread::yield(); + } + return predicate(); +} + +MySrvC* add_server(const char* hostname, int64_t weight) { + srv_info_t info; + info.addr = hostname; + info.port = 3306; + info.kind = "AWS locality selection test"; + srv_opts_t options; + options.weigth = weight; + options.max_conns = 100; + options.use_ssl = 1; + + MyHGM->wrlock(); + const int result = MyHGM->create_new_server_in_hg(kHostgroup, info, options); + MyHGC* hostgroup = MyHGM->MyHGC_find(kHostgroup); + MySrvC* server = nullptr; + if (hostgroup != nullptr) { + for (unsigned int i = 0; i < hostgroup->mysrvs->cnt(); ++i) { + MySrvC* candidate = hostgroup->mysrvs->idx(i); + if (strcmp(candidate->address, hostname) == 0) { + server = candidate; + break; + } + } + } + MyHGM->wrunlock(); + if (result != 0 || server == nullptr) { + BAIL_OUT("failed to add locality server %s", hostname); + } + return server; +} + +class SessionFixture { +public: + explicit SessionFixture(MySQL_Thread& worker) { + session = new MySQL_Session(); + session->thread = &worker; + session->connections_handler = true; + frontend_stream = new MySQL_Data_Stream(); + frontend_stream->init(MYDS_FRONTEND, session, -1); + frontend = new MySQL_Connection(); + frontend_stream->attach_connection(frontend); + frontend_stream->myprot.init(&frontend_stream, frontend->userinfo, session); + session->client_myds = frontend_stream; + frontend->userinfo->set( + const_cast(kUser), const_cast("password"), + const_cast(kSchema), nullptr); + frontend->set_backend_auth_type(MySQLBackendAuthType::PASSWORD); + } + + ~SessionFixture() { delete session; } + + MySQL_Session* session { nullptr }; + MySQL_Data_Stream* frontend_stream { nullptr }; + MySQL_Connection* frontend { nullptr }; +}; + +MySQL_Connection* make_connection(MySrvC* server, int fd) { + MySQL_Connection* connection = new MySQL_Connection(); + connection->mysql = mysql_init(nullptr); + if (connection->mysql == nullptr) BAIL_OUT("mysql_init failed"); + connection->ret_mysql = connection->mysql; + connection->mysql->charset = mariadb_get_charset_by_name("utf8mb4"); + connection->parent = server; + connection->userinfo->set( + const_cast(kUser), const_cast("password"), + const_cast(kSchema), nullptr); + connection->set_backend_auth_type(MySQLBackendAuthType::PASSWORD); + connection->healthy = true; + connection->reusable = true; + connection->send_quit = false; + connection->fd = fd; + connection->async_state_machine = ASYNC_IDLE; + server->ConnectionsUsed->add(connection); + return connection; +} + +} // namespace + +int main() { + plan(30); + ok(aws_locality_saturating_add( + std::numeric_limits::max() - 2, 5) == + std::numeric_limits::max(), + "locality weight sums saturate instead of overflowing"); + const uint64_t lottery_weights[] = {40, 40, 30}; + ok(aws_locality_weighted_index(lottery_weights, 3, 0) == 0, + "locality lottery selects the first candidate at its lower boundary"); + ok(aws_locality_weighted_index(lottery_weights, 3, 39) == 0, + "locality lottery selects the first candidate at its upper boundary"); + ok(aws_locality_weighted_index(lottery_weights, 3, 40) == 1, + "locality lottery selects the second candidate at its lower boundary"); + ok(aws_locality_weighted_index(lottery_weights, 3, 79) == 1, + "locality lottery selects the second candidate at its upper boundary"); + ok(aws_locality_weighted_index(lottery_weights, 3, 80) == 2, + "locality lottery selects the final candidate at its lower boundary"); + const uint64_t zero_weights[] = {0, 0}; + ok(aws_locality_weighted_index(zero_weights, 2, 17) == 2, + "shared locality lottery rejects an all-zero candidate set"); + ok(test_init_minimal() == 0 && test_init_auth() == 0 && test_init_query_processor() == 0 && + test_init_hostgroups() == 0, + "minimal worker and Hostgroup Manager fixtures initialize"); + GloMyLogger = new MySQL_Logger(); + + auto provider_state = std::make_shared(); + ok(install_global_aws_metadata_provider( + new FakeProvider(provider_state), destroy_provider, nullptr), + "fake metadata provider installs through the production lease registry"); + + MySrvC* local = add_server(kLocal, 10); + MySrvC* regional = add_server(kRegional, 20); + MySrvC* remote = add_server(kRemote, 30); + MyHGC* hostgroup = MyHGM->MyHGC_find(kHostgroup); + init_myhgc_hostgroup_settings( + R"({"aws":{"locality_awareness":{"same_region_multiplier":2.0,"same_az_multiplier":4.0}}})", + hostgroup); + MyHGM->refresh_aws_locality_configuration(); + MyHGM->set_aws_locality_awareness_enabled(true); + ok(wait_for_requests(provider_state, 3), + "selection fixture requests local identity and both backend Regions"); + + AwsMetadataResult local_result; + local_result.status = AwsMetadataStatus::ok; + local_result.local = {"us-east-1", "us-east-1a", "111122223333"}; + ok(complete(provider_state, AwsMetadataRequestKind::local_location, "", + std::move(local_result)), "local identity completion is accepted"); + + AwsMetadataResult east; + east.status = AwsMetadataStatus::ok; + east.endpoints.push_back({kLocal, 3306, AwsEndpointType::instance, + "us-east-1", "us-east-1a", "111122223333"}); + east.endpoints.push_back({kRegional, 3306, AwsEndpointType::cluster, + "us-east-1", "us-east-1a", "111122223333"}); + ok(complete(provider_state, AwsMetadataRequestKind::rds_region, "us-east-1", + std::move(east)), "same-Region endpoint completion is accepted"); + + AwsMetadataResult west; + west.status = AwsMetadataStatus::ok; + west.endpoints.push_back({kRemote, 3306, AwsEndpointType::instance, + "eu-west-1", "eu-west-1a", "111122223333"}); + ok(complete(provider_state, AwsMetadataRequestKind::rds_region, "eu-west-1", + std::move(west)), "remote endpoint completion is accepted"); + + ok(wait_until([&] { + auto snapshot = MyHGM->aws_locality_manager()->snapshot(); + const auto* entry = snapshot->find(kHostgroup, kRemote, 3306); + return entry != nullptr && entry->status == AwsLocalityMetadataStatus::fresh; + }), "fresh immutable selection snapshot is published"); + + auto snapshot = MyHGM->aws_locality_manager()->snapshot(); + ok(snapshot->effective_weight(kHostgroup, kLocal, 3306, 10) == 40 && + snapshot->effective_weight(kHostgroup, kRegional, 3306, 20) == 40 && + snapshot->effective_weight(kHostgroup, kRemote, 3306, 30) == 30, + "literal same-AZ, same-Region, and remote weights evaluate to 40/40/30"); + const auto* regional_entry = snapshot->find(kHostgroup, kRegional, 3306); + ok(regional_entry != nullptr && + regional_entry->locality == AwsLocalityClass::same_region && + regional_entry->multiplier == 2.0, + "cluster endpoints receive Region bias only, even when their reported AZ matches"); + + GloMTH->set_variable("aws_locality_awareness", "true"); + GloMTH->num_threads = 1; + { + MySQL_Thread worker; + if (!worker.init()) BAIL_OUT("worker init failed"); + worker.curtime = 10000000; + SessionFixture session(worker); + + int global_counts[3] = {0, 0, 0}; + for (int i = 0; i < 12000; ++i) { + MySrvC* selected = hostgroup->get_random_MySrvC(nullptr, 0, -1, session.session); + if (selected == local) ++global_counts[0]; + else if (selected == regional) ++global_counts[1]; + else if (selected == remote) ++global_counts[2]; + } + const double local_share = static_cast(global_counts[0]) / 12000.0; + const double regional_share = static_cast(global_counts[1]) / 12000.0; + const double remote_share = static_cast(global_counts[2]) / 12000.0; + ok(local_share > 0.31 && local_share < 0.42 && + regional_share > 0.31 && regional_share < 0.42 && + remote_share > 0.20 && remote_share < 0.34, + "global server lottery follows effective 40:40:30 weights (%.3f/%.3f/%.3f)", + local_share, regional_share, remote_share); + ok(local->weight == 10 && regional->weight == 20 && remote->weight == 30, + "global locality selection never mutates configured server weights"); + local->weight = 0; + regional->weight = 0; + remote->weight = 0; + ok(hostgroup->get_random_MySrvC(nullptr, 0, -1, session.session) != nullptr, + "global selection retains an eligible fallback when all configured weights are zero"); + local->weight = 10; + regional->weight = 20; + remote->weight = 30; + + MySQL_Connection* local_connection = make_connection(local, 100); + std::vector remote_connections; + for (int i = 0; i < 5; ++i) { + remote_connections.push_back(make_connection(remote, 200 + i)); + } + for (int i = 5; i < 33; ++i) { + remote_connections.push_back(make_connection(remote, 200 + i)); + } + worker.push_MyConn_local(local_connection); + for (auto* connection : remote_connections) worker.push_MyConn_local(connection); + + hot_path_allocations = 0; + track_hot_path_allocations = true; + MySQL_Connection* allocation_probe = worker.get_MyConn_local( + kHostgroup, session.session, nullptr, 0, -1, + MySQLBackendAuthType::PASSWORD); + track_hot_path_allocations = false; + ok(allocation_probe != nullptr && hot_path_allocations == 0, + "locality selection performs no heap allocation with more than 32 cached connections"); + if (allocation_probe != nullptr) worker.push_MyConn_local(allocation_probe); + + int local_parent = 0; + int remote_parent = 0; + for (int i = 0; i < 6000; ++i) { + MySQL_Connection* selected = worker.get_MyConn_local( + kHostgroup, session.session, nullptr, 0, -1, + MySQLBackendAuthType::PASSWORD); + if (selected == nullptr) BAIL_OUT("local selection returned no candidate"); + if (selected->parent == local) ++local_parent; + else if (selected->parent == remote) ++remote_parent; + worker.push_MyConn_local(selected); + } + const double local_parent_share = static_cast(local_parent) / 6000.0; + ok(local_parent_share > 0.53 && local_parent_share < 0.62 && + local_parent + remote_parent == 6000, + "local cache chooses 40:30 weighted parents, not the 1:5 connection count (%.3f local)", + local_parent_share); + local->weight = 0; + remote->weight = 0; + MySQL_Connection* zero_weight_selected = worker.get_MyConn_local( + kHostgroup, session.session, nullptr, 0, -1, + MySQLBackendAuthType::PASSWORD); + ok(zero_weight_selected != nullptr, + "local cache reuses an eligible connection when all parent weights are zero"); + if (zero_weight_selected != nullptr) worker.push_MyConn_local(zero_weight_selected); + local->weight = 10; + remote->weight = 30; + + GTID_Server_Data local_gtid(nullptr, const_cast(kLocal), 0, 3306); + local_gtid.add_gtid_from_ok("aaaaaaaa-0000-1111-2222-aaaaaaaaaaaa:42"); + MyHGM->gtid_map.emplace(std::string(kLocal) + ":3306", &local_gtid); + local->aws_aurora_current_lag_us = 5000; + char gtid_uuid[] = "aaaaaaaa000011112222aaaaaaaaaaaa"; + MySQL_Connection* gtid_selected = worker.get_MyConn_local( + kHostgroup, session.session, gtid_uuid, 42, 1, + MySQLBackendAuthType::PASSWORD); + ok(gtid_selected != nullptr && gtid_selected->parent == local, + "GTID-qualified local reuse preserves the legacy max-lag exemption"); + if (gtid_selected != nullptr) worker.push_MyConn_local(gtid_selected); + local->aws_aurora_current_lag_us = 0; + MyHGM->gtid_map.erase(std::string(kLocal) + ":3306"); + + remote->aws_aurora_current_lag_us = 5000; + MySQL_Connection* selected = worker.get_MyConn_local( + kHostgroup, session.session, nullptr, 0, 1, + MySQLBackendAuthType::PASSWORD); + ok(selected != nullptr && selected->parent == local, + "replication-lag eligibility excludes a remote parent before locality weighting"); + worker.push_MyConn_local(selected); + remote->aws_aurora_current_lag_us = 0; + + MySQL_Connection* incompatible = make_connection(regional, 300); + incompatible->userinfo->set( + const_cast("other_user"), const_cast("password"), + const_cast(kSchema), nullptr); + worker.push_MyConn_local(incompatible); + selected = worker.get_MyConn_local( + kHostgroup, session.session, nullptr, 0, -1, + MySQLBackendAuthType::PASSWORD); + ok(selected != nullptr && selected->parent != regional, + "authentication incompatibility excludes a parent before locality weighting"); + worker.push_MyConn_local(selected); + + local_connection->healthy = false; + selected = worker.get_MyConn_local( + kHostgroup, session.session, nullptr, 0, -1, + MySQLBackendAuthType::PASSWORD); + ok(selected != nullptr && selected->parent == remote, + "health eligibility excludes a preferred local connection before locality weighting"); + worker.push_MyConn_local(selected); + local_connection->healthy = true; + + local->session_track_backoff_until.store(worker.curtime + 1, std::memory_order_relaxed); + mysql_thread___session_track_variables = session_track_variables::ENFORCED; + selected = worker.get_MyConn_local( + kHostgroup, session.session, nullptr, 0, -1, + MySQLBackendAuthType::PASSWORD); + ok(selected != nullptr && selected->parent == remote, + "session-capability backoff excludes a local parent before locality weighting"); + worker.push_MyConn_local(selected); + mysql_thread___session_track_variables = session_track_variables::DISABLED; + local->session_track_backoff_until.store(0, std::memory_order_relaxed); + + local_connection->options.client_flag |= CLIENT_FOUND_ROWS; + selected = worker.get_MyConn_local( + kHostgroup, session.session, nullptr, 0, -1, + MySQLBackendAuthType::PASSWORD); + ok(selected != nullptr && selected->parent == remote, + "session option incompatibility is enforced before locality weighting"); + worker.push_MyConn_local(selected); + local_connection->options.client_flag &= ~CLIENT_FOUND_ROWS; + } + MyHGM->set_aws_locality_awareness_enabled(false); + shutdown_global_aws_metadata_provider(); + MyHGM->set_aws_locality_awareness_enabled(true); + ok(wait_until([&] { + const auto unavailable = MyHGM->aws_locality_manager()->snapshot(); + const auto* entry = unavailable->find(kHostgroup, kLocal, 3306); + return entry != nullptr && + entry->status == AwsLocalityMetadataStatus::error && + entry->failure_category == "provider_unavailable" && + unavailable->effective_weight(kHostgroup, kLocal, 3306, 10) == 10; + }), "missing provider publishes fixed provider_unavailable with neutral weight"); + { + MySQL_Thread neutral_worker; + if (!neutral_worker.init()) BAIL_OUT("neutral worker init failed"); + SessionFixture neutral_session(neutral_worker); + int neutral_counts[3] = {0, 0, 0}; + for (int i = 0; i < 12000; ++i) { + MySrvC* selected = hostgroup->get_random_MySrvC( + nullptr, 0, -1, neutral_session.session); + if (selected == local) ++neutral_counts[0]; + else if (selected == regional) ++neutral_counts[1]; + else if (selected == remote) ++neutral_counts[2]; + } + const double neutral_local_share = + static_cast(neutral_counts[0]) / 12000.0; + const double neutral_regional_share = + static_cast(neutral_counts[1]) / 12000.0; + const double neutral_remote_share = + static_cast(neutral_counts[2]) / 12000.0; + ok(neutral_local_share > 0.13 && neutral_local_share < 0.20 && + neutral_regional_share > 0.29 && neutral_regional_share < 0.38 && + neutral_remote_share > 0.45 && neutral_remote_share < 0.55 && + local->weight == 10 && regional->weight == 20 && remote->weight == 30, + "provider absence selects by unchanged configured 10:20:30 weights (%.3f/%.3f/%.3f)", + neutral_local_share, neutral_regional_share, neutral_remote_share); + } + MyHGM->set_aws_locality_awareness_enabled(false); + test_cleanup_hostgroups(); + delete GloMyLogger; + GloMyLogger = nullptr; + test_cleanup_query_processor(); + test_cleanup_auth(); + test_cleanup_minimal(); + return exit_status(); +} diff --git a/test/tap/tests/unit/aws_locality_stats_unit-t.cpp b/test/tap/tests/unit/aws_locality_stats_unit-t.cpp new file mode 100644 index 0000000000..5a1514ef27 --- /dev/null +++ b/test/tap/tests/unit/aws_locality_stats_unit-t.cpp @@ -0,0 +1,269 @@ +#include "Aws_Locality_Manager.h" +#include "MySQL_HostGroups_Manager.h" +#include "ProxySQL_PluginManager.h" +#include "sqlite3db.h" +#include "tap.h" +#include "test_globals.h" + +#include +#include +#include +#include +#include + +extern MySQL_HostGroups_Manager* MyHGM; + +namespace { + +constexpr const char* kLocalityStatsProjectionFixture = + "CREATE TABLE stats_mysql_aws_locality (" + "hostgroup_id INT NOT NULL, hostname VARCHAR NOT NULL, port INT NOT NULL, " + "endpoint_type VARCHAR NOT NULL, configured_weight INT NOT NULL, " + "effective_weight INT NOT NULL, local_region VARCHAR NOT NULL, " + "local_az VARCHAR NOT NULL, backend_region VARCHAR NOT NULL, " + "backend_az VARCHAR NOT NULL, account_match VARCHAR NOT NULL, " + "locality VARCHAR NOT NULL, active_multiplier REAL NOT NULL, " + "metadata_status VARCHAR NOT NULL, last_success_timestamp INT NOT NULL, " + "last_attempt_timestamp INT NOT NULL, last_error_category VARCHAR NOT NULL, " + "PRIMARY KEY(hostgroup_id, hostname, port))"; + +class CountingProvider final : public AwsMetadataProvider { +public: + AwsMetadataRequestHandle request( + const AwsMetadataRequest&, + std::weak_ptr) override { + ++requests; + return {requests.load()}; + } + void cancel(AwsMetadataRequestHandle) override {} + void shutdown() override {} + std::atomic requests {0}; +}; + +CountingProvider* counting_provider = nullptr; + +void destroy_counting_provider(AwsMetadataProvider* provider) { + delete static_cast(provider); + counting_provider = nullptr; +} + +AwsLocalityHostgroupConfig disabled_hostgroup( + uint32_t hostgroup_id, const std::string& hostname, int64_t weight) { + AwsLocalityHostgroupConfig config; + config.hostgroup_id = hostgroup_id; + config.policy.valid = true; + AwsEndpointCandidate endpoint; + endpoint.recognized = true; + endpoint.hostgroup_id = hostgroup_id; + endpoint.hostname = hostname; + endpoint.port = 3306; + endpoint.region = "us-east-1"; + endpoint.partition = "aws"; + config.backends.emplace_back(std::move(endpoint), weight); + return config; +} + +AwsLocalitySnapshotEntry diagnostic_row( + uint32_t hostgroup_id, + AwsLocalityMetadataStatus status, + double multiplier, + int64_t weight) { + AwsLocalitySnapshotEntry row; + row.hostgroup_id = hostgroup_id; + row.hostname = hostgroup_id == 6 + ? "db'quoted.abcdefghijkl.us-east-1.rds.amazonaws.com" + : "db-" + std::to_string(hostgroup_id) + + ".abcdefghijkl.us-east-1.rds.amazonaws.com"; + row.port = 3306; + row.endpoint_type = hostgroup_id == 1 ? AwsEndpointType::unknown + : hostgroup_id == 2 ? AwsEndpointType::instance + : hostgroup_id == 3 ? AwsEndpointType::cluster + : hostgroup_id == 4 ? AwsEndpointType::reader + : AwsEndpointType::custom; + row.configured_weight = weight; + row.local = {"us-east-1", "us-east-1a", "111122223333"}; + row.backend.region = hostgroup_id == 4 ? "eu-west-1" : "us-east-1"; + row.backend.availability_zone = hostgroup_id == 2 ? "us-east-1a" : ""; + row.backend.account_id = hostgroup_id == 1 ? "" + : hostgroup_id == 3 ? "444455556666" : "111122223333"; + row.locality = hostgroup_id == 2 ? AwsLocalityClass::same_az + : hostgroup_id == 3 ? AwsLocalityClass::same_region + : hostgroup_id == 4 ? AwsLocalityClass::remote + : AwsLocalityClass::unknown; + row.multiplier = multiplier; + row.status = status; + row.last_success_timestamp = 1700000000 + hostgroup_id; + row.last_attempt_timestamp = 1700000100 + hostgroup_id; + row.failure_category = status == AwsLocalityMetadataStatus::stale + ? "throttled" : status == AwsLocalityMetadataStatus::error + ? "access_denied" : ""; + return row; +} + +} // namespace + +int main() { + plan(22); + if (test_globals_init() != 0) { + BAIL_OUT("test global initialization failed"); + } + + SQLite3DB statsdb; + statsdb.open((char*)":memory:", + SQLITE_OPEN_READWRITE | SQLITE_OPEN_CREATE | SQLITE_OPEN_FULLMUTEX); + ok(statsdb.return_one_int( + "SELECT count(*) FROM sqlite_master WHERE name='stats_mysql_aws_locality'") == 0, + "public core does not register an AWS locality stats table"); + + std::unique_ptr manager; + std::string error; + ok(proxysql_load_configured_plugins(manager, {}, error) && manager == nullptr, + "provider-neutral plugin services expose no always-present locality schema"); + if (!error.empty()) diag("plugin error: %s", error.c_str()); + + ok(statsdb.execute(kLocalityStatsProjectionFixture), + "test-owned schema fixture accepts the public projection callback"); + ok(statsdb.return_one_int( + "SELECT count(*) FROM pragma_table_info('stats_mysql_aws_locality')") == 17, + "projection fixture has the external-provider contract's 17 columns"); + ok(statsdb.return_one_int( + "SELECT count(*) FROM pragma_table_info('stats_mysql_aws_locality') " + "WHERE name IN ('hostgroup_id','hostname','port','endpoint_type'," + "'configured_weight','effective_weight','local_region','local_az'," + "'backend_region','backend_az','account_match','locality'," + "'active_multiplier','metadata_status','last_success_timestamp'," + "'last_attempt_timestamp','last_error_category')") == 17, + "projection callback targets the documented external schema columns"); + + std::vector rows; + rows.push_back(diagnostic_row(1, AwsLocalityMetadataStatus::pending, 4.0, 10)); + rows.push_back(diagnostic_row(2, AwsLocalityMetadataStatus::fresh, 2.5, 11)); + rows.push_back(diagnostic_row(3, AwsLocalityMetadataStatus::stale, 4.0, 12)); + rows.push_back(diagnostic_row(4, AwsLocalityMetadataStatus::expired, 5.0, 13)); + rows.push_back(diagnostic_row(5, AwsLocalityMetadataStatus::error, 6.0, 14)); + rows.push_back(diagnostic_row(6, AwsLocalityMetadataStatus::disabled, 7.0, 15)); + ok(MySQL_HostGroups_Manager::project_aws_locality_stats(&statsdb, rows), + "one retained diagnostics snapshot projects transactionally"); + ok(statsdb.return_one_int("SELECT count(*) FROM stats_mysql_aws_locality") == 6, + "projection emits one row per configured backend"); + ok(statsdb.return_one_int( + "SELECT count(DISTINCT metadata_status) FROM stats_mysql_aws_locality") == 6, + "pending, fresh, stale, expired, error, and disabled are explicit"); + ok(statsdb.return_one_int( + "SELECT count(*) FROM stats_mysql_aws_locality WHERE " + "(hostgroup_id=2 AND effective_weight=27 AND active_multiplier=2.5) OR " + "(hostgroup_id=3 AND effective_weight=48 AND active_multiplier=4.0)") == 2, + "fresh/stale rows expose integer-cast weighted multipliers"); + ok(statsdb.return_one_int( + "SELECT count(*) FROM stats_mysql_aws_locality WHERE hostgroup_id IN (1,4,5,6) " + "AND effective_weight=configured_weight AND active_multiplier=1.0") == 4, + "pending/expired/error/disabled rows force neutral effective weights"); + ok(statsdb.return_one_int( + "SELECT count(*) FROM stats_mysql_aws_locality WHERE " + "(hostgroup_id=2 AND endpoint_type='instance' AND locality='same_az') OR " + "(hostgroup_id=3 AND endpoint_type='cluster' AND locality='same_region') OR " + "(hostgroup_id=4 AND endpoint_type='reader' AND locality='remote')") == 3, + "endpoint and locality classifications use stable strings"); + ok(statsdb.return_one_int( + "SELECT count(*) FROM stats_mysql_aws_locality WHERE " + "(hostgroup_id=1 AND account_match='unknown') OR " + "(hostgroup_id=2 AND account_match='same') OR " + "(hostgroup_id=3 AND account_match='different')") == 3, + "account comparison exposes unknown/same/different without identifiers"); + ok(statsdb.return_one_int( + "SELECT count(*) FROM stats_mysql_aws_locality WHERE hostgroup_id=3 " + "AND last_success_timestamp=1700000003 AND last_attempt_timestamp=1700000103 " + "AND last_error_category='throttled'") == 1, + "timestamps and fixed failure category survive projection"); + ok(statsdb.return_one_int( + "SELECT count(*) FROM stats_mysql_aws_locality WHERE hostname=" + "'db''quoted.abcdefghijkl.us-east-1.rds.amazonaws.com'") == 1, + "projection safely quotes endpoint text"); + + std::vector concurrent_rows_a; + std::vector concurrent_rows_b; + for (uint32_t index = 0; index < 200; ++index) { + concurrent_rows_a.push_back(diagnostic_row(1000 + index, + AwsLocalityMetadataStatus::fresh, 2.0, 10)); + concurrent_rows_b.push_back(diagnostic_row(2000 + index, + AwsLocalityMetadataStatus::stale, 3.0, 10)); + } + std::atomic start_concurrent_projection { false }; + std::atomic successful_projections { 0 }; + auto project_repeatedly = [&](const std::vector& projection) { + while (!start_concurrent_projection.load(std::memory_order_acquire)) { + std::this_thread::yield(); + } + for (unsigned int iteration = 0; iteration < 10; ++iteration) { + if (MySQL_HostGroups_Manager::project_aws_locality_stats(&statsdb, projection)) { + successful_projections.fetch_add(1, std::memory_order_relaxed); + } + } + }; + std::thread projection_a(project_repeatedly, std::cref(concurrent_rows_a)); + std::thread projection_b(project_repeatedly, std::cref(concurrent_rows_b)); + start_concurrent_projection.store(true, std::memory_order_release); + projection_a.join(); + projection_b.join(); + ok(successful_projections.load() == 20 && + statsdb.return_one_int("SELECT count(*) FROM stats_mysql_aws_locality") == 200, + "concurrent runtime-view refreshes serialize complete replacement transactions"); + + GloVars.prometheus_registry = std::make_shared(); + { + MySQL_HostGroups_Manager hostgroups; + MyHGM = &hostgroups; + counting_provider = new CountingProvider(); + ok(install_global_aws_metadata_provider( + counting_provider, &destroy_counting_provider, nullptr), + "network-request counter installs through the production registry"); + + hostgroups.aws_locality_manager()->configure({disabled_hostgroup( + 101, "first.abcdefghijkl.us-east-1.rds.amazonaws.com", 7)}); + hostgroups.refresh_aws_locality_stats(&statsdb); + ok(statsdb.return_one_int( + "SELECT count(*) FROM stats_mysql_aws_locality WHERE hostgroup_id=101 " + "AND metadata_status='disabled' AND configured_weight=7 " + "AND effective_weight=7") == 1, + "public callback projects the MySQL manager's current snapshot"); + ok(counting_provider->requests.load() == 0, + "query-time refresh issues no metadata-provider request"); + + hostgroups.aws_locality_manager()->configure({disabled_hostgroup( + 202, "second.abcdefghijkl.us-east-1.rds.amazonaws.com", 9)}); + hostgroups.refresh_aws_locality_stats(&statsdb); + ok(statsdb.return_one_int("SELECT count(*) FROM stats_mysql_aws_locality") == 1 && + statsdb.return_one_int( + "SELECT count(*) FROM stats_mysql_aws_locality WHERE hostgroup_id=202") == 1, + "generation swap replaces the prior projection without mixed rows"); + + hostgroups.aws_locality_manager()->configure({}); + hostgroups.refresh_aws_locality_stats(&statsdb); + ok(statsdb.return_one_int("SELECT count(*) FROM stats_mysql_aws_locality") == 0, + "no valid locality policy produces zero rows"); + ok(counting_provider->requests.load() == 0, + "repeated generation queries remain network-free"); + MyHGM = nullptr; + } + shutdown_global_aws_metadata_provider(); + GloVars.prometheus_registry.reset(); + + statsdb.execute("PRAGMA query_only = ON"); + char* write_error = nullptr; + SQLite3_result* write_result = statsdb.execute_statement( + "INSERT INTO stats_mysql_aws_locality " + "(hostgroup_id,hostname,port,endpoint_type,configured_weight,effective_weight," + "local_region,local_az,backend_region,backend_az,account_match,locality," + "active_multiplier,metadata_status,last_success_timestamp,last_attempt_timestamp," + "last_error_category) VALUES (1,'x',3306,'unknown',1,1,'','','',''," + "'unknown','unknown',1.0,'disabled',0,0,'')", &write_error); + ok(write_error != nullptr, + "stats listener query-only mode rejects writes to the projection"); + free(write_error); + delete write_result; + statsdb.execute("PRAGMA query_only = OFF"); + + proxysql_stop_configured_plugins(manager, error); + test_globals_cleanup(); + return exit_status(); +} diff --git a/test/tap/tests/unit/connection_pool_unit-t.cpp b/test/tap/tests/unit/connection_pool_unit-t.cpp index 21b7a13644..52ac64c85e 100644 --- a/test/tap/tests/unit/connection_pool_unit-t.cpp +++ b/test/tap/tests/unit/connection_pool_unit-t.cpp @@ -96,6 +96,8 @@ static void test_pool_quality_1_reuse() { auto d = evaluate_pool_state(10, 5, 100, 1, false, 0); ok(d.create_new_connection == false, "pool q=1: reuses when free >= used"); + ok(d.evict_connections == false, + "pool q=1: ordinary CHANGE_USER candidate keeps legacy reuse behavior"); } static void test_pool_quality_2_3() { @@ -124,7 +126,7 @@ static void test_pool_empty() { // ============================================================================ int main() { - plan(23); + plan(24); int rc = test_init_minimal(); ok(rc == 0, "test_init_minimal() succeeds"); @@ -138,7 +140,7 @@ int main() { test_pool_quality_0(); // 1 test_pool_quality_0_evict(); // 3 test_pool_quality_1_create(); // 1 - test_pool_quality_1_reuse(); // 1 + test_pool_quality_1_reuse(); // 2 test_pool_quality_2_3(); // 2 test_pool_warming(); // 3 test_pool_empty(); // 1 diff --git a/test/tap/tests/unit/mariadb_tls_server_name_unit-t.cpp b/test/tap/tests/unit/mariadb_tls_server_name_unit-t.cpp new file mode 100644 index 0000000000..a7051eb9bb --- /dev/null +++ b/test/tap/tests/unit/mariadb_tls_server_name_unit-t.cpp @@ -0,0 +1,49 @@ +/** + * @file mariadb_tls_server_name_unit-t.cpp + * @brief Verify Connector/C separates the TLS identity from the TCP host. + */ + +#include "tap.h" + +#include + +#include + +extern "C" const char *ma_tls_get_server_name(MYSQL *mysql); + +int main() { + plan(8); + + MYSQL *mysql = mysql_init(nullptr); + if (mysql == nullptr) { + BAIL_OUT("mysql_init() failed"); + } + + const char *const transport_host = "127.0.0.1"; + mysql->host = strdup(transport_host); + ok(mysql->host != nullptr, "transport host is available for TLS fallback"); + ok(std::strcmp(ma_tls_get_server_name(mysql), transport_host) == 0, + "unset TLS server-name option falls back to the transport host"); + + char tls_server_name[] = "db.cluster-abc.us-east-1.rds.amazonaws.com"; + ok(mysql_options(mysql, MARIADB_OPT_TLS_SERVER_NAME, tls_server_name) == 0, + "sets TLS server-name option"); + tls_server_name[0] = 'x'; + + char *configured_name = nullptr; + ok(mysql_get_optionv(mysql, MARIADB_OPT_TLS_SERVER_NAME, &configured_name) == 0, + "gets TLS server-name option"); + ok(configured_name != tls_server_name && + std::strcmp(configured_name, "db.cluster-abc.us-east-1.rds.amazonaws.com") == 0, + "TLS server-name option owns a copy of the caller buffer"); + ok(std::strcmp(ma_tls_get_server_name(mysql), configured_name) == 0, + "TLS server-name override takes precedence over the transport host"); + + ok(mysql_options(mysql, MARIADB_OPT_TLS_SERVER_NAME, "") == 0, + "sets an empty TLS server-name option"); + ok(std::strcmp(ma_tls_get_server_name(mysql), transport_host) == 0, + "empty TLS server-name option falls back to the transport host"); + + mysql_close(mysql); + return exit_status(); +}