Adding OAuth2 to Jupyter Notebooks on Kubernetes

February 1, 2024
Share this post

TrueFoundry users can deploy Jupyter Notebooks on their personal cloud accounts, such as AWS, Azure, or GCP. This feature allows them to conduct machine learning experiments and training jobs on their own machines with ease. Initially, notebooks deployed through TrueFoundry were secured using a username-password combination. However, in response to widespread client requests, we have integrated Single Sign-On. This means users can now conveniently access their notebooks with the same login they use for TrueFoundry. This blog post delves into the specifics of how we implemented this feature.

Launching a Jupyter Notebook on TrueFoundry

Notebooks on TrueFoundry

TrueFoundry internally uses a fork of Kubeflow Notebook Controller to orchestrate the deployment of notebooks. The controller provides various features that we leverage, like:

  1. Simplified notebook spec: The Kubeflow Notebook APIs are simple and the controller orchestrates the creation of the Jupyter Notebook deployments.
  2. Automatic Culling: The controller automatically shuts down the notebook after a certain period of inactivity. This is incredibly useful to our clients who run experiments on notebooks backed by GPU machines.
  3. Persistent Home Directory: The controller takes care of creating a persistent volume that saves user progress on the notebook across sessions.
  4. Extensible base images: The controller supports a suite of base notebook images of Jupyter Notebook and VS Code maintained by TrueFoundry. The user can extend the features on these Docker images by adding a startup script or installing specific libraries.

For context, here’s what a simple Kubeflow Notebook object looks like:

kind: Notebook
 name: my-notebook
       - name: my-notebook
         image: kubeflownotebookswg/jupyter:master

Basic Auth for Notebooks

Before implementing OAuth2, TrueFoundry provided users with the option to enhance the security of their public notebooks by integrating basic authentication. This added layer of security was crucial to ensure that only authorized individuals could access the sensitive content of these notebooks. To implement this feature, TrueFoundry utilized the capabilities of WebAssembly (Wasm) plugins within the Istio proxy, specifically the Envoy proxy.

Istio, an open-source service mesh, offers a framework for managing network communications between various service workloads. With Istio, TrueFoundry was empowered to inject custom logic directly into the network layer, which is managed by the Envoy proxy. This approach allowed for effective control and security of the traffic flowing to and from their Jupyter Notebooks. The key to the implementation of basic auth was the WasmPlugin, a feature of Istio that facilitates the deployment of WebAssembly modules within the Envoy proxy.

This basic authentication WasmPlugin is integrated into a sequence of network filters within the Envoy proxy. These filters enable the execution of higher-level functions related to access control, transformation, data enrichment, auditing, and more, thereby enhancing the overall security and functionality of the service mesh. Here’s a simplified version of the spec for adding basic auth filter to the Envoy filter chain:

kind: WasmPlugin
 name: basic-auth
 namespace: istio-ingress
 phase: AUTHN
     - credentials:
         - user:pass
       prefix: /secret/
     istio: ingressgateway
 url: oci://

OAuth2 for Notebook

For implementing OAuth2 in our notebooks, we utilized an Envoy filter, but the approach differed from that of basic authentication. Unlike the basic auth where we could conveniently insert a pre-built WasmPlugin into the filter chain, OAuth2 required a more tailored solution. To achieve this, we employed an HTTP filter specifically designed for OAuth. At TrueFoundry, our Single Sign-On system integrates with FusionAuth, serving as our OAuth provider.

Here’s how the Envoy Filter spec looks like – refer to the comments in the file for more details:

kind: EnvoyFilter
 name: truefoundry-notebook-tfy-oauth2  # Name of the EnvoyFilter
 namespace: auth-test  # Namespace where the EnvoyFilter is deployed
   labels: truefoundry-notebook  # Selector targeting workloads with specific labels
 - applyTo: CLUSTER
     context: SIDECAR_OUTBOUND
     operation: ADD
       name: tfy-oauth2  # Name of the cluster for OAuth2 authentication service
       type: LOGICAL_DNS  # Type of service discovery (DNS)
       connect_timeout: 5s  # Timeout for establishing a connection
       lb_policy: ROUND_ROBIN  # Load balancing policy
       # other load balancing config
 - applyTo: HTTP_FILTER
     context: SIDECAR_INBOUND
           name: ""
             name: envoy.filters.http.jwt_authn
     operation: INSERT_BEFORE  # Inserting this filter before the JWT auth filter
      name: envoy.filters.http.tfy-oauth  # Name of the OAuth filter
          use_refresh_token: false  # Whether to use a refresh token
            - name: Authorization
              present_match: true  # Pass through if Authorization header is present
          forward_bearer_token: true  # Forward bearer token to upstream
          auth_type: BASIC_AUTH  # Type of authentication used
            cluster: tfy-oauth2  # Cluster for token endpoint
            uri: <token-endpoint-uri-of-oauth-provider>
            timeout: 5s  # Timeout for token endpoint
          authorization_endpoint: <authorization-endpoint-uri-of-oauth-provider>
          redirect_uri: https://%REQ(:authority)%/truefoundry-notebook/_auth/callback  # Redirect URI for callback
              exact: /truefoundry-notebook/_auth/callback  # Path for redirect URI
              exact: /truefoundry-notebook/_auth/signout  # Path for signout
            client_id: <client-id-for-oauth>
               # configuration to fetch token secret
               # read more about how we fetch secrets here:
               # configuration to fetch hmac

When a user attempts to access a service protected by the OAuth2 filter for the first time, they are redirected to the authorization_endpoint. This endpoint is the URL of our external OAuth Provider, which, in our implementation, is the FusionAuth-based TrueFoundry login modal. This redirection is a critical step in the OAuth process, guiding users to a secure location where they can authenticate and consequently grant the necessary permissions for access to the service.

Once the login is complete, FusionAuth will redirect you to the redirect_uri (configured in the filter specification), adding a secret, temporary authorization code there. This request is intercepted by the filter and it makes a request to token_endpoint, exchanging the code for a JWT token. Finally, the filter sets cookies with the JWT token.

Subsequent accesses to the service are passed through the HTTP Filter since the cookie sets the Authorization header with JWT as the value. The filter is configured to pass through such requests (refer pass_through_matcher in the spec). To validate that the JWT is a valid token, we create a RequestAuthentication policy that will check with the OAuth provider:

kind: RequestAuthentication
 # ...
   # ...
 - issuer: ""
   - name: Authorization
     prefix: "Bearer "
     - <client-id>
   jwksUri: <oauth-provider-jwks-uri>
   forwardOriginalToken: true

Finally, we add the Authorization Policy that specify what requests to apply RequestAuthentication to. We want to apply authorization to all requests on port 8888:

kind: AuthorizationPolicy
 name: best-notebook-tfy-oauth2
 namespace: auth-test
   matchLabels: best-notebook
 action: DENY
 - from:
   - source:
       notRequestPrincipals: ["*"]
     - operation:
           - "8888"

Build, Train, and Deploy LLM/ML Faster
Start Your Free 7-Day Trial Now!

Discover More

June 13, 2024

Enhancing Customer Support with Real-Time AI Assistance Using Cognita

Engineering and Product
May 3, 2024

Top Tools for Fine-tuning

Engineering and Product
February 29, 2024

SSH Server Containers For Development on Kubernetes

Engineering and Product
March 6, 2024

Prompting, RAG or Fine-tuning - the right choice?

Engineering and Product

Related Blogs

June 1, 2024
5 min read

A Guide to Cloud Node Auto-Provisioning

Blazingly fast way to build, track and deploy your models!