From 95c3a6c14be9efb6f6edfb765da336c0d4c7f5ef Mon Sep 17 00:00:00 2001
From: Pascal Fischer <pascal@netbird.io>
Date: Thu, 8 Feb 2024 12:48:19 +0100
Subject: [PATCH] add IDMSv2 to AWS detection

---
 aws.go | 40 ++++++++++++++++++++++++++++++++++++++--
 1 file changed, 38 insertions(+), 2 deletions(-)

diff --git a/aws.go b/aws.go
index 1848e3f..2c21387 100644
--- a/aws.go
+++ b/aws.go
@@ -5,9 +5,45 @@ import (
 )
 
 func detectAWS() string {
-	resp, err := hc.Get("http://169.254.169.254/latest/")
-	if err == nil && resp.StatusCode == http.StatusOK {
+	v1ResultChan := make(chan bool, 1)
+	v2ResultChan := make(chan bool, 1)
+
+	go func() {
+		v1ResultChan <- detectIDMSv1()
+	}()
+
+	go func() {
+		v2ResultChan <- detectIDMSv2()
+	}()
+
+	v1Result, v2Result := <-v1ResultChan, <-v2ResultChan
+
+	if v1Result || v2Result {
 		return "Amazon Web Services"
 	}
 	return ""
 }
+
+func detectIDMSv1() bool {
+	resp, err := hc.Get("http://169.254.169.254/latest/")
+	if err != nil {
+		return false
+	}
+	defer resp.Body.Close()
+	return resp.StatusCode == http.StatusOK
+}
+
+func detectIDMSv2() bool {
+	req, err := http.NewRequest("PUT", "http://169.254.169.254/latest/api/token", nil)
+	if err != nil {
+		return false
+	}
+	req.Header.Set("X-aws-ec2-metadata-token-ttl-seconds", "21600")
+
+	resp, err := hc.Do(req)
+	if err != nil {
+		return false
+	}
+	defer resp.Body.Close()
+	return resp.StatusCode == http.StatusOK
+}