Refactor AuthContext creation into standalone AuthContextMiddleware.

tensorboard/backend/BUILD

@@ -53,15 +53,14 @@ py_library(
     srcs_version = "PY3",
     visibility = ["//visibility:public"],
     deps = [
+        ":auth_context",
         ":client_feature_flags",
         ":empty_path_redirect",
         ":experiment_id",
         ":experimental_plugin",
         ":http_util",
         ":path_prefix",
         ":security_validator",
-        "//tensorboard:auth",
-        "//tensorboard:context",
         "//tensorboard:errors",
         "//tensorboard:plugin_util",
         "//tensorboard/plugins/core:core_plugin",
@@ -88,6 +87,31 @@ py_test(
     ],
 )
 
+py_library(
+    name = "auth_context",
+    srcs = ["auth_context.py"],
+    srcs_version = "PY3",
+    deps = [
+        "//tensorboard:auth",
+        "//tensorboard:context",
+    ],
+)
+
+py_test(
+    name = "auth_context_test",
+    size = "small",
+    srcs = ["auth_context_test.py"],
+    srcs_version = "PY3",
+    tags = ["support_notf"],
+    deps = [
+        ":auth_context",
+        "//tensorboard:auth",
+        "//tensorboard:context",
+        "//tensorboard:test",
+        "@org_pocoo_werkzeug",
+    ],
+)
+
 py_library(
     name = "empty_path_redirect",
     srcs = ["empty_path_redirect.py"],

tensorboard/backend/application.py

@@ -32,8 +32,7 @@
 
 from tensorboard import errors
 from tensorboard import plugin_util
-from tensorboard import auth
-from tensorboard import context
+from tensorboard.backend import auth_context
 from tensorboard.backend import client_feature_flags
 from tensorboard.backend import empty_path_redirect
 from tensorboard.backend import experiment_id
@@ -326,7 +325,7 @@ def _create_wsgi_app(self):
         app = self._route_request
         for middleware in self._extra_middlewares:
             app = middleware(app)
-        app = _auth_context_middleware(app, self._auth_providers)
+        app = auth_context.AuthContextMiddleware(app, self._auth_providers)
         app = client_feature_flags.ClientFeatureFlagsMiddleware(app)
         app = empty_path_redirect.EmptyPathRedirectMiddleware(app)
         app = experiment_id.ExperimentIdMiddleware(app)
@@ -582,17 +581,6 @@ def wrapper(environ, start_response):
     return wrapper
 
 
-def _auth_context_middleware(wsgi_app, auth_providers):
-    def wrapper(environ, start_response):
-        environ = dict(environ)
-        auth_ctx = auth.AuthContext(auth_providers, environ)
-        ctx = context.from_environ(environ).replace(auth=auth_ctx)
-        context.set_in_environ(environ, ctx)
-        return wsgi_app(environ, start_response)
-
-    return wrapper
-
-
 def _clean_path(path):
     """Removes a trailing slash from a non-root path.
 

tensorboard/backend/auth_context.py

@@ -0,0 +1,38 @@
+# Copyright 2023 The TensorFlow Authors. All Rights Reserved.
+#
+# 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.
+# ==============================================================================
+from tensorboard import auth
+from tensorboard import context
+
+
+class AuthContextMiddleware:
+    """WSGI middleware to inject an AuthContext into the RequestContext."""
+
+    def __init__(self, application, auth_providers):
+        """Initializes this middleware.
+
+        Args:
+          application: A WSGI application to delegate to.
+          auth_providers: The auth_providers to provide to the AuthContext.
+        """
+        self._application = application
+        self._auth_providers = auth_providers
+
+    def __call__(self, environ, start_response):
+        """Invoke this WSGI application."""
+        environ = dict(environ)
+        auth_ctx = auth.AuthContext(self._auth_providers, environ)
+        ctx = context.from_environ(environ).replace(auth=auth_ctx)
+        context.set_in_environ(environ, ctx)
+        return self._application(environ, start_response)

tensorboard/backend/auth_context_test.py

@@ -0,0 +1,73 @@
+# Copyright 2023 The TensorFlow Authors. All Rights Reserved.
+#
+# 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.
+# ==============================================================================
+"""Tests for `tensorboard.backend.auth_context`."""
+from werkzeug import test as werkzeug_test
+from werkzeug import wrappers
+
+from tensorboard import auth
+from tensorboard import context
+from tensorboard import test as tb_test
+from tensorboard.backend import auth_context
+
+
+class SimpleAuthProvider(auth.AuthProvider):
+    """Simple AuthProvider that returns the value it is initialized with."""
+
+    def __init__(self, credential):
+        self._credential = credential
+
+    def authenticate(self, environ):
+        return self._credential
+
+
+def _create_auth_provider_verifier_app(expected_auth_key):
+    """Generates a WSGI application for verifying AuthContextMiddleware.
+
+    It should be placed downstream from the AuthContextMiddleware. It will
+    generate a credential using the AuthContext and the given key.
+
+    Args:
+      expected_auth_key: The auth key to be used for invoking the AuthContext.
+        This key should correspond to the auth_providers used to configure
+        the AuthContextMiddleware .
+    """
+
+    def _app(environ, start_response):
+        ctx = context.from_environ(environ)
+        credential = ctx.auth.get(expected_auth_key)
+        start_response("200 OK", [("Content-Type", "text\plain")])
+        return f"credential: {credential}"
+
+    return _app
+
+
+class AuthContextMiddlewareTest(tb_test.TestCase):
+    """Tests for `AuthContextMiddleware`"""
+
+    def test_populates_auth_context(self):
+        app = auth_context.AuthContextMiddleware(
+            _create_auth_provider_verifier_app("my_key"),
+            {"my_key": SimpleAuthProvider("my_credential")},
+        )
+
+        server = werkzeug_test.Client(app, wrappers.Response)
+        response = server.get("")
+        self.assertEqual(
+            "credential: my_credential", response.get_data().decode()
+        )
+
+
+if __name__ == "__main__":
+    tb_test.main()