Added logic to validate token and check priority permissions

server/src/api/schemas.py

@@ -115,6 +115,7 @@ class Job(Schema):
     allocate_data = fields.Dict(required=False)
     reserve_data = fields.Dict(required=False)
     job_status_webhook = fields.String(required=False)
+    job_priority = fields.Integer(required=False)
 
 
 class JobId(Schema):

server/src/api/v1.py

@@ -76,8 +76,9 @@ def job_post(json_data: dict):
     if not job_queue:
         abort(422, message="Invalid data or no job_queue specified")
 
+    auth_token = request.headers.get("Authorization")
     try:
-        job = job_builder(json_data)
+        job = job_builder(json_data, auth_token)
     except ValueError:
         abort(400, message="Invalid job_id specified")
 
@@ -101,7 +102,40 @@ def has_attachments(data: dict) -> bool:
     )
 
 
-def job_builder(data):
+def check_token_permission(
+    auth_token: str, secret_key: str, priority: int, queue: str
+) -> bool:
+    """
+    Validates token received from client and checks if it can
+    push a job to the queue with the requested priority
+    """
+    if auth_token is None:
+        abort(401, "No authentication token specified")
+    else:
+        try:
+            decoded_jwt = jwt.decode(
+                auth_token,
+                secret_key,
+                algorithms="HS256",
+                options={"require": ["exp", "iat", "sub", "permissions"]},
+            )
+        except jwt.exceptions.ExpiredSignatureError:
+            abort(403, "Token has expired")
+        # except jwt.exceptions.InvalidTokenError:
+        #   abort(403, "Invalid Token")
+
+        give_permission = False
+        for permission in decoded_jwt["permissions"]:
+            if (
+                permission["queue_name"] == queue
+                and permission["max_priority"] >= priority
+            ):
+                give_permission = True
+                break
+        return give_permission
+
+
+def job_builder(data: dict, auth_token: str):
     """Build a job from a dictionary of data"""
     job = {
         "created_at": datetime.now(timezone.utc),
@@ -124,6 +158,27 @@ def job_builder(data):
     if has_attachments(data):
         data["attachments_status"] = "waiting"
 
+    if "job_priority" in data:
+        priority_level = data["job_priority"]
+        job_queue = data["job_queue"]
+        allowed = check_token_permission(
+            auth_token,
+            SECRET_KEY,
+            priority_level,
+            job_queue,
+        )
+        if not allowed:
+            abort(
+                403,
+                (
+                    f"Not enough permissions to push to {job_queue}",
+                    f"with priority {priority_level}",
+                ),
+            )
+        job["job_priority"] = priority_level
+        data.pop("job_priority")
+    else:
+        job["job_priority"] = 0
     job["job_id"] = job_id
     job["job_data"] = data
     return job

server/tests/test_v1.py

@@ -777,3 +777,21 @@ def test_authenticate_invalid_client_key(mongo_app_with_permissions):
         headers={"client-key": client_key},
     )
     assert output.status_code == 401
+
+
+def test_job_with_priority(mongo_app_with_permissions):
+    """Tests authentication endpoint which returns JWT with permissions"""
+    app, _, client_id, client_key, _ = mongo_app_with_permissions
+    v1.SECRET_KEY = "my_secret_key"
+
+    authenticate_output = app.get(
+        f"/v1/authenticate/token/{client_id}",
+        headers={"client-key": client_key},
+    )
+    token = authenticate_output.data.decode("utf-8")
+    job = {"job_queue": "myqueue2", "job_priority": 200}
+    job_response = app.post(
+        "/v1/job", json=job, headers={"Authorization": token}
+    )
+    print(job_response.text)
+    assert 200 == job_response.status_code