Add built-in support for generating kernel gateway auth tokens using an external command.

[ojarjur]

Problem

I want to use an external command to generate the token used by the GatewayClient for connecting to kernel gateways.

The new GatewayTokenRenewerBase functionality provides a great starting point for that, but it still requires a custom implementation of a class extending that base class in order to work, and ideally this would just require specifying the external command to run as part of my .jupyter/jupyter_lab_config.py config file.

I think this is a generic enough feature that it makes sense to include in the jupyter_server repository rather than being some sort of an external extension, but I'm also open to any suggestions about better homes for it.

Proposed Solution

I have prototyped a proof-of-concept of this by adding the below classes to my ~/.jupyter/jupyter_lab.config.py config file, and confirmed that it was working.

I'm happy to also put together a PR with equivalent changes if you all agree that this is a good addition to the codebase.

# Proof of concept snippet for a generic command-based token renewer
#
# Copyright 2023 Google LLC.
# SPDX-License-Identifier: Apache-2.0
import datetime
import subprocess
import typing


from abc import abstractmethod
from traitlets import Int, List
from jupyter_server.gateway.gateway_client import GatewayTokenRenewerBase


class CachedTokenRenewerBase(GatewayTokenRenewerBase):
  """Token renewer base class that only renews the token after a specified timeout.""" 

  token_lifetime_seconds = Int(
    default_value=300,
    config=True,
    help="""Time (in seconds) to wait between successive token renewals.""",
  )

  @abstractmethod
  def force_new_token(
      self,
      auth_header_key: str,
      auth_scheme: typing.Union[str, None],
      **kwargs: typing.Any,
  ):
    pass

  _created = datetime.datetime.min

  def get_token(
      self,
      auth_header_key: str,
      auth_scheme: typing.Union[str, None],
      auth_token: str,
      **kwargs: typing.Any,
  ):
    current_time = datetime.datetime.now()
    duration = (current_time - self._created).total_seconds()
    if (not auth_token) or (duration > self.token_lifetime_seconds):
      auth_token = self.force_new_token(auth_header_key, auth_scheme, **kwargs)
      self._created = datetime.datetime.now()

    return auth_token


class CommandTokenRenewer(CachedTokenRenewerBase):
  """Token renewer that invokes an external command to generate the token."""

  token_command = List(
    default_value=[],
    config=True,
    help="""External command run to generate auth tokens.""",
  )

  def force_new_token(
      self,
      auth_header_key: str,
      auth_scheme: typing.Union[str, None],
      **kwargs: typing.Any,
  ):
    p = subprocess.run(
      self.token_command,
      stdin=subprocess.DEVNULL,
      capture_output=True,
      check=True,
      encoding='UTF-8')
    return p.stdout.strip()

Additional context

Note that the reason for including the caching behavior in the P.o.C. is because invoking an external command is slow enough that you do not want to do it on every request to the kernel gateway.

Also, in my particular case, the kernel gateway is running in GCP, so the external command is Google's gcloud command line tool, but I believe this is a generic use case and this proposal is not in any way specific to Google.