From 760661fd454bd39de8c95996abd12173eeaf8b1b Mon Sep 17 00:00:00 2001
From: Nicholas Devenish <ndevenish@gmail.com>
Date: Thu, 13 Jun 2024 13:49:04 +0100
Subject: [PATCH] Slurm JWT Checks: Fix tests

The previous change in a1a9e9d1 meant that the test data now failed.
---
 src/zocalo/util/slurm/__init__.py | 19 +++++++++++++++++--
 tests/util/test_slurm.py          |  7 +++++--
 2 files changed, 22 insertions(+), 4 deletions(-)

diff --git a/src/zocalo/util/slurm/__init__.py b/src/zocalo/util/slurm/__init__.py
index bc8b51a..7b07dd0 100644
--- a/src/zocalo/util/slurm/__init__.py
+++ b/src/zocalo/util/slurm/__init__.py
@@ -21,9 +21,24 @@ def validate_is_jwt(token: str) -> bool:
     header, payload, _ = token.split(".")
     try:
         # Check both header and payload are valid base64-encoded json objects
+        # Note that JWT are Base64URL, which might not have padding.
         if not (
-            isinstance(json.loads(base64.b64decode(header, validate=True)), dict)
-            and isinstance(json.loads(base64.b64decode(payload, validate=True)), dict)
+            isinstance(
+                json.loads(
+                    base64.urlsafe_b64decode(
+                        header + "=" * (4 - len(header) % 4)
+                    ).decode()
+                ),
+                dict,
+            )
+            and isinstance(
+                json.loads(
+                    base64.urlsafe_b64decode(
+                        payload + "=" * (4 - len(payload) % 4)
+                    ).decode()
+                ),
+                dict,
+            )
         ):
             return False
     except (binascii.Error, json.JSONDecodeError):
diff --git a/tests/util/test_slurm.py b/tests/util/test_slurm.py
index 8e79f5c..14215cc 100644
--- a/tests/util/test_slurm.py
+++ b/tests/util/test_slurm.py
@@ -4,6 +4,9 @@
 import zocalo.configuration
 from zocalo.util import slurm
 
+# A sample (valid but not useful) JWT token
+SAMPLE_JWT_TOKEN = "eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJzdWIiOiIxMjM0NTY3ODkwIiwibmFtZSI6IkpvaG4gRG9lIiwiaWF0IjoxNTE2MjM5MDIyfQ.SflKxwRJSMeKKF2QT4fwpMeJf36POk6yJV_adQssw5c"
+
 
 @pytest.fixture
 def zocalo_configuration(mocker):
@@ -11,7 +14,7 @@ def zocalo_configuration(mocker):
     zc.slurm = {
         "url": "http://slurm.example.com:1234",
         "user": "foo",
-        "user_token": "sometoken",
+        "user_token": SAMPLE_JWT_TOKEN,
         "api_version": "v0.0.40",
     }
     return zc
@@ -229,7 +232,7 @@ def test_get_slurm_api_from_zocalo_configuration(slurm_api):
     assert slurm_api.url == "http://slurm.example.com:1234"
     assert slurm_api.version == "v0.0.40"
     assert slurm_api.user_name == "foo"
-    assert slurm_api.user_token == "sometoken"
+    assert slurm_api.user_token == SAMPLE_JWT_TOKEN
 
 
 def test_get_slurm_api_user_token_external_file(tmp_path):