Adds hmac signing of the invocation result file (#4366)

See DIAGNijmegen/rse-roadmap#443

Author: jmsmkn

app/grandchallenge/algorithms/migrations/0081_job_signing_key.py

@@ -0,0 +1,22 @@
+# Generated by Django 4.2.25 on 2025-10-20 14:20
+
+from django.db import migrations, models
+
+
+class Migration(migrations.Migration):
+
+    dependencies = [
+        ("algorithms", "0080_alter_algorithmimage_options_and_more"),
+    ]
+
+    operations = [
+        migrations.AddField(
+            model_name="job",
+            name="signing_key",
+            field=models.BinaryField(
+                help_text="The key used to sign the inference result file",
+                max_length=32,
+                null=True,
+            ),
+        ),
+    ]

app/grandchallenge/algorithms/migrations/0082_auto_20251020_1421.py

@@ -0,0 +1,25 @@
+# Generated by Django 4.2.25 on 2025-10-20 14:21
+
+from django.db import migrations
+
+
+class Migration(migrations.Migration):
+
+    dependencies = [
+        ("algorithms", "0081_job_signing_key"),
+    ]
+
+    operations = [
+        migrations.RunSQL(
+            sql="CREATE EXTENSION IF NOT EXISTS pgcrypto;",
+            reverse_sql="",
+            elidable=True,
+        ),
+        migrations.RunSQL(
+            sql=(
+                "UPDATE algorithms_job SET signing_key = gen_random_bytes(32) WHERE signing_key IS NULL;"
+            ),
+            reverse_sql="UPDATE algorithms_job SET signing_key = NULL;",
+            elidable=True,
+        ),
+    ]

app/grandchallenge/algorithms/migrations/0083_alter_job_signing_key.py

@@ -0,0 +1,25 @@
+# Generated by Django 4.2.25 on 2025-10-20 14:21
+
+import secrets
+
+from django.db import migrations, models
+
+
+class Migration(migrations.Migration):
+
+    dependencies = [
+        ("algorithms", "0082_auto_20251020_1421"),
+    ]
+
+    operations = [
+        migrations.AlterField(
+            model_name="job",
+            name="signing_key",
+            field=models.BinaryField(
+                default=secrets.token_bytes,
+                help_text="The key used to sign the inference result file",
+                max_length=32,
+                unique=True,
+            ),
+        ),
+    ]

app/grandchallenge/challenges/models.py

@@ -1292,6 +1292,7 @@ def compute_costs_euro_cents_per_hour(self):
                 time_limit=self.inference_time_limit_in_minutes,
                 requires_gpu_type=gpu_type,
                 use_warm_pool=False,
+                signing_key=b"",
             )
             for gpu_type in self.algorithm_selectable_gpu_type_choices
         ]

app/grandchallenge/components/backends/base.py

@@ -1,9 +1,13 @@
 import asyncio
+import binascii
 import functools
+import hashlib
+import hmac
 import io
 import json
 import logging
 import os
+import secrets
 from abc import ABC, abstractmethod
 from json import JSONDecodeError
 from math import ceil
@@ -272,6 +276,7 @@ def __init__(
         time_limit: int,
         requires_gpu_type: GPUTypeChoices,
         use_warm_pool: bool,
+        signing_key: bytes,
         algorithm_model=None,
         ground_truth=None,
         **kwargs,
@@ -285,6 +290,7 @@ def __init__(
         self._use_warm_pool = (
             use_warm_pool and settings.COMPONENTS_USE_WARM_POOL
         )
+        self._signing_key = signing_key
         self._stdout = []
         self._stderr = []
         self.__s3_client = None
@@ -388,15 +394,21 @@ def invocation_environment(self):
             "GRAND_CHALLENGE_COMPONENT_MAX_MEMORY_MB": str(
                 self._max_memory_mb
             ),
+            "GRAND_CHALLENGE_COMPONENT_SIGNING_KEY_HEX": binascii.hexlify(
+                self._signing_key
+            ).decode("ascii"),
         }
+
         if self._algorithm_model:
             env["GRAND_CHALLENGE_COMPONENT_MODEL"] = (
                 f"s3://{settings.COMPONENTS_INPUT_BUCKET_NAME}/{self._algorithm_model_key}"
             )
+
         if self._ground_truth:
             env["GRAND_CHALLENGE_COMPONENT_GROUND_TRUTH"] = (
                 f"s3://{settings.COMPONENTS_INPUT_BUCKET_NAME}/{self._ground_truth_key}"
             )
+
         return env
 
     @property
@@ -770,43 +782,57 @@ def _get_upload_input_content_task(*, content, key):
         )
 
     def _get_task_return_code(self):
-        with io.BytesIO() as fileobj:
-            try:
-                self._s3_client.download_fileobj(
-                    Fileobj=fileobj,
-                    Bucket=settings.COMPONENTS_OUTPUT_BUCKET_NAME,
-                    Key=self._result_key,
-                )
-            except botocore.exceptions.ClientError as error:
-                if error.response["Error"]["Code"] == "404":
-                    raise UncleanExit(
-                        "The invocation request did not return a result"
-                    ) from error
-                else:
-                    raise
+        try:
+            response = self._s3_client.get_object(
+                Bucket=settings.COMPONENTS_OUTPUT_BUCKET_NAME,
+                Key=self._result_key,
+            )
+        except botocore.exceptions.ClientError as error:
+            if error.response["Error"]["Code"] == "404":
+                raise UncleanExit(
+                    "The invocation request did not return a result"
+                ) from error
+            else:
+                raise
 
