auth: Add native flow support for Apple authentication.

Overrides some of internal functions of python-social-auth
to handle native flow.

Credits to Mateusz Mandera for the overridden functions.

Co-authored-by: Mateusz Mandera <mateusz.mandera@zulip.com>

zerver/tests/test_auth_backends.py

@@ -5,7 +5,8 @@ import json
 import re
 import time
 import urllib
-from typing import Any, Callable, Dict, List, Optional, Sequence, Tuple
+from contextlib import contextmanager
+from typing import Any, Callable, Dict, Iterator, List, Optional, Sequence, Tuple
 from unittest import mock
 
 import jwt
@@ -1919,12 +1920,16 @@ class AppleAuthMixin:
     AUTH_FINISH_URL = "/complete/apple/"
     CONFIG_ERROR_URL = "/config-error/apple"
 
-    def generate_id_token(self, account_data_dict: Dict[str, str]) -> str:
+    def generate_id_token(self, account_data_dict: Dict[str, str], audience: Optional[str]=None) -> str:
         payload = account_data_dict
 
         # This setup is important because python-social-auth decodes `id_token`
         # with `SOCIAL_AUTH_APPLE_CLIENT` as the `audience`
         payload['aud'] = settings.SOCIAL_AUTH_APPLE_CLIENT
+
+        if audience is not None:
+            payload['aud'] = audience
+
         headers = {"kid": "SOMEKID"}
         private_key = settings.APPLE_ID_TOKEN_GENERATION_KEY
 
@@ -2059,6 +2064,153 @@ class AppleIdAuthBackendTest(AppleAuthMixin, SocialAuthBase):
             self.logger_output("Wrong state parameter given.", "warning"),
         ])
 
