Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
11 changes: 9 additions & 2 deletions src/promptflow-devkit/promptflow/_cli/_pf/_service.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand All @@ -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:
Expand Down Expand Up @@ -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:
Expand Down
Original file line number Diff line number Diff line change
@@ -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)
Original file line number Diff line number Diff line change
@@ -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
Loading