Add GatewayIdentityProvider, get tests working

Author: kevin-bates

conftest.py

@@ -1,15 +1,13 @@
 # Copyright (c) Jupyter Development Team.
 # Distributed under the terms of the Modified BSD License.
 
+import os
 import logging
 import pytest
-
-from jupyter_server import version_info as jupyter_server_version
-
+from binascii import hexlify
+from traitlets.config import Config
 from kernel_gateway.gatewayapp import KernelGatewayApp
 
-is_v2 = jupyter_server_version[0] == 2
-
 pytest_plugins = ["pytest_jupyter.jupyter_core", "pytest_jupyter.jupyter_server"]
 
 
@@ -45,20 +43,22 @@ def _configurable_serverapp(
         config=jp_server_config,
         base_url=jp_base_url,
         argv=jp_argv,
-        environ=jp_environ,
         http_port=jp_http_port,
-        tmp_path=tmp_path,
-        io_loop=io_loop,
-        root_dir=jp_root_dir,
         **kwargs,
     ):
+        c = Config(config)
+
+        if "auth_token" not in c.KernelGatewayApp and not c.IdentityProvider.token:
+            default_token = hexlify(os.urandom(4)).decode("ascii")
+            c.IdentityProvider.token = default_token
+
         app = KernelGatewayApp.instance(
             # Set the log level to debug for testing purposes
             log_level="DEBUG",
-            port=jp_http_port,
+            port=http_port,
             port_retries=0,
             base_url=base_url,
-            config=config,
+            config=c,
             **kwargs,
         )
         app.log.propagate = True
@@ -101,6 +101,4 @@ def jp_server_cleanup(jp_asyncio_loop):
 @pytest.fixture
 def jp_auth_header(jp_serverapp):
     """Configures an authorization header using the token from the serverapp fixture."""
-    if not is_v2:
-        return {"Authorization": f"token {jp_serverapp.auth_token}"}
     return {"Authorization": f"token {jp_serverapp.identity_provider.token}"}

kernel_gateway/auth/identity.py

@@ -0,0 +1,38 @@
+# Copyright (c) Jupyter Development Team.
+# Distributed under the terms of the Modified BSD License.
+"""Gateway Identity Provider interface
+
+This defines the _authentication_ layer of Jupyter Server,
+to be used in combination with Authorizer for _authorization_.
+"""
+from traitlets import default
+
+from jupyter_server.auth.identity import IdentityProvider
+from jupyter_server.base.handlers import JupyterHandler
+
+
+class GatewayIdentityProvider(IdentityProvider):
+    """
+    Interface for providing identity management and authentication for a Gateway server.
+    """
+    @default("token")
+    def _token_default(self):
+        # if the superclass generated a token, but auth_token is configured on
+        # the Gateway server, reset token_generated and use the configured value.
+        token_default = super()._token_default()
+        if self.token_generated and self.parent.auth_token:
+            self.token_generated = False
+            return self.parent.auth_token
+        return token_default
+
+    def should_check_origin(self, handler: JupyterHandler) -> bool:
+        """Should the Handler check for CORS origin validation?
+
+        Origin check should be skipped for token-authenticated requests.
+
+        Returns:
+        - True, if Handler must check for valid CORS origin.
+        - False, if Handler should skip origin check since requests are token-authenticated.
+        """
+        # Always check the origin unless operator configured gateway to allow any
+        return handler.settings["kg_allow_origin"] != "*"

kernel_gateway/gatewayapp.py

@@ -6,42 +6,46 @@
 import errno
 import importlib
 import logging
+import hashlib
+import hmac
 import os
 import sys
 import signal
 import select
 import socket
 import ssl
 import threading
+from base64 import encodebytes
 from distutils.util import strtobool
 
 import nbformat
 from jupyter_server.services.kernels.kernelmanager import MappingKernelManager
 
 from urllib.parse import urlparse
-from traitlets import Unicode, Integer, default, observe, Type, Instance, List, CBool
+from traitlets import Unicode, Integer, Bytes, default, observe, Type, Instance, List, CBool
 
 from jupyter_core.application import JupyterApp, base_aliases
-from jupyter_core.utils import run_sync
 from jupyter_client.kernelspec import KernelSpecManager
 
 from tornado import httpserver
 from tornado import web, ioloop
 from tornado.log import enable_pretty_logging, LogFormatter
 
+from jupyter_core.paths import secure_write
 from jupyter_server.serverapp import random_ports
 from ._version import __version__
 from .services.sessions.sessionmanager import SessionManager
 from .services.kernels.manager import SeedingMappingKernelManager
 
+from .auth.identity import GatewayIdentityProvider
+
 # Only present for generating help documentation
 from .notebook_http import NotebookHTTPPersonality
 from .jupyter_websocket import JupyterWebsocketPersonality
 
-try:
-    from jupyter_server.auth.authorizer import AllowAllAuthorizer
-except ImportError:
-    AllowAllAuthorizer = object
+from jupyter_server.auth.authorizer import AllowAllAuthorizer, Authorizer
+from jupyter_server.services.kernels.connection.base import BaseKernelWebsocketConnection
+from jupyter_server.services.kernels.connection.channels import ZMQChannelsWebsocketConnection
 
 
 # Add additional command line aliases
@@ -311,6 +315,51 @@ def ssl_version_default(self):
         ssl_from_env = os.getenv(self.ssl_version_env)
         return ssl_from_env if ssl_from_env is None else int(ssl_from_env)
 
