Repository: dwavesystems/dwave-cloud-client
Pull request #593 (2023-10-27 20:04:16)

Update auth flow pages
Redirect to Leap success page (or Leap error page) upon redirect auth flow completion.


Patch:
diff --git a/dwave/cloud/auth/flows.py b/dwave/cloud/auth/flows.py
index 52b7222..c3ce04b 100644
--- a/dwave/cloud/auth/flows.py
+++ b/dwave/cloud/auth/flows.py
@@ -19,16 +19,17 @@ import time
 import webbrowser
 from operator import sub
 from typing import Any, Callable, Dict, Optional, Union, Sequence
-from urllib.parse import urljoin, parse_qsl
+from urllib.parse import urljoin
 
 import click
 from authlib.integrations.requests_client import OAuth2Session, OAuthError
 from authlib.oauth2.rfc6749 import OAuth2Token
 from authlib.common.security import generate_token
+from authlib.common.urls import add_params_to_uri
 
 from dwave.cloud.auth.config import OCEAN_SDK_CLIENT_ID, OCEAN_SDK_SCOPES
 from dwave.cloud.auth.creds import Credentials
-from dwave.cloud.auth.server import SingleRequestAppServer, RequestCaptureApp
+from dwave.cloud.auth.server import SingleRequestAppServer, RequestCaptureAndRedirectApp
 from dwave.cloud.config.models import ClientConfig
 from dwave.cloud.regions import resolve_endpoints
 from dwave.cloud.utils import pretty_argvalues
@@ -252,6 +253,14 @@ class LeapAuthFlow(AuthFlow):
     def _infer_token_endpoint(leap_api_endpoint: str) -> str:
         return urljoin(leap_api_endpoint, '/leap/openid/token')
 
+    @staticmethod
+    def _infer_leap_success_uri(leap_api_endpoint: str) -> str:
+        return urljoin(leap_api_endpoint, '/leap/openid/success')
+
+    @staticmethod
+    def _infer_leap_error_uri(leap_api_endpoint: str) -> str:
+        return urljoin(leap_api_endpoint, '/leap/openid/error')
+
     @classmethod
     def from_config_model(cls, config: ClientConfig, **kwargs) -> LeapAuthFlow:
         """Construct a :class:`.LeapAuthModel` initialized with Ocean SDK's
@@ -353,7 +362,47 @@ class LeapAuthFlow(AuthFlow):
             flow.run_redirect_flow(open_browser=True)
 
         """
