From ab46b35b7209c4e0404d5d540a602635f734403b Mon Sep 17 00:00:00 2001 From: Varun Valada Date: Mon, 19 Aug 2024 17:47:48 -0500 Subject: [PATCH 1/3] Added logic to validate token and check priority permissions --- docs/reference/job-schema.rst | 3 ++ server/src/api/schemas.py | 1 + server/src/api/v1.py | 56 ++++++++++++++++++++++++++++++++--- server/tests/test_v1.py | 39 +++++++++++++++++++++++- 4 files changed, 94 insertions(+), 5 deletions(-) diff --git a/docs/reference/job-schema.rst b/docs/reference/job-schema.rst index f5d57b46..acd11432 100644 --- a/docs/reference/job-schema.rst +++ b/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 GET request to /v1/authenticate/ with a valid client_id and client-key. Example jobs in YAML ---------------------------- diff --git a/server/src/api/schemas.py b/server/src/api/schemas.py index 411c381b..13722f76 100644 --- a/server/src/api/schemas.py +++ b/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): diff --git a/server/src/api/v1.py b/server/src/api/v1.py index 7068c94b..5e4bb89f 100644 --- a/server/src/api/v1.py +++ b/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,33 @@ 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", "max_priority"]}, + ) + except jwt.exceptions.ExpiredSignatureError: + abort(403, "Token has expired") + except jwt.exceptions.InvalidTokenError: + abort(403, "Invalid Token") + + max_priority = decoded_jwt["max_priority"].get(queue, 0) + 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 +149,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, + 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 +722,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( { diff --git a/server/tests/test_v1.py b/server/tests/test_v1.py index 1ac40bf0..e4e4aab2 100644 --- a/server/tests/test_v1.py +++ b/server/tests/test_v1.py @@ -781,9 +781,46 @@ 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_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 From b5988fb5babc18ec224c0019922a3c0ea4780932 Mon Sep 17 00:00:00 2001 From: Varun Valada Date: Mon, 7 Oct 2024 11:15:04 -0500 Subject: [PATCH 2/3] Addressed comments and added star priority option --- docs/reference/job-schema.rst | 2 +- server/src/api/v1.py | 36 ++++++++++++++++++---------------- server/tests/test_v1.py | 37 ++++++++++++++++++++++++++++++++++- 3 files changed, 56 insertions(+), 19 deletions(-) diff --git a/docs/reference/job-schema.rst b/docs/reference/job-schema.rst index acd11432..f02a2ddd 100644 --- a/docs/reference/job-schema.rst +++ b/docs/reference/job-schema.rst @@ -54,7 +54,7 @@ The following table lists the key elements that a job definition file should con - | (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 GET request to /v1/authenticate/ with a valid client_id and client-key. + - | (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 ---------------------------- diff --git a/server/src/api/v1.py b/server/src/api/v1.py index 5e4bb89f..8f3ec5f7 100644 --- a/server/src/api/v1.py +++ b/server/src/api/v1.py @@ -100,7 +100,7 @@ def has_attachments(data: dict) -> bool: ) -def check_token_permission( +def check_token_priority_permission( auth_token: str, secret_key: str, priority: int, queue: str ) -> bool: """ @@ -108,22 +108,24 @@ def check_token_permission( 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", "max_priority"]}, - ) - except jwt.exceptions.ExpiredSignatureError: - abort(403, "Token has expired") - except jwt.exceptions.InvalidTokenError: - abort(403, "Invalid Token") + 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 = decoded_jwt["max_priority"].get(queue, 0) - return max_priority >= priority + 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): @@ -152,7 +154,7 @@ def job_builder(data: dict, auth_token: str): if "job_priority" in data: priority_level = data["job_priority"] job_queue = data["job_queue"] - allowed = check_token_permission( + allowed = check_token_priority_permission( auth_token, os.environ.get("JWT_SIGNING_KEY"), priority_level, diff --git a/server/tests/test_v1.py b/server/tests/test_v1.py index e4e4aab2..ee3c2aa8 100644 --- a/server/tests/test_v1.py +++ b/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 @@ -824,3 +824,38 @@ def test_priority_invalid_queue(mongo_app_with_permissions): "/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 From 67a725acc34fe90b9ebfc0c0ffb4fbe5834338a5 Mon Sep 17 00:00:00 2001 From: Varun Valada Date: Mon, 7 Oct 2024 11:29:44 -0500 Subject: [PATCH 3/3] Added test for star priority --- server/tests/conftest.py | 1 + server/tests/test_v1.py | 18 ++++++++++++++++++ 2 files changed, 19 insertions(+) diff --git a/server/tests/conftest.py b/server/tests/conftest.py index b44709f1..2b298b25 100644 --- a/server/tests/conftest.py +++ b/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, } diff --git a/server/tests/test_v1.py b/server/tests/test_v1.py index ee3c2aa8..3cc6dde8 100644 --- a/server/tests/test_v1.py +++ b/server/tests/test_v1.py @@ -803,6 +803,24 @@ def test_job_with_priority(mongo_app_with_permissions): 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