From eb0972f9005b9013d057ed9451b4b0f8352a316c Mon Sep 17 00:00:00 2001 From: Varun Valada Date: Tue, 10 Dec 2024 12:21:41 -0600 Subject: [PATCH] Add extended reservation checks in the server --- server/src/api/v1.py | 36 ++++++++++++++++- server/tests/conftest.py | 2 + server/tests/test_v1_authorization.py | 58 +++++++++++++++++++++++++++ 3 files changed, 94 insertions(+), 2 deletions(-) diff --git a/server/src/api/v1.py b/server/src/api/v1.py index 346c8b1c..b8255c94 100644 --- a/server/src/api/v1.py +++ b/server/src/api/v1.py @@ -157,8 +157,29 @@ def check_token_queue(auth_token: str, secret_key: str, queue: str) -> bool: return queue in allowed_queues +def check_token_reservation_timeout( + auth_token: str, secret_key: str, reservation_timeout: int, queue: str +) -> bool: + """ + Checks if the requested reservation is either less than the max + or that their token gives them the permission to use a higher one + """ + # Max reservation time defaults to 6 hours + max_reservation_time = 6 * 60 * 60 + if reservation_timeout <= max_reservation_time: + return True + decoded_jwt = decode_jwt_token(auth_token, secret_key) + max_reservation_time_dict = decoded_jwt.get("max_reservation_time", {}) + max_reservation_time = max_reservation_time_dict.get(queue, 0) + return reservation_timeout <= max_reservation_time + + def check_token_permissions( - auth_token: str, secret_key: str, priority: int, queue: str + auth_token: str, + secret_key: str, + priority: int, + queue: str, + reservation_timeout: int, ) -> bool: """ Validates token received from client and checks if it can @@ -168,7 +189,10 @@ def check_token_permissions( auth_token, secret_key, queue, priority ) queue_allowed = check_token_queue(auth_token, secret_key, queue) - return priority_allowed and queue_allowed + reservation_time_allowed = check_token_reservation_timeout( + auth_token, secret_key, reservation_timeout, queue + ) + return priority_allowed and queue_allowed and reservation_time_allowed def job_builder(data: dict, auth_token: str): @@ -196,11 +220,15 @@ def job_builder(data: dict, auth_token: str): priority_level = data.get("job_priority", 0) job_queue = data["job_queue"] + reserve_data = data.get("reserve_data", {}) + # default reservation timeout is 1 hour + reservation_timeout = reserve_data.get("timeout", 3600) allowed = check_token_permissions( auth_token, os.environ.get("JWT_SIGNING_KEY"), priority_level, job_queue, + reservation_timeout, ) if not allowed: abort( @@ -766,6 +794,10 @@ def generate_token(allowed_resources, secret_key): token_payload["max_priority"] = allowed_resources["max_priority"] if "allowed_queues" in allowed_resources: token_payload["allowed_queues"] = allowed_resources["allowed_queues"] + if "max_reservation_time" in allowed_resources: + token_payload["max_reservation_time"] = allowed_resources[ + "max_reservation_time" + ] token = jwt.encode(token_payload, secret_key, algorithm="HS256") return token diff --git a/server/tests/conftest.py b/server/tests/conftest.py index 849ec9a3..0a2f5898 100644 --- a/server/tests/conftest.py +++ b/server/tests/conftest.py @@ -84,12 +84,14 @@ def mongo_app_with_permissions(mongo_app): "myqueue2": 200, } allowed_queues = ["rqueue1", "rqueue2"] + max_reservation_time = {"myqueue": 30000} mongo.client_permissions.insert_one( { "client_id": client_id, "client_secret_hash": client_key_hash, "max_priority": max_priority, "allowed_queues": allowed_queues, + "max_reservation_time": max_reservation_time, } ) restricted_queues = [ diff --git a/server/tests/test_v1_authorization.py b/server/tests/test_v1_authorization.py index b05ebd08..7dddd349 100644 --- a/server/tests/test_v1_authorization.py +++ b/server/tests/test_v1_authorization.py @@ -316,3 +316,61 @@ def test_restricted_queue_reject_no_token(mongo_app_with_permissions): job = {"job_queue": "rqueue1"} job_response = app.post("/v1/job", json=job) assert 401 == job_response.status_code + + +def test_extended_reservation_allowed(mongo_app_with_permissions): + """ + Tests that jobs that include extended reservation are accepted when + the token gives them permission + """ + app, _, client_id, client_key, _ = mongo_app_with_permissions + authenticate_output = app.post( + "/v1/oauth2/token", + headers=create_auth_header(client_id, client_key), + ) + token = authenticate_output.data.decode("utf-8") + job = {"job_queue": "myqueue", "reserve_data": {"timeout": 30000}} + job_response = app.post( + "/v1/job", json=job, headers={"Authorization": token} + ) + assert 200 == job_response.status_code + + +def test_extended_reservation_rejected(mongo_app_with_permissions): + """ + Tests that jobs that include extended reservation are rejected when + the token does not give them permission + """ + app, _, client_id, client_key, _ = mongo_app_with_permissions + authenticate_output = app.post( + "/v1/oauth2/token", + headers=create_auth_header(client_id, client_key), + ) + token = authenticate_output.data.decode("utf-8") + job = {"job_queue": "myqueue2", "reserve_data": {"timeout": 21601}} + job_response = app.post( + "/v1/job", json=job, headers={"Authorization": token} + ) + assert 403 == job_response.status_code + + +def test_extended_reservation_reject_no_token(mongo_app_with_permissions): + """ + Tests that jobs that included extended reservation are rejected + when no token is included + """ + app, _, _, _, _ = mongo_app_with_permissions + job = {"job_queue": "myqueue", "reserve_data": {"timeout": 21601}} + job_response = app.post("/v1/job", json=job) + assert 401 == job_response.status_code + + +def test_normal_reservation_no_token(mongo_app): + """ + Tests that jobs that include reservation times less than the maximum + are accepted when no token is included + """ + app, _ = mongo_app + job = {"job_queue": "myqueue", "reserve_data": {"timeout": 21600}} + job_response = app.post("/v1/job", json=job) + assert 200 == job_response.status_code