-        app = RequestCaptureApp(message=self._REDIRECT_DONE_MSG)
+        error_uri = self._infer_leap_error_uri(self.leap_api_endpoint)
+        success_uri = self._infer_leap_success_uri(self.leap_api_endpoint)
+
+        def exchange_code(app: RequestCaptureAndRedirectApp) -> str:
+            # error responses redirect immediately
+            if 'error' in app.query:
+                app.set_exception(
+                    OAuthError(error=app.query['error'],
+                               description=app.query.get('error_description')))
+                return add_params_to_uri(error_uri, app.query)
+
+            # when code received, exchange it for token
+            try:
+                self.fetch_token(code=app.query.get('code'), state=app.query.get('state'))
+
+            except OAuthError as exc:
+                # store for main thread
+                app.set_exception(exc)
+
+                # redirect to leap error page
+                query = dict(error=exc.error, error_description=exc.description,
+                             state=app.query.get('state'))
+                return add_params_to_uri(error_uri, query)
+
+            except Exception as exc:
+                # store for main thread
+                app.set_exception(exc)
+
+                # redirect to leap error page
+                query = dict(error=type(exc).__name__, error_description=str(exc),
+                             state=app.query.get('state'))
+                return add_params_to_uri(error_uri, query)
+
+            # redirect to leap success page
+            return add_params_to_uri(success_uri, app.query)
+
+        app = RequestCaptureAndRedirectApp(
+            message=self._REDIRECT_DONE_MSG,
+            redirect_uri=exchange_code,
+            include_query=False)
+
         srv = SingleRequestAppServer(
             host=self._REDIRECT_HOST,
             base_port=self._REDIRECT_PORT_RANGE[0],
@@ -380,5 +429,7 @@ class LeapAuthFlow(AuthFlow):
             print(self._AUTH_TIMEOUT_MSG)
             return
 
-        q = dict(parse_qsl(app.query))
-        return self.fetch_token(code=q.get('code'), state=q.get('state'))
+        # now raise auth exception if set by the app during code exchange
+        app.exception()
+
+        return self.token
diff --git a/dwave/cloud/auth/server.py b/dwave/cloud/auth/server.py
index e45a7fe..778d5d3 100644
--- a/dwave/cloud/auth/server.py
+++ b/dwave/cloud/auth/server.py
@@ -18,8 +18,8 @@ import sys
 import threading
 import traceback
 from socketserver import ThreadingMixIn
-from typing import Optional, Callable, Iterator
-from urllib.parse import urlsplit
+from typing import Callable, Iterator, Optional, Union
+from urllib.parse import urljoin, urlsplit, parse_qsl, SplitResult
 from wsgiref.simple_server import make_server, WSGIRequestHandler, WSGIServer
 from wsgiref.util import request_uri
 
@@ -250,10 +250,11 @@ class SingleRequestAppServer(BackgroundAppServer):
         app = kwargs.pop('app')
 
         def single_request_app(*args, **kwargs):
-            ret = app(*args, **kwargs)
-            # we're done after the first successful request
-            self.server.shutdown()
-            return ret
+            try:
+                return app(*args, **kwargs)
+            finally:
+                # we're done after the first request
+                self.server.shutdown()
 
         super().__init__(app=single_request_app, **kwargs)
 
@@ -265,15 +266,60 @@ class RequestCaptureApp:
 
     def __init__(self, message: str):
         self.message = message
-        self.uri = None
-        self.query = None
+        self.uri: str = None
+        self.parts: SplitResult = None
+        self.query: dict = None
 
-    def __call__(self, environ: dict, start_response: Callable):
+    def store_request(self, environ: dict):
         # store the URI accessed
         self.uri = request_uri(environ, include_query=True)
-        self.query = urlsplit(self.uri).query
+        self.parts = urlsplit(self.uri)
+        self.query = dict(parse_qsl(self.parts.query))
         # in the future, we might also store: method, body data, headers, etc,
         # but we don't need that for now
 
+    def set_exception(self, exc):
+        self._exc = exc
+
+    def exception(self):
+        if hasattr(self, '_exc') and self._exc:
+            raise self._exc
+
+    def __call__(self, environ: dict, start_response: Callable):
+        self.store_request(environ)
+
         start_response("200 OK", [('Content-Type', 'text/plain; charset=utf-8')])
         return [self.message.encode('utf-8')]
+
+
+class RequestCaptureAndRedirectApp(RequestCaptureApp):
+    """A simple WSGI application that stores request data (currently only URL),
+    and redirects to `redirect_uri`, possibly dynamically generated.
+    """
+
+    def __init__(self, message: str,
+                 redirect_uri: Union[str, Callable],
+                 include_query: bool = True):
+        super().__init__(message)
+
+        self.redirect_uri = redirect_uri
+        self.include_query = include_query
+
+    def __call__(self, environ: dict, start_response: Callable):
+        self.store_request(environ)
+
+        if callable(self.redirect_uri):
+            uri = self.redirect_uri(self)
+        else:
+            uri = self.redirect_uri
+
+        if self.include_query:
+            uri = urljoin(uri, f'?{self.parts.query}')
+
+        start_response('302 Found', [
+            ('Content-Type', 'text/plain; charset=utf-8'),
+            ('Location', uri),
+        ])
+
+        # just in case, include the message in body
+        return [self.message.encode('utf-8')]
diff --git a/dwave/cloud/cli.py b/dwave/cloud/cli.py
index b4000af..c2ff0f9 100644
--- a/dwave/cloud/cli.py
+++ b/dwave/cloud/cli.py
@@ -951,8 +951,6 @@ def login(*, config_file, profile, oob, output):
     flow = LeapAuthFlow.from_config_model(config)
 
     if oob:
-        # XXX: until the standard OOB URI is supported on Leap
-        flow._OOB_REDIRECT_URI = 'oob'
         flow.run_oob_flow(open_browser=True)
     else:
         flow.run_redirect_flow(open_browser=True)
diff --git a/releasenotes/notes/nicer-auth-flow-redirect-pages-088cc06fde14d3ca.yaml b/releasenotes/notes/nicer-auth-flow-redirect-pages-088cc06fde14d3ca.yaml
new file mode 100644
index 0000000..2d34b3b
--- /dev/null
+++ b/releasenotes/notes/nicer-auth-flow-redirect-pages-088cc06fde14d3ca.yaml
@@ -0,0 +1,6 @@
+---
+features:
+  - |
+    Redirect user back to a Leap info page when auth redirect flow successfully
+    completes. Similarly, redirect user to a Leap error page in case of an auth
+    error.


Test patch:
diff --git a/tests/auth/test_flows.py b/tests/auth/test_flows.py
index 1831331..1506b19 100644
--- a/tests/auth/test_flows.py
+++ b/tests/auth/test_flows.py
@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import sys
 import time
 import threading
 import unittest
@@ -22,7 +23,7 @@ from urllib.parse import urlsplit, parse_qsl, urljoin
 import requests
 import requests_mock
 
-from dwave.cloud.auth.flows import AuthFlow, LeapAuthFlow
+from dwave.cloud.auth.flows import AuthFlow, LeapAuthFlow, OAuthError
 from dwave.cloud.auth.config import OCEAN_SDK_CLIENT_ID, OCEAN_SDK_SCOPES
 from dwave.cloud.auth.creds import Credentials
 from dwave.cloud.config import ClientConfig
@@ -242,7 +243,7 @@ class TestLeapAuthFlow(unittest.TestCase):
         self.assertEqual(flow.client_id, client_id)
 
 
-class TestLeapAuthFlowRunners(unittest.TestCase):
+class TestLeapAuthFlowOOB(unittest.TestCase):
 
     @mock.patch('click.echo', return_value=None)
     def test_oob(self, m):
@@ -256,8 +257,11 @@ class TestLeapAuthFlowRunners(unittest.TestCase):
                 flow.run_oob_flow()
                 fetch_token.assert_called_once_with(code=mock_code)
 
+class TestLeapAuthFlowRedirect(unittest.TestCase):
+
     @mock.patch('click.echo', return_value=None)
-    def test_redirect(self, m):
+    def test_success(self, m):
+        # success case: access authorized, token fetched
         config = ClientConfig(leap_api_endpoint='https://example.com/leap')
         flow = LeapAuthFlow.from_config_model(config)
 
@@ -276,8 +280,173 @@ class TestLeapAuthFlowRunners(unittest.TestCase):
 
             ready.wait()
             response = requests.get(ctx['redirect_uri'],
-                                    params=dict(code=mock_code, state=ctx['state'])).text
-            self.assertEqual(response, flow._REDIRECT_DONE_MSG)
+                                    params=dict(code=mock_code, state=ctx['state']))
+            self.assertEqual(len(response.history), 1)
+            self.assertEqual(response.history[0].status_code, 302)
+            location = response.history[0].headers.get('Location')
+            self.assertTrue(location.startswith(config.leap_api_endpoint))
+            self.assertIn('/success', location)
 
             f.join()
             fetch_token.assert_called_once_with(code=mock_code, state=ctx['state'])