+class AppleAuthBackendNativeFlowTest(AppleAuthMixin, SocialAuthBase):
+    __unittest_skip__ = False
+
+    SIGNUP_URL = '/complete/apple/'
+    LOGIN_URL = '/complete/apple/'
+
+    def prepare_login_url_and_headers(
+        self,
+        subdomain: Optional[str]=None,
+        mobile_flow_otp: Optional[str]=None,
+        desktop_flow_otp: Optional[str]=None,
+        is_signup: bool=False,
+        next: str='',
+        multiuse_object_key: str='',
+        alternative_start_url: Optional[str]=None,
+        id_token: Optional[str]=None,
+        *,
+        user_agent: Optional[str]=None,
+    ) -> Tuple[str, Dict[str, Any]]:
+        url, headers = super().prepare_login_url_and_headers(
+            subdomain, mobile_flow_otp, desktop_flow_otp, is_signup, next,
+            multiuse_object_key, alternative_start_url=alternative_start_url,
+            user_agent=user_agent,
+        )
+
+        params = {'native_flow': 'true'}
+
+        if id_token is not None:
+            params['id_token'] = id_token
+
+        if is_signup:
+            params['is_signup'] = '1'
+
+        if subdomain:
+            params['subdomain'] = subdomain
+
+        url += "&%s" % (urllib.parse.urlencode(params),)
+        return url, headers
+
+    def social_auth_test(self, account_data_dict: Dict[str, str],
+                         *, subdomain: Optional[str]=None,
+                         mobile_flow_otp: Optional[str]=None,
+                         desktop_flow_otp: Optional[str]=None,
+                         is_signup: bool=False,
+                         next: str='',
+                         multiuse_object_key: str='',
+                         alternative_start_url: Optional[str]=None,
+                         skip_id_token: bool=False,
+                         user_agent: Optional[str]=None,
+                         **extra_data: Any) -> HttpResponse:
+        """In Apple's native authentication flow, the client app authenticates
+        with Apple and receives the JWT id_token, before contacting
+        the Zulip server.  The app sends an appropriate request with
+        it to /complete/apple/ to get logged in.  See the backend
+        class for details.
+
+        As a result, we need a custom social_auth_test function that
+        effectively just does the second half of the flow (i.e. the
+        part after the redirect from this third-party authentication
+        provider) with a properly generated id_token.
+        """
+
+        if not skip_id_token:
+            id_token = self.generate_id_token(account_data_dict, settings.SOCIAL_AUTH_APPLE_BUNDLE_ID)
+        else:
+            id_token = None
+
+        url, headers = self.prepare_login_url_and_headers(
+            subdomain, mobile_flow_otp, desktop_flow_otp, is_signup, next,
+            multiuse_object_key, alternative_start_url=self.AUTH_FINISH_URL,
+            user_agent=user_agent, id_token=id_token,
+        )
+
+        with self.apple_jwk_url_mock():
+            result = self.client_get(url, **headers)
+
+        return result
+
+    @contextmanager
+    def apple_jwk_url_mock(self) -> Iterator[None]:
+        with responses.RequestsMock(assert_all_requests_are_fired=False) as requests_mock:
+            # The server fetches public keys for validating the id_token
+            # from Apple servers. We need to mock that URL to return our key,
+            # created for these tests.
+            requests_mock.add(
+                requests_mock.GET,
+                self.BACKEND_CLASS.JWK_URL,
+                status=200,
+                json=json.loads(settings.APPLE_JWK),
+            )
+            yield
+
+    def test_no_id_token_sent(self) -> None:
+        account_data_dict = self.get_account_data_dict(email=self.email, name=self.name)
+        result = self.social_auth_test(account_data_dict,
+                                       expect_choose_email_screen=False,
+                                       subdomain='zulip', next='/user_uploads/image',
+                                       skip_id_token=True)
+        self.assert_json_error(result, "Missing id_token parameter")
+
+    def test_social_auth_session_fields_cleared_correctly(self) -> None:
+        mobile_flow_otp = '1234abcd' * 8
+
+        def initiate_auth(mobile_flow_otp: Optional[str]=None) -> None:
+            url, headers = self.prepare_login_url_and_headers(subdomain='zulip',
+                                                              id_token='invalid',
+                                                              mobile_flow_otp=mobile_flow_otp)
+            result = self.client_get(url, **headers)
+            self.assertEqual(result.status_code, 302)
+
+        # Start Apple auth with mobile_flow_otp param. It should get saved into the session
+        # on SOCIAL_AUTH_SUBDOMAIN.
+        initiate_auth(mobile_flow_otp)
+        self.assertEqual(self.client.session['mobile_flow_otp'], mobile_flow_otp)
+
+        # Make a request without mobile_flow_otp param and verify the field doesn't persist
+        # in the session from the previous request.
+        initiate_auth()
+        self.assertEqual(self.client.session.get('mobile_flow_otp'), None)
+
+    def test_id_token_with_invalid_aud_sent(self) -> None:
+        account_data_dict = self.get_account_data_dict(email=self.email, name=self.name)
+        url, headers = self.prepare_login_url_and_headers(
+            subdomain='zulip', alternative_start_url=self.AUTH_FINISH_URL,
+            id_token=self.generate_id_token(account_data_dict, audience='com.different.app'),
+        )
+
+        with self.apple_jwk_url_mock(), mock.patch('logging.info') as mock_info:
+            result = self.client_get(url, **headers)
+            mock_info.assert_called_once_with('/complete/apple/: %s',
+                                              'Authentication failed: Token validation failed')
+        return result
+
+    def test_social_auth_desktop_success(self) -> None:
+        """
+        The desktop app doesn't use the native flow currently and the desktop app flow in its
+        current form happens in the browser, thus only the webflow is viable there.
+        """
+        pass
+
+    def test_social_auth_no_key(self) -> None:
+        """
+        The basic validation of server configuration is handled on the
+        /login/social/apple/ endpoint which isn't even a part of the native flow.
+        """
+        pass
+
 class GitHubAuthBackendTest(SocialAuthBase):
     __unittest_skip__ = False
 

zproject/backends.py