-            fileobj.seek(0)
+        body = response["Body"].read()
 
-            try:
-                result = json.loads(
-                    fileobj.read().decode("utf-8"),
-                )
-            except JSONDecodeError:
-                raise ComponentException(
-                    "The invocation request did not return valid json"
-                )
+        signature_hmac_sha256 = response["Metadata"].get(
+            "signature_hmac_sha256"
+        )
+        body_signature_hmac_sha256 = hmac.new(
+            key=self._signing_key, msg=body, digestmod=hashlib.sha256
+        ).hexdigest()
+
+        if signature_hmac_sha256 and not secrets.compare_digest(
+            body_signature_hmac_sha256, signature_hmac_sha256
+        ):
+            # TODO The signature should always be present when all images use sagemaker shim >= 0.5.0
+            logger.error(
+                "The invocation response object has been tampered with"
+            )
+            raise ComponentException(
+                "The invocation response object has been tampered with"
+            )
 
-            logger.info(f"{result=}")
+        try:
+            result = json.loads(body.decode("utf-8"))
+        except JSONDecodeError:
+            raise ComponentException(
+                "The invocation request did not return valid json"
+            )
 
-            if result["pk"] != self._job_id:
-                raise RuntimeError("Wrong result key for this job")
+        logger.info(f"{result=}")
 
-            try:
-                return int(result["return_code"])
-            except (KeyError, ValueError):
-                raise ComponentException(
-                    "The invocation response object is not valid"
-                )
+        if result["pk"] != self._job_id:
+            raise RuntimeError("Wrong result key for this job")
+
+        try:
+            return int(result["return_code"])
+        except (KeyError, ValueError):
+            raise ComponentException(
+                "The invocation response object is not valid"
+            )
 
     def _handle_completed_job(self):
         users_process_exit_code = self._get_task_return_code()

app/grandchallenge/components/models.py

@@ -1,6 +1,7 @@
 import json
 import logging
 import re
+import secrets
 from enum import Enum
 from json import JSONDecodeError
 from pathlib import Path
@@ -1628,6 +1629,13 @@ class ComponentJob(FieldChangeMixin, UUIDModel):
             "/input/foo/bar/<relative_path>"
         ),
     )
+    signing_key = models.BinaryField(
+        max_length=32,
+        default=secrets.token_bytes,
+        unique=True,
+        editable=False,
+        help_text="The key used to sign the inference result file",
+    )
     task_on_success = models.JSONField(
         default=None,
         null=True,
@@ -1688,6 +1696,7 @@ def save(self, *args, **kwargs):
                 "requires_gpu_type",
                 "requires_memory_gb",
                 "time_limit",
+                "signing_key",
             ):
                 if self.has_changed(field):
                     raise ValueError(f"{field} cannot be changed")
@@ -1757,6 +1766,7 @@ def executor_kwargs(self):
             "requires_gpu_type": self.requires_gpu_type,
             "memory_limit": self.requires_memory_gb,
             "use_warm_pool": self.use_warm_pool,
+            "signing_key": self.signing_key,
         }
 
     def get_executor(self, *, backend):

app/grandchallenge/evaluation/migrations/0098_evaluation_signing_key.py

@@ -0,0 +1,22 @@
+# Generated by Django 4.2.25 on 2025-10-20 14:20
+
+from django.db import migrations, models
+
+
+class Migration(migrations.Migration):
+
+    dependencies = [
+        ("evaluation", "0097_alter_evaluationgroundtruth_options"),
+    ]
+
+    operations = [
+        migrations.AddField(
+            model_name="evaluation",
+            name="signing_key",
+            field=models.BinaryField(
+                help_text="The key used to sign the inference result file",
+                max_length=32,
+                null=True,
+            ),
+        ),
+    ]

app/grandchallenge/evaluation/migrations/0099_auto_20251020_1420.py

@@ -0,0 +1,25 @@
+# Generated by Django 4.2.25 on 2025-10-20 14:20
+
+from django.db import migrations
+
+
+class Migration(migrations.Migration):
+
+    dependencies = [
+        ("evaluation", "0098_evaluation_signing_key"),
+    ]
+
+    operations = [
+        migrations.RunSQL(
+            sql="CREATE EXTENSION IF NOT EXISTS pgcrypto;",
+            reverse_sql="",
+            elidable=True,
+        ),
+        migrations.RunSQL(
+            sql=(
+                "UPDATE evaluation_evaluation SET signing_key = gen_random_bytes(32) WHERE signing_key IS NULL;"
+            ),
+            reverse_sql="UPDATE evaluation_evaluation SET signing_key = NULL;",
+            elidable=True,
+        ),
+    ]

app/grandchallenge/evaluation/migrations/0100_alter_evaluation_signing_key.py

@@ -0,0 +1,25 @@
+# Generated by Django 4.2.25 on 2025-10-20 14:21
+
+import secrets
+
+from django.db import migrations, models
+
+
+class Migration(migrations.Migration):
+
+    dependencies = [
+        ("evaluation", "0099_auto_20251020_1420"),
+    ]
+
+    operations = [
+        migrations.AlterField(
+            model_name="evaluation",
+            name="signing_key",
+            field=models.BinaryField(
+                default=secrets.token_bytes,
+                help_text="The key used to sign the inference result file",
+                max_length=32,
+                unique=True,
+            ),
+        ),
+    ]