aws: support AWS_EC2_METADATA_SERVICE_ENDPOINT for custom IMDS endpoints

include/fluent-bit/aws/flb_aws_imds.h

@@ -23,7 +23,11 @@
 #define FLB_AWS_IMDS_HOST "169.254.169.254"
 #define FLB_AWS_IMDS_HOST_LEN 15
 #define FLB_AWS_IMDS_PORT 80
-#define FLB_AWS_IMDS_TIMEOUT 1  /* 1 second */
+#define FLB_AWS_IMDS_TIMEOUT 1  /* 1 second - for standard AWS IMDS */
+#define FLB_AWS_IMDS_TIMEOUT_CUSTOM 10  /* 10 seconds - for custom IMDS endpoints like IAM Roles Anywhere */
+
+/* Environment variable for custom IMDS endpoint */
+#define AWS_EC2_METADATA_SERVICE_ENDPOINT_ENV "AWS_EC2_METADATA_SERVICE_ENDPOINT"
 
 #define FLB_AWS_IMDS_VERSION_EVALUATE 0
 #define FLB_AWS_IMDS_VERSION_1 1

src/aws/flb_aws_credentials_ec2.c

@@ -232,6 +232,14 @@ struct flb_aws_provider *flb_ec2_provider_create(struct flb_config *config,
     struct flb_aws_provider_ec2 *implementation;
     struct flb_aws_provider *provider;
     struct flb_upstream *upstream;
+    char *endpoint;
+    flb_sds_t host = NULL;
+    flb_sds_t port_str = NULL;
+    flb_sds_t protocol = NULL;
+    flb_sds_t path = NULL;
+    const char *use_host;
+    int use_port;
+    int ret;
 
     provider = flb_calloc(1, sizeof(struct flb_aws_provider));
 
@@ -253,17 +261,57 @@ struct flb_aws_provider *flb_ec2_provider_create(struct flb_config *config,
     provider->provider_vtable = &ec2_provider_vtable;
     provider->implementation = implementation;
 
-    upstream = flb_upstream_create(config, FLB_AWS_IMDS_HOST, FLB_AWS_IMDS_PORT,
-                                   FLB_IO_TCP, NULL);
+    /* Check for custom IMDS endpoint */
+    use_host = FLB_AWS_IMDS_HOST;
+    use_port = FLB_AWS_IMDS_PORT;
+    int use_custom_endpoint = FLB_FALSE;
+
+    endpoint = getenv(AWS_EC2_METADATA_SERVICE_ENDPOINT_ENV);
+    if (endpoint && strlen(endpoint) > 0) {
+        ret = flb_utils_url_split_sds(endpoint, &protocol, &host, &port_str, &path);
+        if (ret >= 0 && host) {
+            use_host = host;
+            use_custom_endpoint = FLB_TRUE;
+            if (port_str) {
+                use_port = atoi(port_str);
+                if (use_port <= 0 || use_port > 65535) {
+                    use_port = FLB_AWS_IMDS_PORT;
+                }
+            }
+            flb_info("[aws_credentials] Using custom IMDS endpoint: %s:%d", 
+                     use_host, use_port);
+        }
+        flb_sds_destroy(protocol);
+        flb_sds_destroy(port_str);
+        flb_sds_destroy(path);
+    }
+
+    upstream = flb_upstream_create(config, use_host, use_port, FLB_IO_TCP, NULL);
+
+    if (host) {
+        flb_sds_destroy(host);
+    }
+
     if (!upstream) {
         flb_aws_provider_destroy(provider);
         flb_debug("[aws_credentials] unable to connect to EC2 IMDS.");
         return NULL;
     }
 
-    /* IMDSv2 token request will timeout if hops = 1 and running within container */
-    upstream->base.net.connect_timeout = FLB_AWS_IMDS_TIMEOUT;
-    upstream->base.net.io_timeout = FLB_AWS_IMDS_TIMEOUT;
+    /* 
+     * Set timeout based on endpoint type:
+     * - Standard AWS IMDS: 1 second (fast local endpoint)
+     * - Custom IMDS endpoints (e.g., IAM Roles Anywhere): 10 seconds (certificate auth takes longer)
+     */
+    if (use_custom_endpoint) {
+        upstream->base.net.connect_timeout = FLB_AWS_IMDS_TIMEOUT_CUSTOM;
+        upstream->base.net.io_timeout = FLB_AWS_IMDS_TIMEOUT_CUSTOM;
+        flb_info("[aws_credentials] Using extended timeout (%d seconds) for custom IMDS endpoint", 
+                 FLB_AWS_IMDS_TIMEOUT_CUSTOM);
+    } else {
+        upstream->base.net.connect_timeout = FLB_AWS_IMDS_TIMEOUT;
+        upstream->base.net.io_timeout = FLB_AWS_IMDS_TIMEOUT;
+    }
     upstream->base.net.keepalive = FLB_FALSE; /* On timeout, the connection is broken */
 
     implementation->client = generator->create();
@@ -278,7 +326,7 @@ struct flb_aws_provider *flb_ec2_provider_create(struct flb_config *config,
     implementation->client->provider = NULL;
     implementation->client->region = NULL;
     implementation->client->service = NULL;
-    implementation->client->port = 80;
+    implementation->client->port = use_port;
     implementation->client->flags = 0;
     implementation->client->proxy = NULL;
     implementation->client->upstream = upstream;

src/aws/flb_aws_imds.c

@@ -80,17 +80,9 @@ struct flb_aws_imds *flb_aws_imds_create(const struct flb_aws_imds_config *imds_
         flb_aws_imds_destroy(ctx);
         return NULL;
     }
-    if (0 != strncmp(ec2_imds_client->upstream->tcp_host, FLB_AWS_IMDS_HOST,
-                     FLB_AWS_IMDS_HOST_LEN)) {
-        flb_debug("[imds] ec2_imds_client tcp host must be set to %s", FLB_AWS_IMDS_HOST);
-        flb_aws_imds_destroy(ctx);
-        return NULL;
-    }
-    if (ec2_imds_client->upstream->tcp_port != FLB_AWS_IMDS_PORT) {
-        flb_debug("[imds] ec2_imds_client tcp port must be set to %i", FLB_AWS_IMDS_PORT);
-        flb_aws_imds_destroy(ctx);
-        return NULL;
-    }
+
+    /* Allow custom IMDS endpoints via AWS_EC2_METADATA_SERVICE_ENDPOINT */
+    /* The hardcoded host/port checks have been removed to support custom endpoints */
 
     /* Connect client */
     ctx->ec2_imds_client = ec2_imds_client;
@@ -250,7 +242,6 @@ static int get_imds_version(struct flb_aws_imds *ctx)
 {
     int ret;
     struct flb_aws_client *client = ctx->ec2_imds_client;
-    struct flb_aws_header invalid_token_header;
     struct flb_http_client *c = NULL;
 
     if (ctx->imds_version != FLB_AWS_IMDS_VERSION_EVALUATE) {
@@ -259,65 +250,37 @@ static int get_imds_version(struct flb_aws_imds *ctx)
 
     /*
      * Evaluate version
-     * To evaluate wether IMDSv2 is available, send an invalid token
-     * in IMDS request. If response status is 'Unauthorized', then IMDSv2
-     * is available.
+     * Try to get an IMDSv2 token first. If that fails, fall back to IMDSv1.
+     * This approach is more compatible with custom IMDS implementations like IAM Roles Anywhere.
      */
-    invalid_token_header = imds_v2_token_token_header_template;
-    invalid_token_header.val = "INVALID";
-    invalid_token_header.val_len = 7;
-    c = client->client_vtable->request(client, FLB_HTTP_GET, FLB_AWS_IMDS_ROOT, NULL, 0,
-                                       &invalid_token_header, 1);
-
+    ctx->imds_version = FLB_AWS_IMDS_VERSION_2;
+    ret = refresh_imds_v2_token(ctx);
+    if (ret == 0) {
+        /* Successfully got IMDSv2 token */
+        flb_info("[imds] using IMDSv2");
+        return FLB_AWS_IMDS_VERSION_2;
+    }
+
+    /* IMDSv2 token request failed, try IMDSv1 */
+    flb_debug("[imds] IMDSv2 token request failed, testing IMDSv1");
+    ctx->imds_version = FLB_AWS_IMDS_VERSION_EVALUATE;
+    c = client->client_vtable->request(client, FLB_HTTP_GET, FLB_AWS_IMDS_ROOT,
+                                       NULL, 0, NULL, 0);
+
     if (!c) {
         flb_debug("[imds] imds endpoint unavailable");
         return FLB_AWS_IMDS_VERSION_EVALUATE;
     }
-
-    /* Unauthorized response means that IMDS version 2 is in use */
-    if (c->resp.status == 401) {
-        ctx->imds_version = FLB_AWS_IMDS_VERSION_2;
-        ret = refresh_imds_v2_token(ctx);
-        if (ret == -1) {
-            /*
-             * Token cannot be refreshed, test IMDSv1
-             * If IMDSv1 cannot be used, response will be status 401
-             */
-            flb_http_client_destroy(c);
-            ctx->imds_version = FLB_AWS_IMDS_VERSION_EVALUATE;
-            c = client->client_vtable->request(client, FLB_HTTP_GET, FLB_AWS_IMDS_ROOT,
-                                               NULL, 0, NULL, 0);
-            if (!c) {
-                flb_debug("[imds] imds v1 attempt, endpoint unavailable");
-                return FLB_AWS_IMDS_VERSION_EVALUATE;
-            }
-
-            if (c->resp.status == 200) {
-                flb_info("[imds] to use IMDSv2, set --http-put-response-hop-limit to 2");
-            }
-            else {
-                /* IMDSv1 unavailable. IMDSv2 beyond network hop count */
-                flb_warn("[imds] failed to retrieve IMDSv2 token and IMDSv1 unavailable. "
-                        "This is likely due to instance-metadata-options "
-                        "--http-put-response-hop-limit being set to 1 and --http-tokens "
-                        "set to required. "
-                        "To use IMDSv2, please set --http-put-response-hop-limit to 2 as "
-                        "described https://docs.aws.amazon.com/AWSEC2/latest/UserGuide/"
-                        "configuring-instance-metadata-options.html");
-            }
-        }
-    }
-
-    /*
-     * Success means that IMDS version 1 is in use
-     */
+
     if (c->resp.status == 200) {
-        flb_warn("[imds] falling back on IMDSv1");
+        flb_info("[imds] falling back to IMDSv1");
         ctx->imds_version = FLB_AWS_IMDS_VERSION_1;
+        flb_http_client_destroy(c);
+        return FLB_AWS_IMDS_VERSION_1;
     }
-
+    
     flb_http_client_destroy(c);
-    return ctx->imds_version;
+    return FLB_AWS_IMDS_VERSION_EVALUATE;
 }
 
 /*

tests/internal/aws_credentials_ec2.c

@@ -1025,6 +1025,44 @@ static void test_ec2_imds_create_and_destroy()
     flb_config_exit(config_fluent);
 }
 
+static void test_ec2_provider_custom_endpoint()
+{
+    setenv("AWS_EC2_METADATA_SERVICE_ENDPOINT", "http://127.0.0.1:9911", 1);
+
+    setup_test(FLB_AWS_CLIENT_MOCK(
+        response(
+            expect(URI, "/latest/api/token"),
+            expect(METHOD, FLB_HTTP_PUT),
+            set(STATUS, 200),
+            set(PAYLOAD, "TESTTOKEN")
+        ),
+        response(
+            expect(URI, "/latest/meta-data/iam/security-credentials/"),
+            expect(HEADER, "X-aws-ec2-metadata-token", "TESTTOKEN"),
+            expect(METHOD, FLB_HTTP_GET),
+            set(STATUS, 200),
+            set(PAYLOAD, "test-role")
+        ),
+        response(
+            expect(URI, "/latest/meta-data/iam/security-credentials/test-role"),
+            expect(HEADER, "X-aws-ec2-metadata-token", "TESTTOKEN"),
+            expect(METHOD, FLB_HTTP_GET),
+            set(STATUS, 200),
+            set(PAYLOAD, "{\"AccessKeyId\":\"AKIATEST\",\"SecretAccessKey\":\"SECRET\",\"Token\":\"TOKEN\"}")
+        )
+    ));
+
+    creds = provider->provider_vtable->get_credentials(provider);
+    TEST_CHECK(creds != NULL);
+    TEST_CHECK(strcmp("AKIATEST", creds->access_key_id) == 0);
+    TEST_CHECK(strcmp("SECRET", creds->secret_access_key) == 0);
+    TEST_CHECK(strcmp("TOKEN", creds->session_token) == 0);
+
+    flb_aws_credentials_destroy(creds);
+    unsetenv("AWS_EC2_METADATA_SERVICE_ENDPOINT");
+    cleanup_test();
+}
+
 TEST_LIST = {
     { "test_ec2_provider_v2" , test_ec2_provider_v2},
     { "test_ec2_provider_v1" , test_ec2_provider_v1},
@@ -1033,5 +1071,6 @@ TEST_LIST = {
     { "test_ec2_provider_acquire_token_error" , test_ec2_provider_acquire_token_error},
     { "test_ec2_provider_metadata_request_error" , test_ec2_provider_metadata_request_error},
     { "test_ec2_imds_create_and_destroy" , test_ec2_imds_create_and_destroy},
+    { "test_ec2_provider_custom_endpoint" , test_ec2_provider_custom_endpoint},
     { 0 }
 };