From 033531df79670dedcd3d726b06e385867d1eb301 Mon Sep 17 00:00:00 2001
From: Moritz Sanft <58110325+msanft@users.noreply.github.com>
Date: Mon, 11 Nov 2024 08:41:53 +0100
Subject: [PATCH 1/2] packages/nixos: add IMDS setup script

Azure needs special care for enabling IMDS within Peerpods. This adds a script to setup IMDS through Proxy ARP (from Peerpods upstream, see https://github.com/confidential-containers/cloud-api-adaptor/blob/main/src/cloud-api-adaptor/podvm/files/usr/local/bin/setup-nat-for-imds.sh), so that all requests to the IMDS from within the pod are routed through an interface that is peered to the Pod VM. Verified to work in 2 distinct Azure peer pods.
---
 .../by-name/cloud-api-adaptor/package.nix     | 19 ++++++++++++++++++
 packages/nixos/azure.nix                      | 20 +++++++++++++++++++
 2 files changed, 39 insertions(+)

diff --git a/packages/by-name/cloud-api-adaptor/package.nix b/packages/by-name/cloud-api-adaptor/package.nix
index 1e1e915201..c487317241 100644
--- a/packages/by-name/cloud-api-adaptor/package.nix
+++ b/packages/by-name/cloud-api-adaptor/package.nix
@@ -10,6 +10,9 @@
   writeShellApplication,
   gnugrep,
   iptables,
+  iproute2,
+  sysctl,
+  gawk,
   runCommand,
   applyPatches,
   makeWrapper,
@@ -102,6 +105,22 @@ buildGoModule rec {
         "SC2153"
       ];
     };
+
+    setup-nat-for-imds = writeShellApplication {
+      name = "setup-nat-for-imds";
+      runtimeInputs = [
+        iproute2
+        iptables
+        sysctl
+        gawk
+      ];
+      # TODO(burgerdev): generalize for all link-local IPs and investigate routing simplification
+      text = builtins.readFile "${cloud-api-adaptor.src}/src/cloud-api-adaptor/podvm/files/usr/local/bin/setup-nat-for-imds.sh";
+      meta = {
+        mainProgram = "setup-nat-for-imds";
+        homepage = "https://github.com/confidential-containers/cloud-api-adaptor/blob/main/src/cloud-api-adaptor/podvm/files/usr/local/bin/setup-nat-for-imds.sh";
+      };
+    };
   };
 
   meta = {
diff --git a/packages/nixos/azure.nix b/packages/nixos/azure.nix
index 612b48e767..4f2a70a94f 100644
--- a/packages/nixos/azure.nix
+++ b/packages/nixos/azure.nix
@@ -85,5 +85,25 @@ in
         ExecStart = "${lib.getExe pkgs.azure-no-agent}";
       };
     };
+
+    systemd.services.setup-nat-for-imds = {
+      wantedBy = [ "multi-user.target" ];
+      requires = [ "netns@podns.service" ];
+      wants = [ "network-online.target" ];
+      after = [
+        "network-online.target"
+        "netns@podns.service"
+      ];
+      description = "Setup NAT for IMDS";
+      serviceConfig = {
+        Type = "oneshot";
+        RemainAfterExit = "yes";
+        # TODO(msanft): Find out why just ordering this after network-online.target
+        # isn't sufficient. (Errors with saying that the network is unreachable)
+        Restart = "on-failure";
+        RestartSec = "5s";
+        ExecStart = "${lib.getExe pkgs.cloud-api-adaptor.setup-nat-for-imds}";
+      };
+    };
   };
 }

From b6abc8d51e3d94526980ad8f5808da7d6a4cb27b Mon Sep 17 00:00:00 2001
From: Moritz Sanft <58110325+msanft@users.noreply.github.com>
Date: Mon, 2 Dec 2024 12:16:01 +0100
Subject: [PATCH 2/2] packages/test-peerpods: test IMDS functionality

This adds a verification of IMDS functionality to the peer-pods smoke test.
---
 packages/test-peerpods.sh | 27 +++++++++++++++++++--------
 1 file changed, 19 insertions(+), 8 deletions(-)

diff --git a/packages/test-peerpods.sh b/packages/test-peerpods.sh
index c2e3dc15bd..d28f6f8f8c 100644
--- a/packages/test-peerpods.sh
+++ b/packages/test-peerpods.sh
@@ -39,9 +39,17 @@ if [[ $found != true ]]; then
   exit 1
 fi
 
+run_tests() {
+  pod="$(kubectl get pod -l app=alpine -o jsonpath='{.items[0].metadata.name}')"
+
+  # Check IMDS functionality.
+  # -f makes this fail on a 500 status code.
+  kubectl exec "$pod" -- curl -f -i -H "Metadata: true" http://169.254.169.254/metadata/THIM/amd/certification
+}
+
 cleanup() {
-  kubectl delete deploy nginx
-  kubectl wait --for=delete pod --selector=app=nginx --timeout=5m
+  kubectl delete deploy alpine
+  kubectl wait --for=delete pod --selector=app=alpine --timeout=5m
 }
 
 trap cleanup EXIT
@@ -52,26 +60,29 @@ kubectl apply -f - <<EOF
 apiVersion: apps/v1
 kind: Deployment
 metadata:
-  name: nginx
+  name: alpine
 spec:
   selector:
     matchLabels:
-      app: nginx
+      app: alpine
   replicas: 1
   template:
     metadata:
       labels:
-        app: nginx
+        app: alpine
     spec:
       runtimeClassName: kata-remote
       containers:
-      - name: nginx
-        image: nginx
+      - name: alpine
+        image: alpine/curl
         imagePullPolicy: Always
+        command: ["sleep", "3600"]
 EOF
 
-if ! kubectl wait --for=condition=available --timeout=5m deployment/nginx; then
+if ! kubectl wait --for=condition=available --timeout=5m deployment/alpine; then
   kubectl describe pods
   kubectl logs -n confidential-containers-system -l app=cloud-api-adaptor --tail=-1 --all-containers
   exit 1
 fi
+
+run_tests