Refactor AuthContext creation into standalone AuthContextMiddleware. (tensorflow#6131)

The TensorBoard application layer creates a middleware that injects
AuthContext into the RequestContext. It would be useful to expose this
middleware for other usage - there are at least a couple test files in
the internal code base that could make use of it.

This change refactors the middleware into its own file for reuse, naming
it AuthContextMiddleware and giving it its own devoted test.

Author: bmd3k

tensorboard/backend/BUILD

@@ -53,15 +53,14 @@ py_library(
     srcs_version = "PY3",
     visibility = ["//visibility:public"],
     deps = [
+        ":auth_context_middleware",
         ":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_middleware",
+    srcs = ["auth_context_middleware.py"],
+    srcs_version = "PY3",
+    deps = [
+        "//tensorboard:auth",
+        "//tensorboard:context",
+    ],
+)
+
+py_test(
+    name = "auth_context_middleware_test",
+    size = "small",
+    srcs = ["auth_context_middleware_test.py"],
+    srcs_version = "PY3",
+    tags = ["support_notf"],
+    deps = [
+        ":auth_context_middleware",
+        "//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_middleware
 from tensorboard.backend import client_feature_flags
 from tensorboard.backend import empty_path_redirect
 from tensorboard.backend import experiment_id
@@ -326,7 +325,9 @@ 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_middleware.AuthContextMiddleware(
+            app, self._auth_providers
+        )
         app = client_feature_flags.ClientFeatureFlagsMiddleware(app)
         app = empty_path_redirect.EmptyPathRedirectMiddleware(app)
         app = experiment_id.ExperimentIdMiddleware(app)
@@ -582,17 +583,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_middleware.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_middleware_test.py

@@ -0,0 +1,74 @@
+# 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_middleware`."""
+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_middleware
+
+
+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 after AuthContextMiddleware in the WSGI handler chain.
+    It will generate a credential using the AuthContext populated by
+    AuthContextMiddleware.
+
+    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_middleware.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()