Support for creating custom tokens without service account credentials (firebase#175)

hiranya911 · web-flow · commit e4950a0e7b6b · 2018-07-12T11:05:31.000-07:00
* Ability to sign custom tokens without a service account

* Fixing tests for py3

* Fix for Python 3.5

* Implemented go/firebase-admin-sign

* Sending the required headers when calling the metadata service

* Renamed service_account option to service_account_+id

* Added discovery failiure test

* Added snippet; Updated error message to be in sync with docs

* Renamed service_account_id option to serviceAccountId for consistency
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -1,5 +1,7 @@
 # Unreleased
 
+- [added] Implemented the ability to create custom tokens without
+  service account credentials.
 - [added] Admin SDK can now read the project ID from both `GCLOUD_PROJECT` and
   `GOOGLE_CLOUD_PROJECT` environment variables.
 
diff --git a/firebase_admin/__init__.py b/firebase_admin/__init__.py
@@ -50,9 +50,9 @@ def initialize_app(credential=None, options=None, name=_DEFAULT_APP_NAME):
       credential: A credential object used to initialize the SDK (optional). If none is provided,
           Google Application Default Credentials are used.
       options: A dictionary of configuration options (optional). Supported options include
-          ``databaseURL``, ``storageBucket``, ``projectId``, ``databaseAuthVariableOverride``
-          and ``httpTimeout``. If ``httpTimeout`` is not set, HTTP connections initiated by client
-          modules such as ``db`` will not time out.
+          ``databaseURL``, ``storageBucket``, ``projectId``, ``databaseAuthVariableOverride``,
+          ``serviceAccountId`` and ``httpTimeout``. If ``httpTimeout`` is not set, HTTP
+          connections initiated by client modules such as ``db`` will not time out.
       name: Name of the app (optional).
     Returns:
       App: A newly initialized instance of App.
diff --git a/firebase_admin/_token_gen.py b/firebase_admin/_token_gen.py
@@ -20,11 +20,13 @@
 import cachecontrol
 import requests
 import six
+from google.auth import credentials
+from google.auth import exceptions
+from google.auth import iam
 from google.auth import jwt
 from google.auth import transport
 import google.oauth2.id_token
-
-from firebase_admin import credentials
+import google.oauth2.service_account
 
 
 # ID token constants
@@ -46,9 +48,12 @@
     'acr', 'amr', 'at_hash', 'aud', 'auth_time', 'azp', 'cnf', 'c_hash',
     'exp', 'firebase', 'iat', 'iss', 'jti', 'nbf', 'nonce', 'sub'
 ])
+METADATA_SERVICE_URL = ('http://metadata/computeMetadata/v1/instance/service-accounts/'
+                        'default/email')
 
 # Error codes
 COOKIE_CREATE_ERROR = 'COOKIE_CREATE_ERROR'
+TOKEN_SIGN_ERROR = 'TOKEN_SIGN_ERROR'
 
 
 class ApiCallError(Exception):
@@ -60,20 +65,81 @@ def __init__(self, code, message, error=None):
         self.detail = error
 
 
+class _SigningProvider(object):
+    """Stores a reference to a google.auth.crypto.Signer."""
+
+    def __init__(self, signer, signer_email):
+        self._signer = signer
+        self._signer_email = signer_email
+
+    @property
+    def signer(self):
+        return self._signer
+
+    @property
+    def signer_email(self):
+        return self._signer_email
+
+    @classmethod
+    def from_credential(cls, google_cred):
+        return _SigningProvider(google_cred.signer, google_cred.signer_email)
+
+    @classmethod
+    def from_iam(cls, request, google_cred, service_account):
+        signer = iam.Signer(request, google_cred, service_account)
+        return _SigningProvider(signer, service_account)
+
+
 class TokenGenerator(object):
     """Generates custom tokens and session cookies."""
 
     def __init__(self, app, client):
