Merge pull request #361 from canonical/CERTTF-371

Add logic to validate token and check priority permissions

docs/reference/job-schema.rst

@@ -52,6 +52,9 @@ The following table lists the key elements that a job definition file should con
     - string
     - /
     - | (Optional) URL to send job status updates to. These updates originate from the agent and get posted to the server which then posts the update to the webhook. If no webhook is specified, these updates will not be generated.
+  * - ``job_priority``
+    - integer
+    - | (Optional) Integer specifying how much priority this job has. Jobs with higher job_priority will be selected by agents before other jobs. Specifying a job priority requires authorization in the form of a JWT obtained by sending a POST request to /v1/oauth2/token with a client id and client key specified in an Authorization header.
 
 Example jobs in YAML
 ----------------------------

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

@@ -20,7 +20,6 @@
 import os
 import uuid
 from datetime import datetime, timezone, timedelta
-
 import pkg_resources
 from apiflask import APIBlueprint, abort
 from flask import jsonify, request, send_file
@@ -75,9 +74,9 @@ def job_post(json_data: dict):
         job_queue = ""
     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 +100,35 @@ def has_attachments(data: dict) -> bool:
     )
 
 
-def job_builder(data):
+def check_token_priority_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, "Unauthorized")
+    try:
+        decoded_jwt = jwt.decode(
+            auth_token,
+            secret_key,
+            algorithms="HS256",
+            options={"require": ["exp", "iat", "sub"]},
+        )
+    except jwt.exceptions.ExpiredSignatureError:
+        abort(403, "Token has expired")
+    except jwt.exceptions.InvalidTokenError:
+        abort(403, "Invalid Token")
+
+    max_priority_dict = decoded_jwt.get("max_priority", {})
+    star_priority = max_priority_dict.get("*", 0)
+    queue_priority = max_priority_dict.get(queue, 0)
+    max_priority = max(star_priority, queue_priority)
+    return max_priority >= priority
+
+
+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 +151,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_priority_permission(
+            auth_token,
+            os.environ.get("JWT_SIGNING_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
@@ -676,6 +724,8 @@ def validate_client_key_pair(client_id: str, client_key: str):
     """
     Checks client_id and key pair for validity and returns their permissions
     """
+    if client_key is None:
+        return None
     client_key_bytes = client_key.encode("utf-8")
     client_permissions_entry = database.mongo.db.client_permissions.find_one(
         {

server/tests/conftest.py

@@ -79,6 +79,7 @@ def mongo_app_with_permissions(mongo_app):
     ).decode("utf-8")
 
     max_priority = {
+        "*": 1,
         "myqueue": 100,
         "myqueue2": 200,
     }

server/tests/test_v1.py

@@ -17,7 +17,7 @@
 Unit tests for Testflinger v1 API
 """
 
-from datetime import datetime
+from datetime import datetime, timedelta
 from io import BytesIO
 import json
 import os
@@ -781,9 +781,99 @@ def test_retrieve_token_invalid_client_key(mongo_app_with_permissions):
     """
     app, _, client_id, _, _ = mongo_app_with_permissions
     client_key = "my_wrong_key"
-
     output = app.post(
         "/v1/oauth2/token",
         headers=create_auth_header(client_id, client_key),
     )
     assert output.status_code == 401
+
+
+def test_job_with_priority(mongo_app_with_permissions):
+    """Tests submission of priority job with valid token"""
+    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", "job_priority": 200}
+    job_response = app.post(
+        "/v1/job", json=job, headers={"Authorization": token}
+    )
+    assert 200 == job_response.status_code
+
+
+def test_star_priority(mongo_app_with_permissions):
+    """
+    Tests submission of priority job for a generic queue
+    with star priority permissions
+    """
+    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": "mygenericqueue", "job_priority": 1}
+    job_response = app.post(
+        "/v1/job", json=job, headers={"Authorization": token}
+    )
+    assert 200 == job_response.status_code
+
+
+def test_priority_no_token(mongo_app_with_permissions):
+    """Tests rejection of priority job with no token"""
+    app, _, _, _, _ = mongo_app_with_permissions
+    job = {"job_queue": "myqueue2", "job_priority": 200}
+    job_response = app.post("/v1/job", json=job)
+    assert 401 == job_response.status_code
+
+
+def test_priority_invalid_queue(mongo_app_with_permissions):
+    """Tests rejection of priority job with invalid queue"""
+    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": "myinvalidqueue", "job_priority": 200}
+    job_response = app.post(
+        "/v1/job", json=job, headers={"Authorization": token}
+    )
+    assert 403 == job_response.status_code
+
+
+def test_priority_expired_token(mongo_app_with_permissions):
+    """Tests rejection of priority job with expired token"""
+    app, _, _, _, _ = mongo_app_with_permissions
+    secret_key = os.environ.get("JWT_SIGNING_KEY")
+    expired_token_payload = {
+        "exp": datetime.utcnow() - timedelta(seconds=2),
+        "iat": datetime.utcnow() - timedelta(seconds=4),
+        "sub": "access_token",
+        "max_priority": {},
+    }
+    token = jwt.encode(expired_token_payload, secret_key, algorithm="HS256")
+    job = {"job_queue": "myqueue", "job_priority": 100}
+    job_response = app.post(
+        "/v1/job", json=job, headers={"Authorization": token}
+    )
+    assert 403 == job_response.status_code
+    assert "Token has expired" in job_response.text
+
+
+def test_missing_fields_in_token(mongo_app_with_permissions):
+    """Tests rejection of priority job with token with missing fields"""
+    app, _, _, _, _ = mongo_app_with_permissions
+    secret_key = os.environ.get("JWT_SIGNING_KEY")
+    incomplete_token_payload = {
+        "max_priority": {},
+    }
+    token = jwt.encode(incomplete_token_payload, secret_key, algorithm="HS256")
+    job = {"job_queue": "myqueue", "job_priority": 100}
+    job_response = app.post(
+        "/v1/job", json=job, headers={"Authorization": token}
+    )
+    assert 403 == job_response.status_code
+    assert "Invalid Token" in job_response.text