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