-        self._app = app
-        self._client = client
+        self.app = app
+        self.client = client
+        self.request = transport.requests.Request()
+        self._signing_provider = None
+
+    def _init_signing_provider(self):
+        """Initializes a signing provider by following the go/firebase-admin-sign protocol."""
+        # If the SDK was initialized with a service account, use it to sign bytes.
+        google_cred = self.app.credential.get_credential()
+        if isinstance(google_cred, google.oauth2.service_account.Credentials):
+            return _SigningProvider.from_credential(google_cred)
+
+        # If the SDK was initialized with a service account email, use it with the IAM service
+        # to sign bytes.
+        service_account = self.app.options.get('serviceAccountId')
+        if service_account:
+            return _SigningProvider.from_iam(self.request, google_cred, service_account)
+
+        # If the SDK was initialized with some other credential type that supports signing
+        # (e.g. GAE credentials), use it to sign bytes.
+        if isinstance(google_cred, credentials.Signing):
+            return _SigningProvider.from_credential(google_cred)
+
+        # Attempt to discover a service account email from the local Metadata service. Use it
+        # with the IAM service to sign bytes.
+        resp = self.request(url=METADATA_SERVICE_URL, headers={'Metadata-Flavor': 'Google'})
+        service_account = resp.data.decode()
+        return _SigningProvider.from_iam(self.request, google_cred, service_account)
+
+    @property
+    def signing_provider(self):
+        """Initializes and returns the SigningProvider instance to be used."""
+        if not self._signing_provider:
+            try:
+                self._signing_provider = self._init_signing_provider()
+            except Exception as error:
+                url = 'https://firebase.google.com/docs/auth/admin/create-custom-tokens'
+                raise ValueError(
+                    'Failed to determine service account: {0}. Make sure to initialize the SDK '
+                    'with service account credentials or specify a service account ID with '
+                    'iam.serviceAccounts.signBlob permission. Please refer to {1} for more '
+                    'details on creating custom tokens.'.format(error, url))
+        return self._signing_provider
 
     def create_custom_token(self, uid, developer_claims=None):
         """Builds and signs a Firebase custom auth token."""
-        if not isinstance(self._app.credential, credentials.Certificate):
-            raise ValueError(
-                'Must initialize Firebase App with a certificate credential '
-                'to call create_custom_token().')
-
         if developer_claims is not None:
             if not isinstance(developer_claims, dict):
                 raise ValueError('developer_claims must be a dictionary')
@@ -93,10 +159,11 @@ def create_custom_token(self, uid, developer_claims=None):
         if not uid or not isinstance(uid, six.string_types) or len(uid) > 128:
             raise ValueError('uid must be a string between 1 and 128 characters.')
 
+        signing_provider = self.signing_provider
         now = int(time.time())
         payload = {
-            'iss': self._app.credential.service_account_email,
-            'sub': self._app.credential.service_account_email,
+            'iss': signing_provider.signer_email,
+            'sub': signing_provider.signer_email,
             'aud': FIREBASE_AUDIENCE,
             'uid': uid,
             'iat': now,
@@ -105,7 +172,12 @@ def create_custom_token(self, uid, developer_claims=None):
 
         if developer_claims is not None:
             payload['claims'] = developer_claims
-        return jwt.encode(self._app.credential.signer, payload)
+        try:
+            return jwt.encode(signing_provider.signer, payload)
+        except exceptions.TransportError as error:
+            msg = 'Failed to sign custom token. {0}'.format(error)
+            raise ApiCallError(TOKEN_SIGN_ERROR, msg, error)
+
 
     def create_session_cookie(self, id_token, expires_in):
         """Creates a session cookie from the provided ID token."""
@@ -131,7 +203,7 @@ def create_session_cookie(self, id_token, expires_in):
             'validDuration': expires_in,
         }
         try:
-            response = self._client.request('post', 'createSessionCookie', json=payload)
+            response = self.client.request('post', 'createSessionCookie', json=payload)
         except requests.exceptions.RequestException as error:
             self._handle_http_error(COOKIE_CREATE_ERROR, 'Failed to create session cookie', error)
         else:
diff --git a/firebase_admin/auth.py b/firebase_admin/auth.py
@@ -109,10 +109,13 @@ def create_custom_token(uid, developer_claims=None, app=None):
 
     Raises:
         ValueError: If input parameters are invalid.
