Warning, /eic-spack/spack_repo/eic/packages/py_rucio_clients/token-upload-client-40.2.0.patch is written in an unsupported language. File is not indexed.
0001 diff --git a/lib/rucio/client/credentialclient.py b/lib/rucio/client/credentialclient.py
0002 index 8e8bdd6f1..52e6257b8 100644
0003 --- a/lib/rucio/client/credentialclient.py
0004 +++ b/lib/rucio/client/credentialclient.py
0005 @@ -12,6 +12,9 @@
0006 # See the License for the specific language governing permissions and
0007 # limitations under the License.
0008
0009 +import json
0010 +from typing import Any, Literal, Optional
0011 +
0012 from requests.status_codes import codes
0013
0014 from rucio.client.baseclient import BaseClient, choice
0015 @@ -31,6 +34,7 @@ class CredentialClient(BaseClient):
0016 """
0017
0018 CREDENTIAL_BASEURL = 'credentials'
0019 + TOKEN_BASEURL = 'tokens'
0020
0021 def get_signed_url(
0022 self,
0023 @@ -108,3 +112,99 @@ class CredentialClient(BaseClient):
0024
0025 exc_cls, exc_msg = self._get_exception(headers=r.headers, status_code=r.status_code, data=r.content)
0026 raise exc_cls(exc_msg)
0027 +
0028 + def _post_token_request(self, data: dict[str, Any]) -> str:
0029 + """Send a sesson token request to the server and return the minted storage token."""
0030 + path = '/'.join([self.TOKEN_BASEURL, 'request'])
0031 + rurl = build_url(choice(self.list_hosts), path=path)
0032 + r = self._send_request(rurl, method=HTTPMethod.POST, data=json.dumps(data))
0033 + if r.status_code == codes.ok:
0034 + return json.loads(r.text)['token']
0035 + exc_cls, exc_msg = self._get_exception(headers=r.headers,
0036 + status_code=r.status_code,
0037 + data=r.content)
0038 + raise exc_cls(exc_msg)
0039 +
0040 +
0041 + def get_upload_token(
0042 + self,
0043 + rse: str,
0044 + scope: str,
0045 + name: str,
0046 + domain: Literal["wan", "lan"] = "wan",
0047 + register_after_upload: bool = False,
0048 + no_register: bool = False,
0049 + scheme: Optional[str] = None,
0050 + ) -> str:
0051 + """
0052 + Request a file-scoped OIDC token for uploading a single DID to an RSE.
0053 +
0054 + The token authorizes writing exactly the ``scope:name`` DID on ``rse``.
0055 + Call it only after the DID is registered — the server refuses tokens for
0056 + unregistered DIDs.
0057 +
0058 + Parameters
0059 + ----------
0060 + rse
0061 + The name of the RSE to upload to.
0062 + scope
0063 + DID scope of the file to upload.
0064 + name
0065 + DID name of the file to upload.
0066 + domain
0067 + Network domain the token applies to: ``"wan"`` or ``"lan"``.
0068 + Defaults to ``"wan"``.
0069 + register_after_upload
0070 + Whether the file is registered after the upload completes.
0071 + Defaults to ``False``.
0072 + no_register
0073 + Whether to skip registration entirely. Defaults to ``False``.
0074 + scheme
0075 + The scheme the upload is forced to (``force_scheme``), or None to
0076 + cover every write protocol of the RSE.
0077 +
0078 + Returns
0079 + -------
0080 + str
0081 + The OIDC access token, valid until it expires.
0082 +
0083 + Raises
0084 + ------
0085 + RucioException
0086 + If the server returns a status code other than ``200 OK``,
0087 + including when the DID is not yet registered.
0088 + """
0089 + return self._post_token_request({
0090 + 'rse': rse,
0091 + 'operation': 'upload',
0092 + 'scope': scope,
0093 + 'name': name,
0094 + 'domain': domain,
0095 + 'register_after_upload': register_after_upload,
0096 + 'no_register': no_register,
0097 + 'scheme': scheme,
0098 + })
0099 +
0100 + def get_download_token(self, rse: str) -> str:
0101 + """
0102 + Request an RSE-scoped OIDC token for downloading from an RSE.
0103 +
0104 + The token authorises reading from ``rse`` and is not tied to a specific
0105 + DID, so a single token can serve many files on that RSE.
0106 +
0107 + Parameters
0108 + ----------
0109 + rse
0110 + The name of the RSE to download from.
0111 +
0112 + Returns
0113 + -------
0114 + str
0115 + The OIDC access token, valid until it expires.
0116 +
0117 + Raises
0118 + ------
0119 + RucioException
0120 + If the server returns a status code other than ``200 OK``.
0121 + """
0122 + return self._post_token_request({'rse': rse, 'operation': 'download'})
0123 diff --git a/lib/rucio/client/uploadclient.py b/lib/rucio/client/uploadclient.py
0124 index 3ca6a0688..f58768a02 100644
0125 --- a/lib/rucio/client/uploadclient.py
0126 +++ b/lib/rucio/client/uploadclient.py
0127 @@ -50,9 +50,23 @@ from rucio.common.utils import execute, generate_uuid, make_valid_did, retry, se
0128 from rucio.rse import rsemanager as rsemgr
0129
0130 if TYPE_CHECKING:
0131 - from collections.abc import Iterable, Mapping
0132 -
0133 - from rucio.common.types import AttachDict, DatasetDict, DIDStringDict, FileToUploadDict, FileToUploadWithCollectedAndDatasetInfoDict, FileToUploadWithCollectedInfoDict, LFNDict, LoggerFunction, PathTypeAlias, RSESettingsDict, TraceBaseDict, TraceDict
0134 + from collections.abc import Iterable, Mapping, Sequence
0135 +
0136 + from rucio.common.types import (
0137 + AttachDict,
0138 + DatasetDict,
0139 + DIDStringDict,
0140 + FileToUploadDict,
0141 + FileToUploadWithCollectedAndDatasetInfoDict,
0142 + FileToUploadWithCollectedInfoDict,
0143 + LFNDict,
0144 + LoggerFunction,
0145 + PathTypeAlias,
0146 + RSEProtocolDict,
0147 + RSESettingsDict,
0148 + TraceBaseDict,
0149 + TraceDict,
0150 + )
0151 from rucio.rse.protocols.protocol import RSEProtocol
0152
0153
0154 @@ -91,10 +105,6 @@ class UploadClient:
0155
0156 self.client: Final[Client] = _client if _client else Client()
0157 self.client_location = detect_client_location()
0158 - # if token should be used, use only JWT tokens
0159 - self.auth_token: Optional[str] = (
0160 - self.client.auth_token if len(self.client.auth_token.split(".")) == 3 else None
0161 - )
0162 self.tracing = tracing
0163 if not self.tracing:
0164 logger(logging.DEBUG, 'Tracing is turned off.')
0165 @@ -438,6 +448,41 @@ class UploadClient:
0166 ignore_availability=ignore_availability,
0167 activity=activity)
0168
0169 + # write protocols in priority order, used by the token existence check
0170 + # and the upload loop below (computed once). Raises RSEProtocolNotSupported
0171 + # if none match, before any token is requested.
0172 + write_protocols = rsemgr.get_protocols_ordered(rse_settings=rse_settings,
0173 + operation='write',
0174 + scheme=force_scheme,
0175 + domain=domain,
0176 + impl=impl)
0177 +
0178 + # Storage bearer for this file. Defaults to None, meaning "no bearer"
0179 + # (the protocol falls back to its own cert/x509 auth) -- never the
0180 + # Rucio session token. On an RSE flagged OIDC_SUPPORT we exchange the
0181 + # Rucio auth for a path-scoped storage token minted by the server.
0182 + # This request must happen after _register_file: the token API
0183 + # refuses tokens for unregistered DIDs. It is skipped for
0184 + # register_after_upload (DID not yet registered) and for explicit
0185 + # pfn uploads.
0186 + upload_auth_token = None
0187 + if rse_attributes.get(RseAttr.OIDC_SUPPORT) \
0188 + and not register_after_upload and not pfn \
0189 + and (not no_register or rse_attributes.get('no_register')):
0190 + try:
0191 + upload_auth_token = self.client.get_upload_token(
0192 + rse, file['did_scope'], file['did_name'], domain=domain, scheme=force_scheme,
0193 + no_register=bool(no_register))
0194 + logger(logging.DEBUG, 'Using scoped upload token for %s:%s at %s'
0195 + % (file['did_scope'], file['did_name'], rse))
0196 + except Exception as error:
0197 + # fall back to the RSE's other auth (e.g. x509) rather than failing the file.
0198 + logger(logging.WARNING,
0199 + 'Could not obtain scoped upload token for %s:%s at %s, '
0200 + 'falling back to non-token auth: %s'
0201 + % (file['did_scope'], file['did_name'], rse, str(error)))
0202 + upload_auth_token = None
0203 +
0204 # if register_after_upload, the file should be overwritten if it is not registered,
0205 # otherwise if the file already exists on RSE we're done
0206 if register_after_upload:
0207 @@ -446,7 +491,7 @@ class UploadClient:
0208 domain=domain,
0209 scheme=force_scheme,
0210 impl=impl,
0211 - auth_token=self.auth_token,
0212 + auth_token=upload_auth_token,
0213 vo=self.client.vo,
0214 logger=logger):
0215 try:
0216 @@ -463,7 +508,7 @@ class UploadClient:
0217 domain=domain,
0218 scheme=force_scheme,
0219 impl=impl,
0220 - auth_token=self.auth_token,
0221 + auth_token=upload_auth_token,
0222 vo=self.client.vo,
0223 logger=logger):
0224 logger(logging.INFO,
0225 @@ -475,32 +520,39 @@ class UploadClient:
0226 domain=domain,
0227 scheme=force_scheme,
0228 impl=impl,
0229 - auth_token=self.auth_token,
0230 + auth_token=upload_auth_token,
0231 vo=self.client.vo,
0232 logger=logger):
0233 logger(logging.INFO, 'File already exists on RSE with different pfn. Skipping upload.')
0234 trace['stateReason'] = 'File already exists'
0235 continue
0236 else:
0237 - if rsemgr.exists(rse_settings,
0238 - pfn if pfn else file_did, # type: ignore (pfn is str)
0239 - domain=domain,
0240 - scheme=force_scheme,
0241 - impl=impl,
0242 - auth_token=self.auth_token,
0243 - vo=self.client.vo,
0244 - logger=logger):
0245 + if upload_auth_token:
0246 + # With a scoped upload token, check existence on the endpoint the token
0247 + # was minted for (the write endpoint). rsemgr.exists() uses the top *read*
0248 + # protocol, which may be a different (e.g. open-access, read-only) host
0249 + # that the token's audience does not cover.
0250 + already_exists = self._exists_on_write_endpoint(rse_settings,
0251 + write_protocols,
0252 + pfn if pfn else file_did,
0253 + domain=domain,
0254 + auth_token=upload_auth_token)
0255 + else:
0256 + already_exists = rsemgr.exists(rse_settings,
0257 + pfn if pfn else file_did, # type: ignore (pfn is str)
0258 + domain=domain,
0259 + scheme=force_scheme,
0260 + impl=impl,
0261 + auth_token=upload_auth_token,
0262 + vo=self.client.vo,
0263 + logger=logger)
0264 + if already_exists:
0265 logger(logging.INFO, 'File already exists on RSE. Skipping upload')
0266 trace['stateReason'] = 'File already exists'
0267 continue
0268
0269 # protocol handling and upload
0270 - protocols = rsemgr.get_protocols_ordered(rse_settings=rse_settings,
0271 - operation='write',
0272 - scheme=force_scheme,
0273 - domain=domain,
0274 - impl=impl)
0275 - protocols.reverse()
0276 + protocols = list(reversed(write_protocols))
0277 success = False
0278 state_reason = ''
0279 logger(logging.DEBUG, str(protocols))
0280 @@ -536,7 +588,8 @@ class UploadClient:
0281 force_pfn=pfn,
0282 transfer_timeout=file.get('transfer_timeout'),
0283 delete_existing=delete_existing,
0284 - sign_service=sign_service)
0285 + sign_service=sign_service,
0286 + auth_token=upload_auth_token)
0287 logger(logging.DEBUG, 'Upload done.')
0288 success = True
0289 file['upload_result'] = {0: True, 1: None, 'success': True, 'pfn': pfn} # TODO: needs to be removed
0290 @@ -996,7 +1049,8 @@ class UploadClient:
0291 force_scheme: Optional[str] = None,
0292 transfer_timeout: Optional[int] = None,
0293 delete_existing: bool = False,
0294 - sign_service: Optional[str] = None
0295 + sign_service: Optional[str] = None,
0296 + auth_token: Optional[str] = None
0297 ) -> Optional[str]:
0298 """
0299 Perform the actual file transfer to an RSE using the appropriate protocol.
0300 @@ -1034,6 +1088,10 @@ class UploadClient:
0301 If True, removes any unregistered or stale file on the storage that matches this PFN.
0302 sign_service
0303 If set, requests a signed URL from the given service (e.g., gcs, s3, swift).
0304 + auth_token
0305 + The storage bearer token to use for this protocol, or None to use the
0306 + protocol's own (e.g. cert/x509) authentication. This is the per-file,
0307 + path-scoped storage token requested from the server
0308
0309 Returns
0310 -------
0311 @@ -1059,7 +1117,8 @@ class UploadClient:
0312 'write',
0313 force_scheme=force_scheme,
0314 domain=domain,
0315 - impl=impl)
0316 + impl=impl,
0317 + auth_token=auth_token)
0318
0319 base_name = lfn.get('filename', lfn['name'])
0320 name = lfn.get('name', base_name)
0321 @@ -1087,7 +1146,8 @@ class UploadClient:
0322 protocol_read = self._create_protocol(rse_settings,
0323 'read',
0324 domain=domain,
0325 - impl=impl)
0326 + impl=impl,
0327 + auth_token=auth_token)
0328 if pfn is not None:
0329 signed_read_pfn = self.client.get_signed_url(rse_settings['rse'], sign_service, 'read', pfn)
0330 pfn = self.client.get_signed_url(rse_settings['rse'], sign_service, 'write', pfn)
0331 @@ -1116,7 +1176,8 @@ class UploadClient:
0332 'delete',
0333 force_scheme=force_scheme,
0334 domain=domain,
0335 - impl=impl)
0336 + impl=impl,
0337 + auth_token=auth_token)
0338 delete_pfn = '%s.rucio.upload' % list(protocol_delete.lfns2pfns(make_valid_did(lfn)).values())[0]
0339 if sign_service:
0340 delete_pfn = self.client.get_signed_url(rse_settings['rse'], sign_service, 'delete', delete_pfn)
0341 @@ -1135,7 +1196,8 @@ class UploadClient:
0342 'delete',
0343 force_scheme=force_scheme,
0344 domain=domain,
0345 - impl=impl)
0346 + impl=impl,
0347 + auth_token=auth_token)
0348 delete_pfn = '%s' % list(protocol_delete.lfns2pfns(make_valid_did(lfn)).values())[0]
0349 if sign_service:
0350 delete_pfn = self.client.get_signed_url(rse_settings['rse'], sign_service, 'delete', delete_pfn)
0351 @@ -1252,13 +1314,67 @@ class UploadClient:
0352 time.sleep(2 ** attempt)
0353 return protocol.stat(pfn)
0354
0355 + def _exists_on_write_endpoint(
0356 + self,
0357 + rse_settings: "RSESettingsDict",
0358 + write_protocols: "Sequence[RSEProtocolDict]",
0359 + pfn_or_did: Union[str, dict[str, str]],
0360 + domain: str = 'wan',
0361 + auth_token: Optional[str] = None
0362 + ) -> bool:
0363 + """
0364 + Check whether a file already exists, using the *write* protocols of the RSE.
0365 +
0366 + Used for token-authenticated uploads: the scoped upload token is minted for
0367 + the write endpoint(s), so the pre-upload existence check must go there too,
0368 + not to the top-priority read protocol (which may be a different, e.g.
0369 + open-access read-only, host the token's audience does not cover).
0370 +
0371 + ``write_protocols`` is the same priority-ordered list the upload loop uses.
0372 + They are tried in that order, mirroring the loop's fallback: if the check
0373 + fails on one protocol (e.g. root), the next one (e.g. https) is tried. Only
0374 + if every protocol fails is the last error raised.
0375 + """
0376 + last_error: Optional[Exception] = None
0377 + for protocol_attr in write_protocols:
0378 + try:
0379 + protocol = rsemgr.create_protocol(rse_settings, 'write', domain=domain,
0380 + auth_token=auth_token, protocol_attr=protocol_attr,
0381 + logger=self.logger)
0382 + protocol.connect()
0383 + except Exception as error:
0384 + self.logger(logging.WARNING, 'Could not connect with %s for the existence check: %s'
0385 + % (protocol_attr['scheme'], error))
0386 + last_error = error
0387 + continue
0388 + try:
0389 + if isinstance(pfn_or_did, str):
0390 + pfn = pfn_or_did
0391 + else:
0392 + pfn = list(protocol.lfns2pfns(make_valid_did(cast("LFNDict", pfn_or_did))).values())[0]
0393 + if isinstance(pfn, RucioException):
0394 + raise pfn
0395 + self.logger(logging.DEBUG, 'Checking on write endpoint if %s exists' % pfn)
0396 + return protocol.exists(pfn)
0397 + except Exception as error:
0398 + self.logger(logging.WARNING, 'Existence check with %s failed, trying next write protocol: %s'
0399 + % (protocol_attr['scheme'], error))
0400 + last_error = error
0401 + finally:
0402 + protocol.close()
0403 + if last_error is not None:
0404 + raise last_error
0405 + raise RSEOperationNotSupported('No write protocol available for the existence check on %s'
0406 + % rse_settings['rse'])
0407 +
0408 def _create_protocol(
0409 self,
0410 rse_settings: "RSESettingsDict",
0411 operation: str,
0412 impl: Optional[str] = None,
0413 force_scheme: Optional[str] = None,
0414 - domain: str = 'wan'
0415 + domain: str = 'wan',
0416 + auth_token: Optional[str] = None
0417 ) -> "RSEProtocol":
0418 """
0419 Creates and returns the protocol object for the requested RSE operation.
0420 @@ -1278,6 +1394,10 @@ class UploadClient:
0421 If provided, forces the protocol to use this scheme.
0422 domain
0423 The network domain to be used, defaulting to 'wan'.
0424 + auth_token
0425 + The storage bearer token to use for this protocol, or None to use the
0426 + protocol's own (e.g. cert/x509) authentication. This is the per-file,
0427 + path-scoped storage token requested from the server
0428
0429 Returns
0430 -------
0431 @@ -1296,7 +1416,7 @@ class UploadClient:
0432 scheme=force_scheme,
0433 domain=domain,
0434 impl=impl,
0435 - auth_token=self.auth_token,
0436 + auth_token=auth_token,
0437 logger=self.logger
0438 )
0439 protocol.connect()
0440 @@ -1430,7 +1550,8 @@ class UploadClient:
0441 def preferred_impl(
0442 self,
0443 rse_settings: "RSESettingsDict",
0444 - domain: str
0445 + domain: str,
0446 + auth_token: Optional[str] = None
0447 ) -> Optional[str]:
0448 """
0449 Select a suitable protocol implementation for read, write, and delete operations on
0450 @@ -1450,6 +1571,10 @@ class UploadClient:
0451 domain
0452 The network domain (e.g., 'lan' or 'wan') in which the protocol must support
0453 all operations.
0454 + auth_token
0455 + The storage bearer token to use for this protocol, or None to use the
0456 + protocol's own (e.g. cert/x509) authentication. This is the per-file,
0457 + path-scoped storage token requested from the server
0458
0459 Returns
0460 -------
0461 @@ -1497,7 +1622,7 @@ class UploadClient:
0462 continue
0463 try:
0464 supported_protocol = rsemgr.create_protocol(rse_settings, 'write', domain=domain, impl=protocol['impl'],
0465 - auth_token=self.auth_token, logger=self.logger)
0466 + auth_token=auth_token, logger=self.logger)
0467 supported_protocol.connect()
0468 except Exception as error:
0469 self.logger(logging.DEBUG, 'Failed to create protocol "%s", exception: %s' % (protocol['impl'], error))