@@ -47,6 +47,7 @@ from social_core.backends.gitlab import GitLabOAuth2
 from social_core.backends.google import GoogleOAuth2
 from social_core.backends.saml import SAMLAuth
 from social_core.exceptions import (
+    AuthCanceled,
     AuthFailed,
     AuthMissingParameter,
     AuthStateForbidden,
@@ -1540,10 +1541,9 @@ class AppleAuthBackend(SocialAuthMixin, AppleIdAuth):
     2. The native flow, intended for users on an Apple device.  In the native flow,
        the device handles authentication of the user with Apple's servers and ends up
        with the JWT id_token (like in the web flow).  The client-side details aren't
-       relevant to us; the app will simply provide the id_token to the server,
+       relevant to us; the app should simply send the id_token as a param to the
-       which can verify its signature and process it like in the web flow.
+       /complete/apple/ endpoint, together with native_flow=true and any other
-
+       appropriate params, such as mobile_flow_otp.
-    So far we only implement the web flow.
     """
     sort_order = 10
     name = "apple"
@@ -1578,6 +1578,9 @@ class AppleAuthBackend(SocialAuthMixin, AppleIdAuth):
                 return locale
         return 'en_US'
 
+    def is_native_flow(self) -> bool:
+        return self.strategy.request_data().get('native_flow', False)
+
     # This method replaces a method from python-social-auth; it is adapted to store
     # the state_token data in redis.
     def get_or_create_state(self) -> str:
@@ -1612,11 +1615,12 @@ class AppleAuthBackend(SocialAuthMixin, AppleIdAuth):
                           token=state)
         return state
 
-    # This method replaces a method from python-social-auth; it is
-    # adapted to retrieve the data stored in redis, transferring to
-    # the session so that it can be accessed by common
-    # python-social-auth code.
     def validate_state(self) -> Optional[str]:
+        """
+        This method replaces a method from python-social-auth; it is
+        adapted to retrieve the data stored in redis, save it in
+        the session so that it can be accessed by the social pipeline.
+        """
         request_state = self.get_request_state()
 
         if not request_state:
@@ -1648,6 +1652,9 @@ class AppleAuthBackend(SocialAuthMixin, AppleIdAuth):
         to make this function a thin wrapper around the upstream
         method; we may want to submit a PR to achieve that.
         '''
+        if self.is_native_flow():
+            audience = self.setting("BUNDLE_ID")
+        else:
             audience = self.setting("SERVICES_ID")
 
         try:
@@ -1663,6 +1670,53 @@ class AppleAuthBackend(SocialAuthMixin, AppleIdAuth):
 
         return decoded
 
+    def auth_complete(self, *args: Any, **kwargs: Any) -> Optional[HttpResponse]:
+        if not self.is_native_flow():
+            # The default implementation in python-social-auth is the browser flow.
+            return super().auth_complete(*args, **kwargs)
+
+        # We handle the Apple's native flow on our own.  In this flow,
+        # before contacting the server, the client obtains an id_token
+        # from Apple directly, and then sends that to /complete/apple/
+        # (the endpoint handled by this function), together with any
+        # other desired parameters from self.standard_relay_params.
+        #
+        # What we'd like to do with the payload is just pass it into
+        # the common code path for the web flow.  In the web flow,
+        # before sending a request to Apple, python-social-auth sets
+        # various values about the intended authentication in the
+        # session, before the redirect.
+        #
+        # Thus, we need to set those session variables here, before
+        # processing the id_token we received using the common do_auth.
+        request_data = self.strategy.request_data()
+        if 'id_token' not in request_data:
+            raise JsonableError(_("Missing id_token parameter"))
+
+        for param in self.standard_relay_params:
+            self.strategy.session_set(param, request_data.get(param))
+
+        # We should get the subdomain from the hostname of the request.
+        self.strategy.session_set('subdomain', get_subdomain(self.strategy.request))
+
+        try:
+            # Things are now ready to be handled by the superclass code. It will
+            # validate the id_token and push appropriate user data to the social pipeline.
+            result = self.do_auth(request_data['id_token'], *args, **kwargs)
+            return result
+        except (AuthFailed, AuthCanceled) as e:
+            # AuthFailed is a general "failure" exception from
+            # python-social-auth that we should convert to None return
+            # value here to avoid getting tracebacks.
+            #
+            # AuthCanceled is raised in the Apple backend
+            # implementation in python-social-auth in certain cases,
+            # though AuthFailed would have been more correct.
+            #
+            # We have an open PR to python-social-auth to clean this up.
+            logging.info("/complete/apple/: %s", str(e))
+            return None
+
 @external_auth_method
 class SAMLAuthBackend(SocialAuthMixin, SAMLAuth):
     auth_backend_name = "SAML"