+        AuthError: If an error occurs while creating the token using the remote IAM service.
     """
     token_generator = _get_auth_service(app).token_generator
-    return token_generator.create_custom_token(uid, developer_claims)
-
+    try:
+        return token_generator.create_custom_token(uid, developer_claims)
+    except _token_gen.ApiCallError as error:
+        raise AuthError(error.code, str(error), error.detail)
 
 def verify_id_token(id_token, app=None, check_revoked=False):
     """Verifies the signature and data for the provided JWT.
diff --git a/integration/test_auth.py b/integration/test_auth.py
@@ -22,8 +22,11 @@
 import pytest
 import requests
 
+import firebase_admin
 from firebase_admin import auth
-
+from firebase_admin import credentials
+import google.oauth2.credentials
+from google.auth import transport
 
 _verify_token_url = 'https://www.googleapis.com/identitytoolkit/v3/relyingparty/verifyCustomToken'
 _verify_password_url = 'https://www.googleapis.com/identitytoolkit/v3/relyingparty/verifyPassword'
@@ -57,6 +60,20 @@ def test_custom_token(api_key):
     claims = auth.verify_id_token(id_token)
     assert claims['uid'] == 'user1'
 
+def test_custom_token_without_service_account(api_key):
+    google_cred = firebase_admin.get_app().credential.get_credential()
+    cred = CredentialWrapper.from_existing_credential(google_cred)
+    custom_app = firebase_admin.initialize_app(cred, {
+        'serviceAccountId': google_cred.service_account_email,
+    }, 'temp-app')
+    try:
+        custom_token = auth.create_custom_token('user1', app=custom_app)
+        id_token = _sign_in(custom_token, api_key)
+        claims = auth.verify_id_token(id_token)
+        assert claims['uid'] == 'user1'
+    finally:
+        firebase_admin.delete_app(custom_app)
+
 def test_custom_token_with_claims(api_key):
     dev_claims = {'premium' : True, 'subscription' : 'silver'}
     custom_token = auth.create_custom_token('user2', dev_claims)
@@ -353,3 +370,20 @@ def test_import_users_with_password(api_key):
         assert len(id_token) > 0
     finally:
         auth.delete_user(uid)
+
+
+class CredentialWrapper(credentials.Base):
+    """A custom Firebase credential that wraps an OAuth2 token."""
+
+    def __init__(self, token):
+        self._delegate = google.oauth2.credentials.Credentials(token)
+
+    def get_credential(self):
+        return self._delegate
+
+    @classmethod
+    def from_existing_credential(cls, google_cred):
+        if not google_cred.token:
+            request = transport.requests.Request()
+            google_cred.refresh(request)
+        return CredentialWrapper(google_cred.token)
diff --git a/snippets/auth/index.py b/snippets/auth/index.py
@@ -50,6 +50,15 @@ def initialize_sdk_with_refresh_token():
     # [END initialize_sdk_with_refresh_token]
     firebase_admin.delete_app(default_app)
 
+def initialize_sdk_with_service_account_id():
+    # [START initialize_sdk_with_service_account_id]
+    options = {
+        'serviceAccountId': 'my-client-id@my-project-id.iam.gserviceaccount.com',
+    }
+    firebase_admin.initialize_app(options=options)
+    # [END initialize_sdk_with_service_account_id]
+    firebase_admin.delete_app(firebase_admin.get_app())
+
 def access_services_default():
     cred = credentials.Certificate('path/to/service.json')
     # [START access_services_default]
diff --git a/tests/test_token_gen.py b/tests/test_token_gen.py
@@ -14,6 +14,7 @@
 
 """Test cases for the firebase_admin._token_gen module."""
 
+import base64
 import datetime
 import json
 import os
@@ -109,8 +110,11 @@ def _instrument_user_manager(app, status, payload):
 
 def _overwrite_cert_request(app, request):
     auth_service = auth._get_auth_service(app)
-    token_verifier = auth_service.token_verifier
-    token_verifier.request = request
+    auth_service.token_verifier.request = request
+
+def _overwrite_iam_request(app, request):
+    auth_service = auth._get_auth_service(app)
+    auth_service.token_generator.request = request
 
 @pytest.fixture(scope='module')
 def auth_app():
@@ -194,6 +198,74 @@ def test_noncert_credential(self, user_mgt_app):
         with pytest.raises(ValueError):
             auth.create_custom_token(MOCK_UID, app=user_mgt_app)
 