+    cookie_secret_file = Unicode(
+        config=True, help="""The file where the cookie secret is stored."""
+    )
+
+    @default("cookie_secret_file")
+    def _default_cookie_secret_file(self):
+        return os.path.join(self.runtime_dir, "jupyter_cookie_secret")
+
+    cookie_secret = Bytes(
+        b"",
+        config=True,
+        help="""The random bytes used to secure cookies.
+        By default this is a new random number every time you start the server.
+        Set it to a value in a config file to enable logins to persist across server sessions.
+
+        Note: Cookie secrets should be kept private, do not share config files with
+        cookie_secret stored in plaintext (you can read the value from a file).
+        """,
+    )
+
+    @default("cookie_secret")
+    def _default_cookie_secret(self):
+        if os.path.exists(self.cookie_secret_file):
+            with open(self.cookie_secret_file, "rb") as f:
+                key = f.read()
+        else:
+            key = encodebytes(os.urandom(32))
+            self._write_cookie_secret_file(key)
+        h = hmac.new(key, digestmod=hashlib.sha256)
+        # h.update(self.password.encode())  # password is deprecated in 2.0
+        return h.digest()
+
+    def _write_cookie_secret_file(self, secret):
+        """write my secret to my secret_file"""
+        self.log.info("Writing Jupyter server cookie secret to %s", self.cookie_secret_file)
+        try:
+            with secure_write(self.cookie_secret_file, True) as f:
+                f.write(secret)
+        except OSError as e:
+            self.log.error(
+                "Failed to write cookie secret to %s: %s",
+                self.cookie_secret_file,
+                e,
+            )
+
     ws_ping_interval_env = "KG_WS_PING_INTERVAL_SECS"
     ws_ping_interval_default_value = 30
     ws_ping_interval = Integer(
@@ -352,6 +401,27 @@ def _default_log_format(self) -> str:
         help="""The kernel manager class to use."""
     )
 
+    kernel_websocket_connection_class = Type(
+        default_value=ZMQChannelsWebsocketConnection,
+        klass=BaseKernelWebsocketConnection,
+        config=True,
+        help="""The kernel websocket connection class to use.""",
+    )
+
+    authorizer_class = Type(
+        default_value=AllowAllAuthorizer,
+        klass=Authorizer,
+        config=True,
+        help="The authorizer class to use.",
+    )
+
+    identity_provider_class = Type(
+        default_value=GatewayIdentityProvider,
+        klass=GatewayIdentityProvider,
+        config=True,
+        help="The identity provider class to use.",
+    )
+
     def _load_api_module(self, module_name):
         """Tries to import the given module name.
 
@@ -467,6 +537,12 @@ def init_configurables(self):
         )
         self.contents_manager = None
 
+        self.identity_provider = self.identity_provider_class(parent=self, log=self.log)
+
+        self.authorizer = self.authorizer_class(
+            parent=self, log=self.log, identity_provider=self.identity_provider
+        )
+
         if self.prespawn_count:
             if self.max_kernels and self.prespawn_count > self.max_kernels:
                 msg = f"Cannot prespawn {self.prespawn_count} kernels; more than max kernels {self.max_kernels}"
@@ -514,13 +590,17 @@ def init_webapp(self):
             allow_origin=self.allow_origin,
             # Set base_url for use in request handlers
             base_url=self.base_url,
+            # Authentication
+            cookie_secret=self.cookie_secret,
             # Always allow remote access (has been limited to localhost >= notebook 5.6)
             allow_remote_access=True,
             # setting ws_ping_interval value that can allow it to be modified for the purpose of toggling ping mechanism
             # for zmq web-sockets or increasing/decreasing web socket ping interval/timeouts.
             ws_ping_interval=self.ws_ping_interval * 1000,
             # Add a pass-through authorizer for now
-            authorizer=AllowAllAuthorizer(),
+            authorizer=self.authorizer_class(parent=self),
+            identity_provider=self.identity_provider,
+            kernel_websocket_connection_class=self.kernel_websocket_connection_class,
         )
 
         # promote the current personality's "config" tagged traitlet values to webapp settings
@@ -684,6 +764,7 @@ def start(self):
     async def _stop(self):
         """Cleanup resources and stop the IO Loop."""
         await self.personality.shutdown()
+        await self.kernel_websocket_connection_class.close_all()
         if getattr(self, "io_loop", None):
             self.io_loop.stop()
 

kernel_gateway/mixins.py

@@ -15,17 +15,25 @@ class CORSMixin(object):
         'kg_allow_headers': 'Access-Control-Allow-Headers',
         'kg_allow_methods': 'Access-Control-Allow-Methods',
         'kg_allow_origin': 'Access-Control-Allow-Origin',
-        'kg_expose_headers' : 'Access-Control-Expose-Headers',
+        'kg_expose_headers': 'Access-Control-Expose-Headers',
         'kg_max_age': 'Access-Control-Max-Age'
     }
 
-    def set_default_headers(self):
+    def set_cors_headers(self):
         """Sets the CORS headers as the default for all responses.
 
         Disables CSP configured by the notebook package. It's not necessary
         for a programmatic API.
+
+        Notes
+        -----
+        This method name was changed from set_default_headers to set_cors_header
+        when adding support for JupyterServer 2.x.  In that release, JS changed the
+        way the headers were implemented due to changes in the way a user is authenticated.
+        See https://github.com/jupyter-server/jupyter_server/pull/671.
         """
-        super(CORSMixin, self).set_default_headers()
+        super().set_cors_headers()
+
         # Add CORS headers after default if they have a non-blank value
         for settings_name, header_name in self.SETTINGS_TO_HEADERS.items():
             header_value = self.settings.get(settings_name)
@@ -75,7 +83,7 @@ def prepare(self):
                     client_token = None
             if client_token != server_token:
                 return self.send_error(401)
-        return super(TokenAuthorizationMixin, self).prepare()
+        return super().prepare()
 
 
 class JSONErrorsMixin(object):