Repository: getmoto/moto
Pull request #8048 (2024-08-28 12:44:29)

[S3] Fix issue with large body not being correctly decoded in put_object


Patch:
diff --git a/moto/cognitoidp/models.py b/moto/cognitoidp/models.py
index 25c8450b3..81a639d8d 100644
--- a/moto/cognitoidp/models.py
+++ b/moto/cognitoidp/models.py
@@ -1389,13 +1389,14 @@ class CognitoIdpBackend(BaseBackend):
         )
 
         return {
+            "ChallengeParameters": {},
             "AuthenticationResult": {
                 "IdToken": id_token,
                 "AccessToken": access_token,
                 "RefreshToken": refresh_token,
                 "ExpiresIn": expires_in,
                 "TokenType": "Bearer",
-            }
+            },
         }
 
     def _validate_auth_flow(
@@ -1592,27 +1593,40 @@ class CognitoIdpBackend(BaseBackend):
             if not timestamp:
                 raise ResourceNotFoundError(timestamp)
 
-            if user.software_token_mfa_enabled:
+            if user.status == UserStatus.FORCE_CHANGE_PASSWORD:
                 return {
-                    "ChallengeName": "SOFTWARE_TOKEN_MFA",
+                    "ChallengeName": "NEW_PASSWORD_REQUIRED",
+                    "ChallengeParameters": {
+                        "USERNAME": username,
+                    },
                     "Session": session,
-                    "ChallengeParameters": {},
                 }
+            if user_pool.mfa_config == "ON" and not user.token_verified:
+                mfas_can_setup = []
+                if user_pool.token_mfa_config == {"Enabled": True}:
+                    mfas_can_setup.append("SOFTWARE_TOKEN_MFA")
+                if user_pool.sms_mfa_config:
+                    mfas_can_setup.append("SMS_MFA")
 
-            if user.sms_mfa_enabled:
                 return {
-                    "ChallengeName": "SMS_MFA",
+                    "ChallengeName": "MFA_SETUP",
+                    "ChallengeParameters": {"MFAS_CAN_SETUP": mfas_can_setup},
+                    "Session": session,
+                }
+            if user.software_token_mfa_enabled or (
+                user_pool.token_mfa_config == {"Enabled": True} and user.token_verified
+            ):
+                return {
+                    "ChallengeName": "SOFTWARE_TOKEN_MFA",
                     "Session": session,
                     "ChallengeParameters": {},
                 }
 
-            if user.status == UserStatus.FORCE_CHANGE_PASSWORD:
+            if user.sms_mfa_enabled:
                 return {
-                    "ChallengeName": "NEW_PASSWORD_REQUIRED",
-                    "ChallengeParameters": {
-                        "USERNAME": username,
-                    },
+                    "ChallengeName": "SMS_MFA",
                     "Session": session,
+                    "ChallengeParameters": {},
                 }
 
             del self.sessions[session]
@@ -1637,8 +1651,7 @@ class CognitoIdpBackend(BaseBackend):
 
         elif challenge_name == "MFA_SETUP":
             username, user_pool = self.sessions[session]
-            auth_params = self._log_user_in(user_pool, client, username)
-            return {"ChallengeParameters": {}, **auth_params}
+            return self._log_user_in(user_pool, client, username)
 
         else:
             return {}
@@ -2012,7 +2025,7 @@ class CognitoIdpBackend(BaseBackend):
             if (
                 user.software_token_mfa_enabled
                 and user.preferred_mfa_setting == "SOFTWARE_TOKEN_MFA"
-            ):
+            ) or (user_pool.mfa_config == "ON" and user.token_verified):
                 return {
                     "ChallengeName": "SOFTWARE_TOKEN_MFA",
                     "ChallengeParameters": {},
diff --git a/moto/s3/responses.py b/moto/s3/responses.py
index 684145cbf..84baf2f20 100644
--- a/moto/s3/responses.py
+++ b/moto/s3/responses.py
@@ -881,9 +881,11 @@ class S3Response(BaseResponse):
             # - you should not use any location constraint
             location_constraint = self._get_location_constraint()
             if self.region == DEFAULT_REGION_NAME:
-                # REGION = us-east-1 - we should never receive a LocationConstraint
-                if location_constraint:
+                # REGION = us-east-1 - we should never receive a LocationConstraint for `us-east-1`
+                # However, you can create buckets in other regions from us-east-1
+                if location_constraint == DEFAULT_REGION_NAME:
                     raise InvalidLocationConstraintException
+
             else:
                 # Non-Standard region - the LocationConstraint must be equal to the actual region the request was send to
                 if not location_constraint:
@@ -893,8 +895,11 @@ class S3Response(BaseResponse):
             if self.body and not location_constraint:
                 raise MalformedXML()
 
+            bucket_region = (
+                location_constraint if location_constraint else DEFAULT_REGION_NAME
+            )
             try:
-                new_bucket = self.backend.create_bucket(bucket_name, self.region)
+                new_bucket = self.backend.create_bucket(bucket_name, bucket_region)
             except BucketAlreadyExists:
                 new_bucket = self.backend.get_bucket(bucket_name)
                 if new_bucket.account_id == self.get_current_account():
@@ -1321,12 +1326,20 @@ class S3Response(BaseResponse):
             line = body_io.readline()
         return bytes(new_body)
 
-    def _handle_encoded_body(self, body: bytes, content_length: int) -> bytes:
+    def _handle_encoded_body(self, body: bytes) -> bytes:
+        decoded_body = b""
         body_io = io.BytesIO(body)
-        # first line should equal '{content_length}\r\n
-        body_io.readline()
-        # Body contains actual data next
-        return body_io.read(content_length)
+        # first line should equal '{content_length}\r\n' while the content_length is a hex number
+        content_length = int(body_io.readline().strip(), 16)
+        while content_length > 0:
+            # read the content_length of the actual data
+            decoded_body += body_io.read(content_length)
+            # next is line with just '\r\n' so we skip it
+            body_io.readline()
+            # read another line with '{content_length}\r\n'
+            content_length = int(body_io.readline().strip(), 16)
+
+        return decoded_body
         # last line should equal
         # amz-checksum-sha256:<..>\r\n
 
@@ -1685,9 +1698,7 @@ class S3Response(BaseResponse):
             request.headers.get("x-amz-content-sha256", None)
             == "STREAMING-UNSIGNED-PAYLOAD-TRAILER"
         ):
