Skip to content

Commit

Permalink
[Auth] Only admins can create new user (#9700)
Browse files Browse the repository at this point in the history
Signed-off-by: Gabriel Fu <hfu.gabriel@gmail.com>
Co-authored-by: Harutaka Kawamura <hkawamura0130@gmail.com>
  • Loading branch information
gabrielfu and harupy committed Oct 3, 2023
1 parent 5a58bae commit 32de215
Show file tree
Hide file tree
Showing 5 changed files with 79 additions and 25 deletions.
3 changes: 3 additions & 0 deletions docs/source/auth/index.rst
Original file line number Diff line number Diff line change
Expand Up @@ -550,6 +550,9 @@ In Python, you can use the ``requests`` library:
Creating a New User
===================

.. important::
To create a new user, you are required to authenticate with admin privileges.

Using MLflow UI
---------------

Expand Down
14 changes: 8 additions & 6 deletions mlflow/server/auth/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -113,13 +113,8 @@
store = SqlAlchemyStore()


UNPROTECTED_ROUTES = [CREATE_USER, SIGNUP]


def is_unprotected_route(path: str) -> bool:
if path.startswith(("/static", "/favicon.ico", "/health")):
return True
return path in UNPROTECTED_ROUTES
return path.startswith(("/static", "/favicon.ico", "/health"))


def make_basic_auth_response() -> Response:
Expand Down Expand Up @@ -315,6 +310,11 @@ def validate_can_read_user():
return username_is_sender()


def validate_can_create_user():
# only admins can create user, but admins won't reach this validator
return False


def validate_can_update_user_password():
return username_is_sender()

Expand Down Expand Up @@ -385,7 +385,9 @@ def get_before_request_handler(request_class):

BEFORE_REQUEST_VALIDATORS.update(
{
(SIGNUP, "GET"): validate_can_create_user,
(GET_USER, "GET"): validate_can_read_user,
(CREATE_USER, "POST"): validate_can_create_user,
(UPDATE_USER_PASSWORD, "PATCH"): validate_can_update_user_password,
(UPDATE_USER_ADMIN, "PATCH"): validate_can_update_user_admin,
(DELETE_USER, "DELETE"): validate_can_delete_user,
Expand Down
10 changes: 9 additions & 1 deletion tests/server/auth/auth_test_utils.py
Original file line number Diff line number Diff line change
@@ -1,20 +1,28 @@
from mlflow.environment_variables import MLFLOW_TRACKING_PASSWORD, MLFLOW_TRACKING_USERNAME
from mlflow.server.auth import auth_config

from tests.helper_functions import random_str
from tests.tracking.integration_test_utils import _send_rest_tracking_post_request

PERMISSION = "READ"
NEW_PERMISSION = "EDIT"
ADMIN_USERNAME = auth_config.admin_username
ADMIN_PASSWORD = auth_config.admin_password


def create_user(tracking_uri):
username = random_str()
password = random_str()
_send_rest_tracking_post_request(
response = _send_rest_tracking_post_request(
tracking_uri,
"/api/2.0/mlflow/users/create",
{
"username": username,
"password": password,
},
auth=(ADMIN_USERNAME, ADMIN_PASSWORD),
)
response.raise_for_status()
return username, password


Expand Down
25 changes: 23 additions & 2 deletions tests/server/auth/test_auth.py
Original file line number Diff line number Diff line change
Expand Up @@ -22,7 +22,8 @@
)
from mlflow.utils.os import is_windows

from tests.server.auth.auth_test_utils import User, create_user
from tests.helper_functions import random_str
from tests.server.auth.auth_test_utils import ADMIN_USERNAME, User, create_user
from tests.tracking.integration_test_utils import (
_init_server,
_send_rest_tracking_post_request,
Expand Down Expand Up @@ -71,6 +72,21 @@ def _mlflow_search_experiments_rest(base_uri, headers):
return response


def _mlflow_create_user_rest(base_uri, headers):
username = random_str()
password = random_str()
response = requests.post(
f"{base_uri}/api/2.0/mlflow/users/create",
headers=headers,
json={
"username": username,
"password": password,
},
)
response.raise_for_status()
return username, password


@pytest.mark.parametrize(
"client",
[
Expand All @@ -88,7 +104,12 @@ def test_authenticate_jwt(client):
assert e.value.response.status_code == 401 # Unauthorized

# authenticated
username, password = create_user(client.tracking_uri)
# we need to use jwt to authenticate as admin so that we can create a new user
bearer_token = jwt.encode({"username": ADMIN_USERNAME}, "secret", algorithm="HS256")
headers = {"Authorization": f"Bearer {bearer_token}"}
username, password = _mlflow_create_user_rest(client.tracking_uri, headers)

# authenticate with the newly created user
headers = {
"Authorization": f'Bearer {jwt.encode({"username": username}, "secret", algorithm="HS256")}'
}
Expand Down
52 changes: 36 additions & 16 deletions tests/server/auth/test_client.py
Original file line number Diff line number Diff line change
Expand Up @@ -17,14 +17,16 @@
from mlflow.utils.os import is_windows

from tests.helper_functions import random_str
from tests.server.auth.auth_test_utils import User, create_user
from tests.server.auth.auth_test_utils import (
ADMIN_PASSWORD,
ADMIN_USERNAME,
NEW_PERMISSION,
PERMISSION,
User,
create_user,
)
from tests.tracking.integration_test_utils import _init_server

PERMISSION = "READ"
NEW_PERMISSION = "EDIT"
ADMIN_USERNAME = auth_config.admin_username
ADMIN_PASSWORD = auth_config.admin_password


@pytest.fixture(autouse=True)
def clear_credentials(monkeypatch):
Expand Down Expand Up @@ -70,18 +72,29 @@ def test_get_client():
assert isinstance(client, AuthServiceClient)


def test_create_user(client):
def test_create_user(client, monkeypatch):
username = random_str()
password = random_str()
user = client.create_user(username, password)

with assert_unauthenticated():
client.create_user(username, password)

with User(ADMIN_USERNAME, ADMIN_PASSWORD, monkeypatch):
user = client.create_user(username, password)
assert user.username == username
assert user.is_admin is False

username2 = random_str()
password2 = random_str()
with User(username, password, monkeypatch), assert_unauthorized():
client.create_user(username2, password2)


def test_get_user(client, monkeypatch):
username = random_str()
password = random_str()
client.create_user(username, password)
with User(ADMIN_USERNAME, ADMIN_PASSWORD, monkeypatch):
client.create_user(username, password)

with User(ADMIN_USERNAME, ADMIN_PASSWORD, monkeypatch):
user = client.get_user(username)
Expand All @@ -92,15 +105,17 @@ def test_get_user(client, monkeypatch):

username2 = random_str()
password2 = random_str()
client.create_user(username2, password2)
with User(ADMIN_USERNAME, ADMIN_PASSWORD, monkeypatch):
client.create_user(username2, password2)
with User(username2, password2, monkeypatch), assert_unauthorized():
client.get_user(username)


def test_update_user_password(client, monkeypatch):
username = random_str()
password = random_str()
client.create_user(username, password)
with User(ADMIN_USERNAME, ADMIN_PASSWORD, monkeypatch):
client.create_user(username, password)

new_password = random_str()
with User(ADMIN_USERNAME, ADMIN_PASSWORD, monkeypatch):
Expand All @@ -117,15 +132,17 @@ def test_update_user_password(client, monkeypatch):

username2 = random_str()
password2 = random_str()
client.create_user(username2, password2)
with User(ADMIN_USERNAME, ADMIN_PASSWORD, monkeypatch):
client.create_user(username2, password2)
with User(username2, password2, monkeypatch), assert_unauthorized():
client.update_user_password(username, new_password)


def test_update_user_admin(client, monkeypatch):
username = random_str()
password = random_str()
client.create_user(username, password)
with User(ADMIN_USERNAME, ADMIN_PASSWORD, monkeypatch):
client.create_user(username, password)

with User(ADMIN_USERNAME, ADMIN_PASSWORD, monkeypatch):
client.update_user_admin(username, True)
Expand All @@ -137,15 +154,17 @@ def test_update_user_admin(client, monkeypatch):

username2 = random_str()
password2 = random_str()
client.create_user(username2, password2)
with User(ADMIN_USERNAME, ADMIN_PASSWORD, monkeypatch):
client.create_user(username2, password2)
with User(username2, password2, monkeypatch), assert_unauthorized():
client.update_user_admin(username, True)


def test_delete_user(client, monkeypatch):
username = random_str()
password = random_str()
client.create_user(username, password)
with User(ADMIN_USERNAME, ADMIN_PASSWORD, monkeypatch):
client.create_user(username, password)

with User(ADMIN_USERNAME, ADMIN_PASSWORD, monkeypatch):
client.update_user_admin(username, True)
Expand All @@ -162,7 +181,8 @@ def test_delete_user(client, monkeypatch):

username2 = random_str()
password2 = random_str()
client.create_user(username2, password2)
with User(ADMIN_USERNAME, ADMIN_PASSWORD, monkeypatch):
client.create_user(username2, password2)
with User(username2, password2, monkeypatch), assert_unauthorized():
client.delete_user(username)

Expand Down

0 comments on commit 32de215

Please sign in to comment.