1# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2# SPDX-License-Identifier: Apache-2.0
3#
4# Licensed under the Apache License, Version 2.0 (the "License");
5# you may not use this file except in compliance with the License.
6# You may obtain a copy of the License at
7#
8# http://www.apache.org/licenses/LICENSE-2.0
9#
10# Unless required by applicable law or agreed to in writing, software
11# distributed under the License is distributed on an "AS IS" BASIS,
12# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13# See the License for the specific language governing permissions and
14# limitations under the License.
15
16from typing import Any, Callable, Optional, Union
17
18# Import botocore patch to handle AIStore redirects.
19# See https://github.com/NVIDIA/aistore/tree/main/python/aistore/botocore_patch
20from aistore.botocore_patch import botocore # noqa: F401
21
22from ..telemetry import Telemetry
23from ..types import AWARE_DATETIME_MIN, CredentialsProvider, ObjectMetadata
24from .s3 import S3StorageProvider, StaticS3CredentialsProvider
25
26PROVIDER = "ais_s3"
27
28# Default dummy credentials for AIStore S3 when auth is disabled
29DEFAULT_ACCESS_KEY = "FAKEKEY"
30DEFAULT_SECRET_KEY = "FAKESECRET"
31
32
[docs]
33class AIStoreS3StorageProvider(S3StorageProvider):
34 """
35 A concrete implementation of the :py:class:`multistorageclient.types.StorageProvider` for interacting with
36 AIStore via its S3-compatible interface.
37 """
38
39 def __init__(
40 self,
41 region_name: str = "",
42 endpoint_url: str = "",
43 base_path: str = "",
44 credentials_provider: Optional[CredentialsProvider] = None,
45 config_dict: Optional[dict[str, Any]] = None,
46 telemetry_provider: Optional[Callable[[], Telemetry]] = None,
47 verify: Optional[Union[bool, str]] = None,
48 **kwargs: Any,
49 ) -> None:
50 """
51 Initializes the :py:class:`AIStoreS3StorageProvider` with AIStore S3 endpoint and optional JWT authentication.
52
53 :param region_name: The AWS region (can be any valid region, AIStore ignores this).
54 :param endpoint_url: The AIStore S3 endpoint (e.g., ``http://localhost:8080/s3`` or ``https://aistore.example.com/s3``).
55 :param base_path: The root prefix path within the bucket where all operations will be scoped.
56 :param credentials_provider: The provider to retrieve AIStore credentials (AISCredentials).
57 If not provided, uses dummy credentials for unauthenticated access.
58 :param config_dict: Resolved MSC config.
59 :param telemetry_provider: A function that provides a telemetry instance.
60 :param verify: Controls SSL certificate verification. Can be ``True`` (verify using system CA bundle, default),
61 ``False`` (skip verification for self-signed certificates), or a string path to a custom CA certificate bundle.
62 :param kwargs: Additional keyword arguments. See :py:class:`S3StorageProvider` for all available options.
63 """
64 self._ais_credentials_provider = credentials_provider
65
66 dummy_s3_credentials = StaticS3CredentialsProvider(access_key=DEFAULT_ACCESS_KEY, secret_key=DEFAULT_SECRET_KEY)
67
68 super().__init__(
69 region_name=region_name or "us-east-1",
70 endpoint_url=endpoint_url,
71 base_path=base_path,
72 credentials_provider=dummy_s3_credentials,
73 config_dict=config_dict,
74 telemetry_provider=telemetry_provider,
75 verify=verify,
76 **kwargs,
77 )
78
79 self._provider_name = PROVIDER
80
81 # Register event handler to inject JWT token if credentials are provided
82 # Use 'before-send' instead of 'before-sign' to inject the header AFTER boto3 signs the request
83 # This prevents boto3 from overwriting our Authorization header with AWS signatures
84 if self._ais_credentials_provider:
85 self._s3_client.meta.events.register("before-send.s3.*", self._inject_auth_header)
86
87 def _get_object_metadata(self, path: str, strict: bool = True) -> ObjectMetadata:
88 """
89 Override to handle AIStore S3 API quirk where HEAD requests on directory-like paths return 400.
90 """
91 try:
92 return super()._get_object_metadata(path, strict=strict)
93 except RuntimeError as error:
94 # AIStore returns 400 for invalid HEAD requests (e.g., directory-like paths)
95 # Treat this the same as FileNotFoundError and check if it's a directory
96 if strict and "status_code: 400" in str(error):
97 path = self._append_delimiter(path)
98 if self._is_dir(path):
99 return ObjectMetadata(
100 key=path,
101 type="directory",
102 content_length=0,
103 last_modified=AWARE_DATETIME_MIN,
104 )
105 raise
106
107 def _inject_auth_header(self, request, **kwargs):
108 """
109 Event handler that injects the JWT Bearer token into the Authorization header.
110
111 This is called after boto3 signs the request but before it's sent over the network.
112 It replaces the AWS signature with the AIStore JWT token.
113 More details: https://github.com/NVIDIA/aistore/tree/main/python/aistore/botocore_patch#boto3-with-aistore-authentication
114
115 :param request: The request object from botocore containing headers, URL, etc.
116 :param kwargs: Additional keyword arguments from the event system.
117 """
118 if self._ais_credentials_provider:
119 credentials = self._ais_credentials_provider.get_credentials()
120 if credentials.token:
121 # Replace the Authorization header with the JWT Bearer token
122 request.headers["Authorization"] = f"Bearer {credentials.token}"