From 167d9f18afee4f2d50610b72dacd1f210c4bc730 Mon Sep 17 00:00:00 2001
From: Ben Whaley <503816+bwhaley@users.noreply.github.com>
Date: Tue, 9 Apr 2024 16:18:25 -0700
Subject: [PATCH] New option to disable IPv6 in the connectivity test

---
 functions/replace-route/app.py                | 27 ++++++++++++++++++-
 .../replace-route/tests/test_replace_route.py | 13 +++++++++
 modules/terraform-aws-alternat/lambda.tf      | 11 +++++---
 modules/terraform-aws-alternat/variables.tf   |  6 +++++
 4 files changed, 53 insertions(+), 4 deletions(-)

diff --git a/functions/replace-route/app.py b/functions/replace-route/app.py
index 2da1461..9e4b09e 100644
--- a/functions/replace-route/app.py
+++ b/functions/replace-route/app.py
@@ -27,9 +27,24 @@
 # Which URLs to check for connectivity
 DEFAULT_CHECK_URLS = ["https://www.example.com", "https://www.google.com"]
 
-
+# The timeout for the connectivity checks.
 REQUEST_TIMEOUT = 5
 
+# Whether or not use IPv6.
+DEFAULT_HAS_IPV6 = True
+
+
+# Overrides socket.getaddrinfo to perform IPv4 lookups
+# See https://github.com/1debit/alternat/issues/87
+def disable_ipv6():
+    prv_getaddrinfo = socket.getaddrinfo
+    def getaddrinfo_ipv4(*args):
+        modified_args = (args[0], args[1], socket.AF_INET) + args[3:]
+        res = prv_getaddrinfo(*modified_args)
+        return res
+    socket.getaddrinfo = getaddrinfo_ipv4
+
+
 def get_az_and_vpc_zone_identifier(auto_scaling_group):
     autoscaling = boto3.client("autoscaling")
 
@@ -156,6 +171,10 @@ def connectivity_test_handler(event, context):
     check_interval = int(os.getenv("CONNECTIVITY_CHECK_INTERVAL", DEFAULT_CONNECTIVITY_CHECK_INTERVAL))
     check_urls = "CHECK_URLS" in os.environ and os.getenv("CHECK_URLS").split(",") or DEFAULT_CHECK_URLS
 
+    has_ipv6 = get_env_bool("HAS_IPV6", DEFAULT_HAS_IPV6)
+    if not has_ipv6:
+        disable_ipv6()
+
     # Run connectivity checks for approximately 1 minute
     run = 0
     num_runs = 60 / check_interval
@@ -167,6 +186,12 @@ def connectivity_test_handler(event, context):
             break
 
 
+def get_env_bool(var_name, default_value=False):
+    value = os.getenv(var_name, default_value)
+    true_values = ["t", "true", "y", "yes", "1"]
+    return str(value).lower() in true_values
+
+
 def handler(event, _):
     try:
         for record in event["Records"]:
diff --git a/functions/replace-route/tests/test_replace_route.py b/functions/replace-route/tests/test_replace_route.py
index 2bc080b..05bc5a2 100644
--- a/functions/replace-route/tests/test_replace_route.py
+++ b/functions/replace-route/tests/test_replace_route.py
@@ -206,3 +206,16 @@ class Context:
     connectivity_test_handler(event=json.loads(cloudwatch_event), context=Context())
 
     verify_nat_gateway_route(mocked_networking)
+
+
+def test_disable_ipv6():
+    results = socket.getaddrinfo("google.com", 80)
+    families = [family for family, _, _, _, _ in results]
+    assert socket.AF_INET6 in families, "Unable to find a non-IPv4 address family in getaddrinfo results before patching getaddrinfo"
+
+    from app import disable_ipv6
+    disable_ipv6()
+    results = socket.getaddrinfo("google.com", 80)
+    for family, _, _, _, _ in results:
+        assert family == socket.AF_INET, "Found a non-IPv4 address family in getaddrinfo results"
+
diff --git a/modules/terraform-aws-alternat/lambda.tf b/modules/terraform-aws-alternat/lambda.tf
index 95723f3..e5c774a 100644
--- a/modules/terraform-aws-alternat/lambda.tf
+++ b/modules/terraform-aws-alternat/lambda.tf
@@ -15,7 +15,7 @@ resource "aws_lambda_function" "alternat_autoscaling_hook" {
   timeout       = var.lambda_timeout
   role          = aws_iam_role.nat_lambda_role.arn
 
-  layers        = var.lambda_layer_arns
+  layers = var.lambda_layer_arns
 
   image_uri = var.lambda_package_type == "Image" ? "${var.alternat_image_uri}:${var.alternat_image_tag}" : null
 
@@ -25,7 +25,11 @@ resource "aws_lambda_function" "alternat_autoscaling_hook" {
   source_code_hash = var.lambda_package_type == "Zip" ? data.archive_file.lambda[0].output_base64sha256 : null
 
   environment {
-    variables = merge(local.autoscaling_func_env_vars, var.lambda_environment_variables)
+    variables = merge(
+      local.autoscaling_func_env_vars,
+      local.has_ipv6_env_var,
+      var.lambda_environment_variables,
+    )
   }
 
   tags = merge({
@@ -39,6 +43,7 @@ locals {
     for obj in var.vpc_az_maps
     : replace(upper(obj.az), "-", "_") => join(",", obj.route_table_ids)
   }
+  has_ipv6_env_var = { "HAS_IPV6" = var.lambda_has_ipv6 }
 }
 
 resource "aws_iam_role" "nat_lambda_role" {
@@ -129,7 +134,7 @@ resource "aws_lambda_function" "alternat_connectivity_tester" {
   timeout       = var.lambda_timeout
   role          = aws_iam_role.nat_lambda_role.arn
 
-  layers        = var.lambda_layer_arns
+  layers = var.lambda_layer_arns
 
   image_uri = var.lambda_package_type == "Image" ? "${var.alternat_image_uri}:${var.alternat_image_tag}" : null
 
diff --git a/modules/terraform-aws-alternat/variables.tf b/modules/terraform-aws-alternat/variables.tf
index bfc959e..42c3c05 100644
--- a/modules/terraform-aws-alternat/variables.tf
+++ b/modules/terraform-aws-alternat/variables.tf
@@ -242,6 +242,12 @@ variable "lambda_environment_variables" {
   default     = null
 }
 
+variable "lambda_has_ipv6" {
+  description = "Controls whether or not the lambda function can use IPv6."
+  type        = bool
+  default     = true
+}
+
 variable "lambda_zip_path" {
   description = "The location where the generated zip file should be stored. Required when `lambda_package_type` is \"Zip\"."
   type        = string