+    def test_sign_with_iam(self):
+        options = {'serviceAccountId': 'test-service-account'}
+        app = firebase_admin.initialize_app(
+            testutils.MockCredential(), name='iam-signer-app', options=options)
+        try:
+            signature = base64.b64encode(b'test').decode()
+            iam_resp = '{{"signature": "{0}"}}'.format(signature)
+            _overwrite_iam_request(app, testutils.MockRequest(200, iam_resp))
+            custom_token = auth.create_custom_token(MOCK_UID, app=app).decode()
+            assert custom_token.endswith('.' + signature)
+            self._verify_signer(custom_token, 'test-service-account')
+        finally:
+            firebase_admin.delete_app(app)
+
+    def test_sign_with_iam_error(self):
+        options = {'serviceAccountId': 'test-service-account'}
+        app = firebase_admin.initialize_app(
+            testutils.MockCredential(), name='iam-signer-app', options=options)
+        try:
+            iam_resp = '{"error": {"code": 403, "message": "test error"}}'
+            _overwrite_iam_request(app, testutils.MockRequest(403, iam_resp))
+            with pytest.raises(auth.AuthError) as excinfo:
+                auth.create_custom_token(MOCK_UID, app=app)
+            assert excinfo.value.code == _token_gen.TOKEN_SIGN_ERROR
+            assert iam_resp in str(excinfo.value)
+        finally:
+            firebase_admin.delete_app(app)
+
+    def test_sign_with_discovered_service_account(self):
+        request = testutils.MockRequest(200, 'discovered-service-account')
+        app = firebase_admin.initialize_app(testutils.MockCredential(), name='iam-signer-app')
+        try:
+            _overwrite_iam_request(app, request)
+            # Force initialization of the signing provider. This will invoke the Metadata service.
+            auth_service = auth._get_auth_service(app)
+            assert auth_service.token_generator.signing_provider is not None
+            # Now invoke the IAM signer.
+            signature = base64.b64encode(b'test').decode()
+            request.response = testutils.MockResponse(
+                200, '{{"signature": "{0}"}}'.format(signature))
+            custom_token = auth.create_custom_token(MOCK_UID, app=app).decode()
+            assert custom_token.endswith('.' + signature)
+            self._verify_signer(custom_token, 'discovered-service-account')
+            assert len(request.log) == 2
+            assert request.log[0][1]['headers'] == {'Metadata-Flavor': 'Google'}
+        finally:
+            firebase_admin.delete_app(app)
+
+    def test_sign_with_discovery_failure(self):
+        request = testutils.MockFailedRequest(Exception('test error'))
+        app = firebase_admin.initialize_app(testutils.MockCredential(), name='iam-signer-app')
+        try:
+            _overwrite_iam_request(app, request)
+            with pytest.raises(ValueError) as excinfo:
+                auth.create_custom_token(MOCK_UID, app=app)
+            assert str(excinfo.value).startswith('Failed to determine service account: test error')
+            assert len(request.log) == 1
+            assert request.log[0][1]['headers'] == {'Metadata-Flavor': 'Google'}
+        finally:
+            firebase_admin.delete_app(app)
+
+    def _verify_signer(self, token, signer):
+        segments = token.split('.')
+        assert len(segments) == 3
+        body = json.loads(base64.b64decode(segments[1]).decode())
+        assert body['iss'] == signer
+        assert body['sub'] == signer
+
 
 class TestCreateSessionCookie(object):
 
diff --git a/tests/testutils.py b/tests/testutils.py
@@ -86,11 +86,25 @@ class MockRequest(transport.Request):
 
     def __init__(self, status, response):
         self.response = MockResponse(status, response)
+        self.log = []
 
     def __call__(self, *args, **kwargs):
+        self.log.append((args, kwargs))
         return self.response
 
 
+class MockFailedRequest(transport.Request):
+    """A mock HTTP request that fails by raising an exception."""
+
+    def __init__(self, error):
+        self.error = error
+        self.log = []
+
+    def __call__(self, *args, **kwargs):
+        self.log.append((args, kwargs))
+        raise self.error
+
+
 class MockGoogleCredential(credentials.Credentials):
     """A mock Google authentication credential."""
     def refresh(self, request):
