fix: inject local user for loopback PFS requests

Author: YizukiAme

src/promptflow-devkit/promptflow/_cli/_pf/_service.py

@@ -23,6 +23,7 @@
     PF_SERVICE_WORKER_NUM,
 )
 from promptflow._sdk._service.app import create_app
+from promptflow._sdk._service.utils.local_user_auth import LocalUserAuthMiddleware
 from promptflow._sdk._service.utils.utils import (
     check_pfs_service_status,
     dump_port_to_config,
@@ -47,10 +48,16 @@
 app = None
 
 
+def _create_app_with_local_user_auth():
+    flask_app, _ = create_app()
+    flask_app.wsgi_app = LocalUserAuthMiddleware(flask_app.wsgi_app)
+    return flask_app
+
+
 def get_app(environ, start_response):
     global app
     if app is None:
-        app, _ = create_app()
+        app = _create_app_with_local_user_auth()
     if os.environ.get(PF_SERVICE_DEBUG) == "true":
         app.logger.setLevel(logging.DEBUG)
     else:
@@ -255,7 +262,7 @@ def _prepare_app_for_foreground_service(port, force_start, service_host):
     port = validate_port(port, force_start, service_host)
     global app
     if app is None:
-        app, _ = create_app()
+        app = _create_app_with_local_user_auth()
     if os.environ.get(PF_SERVICE_DEBUG) == "true":
         app.logger.setLevel(logging.DEBUG)
     else:

src/promptflow-devkit/promptflow/_sdk/_service/utils/local_user_auth.py

@@ -0,0 +1,19 @@
+# ---------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# ---------------------------------------------------------
+
+import getpass
+
+_LOOPBACK_REMOTE_ADDRS = {"127.0.0.1", "::1"}
+
+
+class LocalUserAuthMiddleware:
+    """Inject the current OS user for local PFS loopback requests."""
+
+    def __init__(self, wsgi_app):
+        self._wsgi_app = wsgi_app
+
+    def __call__(self, environ, start_response):
+        if environ.get("REMOTE_ADDR") in _LOOPBACK_REMOTE_ADDRS:
+            environ["REMOTE_USER"] = getpass.getuser()
+        return self._wsgi_app(environ, start_response)

src/promptflow-devkit/tests/sdk_pfs_test/test_local_user_auth_middleware.py

@@ -0,0 +1,61 @@
+# ---------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# ---------------------------------------------------------
+
+import getpass
+
+from promptflow._sdk._service.utils.local_user_auth import LocalUserAuthMiddleware
+
+
+def _remote_user_echo_app(environ, start_response):
+    start_response("200 OK", [])
+    return [environ.get("REMOTE_USER", "").encode()]
+
+
+def _call_middleware(remote_addr, extra_environ=None):
+    environ = {"REMOTE_ADDR": remote_addr}
+    if extra_environ:
+        environ.update(extra_environ)
+
+    response_status = []
+
+    def start_response(status, headers):
+        response_status.append(status)
+
+    response_body = b"".join(LocalUserAuthMiddleware(_remote_user_echo_app)(environ, start_response)).decode()
+    return response_status[0], response_body, environ
+
+
+def test_loopback_request_injects_current_user():
+    status, remote_user, environ = _call_middleware("127.0.0.1")
+
+    assert status == "200 OK"
+    assert remote_user == getpass.getuser()
+    assert environ["REMOTE_USER"] == getpass.getuser()
+
+
+def test_ipv6_loopback_request_injects_current_user():
+    _, remote_user, environ = _call_middleware("::1")
+
+    assert remote_user == getpass.getuser()
+    assert environ["REMOTE_USER"] == getpass.getuser()
+
+
+def test_non_loopback_request_does_not_inject_user():
+    _, remote_user, environ = _call_middleware("192.0.2.10")
+
+    assert remote_user == ""
+    assert "REMOTE_USER" not in environ
+
+
+def test_client_supplied_remote_user_header_is_not_trusted():
+    _, remote_user, environ = _call_middleware(
+        "192.0.2.10",
+        {
+            "HTTP_REMOTE_USER": getpass.getuser(),
+            "HTTP_X_REMOTE_USER": getpass.getuser(),
+        },
+    )
+
+    assert remote_user == ""
+    assert "REMOTE_USER" not in environ