diff --git a/src/promptflow-devkit/promptflow/_cli/_pf/_service.py b/src/promptflow-devkit/promptflow/_cli/_pf/_service.py index d24bad30a91..defe3a0c955 100644 --- a/src/promptflow-devkit/promptflow/_cli/_pf/_service.py +++ b/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: diff --git a/src/promptflow-devkit/promptflow/_sdk/_service/utils/local_user_auth.py b/src/promptflow-devkit/promptflow/_sdk/_service/utils/local_user_auth.py new file mode 100644 index 00000000000..285b67676b2 --- /dev/null +++ b/src/promptflow-devkit/promptflow/_sdk/_service/utils/local_user_auth.py @@ -0,0 +1,23 @@ +# --------------------------------------------------------- +# 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): + # Strip any client-supplied identity headers to prevent spoofing. + environ.pop("HTTP_REMOTE_USER", None) + environ.pop("HTTP_X_REMOTE_USER", None) + + if environ.get("REMOTE_ADDR") in _LOOPBACK_REMOTE_ADDRS: + environ["REMOTE_USER"] = getpass.getuser() + return self._wsgi_app(environ, start_response) diff --git a/src/promptflow-devkit/tests/sdk_pfs_test/test_local_user_auth_middleware.py b/src/promptflow-devkit/tests/sdk_pfs_test/test_local_user_auth_middleware.py new file mode 100644 index 00000000000..7531e1970a2 --- /dev/null +++ b/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