2025-12-01

2026-03-17 14:58:51 -06:00
parent 183e865f8b
commit 4b82b57113
6846 changed files with 954887 additions and 162606 deletions
@@ -0,0 +1,167 @@
+# Copyright 2023 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License"). You
+# may not use this file except in compliance with the License. A copy of
+# the License is located at
+#
+# https://aws.amazon.com/apache2.0/
+#
+# or in the "license" file accompanying this file. This file is
+# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF
+# ANY KIND, either express or implied. See the License for the specific
+# language governing permissions and limitations under the License.
+"""
+This file contains private functionality for interacting with the AWS
+Common Runtime library (awscrt) in boto3.
+
+All code contained within this file is for internal usage within this
+project and is not intended for external consumption. All interfaces
+contained within are subject to abrupt breaking changes.
+"""
+
+import threading
+
+import botocore.exceptions
+from botocore.session import Session
+from s3transfer.crt import (
+    BotocoreCRTCredentialsWrapper,
+    BotocoreCRTRequestSerializer,
+    CRTTransferManager,
+    acquire_crt_s3_process_lock,
+    create_s3_crt_client,
+)
+
+# Singletons for CRT-backed transfers
+CRT_S3_CLIENT = None
+BOTOCORE_CRT_SERIALIZER = None
+
+CLIENT_CREATION_LOCK = threading.Lock()
+PROCESS_LOCK_NAME = 'boto3'
+
+
+def _create_crt_client(session, config, region_name, cred_provider):
+    """Create a CRT S3 Client for file transfer.
+
+    Instantiating many of these may lead to degraded performance or
+    system resource exhaustion.
+    """
+    create_crt_client_kwargs = {
+        'region': region_name,
+        'use_ssl': True,
+        'crt_credentials_provider': cred_provider,
+    }
+    return create_s3_crt_client(**create_crt_client_kwargs)
+
+
+def _create_crt_request_serializer(session, region_name):
+    return BotocoreCRTRequestSerializer(
+        session, {'region_name': region_name, 'endpoint_url': None}
+    )
+
+
+def _create_crt_s3_client(
+    session, config, region_name, credentials, lock, **kwargs
+):
+    """Create boto3 wrapper class to manage crt lock reference and S3 client."""
+    cred_wrapper = BotocoreCRTCredentialsWrapper(credentials)
+    cred_provider = cred_wrapper.to_crt_credentials_provider()
+    return CRTS3Client(
+        _create_crt_client(session, config, region_name, cred_provider),
+        lock,
+        region_name,
+        cred_wrapper,
+    )
+
+
+def _initialize_crt_transfer_primatives(client, config):
+    lock = acquire_crt_s3_process_lock(PROCESS_LOCK_NAME)
+    if lock is None:
+        # If we're unable to acquire the lock, we cannot
+        # use the CRT in this process and should default to
+        # the classic s3transfer manager.
+        return None, None
+
+    session = Session()
+    region_name = client.meta.region_name
+    credentials = client._get_credentials()
+
+    serializer = _create_crt_request_serializer(session, region_name)
+    s3_client = _create_crt_s3_client(
+        session, config, region_name, credentials, lock
+    )
+    return serializer, s3_client
+
+
+def get_crt_s3_client(client, config):
+    global CRT_S3_CLIENT
+    global BOTOCORE_CRT_SERIALIZER
+
+    with CLIENT_CREATION_LOCK:
+        if CRT_S3_CLIENT is None:
+            serializer, s3_client = _initialize_crt_transfer_primatives(
+                client, config
+            )
+            BOTOCORE_CRT_SERIALIZER = serializer
+            CRT_S3_CLIENT = s3_client
+
+    return CRT_S3_CLIENT
+
+
+class CRTS3Client:
+    """
+    This wrapper keeps track of our underlying CRT client, the lock used to
+    acquire it and the region we've used to instantiate the client.
+
+    Due to limitations in the existing CRT interfaces, we can only make calls
+    in a single region and does not support redirects. We track the region to
+    ensure we don't use the CRT client when a successful request cannot be made.
+    """
+
+    def __init__(self, crt_client, process_lock, region, cred_provider):
+        self.crt_client = crt_client
+        self.process_lock = process_lock
+        self.region = region
+        self.cred_provider = cred_provider
+
+
+def is_crt_compatible_request(client, crt_s3_client):
+    """
+    Boto3 client must use same signing region and credentials
+    as the CRT_S3_CLIENT singleton. Otherwise fallback to classic.
+    """
+    if crt_s3_client is None:
+        return False
+
+    boto3_creds = client._get_credentials()
+    if boto3_creds is None:
+        return False
+
+    is_same_identity = compare_identity(
+        boto3_creds.get_frozen_credentials(), crt_s3_client.cred_provider
+    )
+    is_same_region = client.meta.region_name == crt_s3_client.region
+    return is_same_region and is_same_identity
+
+
+def compare_identity(boto3_creds, crt_s3_creds):
+    try:
+        crt_creds = crt_s3_creds()
+    except botocore.exceptions.NoCredentialsError:
+        return False
+
+    is_matching_identity = (
+        boto3_creds.access_key == crt_creds.access_key_id
+        and boto3_creds.secret_key == crt_creds.secret_access_key
+        and boto3_creds.token == crt_creds.session_token
+    )
+    return is_matching_identity
+
+
+def create_crt_transfer_manager(client, config):
+    """Create a CRTTransferManager for optimized data transfer."""
+    crt_s3_client = get_crt_s3_client(client, config)
+    if is_crt_compatible_request(client, crt_s3_client):
+        return CRTTransferManager(
+            crt_s3_client.crt_client, BOTOCORE_CRT_SERIALIZER
+        )
+    return None