Added caching of temporary AWS credentials on EC2 instances with IAM roles

Author: hannes-ucsc

src/bd2k/util/ec2/__init__.py

@@ -0,0 +1,153 @@
+import errno
+import logging
+import threading
+import time
+from datetime import datetime
+
+import os
+from bd2k.util.files import mkdir_p
+
+log = logging.getLogger( __name__ )
+
+cache_path = '~/.cache/aws/cached_temporary_credentials'
+
+datetime_format = "%Y-%m-%dT%H:%M:%SZ"  # incidentally the same as the format used by AWS
+
+
+def datetime_to_str( dt ):
+    """
+    Convert a naive (implicitly UTC) datetime object into a string, explicitly UTC.
+
+    >>> datetime_to_str( datetime( 1970, 1, 1, 0, 0, 0 ) )
+    '1970-01-01T00:00:00Z'
+    """
+    return dt.strftime( datetime_format )
+
+
+def str_to_datetime( s ):
+    """
+    Convert a string, explicitly UTC into a naive (implicitly UTC) datetime object.
+
+    >>> str_to_datetime( '1970-01-01T00:00:00Z' )
+    datetime.datetime(1970, 1, 1, 0, 0)
+
+    Just to show that the constructor args for seconds and microseconds are optional:
+    >>> datetime(1970, 1, 1, 0, 0, 0)
+    datetime.datetime(1970, 1, 1, 0, 0)
+    """
+    return datetime.strptime( s, datetime_format )
+
+
+monkey_patch_lock = threading.RLock( )
+_populate_keys_from_metadata_server_orig = None
+
+
+def enable_metadata_credential_caching( ):
+    """
+    Monkey-patches Boto to allow multiple processes using it to share one set of cached, temporary
+    IAM role credentials. This helps avoid hitting request limits imposed on the metadata service
+    when too many processes concurrently request those credentials. Function is idempotent.
+
+    This function should be called before any AWS connections attempts are made with Boto.
+    """
+    global _populate_keys_from_metadata_server_orig
+    with monkey_patch_lock:
+        if _populate_keys_from_metadata_server_orig is None:
+            from boto.provider import Provider
+            _populate_keys_from_metadata_server_orig = Provider._populate_keys_from_metadata_server
+            Provider._populate_keys_from_metadata_server = _populate_keys_from_metadata_server
+
+
+def disable_metadata_credential_caching( ):
+    """
+    Reverse the effect of enable_metadata_credential_caching()
+    """
+    global _populate_keys_from_metadata_server_orig
+    with monkey_patch_lock:
+        if _populate_keys_from_metadata_server_orig is not None:
+            from boto.provider import Provider
+            Provider._populate_keys_from_metadata_server = _populate_keys_from_metadata_server_orig
+            _populate_keys_from_metadata_server_orig = None
+
+
+def _populate_keys_from_metadata_server( self ):
+    global _populate_keys_from_metadata_server_orig
+    path = os.path.expanduser( cache_path )
+    tmp_path = path + '.tmp'
+    while True:
+        log.debug( 'Attempting to read cached credentials from %s.', path )
+        try:
+            with open( path, 'r' ) as f:
+                record = f.read( ).split( '\n' )
+                if record:
+                    self._access_key = record[ 0 ]
+                    self._secret_key = record[ 1 ]
+                    self._security_token = record[ 2 ]
+                    self._credential_expiry_time = str_to_datetime( record[ 3 ] )
+                else:
+                    log.debug( '%s is empty. Credentials are not temporary.', path )
+                    return
+        except IOError as e:
+            if e.errno == errno.ENOENT:
+                log.debug( 'Cached credentials are missing.' )
+                dir_path = os.path.dirname( path )
+                if not os.path.exists( dir_path ):
+                    log.debug( 'Creating parent directory %s', dir_path )
+                    # A race would be ok at this point
+                    mkdir_p( dir_path )
+            else:
+                raise
+        else:
+            if self._credentials_need_refresh( ):
+                log.debug( 'Cached credentials are expired.' )
+            else:
+                log.debug( 'Cached credentials exist and are still fresh.' )
+                return
+        # We get here if credentials are missing or expired
+        log.debug( 'Racing to create %s.', tmp_path )
+        # Only one process, the winner, will succeed
+        try:
+            fd = os.open( tmp_path, os.O_CREAT | os.O_EXCL | os.O_WRONLY, 0600 )
+        except OSError as e:
+            if e.errno == errno.EEXIST:
+                log.debug( 'Lost the race to create %s. Waiting on winner to remove it.', tmp_path )
+                while os.path.exists( tmp_path ):
+                    time.sleep( .1 )
+                log.debug( 'Winner removed %s. Trying from the top.', tmp_path )
+            else:
+                raise
+        else:
+            try:
+                log.debug( 'Won the race to create %s. '
+                           'Requesting credentials from metadata service.', tmp_path )
+                _populate_keys_from_metadata_server_orig( self )
+            except:
+                os.close( fd )
+                fd = None
+                log.debug( 'Failed to obtain credentials, removing %s.', tmp_path )
+                # This unblocks the loosers.
+                os.unlink( tmp_path )
+                # Bail out. It's too likely to happen repeatedly
+                raise
+            else:
+                if self._credential_expiry_time is None:
+                    os.close( fd )
+                    fd = None
+                    log.debug( 'Credentials are not temporary. '
+                               'Leaving %s empty and renaming it to %s.', tmp_path, path )
+                else:
+                    log.debug( 'Writing credentials to %s.', tmp_path )
+                    with os.fdopen( fd, 'w' ) as fh:
+                        fd = None
+                        fh.write( '\n'.join( [
+                            self._access_key,
+                            self._secret_key,
+                            self._security_token,
+                            datetime_to_str( self._credential_expiry_time ) ] ) )
+                    log.debug( 'Wrote credentials to %s. '
+                               'Renaming it to %s.', tmp_path, path )
+                os.rename( tmp_path, path )
+                return
+            finally:
+                if fd is not None:
+                    os.close( fd )

