Slurm: Validate token is a valid JWT

This avoids a situation we encountered where the filename was being
used as a JWT. The SLURM-REST api server apparently assumed _any_
token being passed was valid, and so forwarded it to the backend.

Author: ndevenish

src/zocalo/util/slurm/__init__.py

@@ -1,5 +1,8 @@
 from __future__ import annotations
 
+import base64
+import binascii
+import json
 import os
 import pathlib
 from typing import Any
@@ -11,6 +14,24 @@
 from . import models
 
 
+def validate_is_jwt(token: str) -> bool:
+    """Checks that a particular string is a JWT token"""
+    if token.count(".") != 2:
+        return False
+    header, payload, _ = token.split(".")
+    try:
+        # Check both header and payload are valid base64-encoded json objects
+        if not (
+            isinstance(json.loads(base64.b64decode(header, validate=True)), dict)
+            and isinstance(json.loads(base64.b64decode(payload, validate=True)), dict)
+        ):
+            return False
+    except (binascii.Error, json.JSONDecodeError):
+        return False
+
+    return True
+
+
 class SlurmRestApi:
     def __init__(
         self,
@@ -25,8 +46,15 @@ def __init__(
         if user_token and os.path.isfile(user_token):
             with open(user_token, "r") as f:
                 self.user_token = f.read().strip()
+        elif isinstance(user_token, pathlib.Path):
+            # We got passed a path, but it isn't a valid one
+            raise RuntimeError(f"SLURM: API token file {user_token} does not exist")
         else:
             assert isinstance(user_token, str)
+            if not validate_is_jwt(user_token):
+                raise RuntimeError(
+                    "SLURM user_token does not appear to be a valid JWT token. Did you pass a nonexistent filename?"
+                )
             self.user_token = user_token
         self.session = requests.Session()
         if self.user_name:

tests/util/test_slurm.py

@@ -244,6 +244,25 @@ def test_get_slurm_api_user_token_external_file(tmp_path):
     assert api.user_token == "foobar"
 
 
+def test_get_slurm_api_invalid_external_file(tmp_path):
+    """Check that passing invalid file doesn't get interpreted as a token"""
+    user_token_file = tmp_path / "bad-slurm-user-token"
+    with pytest.raises(RuntimeError):
+        slurm.SlurmRestApi(
+            url="http://slurm.example.com:1234",
+            version="v0.0.40",
+            user_name="foo",
+            user_token=user_token_file,
+        )
+    with pytest.raises(RuntimeError):
+        slurm.SlurmRestApi(
+            url="http://slurm.example.com:1234",
+            version="v0.0.40",
+            user_name="foo",
+            user_token=str(user_token_file),
+        )
+
+
 def test_get_jobs(requests_mock, slurm_api, jobs_response):
     requests_mock.get(
         "/slurm/v0.0.40/jobs",