From 83ea88e03fe2775601636c5f578b63276910a538 Mon Sep 17 00:00:00 2001
From: Graham Christensen <graham@grahamc.com>
Date: Wed, 18 Nov 2020 11:56:15 -0500
Subject: [PATCH] nixos: ec2 ami: support IMDSv2

AWS's metadata service has two versions. Version 1 allowed plain HTTP
requests to get metadata. However, this was frequently abused when a
user could trick an AWS-hosted server in to proxying requests to the
metadata service. Since the metadata service is frequently used to
generate AWS access keys, this is pretty gnarly. Version two is
identical except it requires the caller to request a token and provide
it on each request.

Today, starting a NixOS AMI in EC2 where the metadata service is
configured to only allow v2 requests fails: the user's SSH key is not
placed, and configuration provided by the user-data is not applied.
The server is useless. This patch addresses that.

Note the dependency on curl is not a joyful one, and it expand the
initrd by 30M. However, see the added comment for more information
about why this is needed. Note the idea of using `echo` and `nc` are
laughable. Don't do that.
---
 nixos/modules/virtualisation/amazon-image.nix |  1 +
 .../virtualisation/ec2-metadata-fetcher.nix   | 45 ++++++++++++++++---
 2 files changed, 41 insertions(+), 5 deletions(-)

diff --git a/nixos/modules/virtualisation/amazon-image.nix b/nixos/modules/virtualisation/amazon-image.nix
index 44cb60809452..819e93a43e57 100644
--- a/nixos/modules/virtualisation/amazon-image.nix
+++ b/nixos/modules/virtualisation/amazon-image.nix
@@ -11,6 +11,7 @@ with lib;
 let
   cfg = config.ec2;
   metadataFetcher = import ./ec2-metadata-fetcher.nix {
+    inherit (pkgs) curl;
     targetRoot = "$targetRoot/";
     wgetExtraOptions = "-q";
   };
diff --git a/nixos/modules/virtualisation/ec2-metadata-fetcher.nix b/nixos/modules/virtualisation/ec2-metadata-fetcher.nix
index b531787c31a2..247bcf513c51 100644
--- a/nixos/modules/virtualisation/ec2-metadata-fetcher.nix
+++ b/nixos/modules/virtualisation/ec2-metadata-fetcher.nix
@@ -1,23 +1,58 @@
-{ targetRoot, wgetExtraOptions }:
+{ curl, targetRoot, wgetExtraOptions }:
+# Note: be very cautious about dependencies, each dependency grows
+# the closure of the initrd. Ideally we would not even require curl,
+# but there is no reasonable way to send an HTTP PUT request without
+# it. Note: do not be fooled: the wget referenced in this script
+# is busybox's wget, not the fully featured one with --method support.
+#
+# Make sure that every package you depend on here is already listed as
+# a channel blocker for both the full-sized and small channels.
+# Otherwise, we risk breaking user deploys in released channels.
 ''
   metaDir=${targetRoot}etc/ec2-metadata
   mkdir -m 0755 -p "$metaDir"
 
+  get_imds_token() {
+    # retry-delay of 1 selected to give the system a second to get going,
+    # but not add a lot to the bootup time
+    ${curl}/bin/curl \
+      -v \
+      --retry 3 \
+      --retry-delay 1 \
+      --fail \
+      -X PUT \
+      --connect-timeout 1 \
+      -H "X-aws-ec2-metadata-token-ttl-seconds: 600" \
+      http://169.254.169.254/latest/api/token
+  }
+
+  try=1
+  while [ $try -le 3 ]; do
+    echo "(attempt $try/3) getting an EC2 instance metadata service v2 token..."
+    IMDS_TOKEN=$(get_imds_token) && break
+    try=$((try + 1))
+    sleep 1
+  done
+
+  if [ "x$IMDS_TOKEN" == "x" ]; then
+    echo "failed to fetch an IMDS2v token."
+  fi
+
   echo "getting EC2 instance metadata..."
 
   if ! [ -e "$metaDir/ami-manifest-path" ]; then
-    wget ${wgetExtraOptions} -O "$metaDir/ami-manifest-path" http://169.254.169.254/1.0/meta-data/ami-manifest-path
+    wget ${wgetExtraOptions} --header "X-aws-ec2-metadata-token: $IMDS_TOKEN" -O "$metaDir/ami-manifest-path" http://169.254.169.254/1.0/meta-data/ami-manifest-path
   fi
 
   if ! [ -e "$metaDir/user-data" ]; then
-    wget ${wgetExtraOptions} -O "$metaDir/user-data" http://169.254.169.254/1.0/user-data && chmod 600 "$metaDir/user-data"
+    wget ${wgetExtraOptions} --header "X-aws-ec2-metadata-token: $IMDS_TOKEN" -O "$metaDir/user-data" http://169.254.169.254/1.0/user-data && chmod 600 "$metaDir/user-data"
   fi
 
   if ! [ -e "$metaDir/hostname" ]; then
-    wget ${wgetExtraOptions} -O "$metaDir/hostname" http://169.254.169.254/1.0/meta-data/hostname
+    wget ${wgetExtraOptions} --header "X-aws-ec2-metadata-token: $IMDS_TOKEN" -O "$metaDir/hostname" http://169.254.169.254/1.0/meta-data/hostname
   fi
 
   if ! [ -e "$metaDir/public-keys-0-openssh-key" ]; then
-    wget ${wgetExtraOptions} -O "$metaDir/public-keys-0-openssh-key" http://169.254.169.254/1.0/meta-data/public-keys/0/openssh-key
+    wget ${wgetExtraOptions} --header "X-aws-ec2-metadata-token: $IMDS_TOKEN" -O "$metaDir/public-keys-0-openssh-key" http://169.254.169.254/1.0/meta-data/public-keys/0/openssh-key
   fi
 ''