src/bd2k/util/ec2/credentials.py

@@ -0,0 +1,83 @@
+import logging
+
+import errno
+
+import os
+import unittest
+
+from bd2k.util.ec2.credentials import (enable_metadata_credential_caching,
+                                       disable_metadata_credential_caching, cache_path)
+
+
+def get_access_key( ):
+    from boto.provider import Provider
+    provider = Provider( 'aws' )
+    return provider.get_access_key( )
+
+
+class CredentialsTest( unittest.TestCase ):
+    def __init__( self, *args, **kwargs ):
+        super( CredentialsTest, self ).__init__( *args, **kwargs )
+        self.cache_path = os.path.expanduser( cache_path )
+
+    @classmethod
+    def setUpClass( cls ):
+        super( CredentialsTest, cls ).setUpClass( )
+        logging.basicConfig( level=logging.DEBUG )
+
+    def setUp( self ):
+        super( CredentialsTest, self ).setUp( )
+        self.cleanUp( )
+
+    def cleanUp( self ):
+        try:
+            os.unlink( self.cache_path )
+        except OSError as e:
+            if e.errno == errno.ENOENT:
+                pass
+            else:
+                raise
+
+    def tearDown( self ):
+        super( CredentialsTest, self ).tearDown( )
+        self.cleanUp( )
+
+    def test_metadata_credential_caching( self ):
+        """
+        Brute forces many concurrent requests for getting temporary credentials. If you comment
+        out the calls to enable_metadata_credential_caching, you should see some failures due to
+        requests timing out. The test will also take much longer in that case.
+        """
+        num_tests = 1000
+        num_processes = 32
+        # Get key without caching
+        access_key = get_access_key( )
+        self.assertFalse( os.path.exists( self.cache_path ) )
+        enable_metadata_credential_caching( )
+        # Again for idempotence
+        enable_metadata_credential_caching( )
+        try:
+            futures = [ ]
+            from multiprocessing import Pool
+            pool = Pool( num_processes )
+            try:
+                for i in range( num_tests ):
+                    futures.append( pool.apply_async( get_access_key ) )
+            except:
+                pool.close( )
+                pool.terminate( )
+                raise
+            else:
+                pool.close( )
+                pool.join( )
+        finally:
+            disable_metadata_credential_caching( )
+            # Again for idempotence
+            disable_metadata_credential_caching( )
+        self.assertTrue( os.path.exists( self.cache_path ) )
+        self.assertEquals( len( futures ), num_tests )
+        access_keys = [ f.get( ) for f in futures ]
+        self.assertEquals( len( access_keys ), num_tests )
+        access_keys = set( access_keys )
+        self.assertEquals( len( access_keys ), 1 )
+        self.assertEquals( access_keys.pop( ), access_key )