+
+    class FlowRunner(threading.Thread):
+        def __init__(self, *, flow, authorize_response_params, **kwargs):
+            super().__init__(**kwargs)
+
+            self.flow = flow
+
+            # set when authorize url is "opened" in browser
+            self.authorize_query = {}
+            self.authorize_ready = threading.Event()
+            self.authorize_response_params = authorize_response_params
+
+        def open_browser(self, url, *args, **kwargs):
+            self.authorize_query.update(parse_qsl(urlsplit(url).query))
+            self.authorize_ready.set()
+
+        def authorize(self):
+            # fake authorize request from Leap to our redirect_uri
+            # `params` determine if auth was denied or approved
+            self.authorize_ready.wait()
+
+            query = self.authorize_response_params.copy()
+            query.update(state=self.authorize_query['state'])
+            return requests.get(self.authorize_query['redirect_uri'], params=query)
+
+        def run(self):
+            try:
+                self.flow.run_redirect_flow(open_browser=self.open_browser)
+            except:
+                # make exception discoverable from the main thread
+                self._exc_info = sys.exc_info()
+
+        def exception(self):
+            # see: https://stackoverflow.com/a/1854263, published under CC BY-SA 4.0.
+            # to be called from main thread
+            if hasattr(self, '_exc_info') and self._exc_info:
+                raise self._exc_info[1].with_traceback(self._exc_info[2])
+
+    def assert_redirect_to_error_page(self, response, flow):
+        # check that `response` is a 302 redirect with redirect location pointing to leap error page
+        error_uri = flow._infer_leap_error_uri(flow.leap_api_endpoint)
+        self.assertEqual(len(response.history), 1)
+        self.assertEqual(response.history[0].status_code, 302)
+        location = response.history[0].headers.get('Location')
+        self.assertTrue(location.startswith(error_uri))
+
+    def assert_redirect_params(self, response, expected_params):
+        # check that redirect location params match `expected params`
+        location = response.history[0].headers.get('Location')
+        location_params = dict(parse_qsl(urlsplit(location).query))
+        self.assertEqual(location_params, expected_params)
+
+    @mock.patch('click.echo', return_value=None)
+    def test_auth_denied(self, m):
+        # error case: access not authorized (e.g. user clicks "Decline")
+        config = ClientConfig(leap_api_endpoint='https://example.com/leap')
+        flow = LeapAuthFlow.from_config_model(config)
+
+        error = 'access_denied'
+        error_description = 'Authorization request denied'
+
+        with mock.patch.object(flow, 'fetch_token') as fetch_token:
+            authorize_response_params = dict(error=error, error_description=error_description)
+
+            t = self.FlowRunner(flow=flow, authorize_response_params=authorize_response_params)
+            t.start()
+
+            response = t.authorize()
+
+            # check response is a redirect to leap error page
+            self.assert_redirect_to_error_page(response, flow)
+
+            # and all request query params are propagated
+            request = response.history[0].request.url
+            request_params = dict(parse_qsl(urlsplit(request).query))
+            self.assert_redirect_params(response, request_params)
+
+            # token fetch shouldn't even be attempted
+            t.join()
+            fetch_token.assert_not_called()
+
+            # oauth exception raised by the redirect flow runner
+            with self.assertRaisesRegex(OAuthError, error_description):
+                t.exception()
+
+    @mock.patch('click.echo', return_value=None)
+    def test_exchange_fails(self, m):
+        # error case: access authorized, but code exchange fails with oauth error
+        config = ClientConfig(leap_api_endpoint='https://example.com/leap')
+        flow = LeapAuthFlow.from_config_model(config)
+
+        code = '1234'
+
+        # mock `fetch_token` to raise an oauth error (e.g. code expired)
+        error = OAuthError(error='invalid_grant',
+                           description='The provided authorization grant is invalid')
+
+        with mock.patch.object(flow, 'fetch_token', side_effect=error) as fetch_token:
+            authorize_response_params = dict(code=code)     # state appended by FlowRunner.authorize()
+
+            t = self.FlowRunner(flow=flow, authorize_response_params=authorize_response_params)
+            t.start()
+
+            response = t.authorize()
+
+            state = t.authorize_query['state']
+
+            # check response is a redirect to leap error page
+            self.assert_redirect_to_error_page(response, flow)
+
+            # and error page params are constructed properly
+            expected_params = dict(error=error.error,
+                                   error_description=error.description,
+                                   state=state)
+            self.assert_redirect_params(response, expected_params)
+
+            # token fetch is attempted (but it fails with `error`)
+            t.join()
+            fetch_token.assert_called_once_with(code=code, state=state)
+
+            # oauth exception raised by the redirect flow runner
+            with self.assertRaisesRegex(OAuthError, str(error)):
+                t.exception()
+
+    @mock.patch('click.echo', return_value=None)
+    def test_non_auth_failure_during_code_exchange(self, m):
+        # error case: access authorized, but code exchange fails with unexpected non-auth error
+        config = ClientConfig(leap_api_endpoint='https://example.com/leap')
+        flow = LeapAuthFlow.from_config_model(config)
+
+        code = '1234'
+
+        # mock `fetch_token` to raise a non-auth error (e.g. file error)
+        error = OSError("credential store unavailable")
+
+        with mock.patch.object(flow, 'fetch_token', side_effect=error) as fetch_token:
+            authorize_response_params = dict(code=code)     # state appended by FlowRunner.authorize()
+
+            t = self.FlowRunner(flow=flow, authorize_response_params=authorize_response_params)
+            t.start()
+
+            response = t.authorize()
+
+            state = t.authorize_query['state']
+
+            # check response is a redirect to leap error page
+            self.assert_redirect_to_error_page(response, flow)
+
+            # and error page params are constructed properly
+            expected_params = dict(error=type(error).__name__,
+                                   error_description=str(error),
+                                   state=state)
+            self.assert_redirect_params(response, expected_params)
+
+            # token fetch is attempted (but it fails with `error`)
+            t.join()
+            fetch_token.assert_called_once_with(code=code, state=state)
+
+            # oauth exception raised by the redirect flow runner
+            with self.assertRaisesRegex(type(error), str(error)):
+                t.exception()
diff --git a/tests/auth/test_server.py b/tests/auth/test_server.py
index c3e7817..cce52ac 100644
--- a/tests/auth/test_server.py
+++ b/tests/auth/test_server.py
@@ -17,13 +17,14 @@ import socket
 import time
 import unittest
 from secrets import token_hex
