Add extended reservation checks in the server

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(
@@ -759,6 +787,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
 

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 = [

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