-            body = self._handle_encoded_body(
-                body, int(request.headers["x-amz-decoded-content-length"])
-            )
+            body = self._handle_encoded_body(body)
 
         bucket_key_enabled = request.headers.get(
             "x-amz-server-side-encryption-bucket-key-enabled", None


Test patch:
diff --git a/tests/test_cognitoidp/test_cognitoidp.py b/tests/test_cognitoidp/test_cognitoidp.py
index 26b0cec20..0e5077a78 100644
--- a/tests/test_cognitoidp/test_cognitoidp.py
+++ b/tests/test_cognitoidp/test_cognitoidp.py
@@ -3071,14 +3071,219 @@ def user_authentication_flow(
     }
 
 
+@cognitoidp_aws_verified(generate_secret=True, with_mfa="ON")
+@pytest.mark.aws_verified
+def test_user_authentication_flow_mfa_on(user_pool=None, user_pool_client=None):
+    conn = boto3.client("cognito-idp", "us-west-2")
+
+    username = str(uuid.uuid4())
+    password = "P2$Sword"
+
+    user_pool_id = user_pool["UserPool"]["Id"]
+    client_id = user_pool_client["UserPoolClient"]["ClientId"]
+    client_secret = user_pool_client["UserPoolClient"]["ClientSecret"]
+
+    secret_hash = pycognito.aws_srp.AWSSRP.get_secret_hash(
+        username=username, client_id=client_id, client_secret=client_secret
+    )
+    conn.sign_up(
+        ClientId=client_id, Username=username, Password=password, SecretHash=secret_hash
+    )
+
+    client_secret = conn.describe_user_pool_client(
+        UserPoolId=user_pool_id, ClientId=client_id
+    )["UserPoolClient"]["ClientSecret"]
+
+    conn.admin_confirm_sign_up(UserPoolId=user_pool_id, Username=username)
+
+    aws_srp = pycognito.aws_srp.AWSSRP(
+        username=username,
+        password=password,
+        pool_id=user_pool_id,
+        client_id=client_id,
+        client_secret=client_secret,
+        client=conn,
+    )
+    auth_params = aws_srp.get_auth_params()
+
+    result = conn.initiate_auth(
+        ClientId=client_id,
+        AuthFlow="USER_SRP_AUTH",
+        AuthParameters=auth_params,
+    )
+
+    challenge_response = aws_srp.process_challenge(
+        result["ChallengeParameters"], auth_params
+    )
+    result = conn.respond_to_auth_challenge(
+        ClientId=client_id,
+        ChallengeName=result["ChallengeName"],
+        ChallengeResponses=challenge_response,
+    )
+    assert result["ChallengeName"] == "MFA_SETUP"
+
+    # add mfa token
+    resp = conn.associate_software_token(Session=result["Session"])
+    secret_code = resp["SecretCode"]
+    totp = pyotp.TOTP(secret_code)
+    user_code = totp.now()
+
+    resp = conn.verify_software_token(Session=resp["Session"], UserCode=user_code)
+    assert resp["Status"] == "SUCCESS"
+
+    # authenticate user once again
+    result = conn.initiate_auth(
+        ClientId=client_id,
+        AuthFlow="USER_SRP_AUTH",
+        AuthParameters={
+            "USERNAME": username,
+            "SRP_A": pycognito.aws_srp.long_to_hex(aws_srp.large_a_value),
+            "SECRET_HASH": secret_hash,
+        },
+    )
+
+    challenge_response = aws_srp.process_challenge(
+        result["ChallengeParameters"], auth_params
+    )
+    result = conn.respond_to_auth_challenge(
+        ClientId=client_id,
+        ChallengeName=result["ChallengeName"],
+        ChallengeResponses=challenge_response,
+    )
+    assert result["ChallengeName"] == "SOFTWARE_TOKEN_MFA"
+
+    # AWS needs a new user code, otherwise they will throw an exception
+    #    ExpiredCodeException: Your software token has already been used once.
+    new_user_code = totp.now()
+    if allow_aws_request():
+        while user_code == new_user_code:
+            new_user_code = totp.now()
+            time.sleep(1)
+    result = conn.respond_to_auth_challenge(
+        ClientId=client_id,
+        Session=result["Session"],
+        ChallengeName=result["ChallengeName"],
+        ChallengeResponses={
+            "SOFTWARE_TOKEN_MFA_CODE": new_user_code,
+            "USERNAME": username,
+            "SECRET_HASH": auth_params["SECRET_HASH"],
+        },
+    )
+    assert result["AuthenticationResult"]["AccessToken"]
+    assert result["AuthenticationResult"]["IdToken"]
+    assert result["AuthenticationResult"]["RefreshToken"]
+
+
 @cognitoidp_aws_verified(generate_secret=True, with_mfa="OPTIONAL")
 @pytest.mark.aws_verified
-def test_user_authentication_flow(user_pool=None, user_pool_client=None):
+def test_user_authentication_flow_mfa_optional(user_pool=None, user_pool_client=None):
     conn = boto3.client("cognito-idp", "us-west-2")
 
-    user_authentication_flow(
-        conn, user_pool=user_pool, user_pool_client=user_pool_client
+    username = str(uuid.uuid4())
+    password = "P2$Sword"
+
+    user_pool_id = user_pool["UserPool"]["Id"]
+    client_id = user_pool_client["UserPoolClient"]["ClientId"]
+    client_secret = user_pool_client["UserPoolClient"]["ClientSecret"]
+
+    secret_hash = pycognito.aws_srp.AWSSRP.get_secret_hash(
+        username=username, client_id=client_id, client_secret=client_secret
+    )
+    conn.sign_up(
+        ClientId=client_id, Username=username, Password=password, SecretHash=secret_hash
+    )
+
+    client_secret = conn.describe_user_pool_client(
+        UserPoolId=user_pool_id, ClientId=client_id
+    )["UserPoolClient"]["ClientSecret"]
+
+    conn.admin_confirm_sign_up(UserPoolId=user_pool_id, Username=username)
+
+    aws_srp = pycognito.aws_srp.AWSSRP(
+        username=username,
+        password=password,
+        pool_id=user_pool_id,
+        client_id=client_id,
+        client_secret=client_secret,
+        client=conn,
     )
+    auth_params = aws_srp.get_auth_params()
+
+    result = conn.initiate_auth(
+        ClientId=client_id,
+        AuthFlow="USER_SRP_AUTH",
+        AuthParameters=auth_params,
+    )
+
+    challenge_response = aws_srp.process_challenge(
+        result["ChallengeParameters"], auth_params
+    )
+    result = conn.respond_to_auth_challenge(
+        ClientId=client_id,
+        ChallengeName=result["ChallengeName"],
+        ChallengeResponses=challenge_response,
+    )
+
+    # add mfa token
+    resp = conn.associate_software_token(
+        AccessToken=result["AuthenticationResult"]["AccessToken"]
+    )
+    secret_code = resp["SecretCode"]
+    totp = pyotp.TOTP(secret_code)
+    user_code = totp.now()
+
+    resp = conn.verify_software_token(
+        AccessToken=result["AuthenticationResult"]["AccessToken"],
+        UserCode=user_code,
+    )
+    assert resp["Status"] == "SUCCESS"
+
+    conn.set_user_mfa_preference(
+        AccessToken=result["AuthenticationResult"]["AccessToken"],
+        SoftwareTokenMfaSettings={"Enabled": True, "PreferredMfa": True},
+    )
+
+    # authenticate user once again
+    result = conn.initiate_auth(
+        ClientId=client_id,
+        AuthFlow="USER_SRP_AUTH",
+        AuthParameters={
+            "USERNAME": username,
+            "SRP_A": pycognito.aws_srp.long_to_hex(aws_srp.large_a_value),
+            "SECRET_HASH": secret_hash,
+        },
+    )
+
+    challenge_response = aws_srp.process_challenge(
+        result["ChallengeParameters"], auth_params
+    )
+    result = conn.respond_to_auth_challenge(
+        ClientId=client_id,
+        ChallengeName=result["ChallengeName"],
+        ChallengeResponses=challenge_response,
+    )
+    assert result["ChallengeName"] == "SOFTWARE_TOKEN_MFA"
+
+    # AWS needs a new user code, otherwise they will throw an exception
+    #    ExpiredCodeException: Your software token has already been used once.
+    new_user_code = totp.now()
+    if allow_aws_request():
+        while user_code == new_user_code:
+            new_user_code = totp.now()
+            time.sleep(1)
+    result = conn.respond_to_auth_challenge(
+        ClientId=client_id,
+        Session=result["Session"],
+        ChallengeName=result["ChallengeName"],
+        ChallengeResponses={
+            "SOFTWARE_TOKEN_MFA_CODE": new_user_code,
+            "USERNAME": username,
+            "SECRET_HASH": auth_params["SECRET_HASH"],
+        },
+    )
+    assert result["AuthenticationResult"]["AccessToken"]
+    assert result["AuthenticationResult"]["IdToken"]
+    assert result["AuthenticationResult"]["RefreshToken"]
 
 
 @mock_aws
@@ -4602,6 +4807,7 @@ def test_initiate_mfa_auth_USER_PASSWORD_AUTH_with_FORCE_CHANGE_PASSWORD_status(
     # 6. Check that the MFA challenge is received
     # 7. Return the MFA challenge
     # 8. Check that the access tokens are received
+    # 9. Login again
 
     client = boto3.client("cognito-idp", "us-west-2")
     username = str(uuid.uuid4())
@@ -4667,6 +4873,31 @@ def test_initiate_mfa_auth_USER_PASSWORD_AUTH_with_FORCE_CHANGE_PASSWORD_status(
     assert result["AuthenticationResult"]["RefreshToken"]
     assert result["AuthenticationResult"]["IdToken"]
 
+    result = client.initiate_auth(
+        ClientId=client_id,
+        AuthFlow="USER_PASSWORD_AUTH",
+        AuthParameters={"USERNAME": username, "PASSWORD": new_password},
+    )
+    assert result["ChallengeName"] == "SOFTWARE_TOKEN_MFA"
+    third_user_code = totp.now()
+    if allow_aws_request():
+        while third_user_code == new_user_code:
+            third_user_code = totp.now()
+            time.sleep(1)
+    result = client.respond_to_auth_challenge(
+        ClientId=client_id,
+        Session=result["Session"],
+        ChallengeName=result["ChallengeName"],
+        ChallengeResponses={
+            "SOFTWARE_TOKEN_MFA_CODE": third_user_code,
+            "USERNAME": username,
+        },
+    )
+    assert result["ChallengeParameters"] == {}
+    assert result["AuthenticationResult"]["AccessToken"]
+    assert result["AuthenticationResult"]["RefreshToken"]
+    assert result["AuthenticationResult"]["IdToken"]
+
 
 @mock_aws
 def test_initiate_auth_USER_PASSWORD_AUTH_user_not_found():
diff --git a/tests/test_s3/test_s3.py b/tests/test_s3/test_s3.py
index c9195d0e9..695f0f766 100644
--- a/tests/test_s3/test_s3.py
+++ b/tests/test_s3/test_s3.py
@@ -292,6 +292,22 @@ def test_create_bucket_with_wrong_location_constraint(constraint):
     )
 
 
+@aws_verified
+@pytest.mark.aws_verified
+@pytest.mark.parametrize("constraint", ["us-east-2", "eu-central-1"])
+def test_create_bucket_in_regions_from_us_east_1(constraint):
+    client = boto3.client("s3", region_name=DEFAULT_REGION_NAME)
+    bucket_name = str(uuid.uuid4())
+    client.create_bucket(
+        Bucket=bucket_name, CreateBucketConfiguration={"LocationConstraint": constraint}
+    )
+    assert (
+        client.get_bucket_location(Bucket=bucket_name)["LocationConstraint"]
+        == constraint
+    )
+    client.delete_bucket(Bucket=bucket_name)
+
+
 @mock_aws
 def test_create_existing_bucket_in_us_east_1():
     """Creating a bucket that already exists in us-east-1 returns the bucket.
@@ -909,6 +925,22 @@ def test_upload_file_with_checksum_algorithm():
     assert random_bytes == actual_content
 
 
+@mock_aws
+def test_put_large_with_checksum_algorithm():
+    s3_client = boto3.client("s3", region_name=DEFAULT_REGION_NAME)
+    bucket = "mybucket"
+    s3_client.create_bucket(Bucket=bucket)
+    s3_client.put_object(
+        Bucket=bucket,
+        Key="the-key",
+        Body=b"test" * 1000000,
+        ChecksumAlgorithm="SHA256",
+    )
+
+    resp = s3_client.get_object(Bucket="mybucket", Key="the-key")
+    assert resp["Body"].read() == b"test" * 1000000
+
+
 @mock_aws
 def test_put_chunked_with_v4_signature_in_body():
     bucket_name = "mybucket"