2828STRICT_ID_PATH = ("datasource" , "Ec2" , "strict_id" )
2929STRICT_ID_DEFAULT = "warn"
3030
31+ API_TOKEN_ROUTE = 'latest/api/token'
32+ API_TOKEN_DISABLED = '_ec2_disable_api_token'
33+ AWS_TOKEN_TTL_SECONDS = '21600'
34+
3135
3236class CloudNames (object ):
3337 ALIYUN = "aliyun"
@@ -62,6 +66,7 @@ class DataSourceEc2(sources.DataSource):
6266 url_max_wait = 120
6367 url_timeout = 50
6468
69+ _api_token = None # API token for accessing the metadata service
6570 _network_config = sources .UNSET # Used to cache calculated network cfg v1
6671
6772 # Whether we want to get network configuration from the metadata service.
@@ -148,11 +153,12 @@ def get_metadata_api_version(self):
148153 min_metadata_version.
149154 """
150155 # Assumes metadata service is already up
156+ url_tmpl = '{0}/{1}/meta-data/instance-id'
157+ headers = self ._get_headers ()
151158 for api_ver in self .extended_metadata_versions :
152- url = '{0}/{1}/meta-data/instance-id' .format (
153- self .metadata_address , api_ver )
159+ url = url_tmpl .format (self .metadata_address , api_ver )
154160 try :
155- resp = uhelp .readurl (url = url )
161+ resp = uhelp .readurl (url = url , headers = headers )
156162 except uhelp .UrlError as e :
157163 LOG .debug ('url %s raised exception %s' , url , e )
158164 else :
@@ -172,12 +178,39 @@ def get_instance_id(self):
172178 # setup self.identity. So we need to do that now.
173179 api_version = self .get_metadata_api_version ()
174180 self .identity = ec2 .get_instance_identity (
175- api_version , self .metadata_address ).get ('document' , {})
181+ api_version , self .metadata_address ,
182+ headers_cb = self ._get_headers ,
183+ exception_cb = self ._refresh_stale_aws_token_cb ).get (
184+ 'document' , {})
176185 return self .identity .get (
177186 'instanceId' , self .metadata ['instance-id' ])
178187 else :
179188 return self .metadata ['instance-id' ]
180189
190+ def _maybe_fetch_api_token (self , mdurls , timeout = None , max_wait = None ):
191+ if self .cloud_name != CloudNames .AWS :
192+ return
193+
194+ urls = []
195+ url2base = {}
196+ url_path = API_TOKEN_ROUTE
197+ request_method = 'PUT'
198+ for url in mdurls :
199+ cur = '{0}/{1}' .format (url , url_path )
200+ urls .append (cur )
201+ url2base [cur ] = url
202+
203+ # use the self._status_cb to check for Read errors, which means
204+ # we can't reach the API token URL, so we should disable IMDSv2
205+ LOG .debug ('Fetching Ec2 IMDSv2 API Token' )
206+ url , response = uhelp .wait_for_url (
207+ urls = urls , max_wait = 1 , timeout = 1 , status_cb = self ._status_cb ,
208+ headers_cb = self ._get_headers , request_method = request_method )
209+
210+ if url and response :
211+ self ._api_token = response
212+ return url2base [url ]
213+
181214 def wait_for_metadata_service (self ):
182215 mcfg = self .ds_cfg
183216
@@ -199,27 +232,39 @@ def wait_for_metadata_service(self):
199232 LOG .warning ("Empty metadata url list! using default list" )
200233 mdurls = self .metadata_urls
201234
202- urls = []
203- url2base = {}
204- for url in mdurls :
205- cur = '{0}/{1}/meta-data/instance-id' .format (
206- url , self .min_metadata_version )
207- urls .append (cur )
208- url2base [cur ] = url
209-
210- start_time = time .time ()
211- url = uhelp .wait_for_url (
212- urls = urls , max_wait = url_params .max_wait_seconds ,
213- timeout = url_params .timeout_seconds , status_cb = LOG .warning )
214-
215- if url :
216- self .metadata_address = url2base [url ]
235+ # try the api token path first
236+ metadata_address = self ._maybe_fetch_api_token (mdurls )
237+ if not metadata_address :
238+ if self ._api_token == API_TOKEN_DISABLED :
239+ LOG .warning ('Retrying with IMDSv1' )
240+ # if we can't get a token, use instance-id path
241+ urls = []
242+ url2base = {}
243+ url_path = '{ver}/meta-data/instance-id' .format (
244+ ver = self .min_metadata_version )
245+ request_method = 'GET'
246+ for url in mdurls :
247+ cur = '{0}/{1}' .format (url , url_path )
248+ urls .append (cur )
249+ url2base [cur ] = url
250+
251+ start_time = time .time ()
252+ url , _ = uhelp .wait_for_url (
253+ urls = urls , max_wait = url_params .max_wait_seconds ,
254+ timeout = url_params .timeout_seconds , status_cb = LOG .warning ,
255+ headers_cb = self ._get_headers , request_method = request_method )
256+
257+ if url :
258+ metadata_address = url2base [url ]
259+
260+ if metadata_address :
261+ self .metadata_address = metadata_address
217262 LOG .debug ("Using metadata source: '%s'" , self .metadata_address )
218263 else :
219264 LOG .critical ("Giving up on md from %s after %s seconds" ,
220265 urls , int (time .time () - start_time ))
221266
222- return bool (url )
267+ return bool (metadata_address )
223268
224269 def device_name_to_device (self , name ):
225270 # Consult metadata service, that has
@@ -376,14 +421,22 @@ def crawl_metadata(self):
376421 return {}
377422 api_version = self .get_metadata_api_version ()
378423 crawled_metadata = {}
424+ if self .cloud_name == CloudNames .AWS :
425+ exc_cb = self ._refresh_stale_aws_token_cb
426+ exc_cb_ud = self ._skip_or_refresh_stale_aws_token_cb
427+ else :
428+ exc_cb = exc_cb_ud = None
379429 try :
380430 crawled_metadata ['user-data' ] = ec2 .get_instance_userdata (
381- api_version , self .metadata_address )
431+ api_version , self .metadata_address ,
432+ headers_cb = self ._get_headers , exception_cb = exc_cb_ud )
382433 crawled_metadata ['meta-data' ] = ec2 .get_instance_metadata (
383- api_version , self .metadata_address )
434+ api_version , self .metadata_address ,
435+ headers_cb = self ._get_headers , exception_cb = exc_cb )
384436 if self .cloud_name == CloudNames .AWS :
385437 identity = ec2 .get_instance_identity (
386- api_version , self .metadata_address )
438+ api_version , self .metadata_address ,
439+ headers_cb = self ._get_headers , exception_cb = exc_cb )
387440 crawled_metadata ['dynamic' ] = {'instance-identity' : identity }
388441 except Exception :
389442 util .logexc (
@@ -393,6 +446,73 @@ def crawl_metadata(self):
393446 crawled_metadata ['_metadata_api_version' ] = api_version
394447 return crawled_metadata
395448
449+ def _refresh_api_token (self , seconds = AWS_TOKEN_TTL_SECONDS ):
450+ """Request new metadata API token.
451+ @param seconds: The lifetime of the token in seconds
452+
453+ @return: The API token or None if unavailable.
454+ """
455+ if self .cloud_name != CloudNames .AWS :
456+ return None
457+ LOG .debug ("Refreshing Ec2 metadata API token" )
458+ request_header = {'X-aws-ec2-metadata-token-ttl-seconds' : seconds }
459+ token_url = '{}/{}' .format (self .metadata_address , API_TOKEN_ROUTE )
460+ try :
461+ response = uhelp .readurl (
462+ token_url , headers = request_header , request_method = "PUT" )
463+ except uhelp .UrlError as e :
464+ LOG .warning (
465+ 'Unable to get API token: %s raised exception %s' ,
466+ token_url , e )
467+ return None
468+ return response .contents
469+
470+ def _skip_or_refresh_stale_aws_token_cb (self , msg , exception ):
471+ """Callback will not retry on SKIP_USERDATA_CODES or if no token
472+ is available."""
473+ retry = ec2 .skip_retry_on_codes (
474+ ec2 .SKIP_USERDATA_CODES , msg , exception )
475+ if not retry :
476+ return False # False raises exception
477+ return self ._refresh_stale_aws_token_cb (msg , exception )
478+
479+ def _refresh_stale_aws_token_cb (self , msg , exception ):
480+ """Exception handler for Ec2 to refresh token if token is stale."""
481+ if isinstance (exception , uhelp .UrlError ) and exception .code == 401 :
482+ # With _api_token as None, _get_headers will _refresh_api_token.
483+ LOG .debug ("Clearing cached Ec2 API token due to expiry" )
484+ self ._api_token = None
485+ return True # always retry
486+
487+ def _status_cb (self , msg , exc = None ):
488+ LOG .warning (msg )
489+ if 'Read timed out' in msg :
490+ LOG .warning ('Cannot use Ec2 IMDSv2 API tokens, using IMDSv1' )
491+ self ._api_token = API_TOKEN_DISABLED
492+
493+ def _get_headers (self , url = '' ):
494+ """Return a dict of headers for accessing a url.
495+
496+ If _api_token is unset on AWS, attempt to refresh the token via a PUT
497+ and then return the updated token header.
498+ """
499+ if self .cloud_name != CloudNames .AWS or (self ._api_token ==
500+ API_TOKEN_DISABLED ):
501+ return {}
502+ # Request a 6 hour token if URL is API_TOKEN_ROUTE
503+ request_token_header = {
504+ 'X-aws-ec2-metadata-token-ttl-seconds' : AWS_TOKEN_TTL_SECONDS }
505+ if API_TOKEN_ROUTE in url :
506+ return request_token_header
507+ if not self ._api_token :
508+ # If we don't yet have an API token, get one via a PUT against
509+ # API_TOKEN_ROUTE. This _api_token may get unset by a 403 due
510+ # to an invalid or expired token
511+ self ._api_token = self ._refresh_api_token ()
512+ if not self ._api_token :
513+ return {}
514+ return {'X-aws-ec2-metadata-token' : self ._api_token }
515+
396516
397517class DataSourceEc2Local (DataSourceEc2 ):
398518 """Datasource run at init-local which sets up network to query metadata.
0 commit comments