Repository: dwavesystems/dwave-cloud-client (python)

Issue:
Implement a utility HTTP server/interface
Local web server is needed to handle redirects from Leap and exchange the authorization code for token (see #561), without further user intervention (like manually pasting the code).

We already implement a simple web server in https://github.com/dwavesystems/dwave-inspector, one that is based on the built-in [`WSGIServer`](https://docs.python.org/3/library/wsgiref.html#wsgiref.simple_server.WSGIServer) and [`WSGIRequestHandler`](https://docs.python.org/3/library/wsgiref.html#wsgiref.simple_server.WSGIRequestHandler), and that it runs Flask.

We might want to either:
- move the dwave-inspector server implementation over to the cloud-client (since it's a dependency anyway); or
- implement a new server based on off-the-shelf `werkzeug.run_simple`, or similar.

Pull request:
Add http server to handle auth callback
Close #563.

Patch:
diff --git a/dwave/cloud/auth/flows.py b/dwave/cloud/auth/flows.py
index f490f50..f7eb4f8 100644
--- a/dwave/cloud/auth/flows.py
+++ b/dwave/cloud/auth/flows.py
@@ -106,6 +106,26 @@ class AuthFlow:
         return url
 
     def fetch_token(self, *, code: str, **kwargs) -> dict:
+        """Exchange ``code`` for access token.
+
+        Args:
+            code:
+                OAuth 2.0 authorization code to be exchanged for access token.
+            state:
+                Authorization state; should match the one generated when creating
+                the authorize URL.
+
+                Note: except for the out-of-band flow, be sure to pass in the
+                ``state`` parameter received on the ``redirect_uri`` for added
+                security.
+        """
+        # make sure states match before proceeding, we shouldn't be tricked
+        # into exchanging arbitrary auth codes.
+        state = kwargs.get('state', None)
+        if state is not None:
+            if state != self.state:
+                raise ValueError('State mismatch')
+
         token = self.session.fetch_token(
             url=self.token_endpoint,
             grant_type='authorization_code',
diff --git a/dwave/cloud/auth/server.py b/dwave/cloud/auth/server.py
new file mode 100644
index 0000000..48085d6
--- /dev/null
+++ b/dwave/cloud/auth/server.py
@@ -0,0 +1,223 @@
+# Copyright 2023 D-Wave Systems Inc.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import logging
+import random
+import sys
+import threading
+import traceback
+from socketserver import ThreadingMixIn
+from typing import Optional, Callable, Iterator
+from wsgiref.simple_server import make_server, WSGIRequestHandler, WSGIServer
+
+logger = logging.getLogger(__name__)
+
+
+class LoggingStream:
+    """Provide file-like interface to a logger."""
+
+    def __init__(self, logger, level):
+        self.logger = logger
+        self.level = level
+
+    def write(self, message):
+        for line in message.split('\n'):
+            if line:
+                self.logger.log(self.level, line)
+
+    def flush(self):
+        pass
+
+# stream interface to our local logger
+request_logging_stream = LoggingStream(logger, logging.DEBUG)
+
+
+class LoggingWSGIRequestHandler(WSGIRequestHandler):
+    """WSGIRequestHandler subclass that logs to our logger, instead of to
+    ``sys.stderr`` (as hardcoded in ``http.server.BaseHTTPRequestHandler``).
+    """
+
+    def log_message(self, format, *args):
+        logger.info(format, *args)
+
+    def get_stderr(self):
+        return request_logging_stream
+
+
+class ErrorLoggingTCPServerMixin:
+    """Extend :class:`http.server.HTTPServer`/:class:`socketserver.TCPServer`
+    to log errors to our logger, instead of to a hard-coded ``sys.stderr`` stream.
+    """
+
+    def handle_error(self, request, client_address):
+        traceback.print_exception(*sys.exc_info(), file=request_logging_stream)
+
+
+class SocketTimeoutTCPServerMixin:
+    """Extend :class:`http.server.HTTPServer`/:class:`socketserver.TCPServer`
+    by setting a timeout on each client connection socket.
+
+    This prevents hanging connections waiting for client data.
+    """
+    connection_timeout = None
+
+    def get_request(self):
+        conn, addr = super().get_request()
+        if self.connection_timeout is not None:
+            conn.settimeout(self.connection_timeout)
+        return conn, addr
+
+
+class NoAddrReuseTCPServerMixin:
+    """Disable address reuse for a TCPServer while socket in TIME_WAIT on Linux,
+    or unless SO_EXCLUSIVEADDRUSE is set on Windows. It's set to truthy value by
+    HTTPServer subclass, so this mixin reverts the TCPServer default.
+
+    See https://learn.microsoft.com/en-us/windows/win32/winsock/so-exclusiveaddruse.
+    """
+    allow_reuse_address = False
+
+
+class ThreadingWSGIServer(ThreadingMixIn, SocketTimeoutTCPServerMixin,
+                          ErrorLoggingTCPServerMixin, NoAddrReuseTCPServerMixin,
+                          WSGIServer):
+    """:class:`~wsgiref.simple_server.WSGIServer` subclass that:
+    - supports multithreading, i.e. handles each request in a new thread,
+    - set a timeout on socket connections, and
+    - logs errors to our configured logger, instead of to stderr.
+    """
+
+
+def iterports(start: int, end: int,
+              n_lin: int, n_rand: Optional[int] = None) -> Iterator[int]:
+    """Server port proposal generator. Starts with a linear search, then
+    switches to a randomized search (random permutation of remaining ports).
+    """
+    # sanity checks
+    if n_rand is None:
+        n_rand = end - start + 1 - n_lin
+    if start < 0 or end < 0 or n_lin < 0 or (n_rand is not None and n_rand < 0):
+        raise ValueError("Non-negative integers required for all parameters")
+    if n_lin + n_rand > end - start + 1:
+        raise ValueError("Sum of tries must be less or equal to population size")
+
+    # linear search
+    yield from range(start, start + n_lin)
+
+    # randomized search
+    yield from random.sample(range(start + n_lin, end + 1), k=n_rand)
+
+
+class BackgroundAppServer(threading.Thread):
+    """WSGI application server container that runs in a background thread,
+    handling each request in its own thread.
+    """
+
+    def _make_server(self) -> ThreadingWSGIServer:
+        """Instantiate a http server, similarly to :func:`~wsgiref.simple_server.make_server`,
+        but bounding it to the first port available, instead of failing if the specified
+        port is unavailable.
+
+        Port search is a combination of linear and randomized search, starting at
+        ``base_port`` and ending with ``max_port`` (both ends inclusive).
+        """
+
+        for port in iterports(start=self.base_port, end=self.max_port,
+                              n_lin=self.linear_tries, n_rand=self.randomized_tries):
+            try:
+                server = make_server(self.host, port, self.app,
+                                     server_class=ThreadingWSGIServer,
+                                     handler_class=LoggingWSGIRequestHandler)
+                server.connection_timeout = self.timeout
+                return server
+            except OSError:
+                # linux: "OSError: [Errno 98] Address already in use"
+                # macos: "OSError: [Errno 48] Address already in use"
+                # win: "OSError: [WinError 10048] Only one usage of each socket address ... permitted"
+                pass
+
+        raise RuntimeError("Unable to find available port in range: "
+                           f"[{self.base_port}, {self.max_port}].")
+
+    @property
+    def server(self) -> ThreadingWSGIServer:
+        """HTTP server accessor that creates the actual server instance
+        (and binds it to host:port) on first access.
+        """
+
+        with self._server_lock:
+            self._server = getattr(self, '_server', None)
+
+            if self._server is None:
+                self._server = self._make_server()
+
+            return self._server
+
+    def __init__(self, *, host: str, base_port: int, max_port: int, app: Callable,
+                 linear_tries: int = 1, randomized_tries: Optional[int] = None,
+                 timeout: Optional[float] = 10):
+        super().__init__(daemon=True)
+
+        # store config, but start the web server (and bind to address) on run()
+        self.host = host
+        self.base_port = base_port
+        self.max_port = max_port
+        self.linear_tries = linear_tries
+        self.randomized_tries = randomized_tries
+        self.app = app
+        self.timeout = timeout
+        self._server_lock = threading.RLock()
+
+    def run(self):
+        """Don't call this method directly. Instead call `.start()`."""
+        logger.debug(f"Running {type(self).__name__} worker thread")
+        try:
+            self.server.serve_forever()
+        finally:
+            self.server.server_close()
+        logger.debug(f"{type(self).__name__} worker thread done.")
+
+    def stop(self):
+        logger.debug(f"{type(self).__name__}.stop()")
+        self.server.shutdown()
+        self.join()
+
+    def root_url(self):
+        return 'http://{}:{}/'.format(*self.server.server_address)
+
+    def wait_shutdown(self, timeout=None):
+        logger.debug(f"{type(self).__name__}.wait_shutdown(timeout={timeout})")
+        self.join(timeout)
+
+
+class SingleRequestAppServer(BackgroundAppServer):
+    """An extension of :class:`.BackgroundAppServer` that terminates after a
+    single request has completed.
+    """
+
+    # note: we can't simply use `server.handle_request()` to handle a single request
+    # because of the problem with some browsers (pre-)opening "ghost" connection (not
+    # sending any data), as it is described in python's `ThreadingHTTPServer` docs
+    # (https://docs.python.org/3/library/http.server.html#http.server.ThreadingHTTPServer)
+
+    def __init__(self, **kwargs):
+        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
+
+        super().__init__(app=single_request_app, **kwargs)
diff --git a/releasenotes/notes/add-auth-server-e9a3a1c71be9fc0e.yaml b/releasenotes/notes/add-auth-server-e9a3a1c71be9fc0e.yaml
new file mode 100644
index 0000000..405e9ab
--- /dev/null
+++ b/releasenotes/notes/add-auth-server-e9a3a1c71be9fc0e.yaml
@@ -0,0 +1,11 @@
+---
+features:
+  - |
+    Add a general-purpose, multi-threaded, WSGI application server,
+    ``dwave.cloud.auth.server.BackgroundAppServer``. This server can handle
+    redirect URIs in the OAuth2 authorization code flow.
+    See `#563 <https://github.com/dwavesystems/dwave-cloud-client/issues/563>`_.
+  - |
+    Add ``dwave.cloud.auth.server.SingleRequestAppServer`` based on
+    ``dwave.cloud.auth.server.BackgroundAppServer`` to specifically handle only
+    one (valid) request to the OAuth redirect URI (and self-terminate after that).


Test patch:
diff --git a/tests/auth/test_flows.py b/tests/auth/test_flows.py
index 5ff5312..6e2121e 100644
--- a/tests/auth/test_flows.py
+++ b/tests/auth/test_flows.py
@@ -54,6 +54,16 @@ class TestAuthFlow(unittest.TestCase):
         self.assertIn('code_challenge', q)
         self.assertEqual(q['code_challenge_method'], 'S256')
 
+    def test_fetch_token_state(self):
+        flow = AuthFlow(**self.test_args)
+
+        # generate auth request (state)
+        _ = flow.get_authorization_url()
+
+        # try exchanging the code with wrong state
+        with self.assertRaisesRegex(ValueError, "State mismatch"):
+            flow.fetch_token(code='not important', state='invalid')
+
     @requests_mock.Mocker()
     def test_fetch_token(self, m):
         # mock the token_endpoint
diff --git a/tests/auth/test_server.py b/tests/auth/test_server.py
new file mode 100644
index 0000000..8647bc0
--- /dev/null
+++ b/tests/auth/test_server.py
@@ -0,0 +1,188 @@
+# Copyright 2023 D-Wave Systems Inc.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import concurrent.futures
+import socket
+import time
+import unittest
+
+import requests
+from parameterized import parameterized
+
+from dwave.cloud.auth.server import (
+    iterports, BackgroundAppServer, SingleRequestAppServer)
+
+
+class TestPortsGenerator(unittest.TestCase):
+
+    @parameterized.expand([
+        (-1, 0, 0, None),
+        (0, -1, 0, None),
+        (0, 0, -1, None),
+        (0, 0, 0, -1),
+        (0, 1, 3, None),
+        (0, 1, 2, 1),
+    ])
+    def test_edge_cases(self, lb, ub, n_lin, n_rand):
+        with self.assertRaises(ValueError):
+            next(iterports(start=lb, end=ub, n_lin=n_lin, n_rand=n_rand))
+
+    @parameterized.expand([
+        (0, 1, 0, None, []),
+        (0, 1, 1, None, [0]),
+        (0, 1, 2, None, [0, 1]),
+        (0, 9, 3, None, [0, 1, 2]),
+        (0, 9, 0, None, []),
+        (0, 9, 0, 10, []),
+    ])
+    def test_sequences(self, lb, ub, n_lin, n_rand, lin):
+        ports = list(iterports(start=lb, end=ub, n_lin=n_lin, n_rand=n_rand))
+        self.assertEqual(ports[:n_lin], lin)
+        self.assertEqual(set(ports[n_lin:]), set(range(lb, ub+1)).difference(lin))
+
+
+class TestBackgroundAppServer(unittest.TestCase):
+
+    def test_port_search(self):
+        base_port = 64000
+
+        first = BackgroundAppServer(
+            host='', base_port=base_port, max_port=base_port, app=None)
+        first.start()
+        self.assertEqual(first.server.server_port, base_port)
+
+        second = BackgroundAppServer(
+            host='', base_port=base_port, max_port=base_port+9, linear_tries=2, app=None)
+        second.start()
+        self.assertEqual(second.server.server_port, base_port+1)
+
+        first.stop()
+        second.stop()
+
+    def test_basics(self):
+        base_port = 64010
+        response = 'It works!'
+
+        def app(environ, start_response):
+            start_response("200 OK", [('Content-Type', 'text/plain; charset=utf-8')])
+            return [response.encode('utf-8')]
+
+        srv = BackgroundAppServer(
+            host='127.0.0.1', base_port=base_port, max_port=base_port+9, app=app)
+        srv.start()
+
+        # test root url
+        self.assertEqual(srv.root_url(), f'http://127.0.0.1:{srv.server.server_port}/')
+
+        # test response
+        self.assertEqual(requests.get(srv.root_url()).text, response)
+
+        srv.stop()
+
+    def test_multithreading(self):
+        base_port = 64020
+        response = 'It works!'
+        response_delay = 2
+        n_requests = 2
+
+        def app(environ, start_response):
+            start_response("200 OK", [('Content-Type', 'text/plain; charset=utf-8')])
+            time.sleep(response_delay)
+            return [response.encode('utf-8')]
+
+        srv = BackgroundAppServer(
+            host='127.0.0.1', base_port=base_port, max_port=base_port+9, app=app)
+        srv.start()
+
+        url = srv.root_url()
+        def get_url(url, timeout=10):
+            return requests.get(url, timeout=timeout).text
+
+        with concurrent.futures.ThreadPoolExecutor(max_workers=n_requests) as executor:
+            fs = {executor.submit(get_url, url): url for _ in range(n_requests)}
+            done, not_done = concurrent.futures.wait(
+                fs, timeout=response_delay*1.3, return_when=concurrent.futures.ALL_COMPLETED)
+
+        # verify all requests finished within (1.3 * response_delay) < (response_delay * n_request)
+        self.assertEqual(len(done), n_requests)
+
+        # all responses correct
+        for f in fs:
+            self.assertEqual(f.result(), response)
+
+        srv.stop()
+
+    def test_timeout(self):
+        # XXX: not the most elegant method of checking if server closes the
+        # connection after `timeout`, so we should revisit this at some point.
+        # I have manually verified the behavior is correct, thought, with
+        # telnet, curl and wireshark.
+
+        base_port = 64030
+        timeout = 0.5
+
+        def app(environ, start_response):
+            start_response("200 OK", [('Content-Type', 'text/plain; charset=utf-8')])
+            return ['It works!'.encode('utf-8')]
+
+        srv = BackgroundAppServer(
+            host='127.0.0.1', base_port=base_port, max_port=base_port+9,
+            app=app, timeout=timeout)
+        srv.start()
+
+        def do_get(timeout=0):
+            sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+            sock.connect(srv.server.server_address)
+
+            time.sleep(timeout)
+
+            try:
+                sock.send('GET /\n\n'.encode('utf-8'))
+                data = sock.recv(1024)
+                sock.close()
+            except ConnectionError:
+                # seems to work only on macos and win
+                data = ''
+
+            return data
+
+        # base case, no timeout, server returns data
+        self.assertTrue(len(do_get(timeout=0)) > 0)
+
+        # client waits 2*connection_timeout, so receives nothing
+        self.assertTrue(len(do_get(timeout=timeout*2)) == 0)
+
+        srv.stop()
+
+
+class TestSingleRequestAppServer(unittest.TestCase):
+
+    def test_function(self):
+        base_port = 64040
+        response = 'It works!'
+
+        def app(environ, start_response):
+            start_response("200 OK", [('Content-Type', 'text/plain; charset=utf-8')])
+            return [response.encode('utf-8')]
+
+        srv = SingleRequestAppServer(
+            host='127.0.0.1', base_port=base_port, max_port=base_port+9, app=app)
+        srv.start()
+
+        # test response
+        self.assertEqual(requests.get(srv.root_url()).text, response)
+
+        # test server shuts down within reasonable time
+        srv.wait_shutdown(timeout=0.5)
+        self.assertFalse(srv.is_alive())