1919
2020import org .ice4j .*;
2121
22- import java .io .*;
2322import java .net .*;
23+ import java .net .http .*;
24+ import java .time .*;
25+ import java .util .*;
2426import java .util .logging .*;
2527
2628/**
@@ -39,6 +41,14 @@ public class AwsCandidateHarvester
3941 private static final Logger logger
4042 = Logger .getLogger (AwsCandidateHarvester .class .getName ());
4143
44+ /**
45+ * The <tt>HttpClient</tt> used by the <tt>AwsCandidateHarvester</tt>
46+ * class and its instances for making requests to IMDS.
47+ */
48+ private static final HttpClient httpClient = HttpClient .newBuilder ()
49+ .connectTimeout (Duration .ofMillis (500 ))
50+ .build ();
51+
4252 /**
4353 * The URL where one obtains AWS public addresses.
4454 */
@@ -52,10 +62,28 @@ public class AwsCandidateHarvester
5262 = "http://169.254.169.254/latest/meta-data/local-ipv4" ;
5363
5464 /**
55- * The URL to use to test whether we are running on Amazon EC2.
65+ * The URL to get IMDSv2 API token for further meta-data requests.
66+ */
67+ private static final String IMDS_API_TOKEN_URL
68+ = "http://169.254.169.254/latest/api/token" ;
69+
70+ /**
71+ * The HTTP header name to provide session token.
5672 */
57- private static final String EC2_TEST_URL
58- = "http://169.254.169.254/latest/meta-data/" ;
73+ private static final String EC2_METADATA_TOKEN_HEADER
74+ = "X-aws-ec2-metadata-token" ;
75+
76+ /**
77+ * The HTTP header name to request session token TTL.
78+ */
79+ private static final String EC2_METADATA_TOKEN_TTL_HEADER
80+ = "X-aws-ec2-metadata-token-ttl-seconds" ;
81+
82+ /**
83+ * The default session token TTL value.
84+ */
85+ private static final String EC2_METADATA_TOKEN_DEFAULT_TTL
86+ = "21600" ;
5987
6088 /**
6189 * Whether we are running on Amazon EC2.
@@ -103,8 +131,11 @@ private static synchronized void obtainEC2Addresses()
103131
104132 try
105133 {
106- localIPStr = fetch (LOCAL_IP_URL );
107- publicIPStr = fetch (PUBLIC_IP_URL );
134+ String metaDataToken = fetch (IMDS_API_TOKEN_URL ,
135+ Collections .singletonMap (EC2_METADATA_TOKEN_TTL_HEADER , EC2_METADATA_TOKEN_DEFAULT_TTL ), "PUT" );
136+ Map <String , String > tokenHeader = Collections .singletonMap (EC2_METADATA_TOKEN_HEADER , metaDataToken );
137+ localIPStr = fetch (LOCAL_IP_URL , tokenHeader );
138+ publicIPStr = fetch (PUBLIC_IP_URL , tokenHeader );
108139
109140 //now let's cross our fingers and hope that what we got above are
110141 //real IP addresses
@@ -183,11 +214,10 @@ private static boolean doTestEc2()
183214 {
184215 try
185216 {
186- URLConnection conn = new URL (EC2_TEST_URL ).openConnection ();
187- conn .setConnectTimeout (500 ); //don't hang for too long
188- conn .getContent ();
217+ String metaDataToken = fetch (IMDS_API_TOKEN_URL ,
218+ Collections .singletonMap (EC2_METADATA_TOKEN_TTL_HEADER , EC2_METADATA_TOKEN_DEFAULT_TTL ), "PUT" );
189219
190- return true ;
220+ return metaDataToken != null ;
191221 }
192222 catch (Exception exc )
193223 {
@@ -201,21 +231,40 @@ private static boolean doTestEc2()
201231 *
202232 * @param url the URL we'd like to open and query.
203233 *
234+ * @param headers the HTTP headers to put into the request.
235+ *
236+ * @throws Exception if anything goes wrong.
237+ */
238+ private static String fetch (String url , Map <String , String > headers )
239+ throws Exception
240+ {
241+ return fetch (url , headers , "GET" );
242+ }
243+
244+ /**
245+ * Retrieves the content at the specified <tt>url</tt>. No more, no less.
246+ *
247+ * @param url the URL we'd like to open and query.
248+ *
249+ * @param headers the HTTP headers to put into the request.
250+ *
251+ * @param method the HTTP method we'd like to use.
252+ *
204253 * @return the String we retrieved from the URL.
205254 *
206255 * @throws Exception if anything goes wrong.
207256 */
208- private static String fetch (String url )
257+ private static String fetch (String url , Map < String , String > headers , String method )
209258 throws Exception
210259 {
211- URLConnection conn = new URL ( url ). openConnection ();
212- BufferedReader in = new BufferedReader (new InputStreamReader (
213- conn . getInputStream (), "UTF-8" ));
214-
215- String retString = in . readLine ();
216-
217- in . close ();
218-
219- return retString ;
260+ HttpRequest . Builder builder = HttpRequest . newBuilder ()
261+ . uri (new URI ( url ))
262+ . method ( method , HttpRequest . BodyPublishers . noBody ( ));
263+ for ( Map . Entry < String , String > header : headers . entrySet ())
264+ {
265+ builder . setHeader ( header . getKey (), header . getValue ());
266+ }
267+ HttpResponse < String > response = httpClient . send ( builder . build (), HttpResponse . BodyHandlers . ofString ());
268+ return response . body () ;
220269 }
221270}
0 commit comments