1 files changed, 370 insertions, 0 deletions
diff --git a/fluent-bit/src/aws/flb_aws_imds.c b/fluent-bit/src/aws/flb_aws_imds.c
new file mode 100644
index 00000000..0e54db16
--- /dev/null
+++ b/fluent-bit/src/aws/flb_aws_imds.c
@@ -0,0 +1,370 @@
+/* -*- Mode: C; tab-width: 4; indent-tabs-mode: nil; c-basic-offset: 4 -*- */
+
+/*  Fluent Bit
+ *  ==========
+ *  Copyright (C) 2015-2022 The Fluent Bit Authors
+ *
+ *  Licensed under the Apache License, Version 2.0 (the "License");
+ *  you may not use this file except in compliance with the License.
+ *  You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ *  Unless required by applicable law or agreed to in writing, software
+ *  distributed under the License is distributed on an "AS IS" BASIS,
+ *  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ *  See the License for the specific language governing permissions and
+ *  limitations under the License.
+ */
+
+#include <fluent-bit/aws/flb_aws_imds.h>
+#include <fluent-bit/flb_aws_credentials.h>
+#include <fluent-bit/flb_aws_util.h>
+#include <fluent-bit/flb_http_client.h>
+#include <fluent-bit/flb_info.h>
+#include <fluent-bit/flb_jsmn.h>
+
+#define FLB_AWS_IMDS_ROOT "/"
+#define FLB_AWS_IMDS_V2_TOKEN_PATH "/latest/api/token"
+
+/* Request headers */
+static struct flb_aws_header imds_v2_token_ttl_header = {
+    .key = "X-aws-ec2-metadata-token-ttl-seconds",
+    .key_len = 36,
+    .val = "21600",  /* 6 hours (ie maximum ttl) */
+    .val_len = 5,
+};
+
+/* Request header templates */
+const static struct flb_aws_header imds_v2_token_token_header_template = {
+    .key = "X-aws-ec2-metadata-token",
+    .key_len = 24,
+    .val = "",     /* Replace with token value */
+    .val_len = 0,  /* Replace with token length */
+};
+
+/* Declarations */
+static int get_imds_version(struct flb_aws_imds *ctx);
+static int refresh_imds_v2_token(struct flb_aws_imds *ctx);
+
+/* Default config values */
+const struct flb_aws_imds_config flb_aws_imds_config_default = {
+    FLB_AWS_IMDS_VERSION_EVALUATE};
+
+/* Create IMDS context */
+struct flb_aws_imds *flb_aws_imds_create(const struct flb_aws_imds_config *imds_config,
+                                         struct flb_aws_client *ec2_imds_client)
+{
+    struct flb_aws_imds *ctx = NULL;
+
+    /* Create context */
+    ctx = flb_calloc(1, sizeof(struct flb_aws_imds));
+    if (!ctx) {
+        flb_errno();
+        return NULL;
+    }
+
+    /*
+     * Set IMDS version to whatever is specified in config
+     * Version may be evaluated later if set to FLB_AWS_IMDS_VERSION_EVALUATE
+     */
+    ctx->imds_version = imds_config->use_imds_version;
+    ctx->imds_v2_token = flb_sds_create_len("INVALID_TOKEN", 13);
+    ctx->imds_v2_token_len = 13;
+
+    /* Detect IMDS support */
+    if (!ec2_imds_client->upstream) {
+        flb_debug(
+            "[imds] unable to connect to EC2 IMDS. ec2_imds_client upstream is null");
+
+        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;
+    }
+
+    /* Connect client */
+    ctx->ec2_imds_client = ec2_imds_client;
+    return ctx;
+}
+
+/* Destroy IMDS context */
+void flb_aws_imds_destroy(struct flb_aws_imds *ctx)
+{
+    if (ctx->imds_v2_token) {
+        flb_sds_destroy(ctx->imds_v2_token);
+    }
+
+    flb_free(ctx);
+}
+
+/* Get IMDS metadata */
+int flb_aws_imds_request(struct flb_aws_imds *ctx, const char *metadata_path,
+                         flb_sds_t *metadata, size_t *metadata_len)
+{
+    return flb_aws_imds_request_by_key(ctx, metadata_path, metadata, metadata_len, NULL);
+}
+
+/* Get IMDS metadata by key */
+int flb_aws_imds_request_by_key(struct flb_aws_imds *ctx, const char *metadata_path,
+                                flb_sds_t *metadata, size_t *metadata_len, char *key)
+{
+    int ret;
+    flb_sds_t tmp;
+
+    struct flb_http_client *c = NULL;
+
+    struct flb_aws_client *ec2_imds_client = ctx->ec2_imds_client;
+    struct flb_aws_header token_header = imds_v2_token_token_header_template;
+
+    /* Get IMDS version */
+    int imds_version = get_imds_version(ctx);
+
+    /* Abort on version detection failure */
+    if (imds_version == FLB_AWS_IMDS_VERSION_EVALUATE) {
+        /* Exit gracefully allowing for retrys */
+        flb_warn("[imds] unable to evaluate IMDS version");
+        return -1;
+    }
+
+    if (imds_version == FLB_AWS_IMDS_VERSION_2) {
+        token_header.val = ctx->imds_v2_token;
+        token_header.val_len = ctx->imds_v2_token_len;
+        flb_debug("[imds] using IMDSv2");
+    }
+    else {
+        flb_debug("[imds] using IMDSv1");
+    }
+
+    c = ec2_imds_client->client_vtable->request(
+        ec2_imds_client, FLB_HTTP_GET, metadata_path, NULL, 0, &token_header,
+        (imds_version == FLB_AWS_IMDS_VERSION_1) ? 0 : 1);
+    if (!c) {
+        /* Exit gracefully allowing for retrys */
+        flb_warn("[imds] failed to retrieve metadata");
+        return -1;
+    }
+
+    /* Detect invalid token */
+    if (imds_version == FLB_AWS_IMDS_VERSION_2 && c->resp.status == 401) {
+        /* Refresh token and retry request */
+        flb_http_client_destroy(c);
+        ret = refresh_imds_v2_token(ctx);
+        if (ret < 0) {
+            flb_debug("[imds] failed to refresh IMDSv2 token");
+            return -1;
+        }
+        token_header.val = ctx->imds_v2_token;
+        token_header.val_len = ctx->imds_v2_token_len;
+        flb_debug("[imds] refreshed IMDSv2 token");
+        c = ec2_imds_client->client_vtable->request(
+            ec2_imds_client, FLB_HTTP_GET, metadata_path, NULL, 0, &token_header, 1);
+        if (!c) {
+            /* Exit gracefully allowing for retries */
+            flb_warn("[imds] failed to retrieve metadata");
+            return -1;
+        }
+    }
+
+    if (c->resp.status != 200) {
+        ret = -1;
+        if (c->resp.status == 404) {
+            ret = -2;
+        }
+        if (c->resp.payload_size > 0) {
+            flb_debug("[imds] metadata request failure response\n%s", c->resp.payload);
+        }
+        flb_http_client_destroy(c);
+        return ret;
+    }
+
+    if (key != NULL) {
+        /* get the value of the key from payload json string */
+        tmp = flb_json_get_val(c->resp.payload, c->resp.payload_size, key);
+        if (!tmp) {
+            tmp = flb_sds_create_len("NULL", 4);
+            flb_error("[imds] %s is undefined in EC2 instance", key);
+        }
+    }
+    else {
+        tmp = flb_sds_create_len(c->resp.payload, c->resp.payload_size);
+    }
+
+    if (!tmp) {
+        flb_errno();
+        flb_http_client_destroy(c);
+        return -1;
+    }
+
+    *metadata = tmp;
+    *metadata_len = key == NULL ? c->resp.payload_size : strlen(tmp);
+
+    flb_http_client_destroy(c);
+    return 0;
+}
+
+/* Get VPC Id */
+flb_sds_t flb_aws_imds_get_vpc_id(struct flb_aws_imds *ctx)
+{
+    int ret;
+    flb_sds_t mac_id = NULL;
+    size_t mac_len = 0;
+    flb_sds_t vpc_id = NULL;
+    size_t vpc_id_len = 0;
+
+    /* get EC2 instance Mac id first before getting VPC id */
+    ret = flb_aws_imds_request(ctx, FLB_AWS_IMDS_MAC_PATH, &mac_id, &mac_len);
+
+    if (ret < 0) {
+        flb_sds_destroy(mac_id);
+        return NULL;
+    }
+
+    /*
+     * the VPC full path should be like:
+     * latest/meta-data/network/interfaces/macs/{mac_id}/vpc-id/"
+     */
+    flb_sds_t vpc_path = flb_sds_create_size(70);
+    vpc_path =
+        flb_sds_printf(&vpc_path, "%s/%s/%s/",
+                       "/latest/meta-data/network/interfaces/macs", mac_id, "vpc-id");
+    ret = flb_aws_imds_request(ctx, vpc_path, &vpc_id, &vpc_id_len);
+
+    flb_sds_destroy(mac_id);
+    flb_sds_destroy(vpc_path);
+
+    return vpc_id;
+}
+
+/* Obtain the IMDS version */
+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) {
+        return ctx->imds_version;
+    }
+
+    /*
+     * Evaluate version
+     * To evaluate wether IMDSv2 is available, send an invalid token
+     * in IMDS request. If response status is 'Unauthorized', then IMDSv2
+     * is available.
+     */
+    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);
+
+    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");
+        ctx->imds_version = FLB_AWS_IMDS_VERSION_1;
+    }
+
+    flb_http_client_destroy(c);
+    return ctx->imds_version;
+}
+
+/*
+ * Get an IMDSv2 token
+ * Token preserved in imds context
+ */
+static int refresh_imds_v2_token(struct flb_aws_imds *ctx)
+{
+    struct flb_http_client *c = NULL;
+    struct flb_aws_client *ec2_imds_client = ctx->ec2_imds_client;
+
+    c = ec2_imds_client->client_vtable->request(ec2_imds_client, FLB_HTTP_PUT,
+                                                FLB_AWS_IMDS_V2_TOKEN_PATH, NULL, 0,
+                                                &imds_v2_token_ttl_header, 1);
+
+    if (!c) {
+        return -1;
+    }
+
+    if (c->resp.status != 200) {
+        if (c->resp.payload_size > 0) {
+            flb_error("[imds] IMDSv2 token retrieval failure response\n%s",
+                      c->resp.payload);
+        }
+
+        flb_http_client_destroy(c);
+        return -1;
+    }
+
+    /* Preserve token information in ctx */
+    if (c->resp.payload_size > 0) {
+        if (ctx->imds_v2_token) {
+            flb_sds_destroy(ctx->imds_v2_token);
+        }
+        ctx->imds_v2_token = flb_sds_create_len(c->resp.payload, c->resp.payload_size);
+        if (!ctx->imds_v2_token) {
+            flb_errno();
+            flb_http_client_destroy(c);
+            return -1;
+        }
+        ctx->imds_v2_token_len = c->resp.payload_size;
+
+        flb_http_client_destroy(c);
+        return 0;
+    }
+
+    flb_debug("[imds] IMDS metadata response was empty");
+    flb_http_client_destroy(c);
+    return -1;
+}