fix: redirect fix auth headers issue (#428)

Co-authored-by: liyang <[email protected]>

Author: bbtfr

megfile/config.py

@@ -60,3 +60,5 @@
 SFTP_MAX_RETRY_TIMES = int(
     os.getenv("MEGFILE_SFTP_MAX_RETRY_TIMES") or DEFAULT_MAX_RETRY_TIMES
 )
+
+HTTP_AUTH_HEADERS = ("Authorization", "Www-Authenticate", "Cookie", "Cookie2")

megfile/lib/s3_prefetch_reader.py

@@ -76,7 +76,11 @@ def _get_content_size(self):
         try:
             start, end = 0, self._block_size - 1
             first_index_response = self._fetch_response(start=start, end=end)
-            content_size = int(first_index_response["ContentRange"].split("/")[-1])
+            if "ContentRange" in first_index_response:
+                content_size = int(first_index_response["ContentRange"].split("/")[-1])
+            else:
+                # usually when read a file only have one block
+                content_size = int(first_index_response["ContentLength"])
         except S3InvalidRangeError:
             # usually when read a empty file
             # can use minio test empty file: https://hub.docker.com/r/minio/minio

megfile/s3_path.py

@@ -7,16 +7,18 @@
 from functools import cached_property, lru_cache, wraps
 from logging import getLogger as get_logger
 from typing import IO, Any, BinaryIO, Callable, Dict, Iterator, List, Optional, Tuple
+from urllib.parse import urlparse
 
 import boto3
 import botocore
-from botocore.awsrequest import AWSResponse
+from botocore.awsrequest import AWSPreparedRequest, AWSResponse
 
 from megfile.config import (
     DEFAULT_BLOCK_SIZE,
     DEFAULT_MAX_BLOCK_SIZE,
     DEFAULT_MIN_BLOCK_SIZE,
     GLOBAL_MAX_WORKERS,
+    HTTP_AUTH_HEADERS,
     S3_CLIENT_CACHE_MODE,
     S3_MAX_RETRY_TIMES,
 )
@@ -76,6 +78,7 @@
     generate_cache_path,
     get_binary_mode,
     get_content_offset,
+    is_domain_or_subdomain,
     is_readable,
     necessary_params,
     process_local,
@@ -162,24 +165,30 @@ def before_callback(operation_model, request_dict, request_context):
         retry_callback=retry_callback,
     )
 
-    def patch_send_request(send_request):
-        def patched_send_request(request_dict, operation_model):
-            http, parsed_response = send_request(request_dict, operation_model)
+    def patch_send(send):
+        def patched_send(request: AWSPreparedRequest) -> AWSResponse:
+            response: AWSResponse = send(request)
             if (
-                request_dict["method"] == "GET"  # only support GET method for now
-                and http.status_code in (301, 302, 307, 308)
-                and "Location" in http.headers
+                request.method == "GET"  # only support GET method for now
+                and response.status_code in (301, 302, 307, 308)
+                and "Location" in response.headers
             ):
-                request_dict["url"] = http.headers["Location"]
-                http, parsed_response = send_request(request_dict, operation_model)
-            return http, parsed_response
-
-        return patched_send_request
+                # Permit sending auth/cookie headers from "foo.com" to "sub.foo.com".
+                # See also: https://go.dev/src/net/http/client.go#L980
+                location = response.headers["Location"]
+                ihost = urlparse(request.url).hostname
+                dhost = urlparse(location).hostname
+                if not is_domain_or_subdomain(dhost, ihost):
+                    for name in HTTP_AUTH_HEADERS:
+                        request.headers.pop(name, None)
+                request.url = location
+                response = send(request)
+            return response
+
+        return patched_send
 
     if redirect:
-        client._endpoint._send_request = patch_send_request(
-            client._endpoint._send_request
-        )
+        client._endpoint._send = patch_send(client._endpoint._send)
 
     return client
 

megfile/utils/__init__.py

@@ -346,3 +346,11 @@ def __get__(  # pyre-ignore[14]
                 val = self.func(cls)
                 setattr(cls, self.attrname, val)
         return val
+
+
+def is_domain_or_subdomain(sub, parent):
+    if sub == parent:
+        return True
+    if sub.endswith(f".{parent}"):
+        return True
+    return False

tests/test_s3.py

@@ -344,8 +344,8 @@ def test_get_s3_client(mocker):
         aws_session_token=session_token,
     )
 
-    # assert _send_request is not patched
-    assert "_send_request" not in client._endpoint.__dict__
+    # assert _send is not patched
+    assert "_send" not in client._endpoint.__dict__
 
     client = s3.get_s3_client(cache_key="test")
     assert client is s3.get_s3_client(cache_key="test")
@@ -383,8 +383,8 @@ def test_get_s3_client_v2():
         == "virtual"
     )
 
-    # assert _send_request is patched
-    assert "_send_request" in client._endpoint.__dict__
+    # assert _send is patched
+    assert "_send" in client._endpoint.__dict__
 
 
 def test_get_s3_client_from_env(mocker):
@@ -409,8 +409,8 @@ def test_get_s3_client_from_env(mocker):
         aws_session_token=session_token,
     )
 
-    # assert _send_request is patched
-    assert "_send_request" in client._endpoint.__dict__
+    # assert _send is patched
+    assert "_send" in client._endpoint.__dict__
 
 
 def test_get_s3_client_with_config(mocker):
@@ -439,8 +439,8 @@ def __eq__(self, other):
         aws_session_token=session_token,
     )
 
-    # assert _send_request is patched
-    assert "_send_request" in client._endpoint.__dict__
+    # assert _send is patched
+    assert "_send" in client._endpoint.__dict__
 
 
 def test_get_s3_session_threading(mocker):

tests/utils/test_init.py

@@ -9,6 +9,7 @@
     cached_classproperty,
     combine,
     get_human_size,
+    is_domain_or_subdomain,
     necessary_params,
     patch_rlimit,
 )
@@ -108,3 +109,12 @@ def test__is_pickle():
     fileObj.name = "test"
     fileObj.mode = "wb"
     assert _is_pickle(fileObj) is False
+
+
+def test_is_domain_or_subdomain():
+    assert is_domain_or_subdomain("test1.com", "test2.com") is False
+    assert is_domain_or_subdomain("test1.test.com", "test2.test.com") is False
+
+    assert is_domain_or_subdomain("test.com", "test.com") is True
+    assert is_domain_or_subdomain("test1.test.com", "test.com") is True
+    assert is_domain_or_subdomain("test.com", "test1.test.com") is False