-from urllib.parse import parse_qsl
+from urllib.parse import urljoin
 
 import requests
 from parameterized import parameterized
 
 from dwave.cloud.auth.server import (
-    iterports, BackgroundAppServer, SingleRequestAppServer, RequestCaptureApp)
+    iterports, BackgroundAppServer, SingleRequestAppServer,
+    RequestCaptureApp, RequestCaptureAndRedirectApp)
 
 
 class TestPortsGenerator(unittest.TestCase):
@@ -233,4 +234,63 @@ class TestRequestCaptureApp(unittest.TestCase):
 
         # test url/query capture
         self.assertTrue(app.uri.startswith(f"{srv.root_url}?"))
-        self.assertEqual(dict(parse_qsl(app.query)), query)
+        self.assertEqual(app.query, query)
+
+
+class TestRequestCaptureAndRedirectApp(unittest.TestCase):
+    base_port = 64070
+
+    def setUp(self):
+        self.success_text = 'success'
+        self.success_app = RequestCaptureApp(message=self.success_text)
+        self.success_srv = SingleRequestAppServer(
+            host='127.0.0.1', base_port=self.base_port, max_port=self.base_port+9,
+            app=self.success_app)
+        self.success_srv.start()
+        self.success_uri = urljoin(self.success_srv.root_url, '/success')
+
+        self.error_text = 'error'
+        self.error_app = RequestCaptureApp(message=self.error_text)
+        self.error_srv = SingleRequestAppServer(
+            host='127.0.0.1', base_port=self.base_port, max_port=self.base_port+9,
+            app=self.error_app)
+        self.error_srv.start()
+        self.error_uri = urljoin(self.error_srv.root_url, '/error')
+
+        def redirect_uri(app):
+            return self.error_uri if 'error' in app.query else self.success_uri
+
+        self.redirect_app = RequestCaptureAndRedirectApp(
+            message='redirecting', redirect_uri=redirect_uri, include_query=True)
+        self.redirect_srv = SingleRequestAppServer(
+            host='127.0.0.1', base_port=self.base_port, max_port=self.base_port+9,
+            app=self.redirect_app)
+        self.redirect_srv.start()
+
+    def test_redirect_on_success(self):
+        query = dict(a=token_hex(8), b=token_hex(16))
+
+        # test response
+        req = requests.get(self.redirect_srv.root_url, params=query)
+        self.assertEqual(req.text, self.success_text)
+
+        # wait for server shutdown to make sure request is done
+        self.redirect_srv.wait_shutdown()
+
+        # test url/query capture
+        self.assertTrue(self.success_app.uri.startswith(self.success_uri))
+        self.assertEqual(self.success_app.query, query)
+
+    def test_redirect_on_error(self):
+        query = dict(error='error-id', error_description='error message')
+
+        # test response
+        req = requests.get(self.redirect_srv.root_url, params=query)
+        self.assertEqual(req.text, self.error_text)
+
+        # wait for server shutdown to make sure request is done
+        self.redirect_srv.wait_shutdown()
+
+        # test url/query capture
+        self.assertTrue(self.error_app.uri.startswith(self.error_uri))
+        self.assertEqual(self.error_app.query, query)