From 770314a829be34e8dab76415e93c2420fc68eced Mon Sep 17 00:00:00 2001 From: Chris Burr Date: Mon, 19 Oct 2026 13:35:03 +0200 Subject: [PATCH 1/3] fix (DMS): presigned URLs from the S3Gateway Accessing an S3 storage element through the S3Gateway was broken with S3 implementations which validate presigned requests: - the gateway looked up urls.get("Fields") instead of the fields of each URL, so the checksum metadata was never stored on upload - boto does not add the Fields of generate_presigned_post to the policy, so once passed S3 rejected them as "Extra input fields" - exists() sent a GET to a URL signed for HEAD, which is rejected as Forbidden --- src/DIRAC/DataManagementSystem/Service/S3GatewayHandler.py | 6 +++--- src/DIRAC/Resources/Storage/S3Storage.py | 7 +++++-- 2 files changed, 8 insertions(+), 5 deletions(-) diff --git a/src/DIRAC/DataManagementSystem/Service/S3GatewayHandler.py b/src/DIRAC/DataManagementSystem/Service/S3GatewayHandler.py index de9cdde6899..3a45262c5c0 100644 --- a/src/DIRAC/DataManagementSystem/Service/S3GatewayHandler.py +++ b/src/DIRAC/DataManagementSystem/Service/S3GatewayHandler.py @@ -150,9 +150,9 @@ def export_createPresignedUrl(self, storageName, s3_method, urls, expiration): failed[url] = "Permission denied" continue - res = returnSingleResult( - s3Plugin.createPresignedUrl({url: urls.get("Fields")}, s3_method, expiration=expiration) - ) + # For put_object, urls maps each url to the fields (metadata) of the file + fields = urls[url] if s3_method == "put_object" else None + res = returnSingleResult(s3Plugin.createPresignedUrl({url: fields}, s3_method, expiration=expiration)) log.debug(f"Presigned URL for {url}: {res}") if res["OK"]: diff --git a/src/DIRAC/Resources/Storage/S3Storage.py b/src/DIRAC/Resources/Storage/S3Storage.py index ac20e294c34..6551e9bfa8e 100644 --- a/src/DIRAC/Resources/Storage/S3Storage.py +++ b/src/DIRAC/Resources/Storage/S3Storage.py @@ -268,7 +268,8 @@ def _presigned_exists(self, urls): # and perform it with requests for url, presignedURL in presignedURLs.items(): try: - response = requests.get(presignedURL, timeout=30) + # The URL is signed for head_object so it must be used with HEAD + response = requests.head(presignedURL, timeout=30) if response.status_code == 200: successful[url] = True elif response.status_code == 404: # not found @@ -768,8 +769,10 @@ def createPresignedUrl(self, urls, s3_method, expiration=3600): fields = keys.get(key) if not isinstance(fields, dict): fields = None + # boto does not add the fields to the policy, without which S3 rejects them + conditions = [{k: v} for k, v in fields.items()] if fields else None response = self.s3_client.generate_presigned_post( - self.bucketName, key, Fields=fields, ExpiresIn=expiration + self.bucketName, key, Fields=fields, Conditions=conditions, ExpiresIn=expiration ) successful[key] = response From cef23fc60aa1e71d06984617121f54b05c15ce25 Mon Sep 17 00:00:00 2001 From: Chris Burr Date: Mon, 19 Oct 2026 13:36:37 +0200 Subject: [PATCH 2/3] test: use SeaweedFS instead of S3Mock for the CI S3 storage S3Mock is pinned to 3.12.0, which does not support presigned POST uploads and does not validate credentials or signatures. Use SeaweedFS, as DiracX does, with real credentials for the S3 storage elements. This also allows uploads through the S3Gateway to be tested again. --- tests/CI/docker-compose.yml | 22 ++++++++++++------- tests/CI/install_server.sh | 4 ++-- .../Resources/Storage/Test_Resources_S3.py | 6 ----- tests/Jenkins/dirac-cfg-update-server.py | 8 +++---- 4 files changed, 20 insertions(+), 20 deletions(-) diff --git a/tests/CI/docker-compose.yml b/tests/CI/docker-compose.yml index 77a076e8bfc..e40a5c9539a 100644 --- a/tests/CI/docker-compose.yml +++ b/tests/CI/docker-compose.yml @@ -100,20 +100,24 @@ services: }" > /jwks/iam-keystore.jwks' pull_policy: always - # Mock of an S3 storage + # S3 storage, used by the S3 storage elements and the DiracX sandbox store s3-direct: platform: linux/amd64 - # Fix the version until https://github.com/adobe/S3Mock/issues/2321 - # is resolved - image: adobe/s3mock:3.12.0 + image: chrislusf/seaweedfs:4.48 container_name: s3-direct hostname: s3-direct + command: ["mini", "-s3.port=9090", "-bucket=my-first-bucket", "-admin.ui=false"] ports: - 9090 - - 9191 environment: - - initialBuckets=my-first-bucket - - debug=true + - AWS_ACCESS_KEY_ID=console + - AWS_SECRET_ACCESS_KEY=console123 + healthcheck: + test: ["CMD", "curl", "-sf", "http://localhost:9090/healthz"] + interval: 5s + timeout: 2s + retries: 15 + start_period: 30s pull_policy: always diracx-wait-for-db: @@ -149,7 +153,7 @@ services: opensearch: condition: service_healthy s3-direct: - condition: service_started + condition: service_healthy iam-login-service: condition: service_healthy dirac-init-certificates: @@ -318,6 +322,8 @@ services: condition: service_completed_successfully diracx-wait-for-db: condition: service_completed_successfully + s3-direct: + condition: service_healthy volumes: - diracx-cs-store:/cs_store/ - diracx-key-store:/keystore/ diff --git a/tests/CI/install_server.sh b/tests/CI/install_server.sh index 1cc754195ab..a15b57ba52d 100755 --- a/tests/CI/install_server.sh +++ b/tests/CI/install_server.sh @@ -65,8 +65,8 @@ Resources { S3 { - Aws_access_key_id = fakeId - Aws_secret_access_key = fakeKey + Aws_access_key_id = console + Aws_secret_access_key = console123 } } } diff --git a/tests/Integration/Resources/Storage/Test_Resources_S3.py b/tests/Integration/Resources/Storage/Test_Resources_S3.py index 42ce7d721a3..18ffcaa7212 100644 --- a/tests/Integration/Resources/Storage/Test_Resources_S3.py +++ b/tests/Integration/Resources/Storage/Test_Resources_S3.py @@ -249,12 +249,6 @@ def test_createDirectory(setuptest): @pytest.mark.order4 def test_putFile(setuptest): """Copy a file""" - # XXX: this is not good ! - # The mock I use for S3 seem to have a bug uploading files - # with presigned URL. So for the time being, I upload directly, - # but this should be checked - # https://github.com/adobe/S3Mock/issues/219 - se = StorageElement("S3-DIRECT") res = se.putFile(putFile) assert res["OK"], res for lfn in putFile: diff --git a/tests/Jenkins/dirac-cfg-update-server.py b/tests/Jenkins/dirac-cfg-update-server.py index bd75aec5b5d..11aad73aa73 100644 --- a/tests/Jenkins/dirac-cfg-update-server.py +++ b/tests/Jenkins/dirac-cfg-update-server.py @@ -202,8 +202,8 @@ # Path = my-first-bucket # Access = remote # SecureConnection = False -# Aws_access_key_id = fakeId #useless -# Aws_secret_access_key = fakeKey #useles +# Aws_access_key_id = console +# Aws_secret_access_key = console123 # } # } # } @@ -227,8 +227,8 @@ csAPI.setOption("Resources/StorageElements/S3-DIRECT/S3/Path", "my-first-bucket") csAPI.setOption("Resources/StorageElements/S3-DIRECT/S3/Access", "remote") csAPI.setOption("Resources/StorageElements/S3-DIRECT/S3/SecureConnection", "False") -csAPI.setOption("Resources/StorageElements/S3-DIRECT/S3/Aws_access_key_id", "FakeId") -csAPI.setOption("Resources/StorageElements/S3-DIRECT/S3/Aws_secret_access_key", "True") +csAPI.setOption("Resources/StorageElements/S3-DIRECT/S3/Aws_access_key_id", "console") +csAPI.setOption("Resources/StorageElements/S3-DIRECT/S3/Aws_secret_access_key", "console123") # Setting up S3 indirect resources for the Test_Resources_S3.py From 34f364dc8604c0b66aaaea97674228dc19c5fef2 Mon Sep 17 00:00:00 2001 From: Chris Burr Date: Mon, 19 Oct 2026 11:43:45 +0200 Subject: [PATCH 3/3] feat (WMS): access the DiracX sandbox store directly from SandboxStoreClient When /DiracX/FutureClientEnabled/WorkloadManagement/SandboxStore is enabled, sandboxes are uploaded to and downloaded from the DiracX S3 backend using presigned URLs instead of being proxied through the SandboxStore service. Transfers are streamed to disk, retried with a jittered backoff and the checksum embedded in the PFN is verified after download. Sandboxes which are not stored in S3 are still downloaded through their storage element. --- .../Client/SandboxStoreClient.py | 47 +++- .../Client/test/Test_SandboxStoreClient.py | 230 ++++++++++++++++++ .../FutureClient/SandboxStoreClient.py | 199 +++++++++++++++ 3 files changed, 463 insertions(+), 13 deletions(-) create mode 100644 src/DIRAC/WorkloadManagementSystem/Client/test/Test_SandboxStoreClient.py create mode 100644 src/DIRAC/WorkloadManagementSystem/FutureClient/SandboxStoreClient.py diff --git a/src/DIRAC/WorkloadManagementSystem/Client/SandboxStoreClient.py b/src/DIRAC/WorkloadManagementSystem/Client/SandboxStoreClient.py index cbad3a3e2a0..9228695cef0 100755 --- a/src/DIRAC/WorkloadManagementSystem/Client/SandboxStoreClient.py +++ b/src/DIRAC/WorkloadManagementSystem/Client/SandboxStoreClient.py @@ -1,5 +1,6 @@ """ Client for the SandboxStore. - Will connect to the WorkloadManagement/SandboxStore service. + Will connect to the WorkloadManagement/SandboxStore service, or to DiracX + if /DiracX/FutureClientEnabled/WorkloadManagement/SandboxStore is enabled. """ from __future__ import annotations @@ -16,11 +17,15 @@ from DIRAC import S_ERROR, S_OK, gLogger from DIRAC.ConfigurationSystem.Client.Helpers.Registry import getVOForGroup +from DIRAC.ConfigurationSystem.Client.PathFinder import useLegacyAdapter from DIRAC.Core.Base.Client import Client from DIRAC.Core.Tornado.Client.ClientSelector import TransferClientSelector as TransferClient from DIRAC.Core.Utilities.File import getGlobbedTotalSize, mkDir from DIRAC.Core.Utilities.ReturnValues import returnSingleResult from DIRAC.Resources.Storage.StorageElement import StorageElement +from DIRAC.WorkloadManagementSystem.FutureClient.SandboxStoreClient import ( + SandboxStoreClient as futureSandboxStoreClient, +) @contextmanager @@ -62,6 +67,11 @@ def __init__(self, rpcClient=None, transferClient=None, smdb=False, **kwargs): self.__vo = None if "delegatedGroup" in kwargs: self.__vo = getVOForGroup(kwargs["delegatedGroup"]) + self.__diracxClient = None + if useLegacyAdapter(self.__serviceName): + self.__diracxClient = futureSandboxStoreClient( + delegatedDN=kwargs.get("delegatedDN"), delegatedGroup=kwargs.get("delegatedGroup") + ) def __getRPCClient(self): """Get an RPC client for SB service""" @@ -144,15 +154,18 @@ def uploadFilesAsSandbox(self, fileList, sizeLimit=0, assignTo=None): result["SandboxFileName"] = tmpFilePath return result - oMD5 = hashlib.md5(usedforsecurity=False) - with open(tmpFilePath, "rb") as fd: - bData = fd.read(10240) - while bData: - oMD5.update(bData) + if self.__diracxClient: + result = self.__diracxClient.uploadSandbox(tmpFilePath, assignTo) + else: + oMD5 = hashlib.md5(usedforsecurity=False) + with open(tmpFilePath, "rb") as fd: bData = fd.read(10240) + while bData: + oMD5.update(bData) + bData = fd.read(10240) - transferClient = self.__getTransferClient() - result = transferClient.sendFile(tmpFilePath, [f"{oMD5.hexdigest()}.tar.bz2", assignTo]) + transferClient = self.__getTransferClient() + result = transferClient.sendFile(tmpFilePath, [f"{oMD5.hexdigest()}.tar.bz2", assignTo]) result["SandboxFileName"] = tmpFilePath try: if result["OK"]: @@ -182,15 +195,20 @@ def downloadSandbox(self, sbLocation, destinationDir="", inMemory=False, unpack= except OSError as e: return S_ERROR(f"Cannot create temporary file: {repr(e)}") - se = StorageElement(seName, vo=self.__vo) - result = returnSingleResult(se.getFile(sePFN, localPath=tmpSBDir)) + sbFileName = os.path.basename(sePFN) + tarFileName = os.path.join(tmpSBDir, sbFileName) + + # Sandboxes stored in S3 are downloaded directly, others still go through their SE + if self.__diracxClient and sePFN.startswith("/S3/"): + result = self.__diracxClient.downloadSandbox(sePFN, tarFileName) + else: + se = StorageElement(seName, vo=self.__vo) + result = returnSingleResult(se.getFile(sePFN, localPath=tmpSBDir)) if not result["OK"]: return result - sbFileName = os.path.basename(sePFN) result = S_OK() - tarFileName = os.path.join(tmpSBDir, sbFileName) if inMemory: try: @@ -239,7 +257,10 @@ def downloadSandbox(self, sbLocation, destinationDir="", inMemory=False, unpack= def downloadSandboxForJob(self, jobId, sbType, destinationPath="", inMemory=False, unpack=True): """Download SB for a job""" - result = self.__getRPCClient().getSandboxesAssignedToEntity(f"Job:{jobId}") + if self.__diracxClient: + result = self.__diracxClient.getSandboxesAssignedToJob(jobId) + else: + result = self.__getRPCClient().getSandboxesAssignedToEntity(f"Job:{jobId}") if not result["OK"]: return result diff --git a/src/DIRAC/WorkloadManagementSystem/Client/test/Test_SandboxStoreClient.py b/src/DIRAC/WorkloadManagementSystem/Client/test/Test_SandboxStoreClient.py new file mode 100644 index 00000000000..8778de8e5f6 --- /dev/null +++ b/src/DIRAC/WorkloadManagementSystem/Client/test/Test_SandboxStoreClient.py @@ -0,0 +1,230 @@ +"""Tests for the SandboxStoreClient when using DiracX""" +# pylint: disable=missing-docstring,redefined-outer-name,unsubscriptable-object + +import hashlib +import io +import tarfile +from http import HTTPStatus +from unittest.mock import MagicMock + +import pytest +import requests +from azure.core.exceptions import HttpResponseError + +from DIRAC.WorkloadManagementSystem.Client import SandboxStoreClient as legacyModule +from DIRAC.WorkloadManagementSystem.FutureClient import SandboxStoreClient as futureModule +from DIRAC.WorkloadManagementSystem.FutureClient.SandboxStoreClient import SandboxStoreClient + + +def makeTarball(content=b"hello"): + buffer = io.BytesIO() + with tarfile.open(fileobj=buffer, mode="w:bz2") as tf: + tarInfo = tarfile.TarInfo(name="hello.txt") + tarInfo.size = len(content) + tf.addfile(tarInfo, io.BytesIO(content)) + return buffer.getvalue() + + +TARBALL = makeTarball() +CHECKSUM = hashlib.sha256(TARBALL).hexdigest() +PFN = f"/S3/sandboxes/vo/vo_user/someone/sha256:{CHECKSUM}.tar.bz2" + + +def fakeResponse(status=200, content=b""): + response = MagicMock() + response.status_code = status + response.iter_content.return_value = [content[i : i + 3] for i in range(0, len(content), 3)] + response.__enter__.return_value = response + if status >= 400: + response.raise_for_status.side_effect = requests.HTTPError(response=response) + return response + + +@pytest.fixture +def api(monkeypatch): + api = MagicMock() + api.jobs.get_sandbox_file.return_value = MagicMock(url="https://s3.invalid/get") + + class FakeDiracXClient: + def __enter__(self): + return api + + def __exit__(self, *args): + pass + + monkeypatch.setattr(futureModule, "DiracXClient", FakeDiracXClient) + monkeypatch.setattr(futureModule.time, "sleep", MagicMock()) + return api + + +def test_download(api, monkeypatch, tmp_path): + get = MagicMock(return_value=fakeResponse(content=TARBALL)) + monkeypatch.setattr(futureModule.requests, "get", get) + + res = SandboxStoreClient().downloadSandbox(PFN, str(tmp_path / "sb.tar.bz2")) + assert res["OK"], res + assert (tmp_path / "sb.tar.bz2").read_bytes() == TARBALL + api.jobs.get_sandbox_file.assert_called_once_with(pfn=PFN) + get.assert_called_once() + + +def test_download_retries(api, monkeypatch, tmp_path): + responses = [fakeResponse(status=503), fakeResponse(content=b"truncated"), fakeResponse(content=TARBALL)] + monkeypatch.setattr(futureModule.requests, "get", MagicMock(side_effect=responses)) + + res = SandboxStoreClient().downloadSandbox(PFN, str(tmp_path / "sb.tar.bz2")) + assert res["OK"], res + assert (tmp_path / "sb.tar.bz2").read_bytes() == TARBALL + # A fresh presigned URL is requested for each attempt + assert api.jobs.get_sandbox_file.call_count == 3 + + +def test_download_gives_up(api, monkeypatch, tmp_path): + get = MagicMock(side_effect=lambda *a, **kw: fakeResponse(status=503)) + monkeypatch.setattr(futureModule.requests, "get", get) + + res = SandboxStoreClient().downloadSandbox(PFN, str(tmp_path / "sb.tar.bz2")) + assert not res["OK"] + assert get.call_count == futureModule.MAX_ATTEMPTS + + +def test_download_not_retryable(api, monkeypatch, tmp_path): + get = MagicMock(return_value=fakeResponse(status=404)) + monkeypatch.setattr(futureModule.requests, "get", get) + + res = SandboxStoreClient().downloadSandbox(PFN, str(tmp_path / "sb.tar.bz2")) + assert not res["OK"] + get.assert_called_once() + + +def test_download_invalid_pfn(api): + res = SandboxStoreClient().downloadSandbox("/SandBox/s/someone.vo_user/abc/def/abcdef.tar.bz2", "/dev/null") + assert not res["OK"] + api.jobs.get_sandbox_file.assert_not_called() + + +@pytest.fixture +def tarballPath(tmp_path): + path = tmp_path / "upload.tar.bz2" + path.write_bytes(TARBALL) + return str(path) + + +def test_upload(api, monkeypatch, tarballPath): + api.jobs.initiate_sandbox_upload.return_value = MagicMock( + url="https://s3.invalid/post", fields={"key": "value"}, pfn=f"SB:SandboxSE|{PFN}" + ) + post = MagicMock(return_value=fakeResponse()) + monkeypatch.setattr(futureModule.requests, "post", post) + + res = SandboxStoreClient().uploadSandbox(tarballPath, {"Job:123": "Output"}) + assert res["OK"], res + assert res["Value"] == f"SB:SandboxSE|{PFN}" + + sandboxInfo = api.jobs.initiate_sandbox_upload.call_args.args[0] + assert sandboxInfo.checksum == CHECKSUM + assert sandboxInfo.size == len(TARBALL) + assert sandboxInfo.format == "tar.bz2" + assert post.call_args.kwargs["data"] == {"key": "value"} + api.jobs.assign_sandbox_to_job.assert_called_once_with(123, PFN) + + +def test_upload_existing(api, monkeypatch, tarballPath): + api.jobs.initiate_sandbox_upload.return_value = MagicMock(url=None, pfn=f"SB:SandboxSE|{PFN}") + post = MagicMock() + monkeypatch.setattr(futureModule.requests, "post", post) + + res = SandboxStoreClient().uploadSandbox(tarballPath, {}) + assert res["OK"], res + post.assert_not_called() + api.jobs.assign_sandbox_to_job.assert_not_called() + + +def test_upload_already_assigned(api, monkeypatch, tarballPath): + api.jobs.initiate_sandbox_upload.return_value = MagicMock(url=None, pfn=f"SB:SandboxSE|{PFN}") + error = HttpResponseError(message="Sandbox already assigned") + error.status_code = HTTPStatus.BAD_REQUEST + api.jobs.assign_sandbox_to_job.side_effect = error + + res = SandboxStoreClient().uploadSandbox(tarballPath, {"Job:123": "Output"}) + assert res["OK"], res + + +def test_upload_too_big(api, monkeypatch, tmp_path): + path = tmp_path / "big.tar.bz2" + path.write_bytes(b"0" * (11 * 1024 * 1024)) + + res = SandboxStoreClient().uploadSandbox(str(path), {}) + assert not res["OK"] + assert "too big" in res["Message"] + api.jobs.initiate_sandbox_upload.assert_not_called() + + +def test_upload_input_assignment(api, tarballPath): + res = SandboxStoreClient().uploadSandbox(tarballPath, {"Job:123": "Input"}) + assert not res["OK"] + api.jobs.initiate_sandbox_upload.assert_not_called() + + +def test_getSandboxesAssignedToJob(api): + api.jobs.get_job_sandboxes.return_value = {"Input": [PFN], "Output": [None]} + + res = SandboxStoreClient().getSandboxesAssignedToJob(123) + assert res["OK"], res + assert res["Value"] == {"Input": [f"SB:SandboxSE|{PFN}"]} + + +def test_getSandboxesAssignedToJob_legacy(api): + legacyPFN = "/SandBox/s/someone.vo_user/abc/def/abcdef.tar.bz2" + api.jobs.get_job_sandboxes.return_value = {"Input": [legacyPFN], "Output": []} + + res = SandboxStoreClient().getSandboxesAssignedToJob(123) + assert res["OK"], res + assert res["Value"] == {"Input": [f"SB:ProductionSandboxSE|{legacyPFN}"]} + + +@pytest.fixture +def legacyClient(monkeypatch): + monkeypatch.setattr(legacyModule, "useLegacyAdapter", lambda *args: True) + futureClient = MagicMock() + monkeypatch.setattr(legacyModule, "futureSandboxStoreClient", MagicMock(return_value=futureClient)) + storageElement = MagicMock() + monkeypatch.setattr(legacyModule, "StorageElement", storageElement) + return legacyModule.SandboxStoreClient(), futureClient, storageElement + + +def test_legacy_download_s3(legacyClient, tmp_path): + client, futureClient, storageElement = legacyClient + + def download(pfn, destination): + with open(destination, "wb") as fh: + fh.write(TARBALL) + return {"OK": True, "Value": None} + + futureClient.downloadSandbox.side_effect = download + + res = client.downloadSandbox(f"SB:SandboxSE|{PFN}", destinationDir=str(tmp_path)) + assert res["OK"], res + assert (tmp_path / "hello.txt").read_bytes() == b"hello" + storageElement.assert_not_called() + + +def test_legacy_download_non_s3(legacyClient, tmp_path): + client, futureClient, storageElement = legacyClient + storageElement.return_value.getFile.return_value = {"OK": True, "Value": {"Successful": {}, "Failed": {"x": "y"}}} + + res = client.downloadSandbox("SB:ProductionSandboxSE|/SandBox/abcdef.tar.bz2", destinationDir=str(tmp_path)) + assert not res["OK"] + storageElement.assert_called_once() + futureClient.downloadSandbox.assert_not_called() + + +def test_legacy_upload(legacyClient, tmp_path): + client, futureClient, _ = legacyClient + futureClient.uploadSandbox.return_value = {"OK": True, "Value": f"SB:SandboxSE|{PFN}"} + (tmp_path / "file.txt").write_text("content") + + res = client.uploadFilesAsSandbox([str(tmp_path / "file.txt")], assignTo={"Job:123": "Output"}) + assert res["OK"], res + assert res["Value"] == f"SB:SandboxSE|{PFN}" + assert futureClient.uploadSandbox.call_args.args[1] == {"Job:123": "Output"} diff --git a/src/DIRAC/WorkloadManagementSystem/FutureClient/SandboxStoreClient.py b/src/DIRAC/WorkloadManagementSystem/FutureClient/SandboxStoreClient.py new file mode 100644 index 00000000000..38a225348e5 --- /dev/null +++ b/src/DIRAC/WorkloadManagementSystem/FutureClient/SandboxStoreClient.py @@ -0,0 +1,199 @@ +"""DiracX implementation of the SandboxStore client. + +Sandboxes are transferred directly between the client and the DiracX S3 +backend using presigned URLs, without going through the SandboxStore service. +""" +from __future__ import annotations + +import hashlib +import os +import random +import re +import time +from collections.abc import Callable, Iterator +from contextlib import contextmanager +from http import HTTPStatus +from typing import TYPE_CHECKING, TypeVar + +import requests +from azure.core.exceptions import HttpResponseError +from diracx.client.models import SandboxInfo + +from DIRAC import gConfig, gLogger +from DIRAC.ConfigurationSystem.Client.Helpers import Registry +from DIRAC.ConfigurationSystem.Client.PathFinder import getServiceSection +from DIRAC.Core.Security.DiracX import DiracXClient, FutureClient +from DIRAC.Core.Utilities.ReturnValues import SErrorException, convertToReturnValue, returnValueOrRaise + +if TYPE_CHECKING: + from diracx.client.sync import SyncDiracClient + +T = TypeVar("T") + +# PFNs generated by DiracX look like /S3/////:. +RE_S3_PFN = re.compile( + r"^/S3/[^/]+/(?:[^/]+/){3}(?P[a-z0-9]+):(?P[0-9a-f]+)\.(?P[a-z0-9.]+)$" +) +# Name of the SE used by DiracX for sandboxes, only used to build "SB:|" URLs +DIRACX_SANDBOX_SE = "SandboxSE" + +MAX_ATTEMPTS = 5 +BACKOFF_BASE_SECONDS = 2 +BACKOFF_MAX_SECONDS = 60 +CHUNK_SIZE = 1024 * 1024 +# (connect, read) timeouts in seconds +DOWNLOAD_TIMEOUT = (30, 60) +UPLOAD_TIMEOUT = (30, 300) + + +class ChecksumMismatchError(Exception): + pass + + +def _isRetryable(exc: Exception) -> bool: + """Whether a transfer to/from the S3 backend is worth retrying""" + if isinstance(exc, (requests.ConnectionError, requests.Timeout, ChecksumMismatchError)): + return True + if isinstance(exc, requests.HTTPError) and exc.response is not None: + # 403 is returned by S3 if the presigned URL expired, a new one is requested on each attempt + status = exc.response.status_code + return status >= 500 or status in (HTTPStatus.FORBIDDEN, HTTPStatus.TOO_MANY_REQUESTS) + return False + + +def _withRetries(description: str, func: Callable[[], T]) -> T: + """Call func, retrying with a jittered exponential backoff if it fails with a retryable error""" + for attempt in range(1, MAX_ATTEMPTS + 1): + try: + return func() + except Exception as e: + if attempt == MAX_ATTEMPTS or not _isRetryable(e): + raise + # Full jitter avoids many clients retrying in lockstep if the backend has a hiccup + delay = random.uniform(0, min(BACKOFF_MAX_SECONDS, BACKOFF_BASE_SECONDS * 2**attempt)) # nosec B311 + gLogger.warn(f"Failed {description}, retrying", f"in {delay:.1f}s (attempt {attempt}): {e!r}") + time.sleep(delay) + raise NotImplementedError("This should never be reached") + + +class SandboxStoreClient(FutureClient): + def __init__(self, delegatedDN: str | None = None, delegatedGroup: str | None = None, **kwargs): + """ + :param delegatedDN: DN of the user to act on behalf of, requires the DiracX legacy exchange key + :param delegatedGroup: group of the user to act on behalf of + """ + self.delegatedDN = delegatedDN + self.delegatedGroup = delegatedGroup + self.serviceSection = getServiceSection("WorkloadManagement/SandboxStore") + + @contextmanager + def _client(self) -> Iterator[SyncDiracClient]: + if self.delegatedGroup: + # Only available on servers which have /DiracX/LegacyExchangeApiKey + from DIRAC.FrameworkSystem.Utilities.diracx import TheImpersonator + + credDict = { + "username": returnValueOrRaise(Registry.getUsernameForDN(self.delegatedDN)), + "group": self.delegatedGroup, + "properties": Registry.getPropertiesForGroup(self.delegatedGroup), + } + with TheImpersonator(credDict, source="SandboxStoreClient") as api: + yield api + else: + with DiracXClient() as api: + yield api + + @convertToReturnValue + def uploadSandbox(self, tarFilePath: str, assignTo: dict[str, str], sbFormat: str = "tar.bz2") -> str: + """Upload a sandbox and optionally assign it to jobs as an output sandbox + + :param tarFilePath: path of the sandbox archive + :param assignTo: dict of {"Job:": "Output"} + :param sbFormat: format of the archive + :returns: the sandbox URL, i.e. "SB:|" + """ + jobIDs = [] + for entity, sbType in assignTo.items(): + if not entity.startswith("Job:") or sbType != "Output": + raise NotImplementedError(f"DiracX can only assign output sandboxes to jobs, got {entity}={sbType}") + jobIDs.append(int(entity.removeprefix("Job:"))) + + size = os.stat(tarFilePath).st_size + maxBytes = gConfig.getValue(f"{self.serviceSection}/MaxSandboxSizeMiB", 10) * 1048576 + if maxBytes and size > maxBytes: + raise SErrorException("Sandbox is too big. Please upload it to a grid storage element") + + hasher = hashlib.sha256() + with open(tarFilePath, "rb") as fh: + while data := fh.read(CHUNK_SIZE): + hasher.update(data) + sandboxInfo = SandboxInfo(checksum_algorithm="sha256", checksum=hasher.hexdigest(), size=size, format=sbFormat) + + def attempt() -> str: + with self._client() as api: + res = api.jobs.initiate_sandbox_upload(sandboxInfo) + if res.url: + with open(tarFilePath, "rb") as fh: + response = requests.post( + res.url, data=res.fields, files={"file": ("file", fh)}, timeout=UPLOAD_TIMEOUT + ) + response.raise_for_status() + gLogger.verbose("Uploaded sandbox", res.pfn) + else: + gLogger.verbose("Sandbox already exists in storage backend", res.pfn) + return res.pfn + + sbURL = _withRetries("sandbox upload", attempt) + + with self._client() as api: + for jobID in jobIDs: + try: + api.jobs.assign_sandbox_to_job(jobID, sbURL.split("|", 1)[-1]) + except HttpResponseError as e: + # Identical sandboxes have the same PFN so a rescheduled job can upload the same one again + if e.status_code != HTTPStatus.BAD_REQUEST or "already assigned" not in str(e): + raise + return sbURL + + @convertToReturnValue + def downloadSandbox(self, pfn: str, destination: str) -> None: + """Download a sandbox from S3 and verify its checksum + + :param pfn: PFN of the sandbox, i.e. "/S3/..." + :param destination: path the sandbox archive is written to + """ + if not (match := RE_S3_PFN.match(pfn)): + raise ValueError(f"Not a DiracX sandbox PFN: {pfn}") + + def attempt() -> None: + with self._client() as api: + res = api.jobs.get_sandbox_file(pfn=pfn) + hasher = hashlib.new(match["algorithm"]) + with requests.get(res.url, stream=True, timeout=DOWNLOAD_TIMEOUT) as response: + response.raise_for_status() + with open(destination, "wb") as fh: + for chunk in response.iter_content(CHUNK_SIZE): + hasher.update(chunk) + fh.write(chunk) + if hasher.hexdigest() != match["checksum"]: + raise ChecksumMismatchError(f"Checksum mismatch for {pfn}: got {hasher.hexdigest()}") + + _withRetries(f"download of {pfn}", attempt) + + @convertToReturnValue + def getSandboxesAssignedToJob(self, jobID: int) -> dict[str, list[str]]: + """Get the sandboxes assigned to a job + + :returns: dict of {"Input"|"Output": ["SB:|", ...]} + """ + with self._client() as api: + sandboxes = api.jobs.get_job_sandboxes(int(jobID)) + # Sandboxes which predate the S3 backend are stored on the LocalSE of the SandboxStore + legacySE = gConfig.getValue(f"{self.serviceSection}/LocalSE", "ProductionSandboxSE") + sbDict = {} + for sbType, pfns in sandboxes.items(): + # DiracX returns [None] if there is no sandbox of this type + for pfn in filter(None, pfns): + seName = DIRACX_SANDBOX_SE if pfn.startswith("/S3/") else legacySE + sbDict.setdefault(sbType, []).append(f"SB:{seName}|{pfn}") + return sbDict