Unverified Commit 998af09c authored by Landon James's avatar Landon James Committed by GitHub
Browse files

Fix S3 ListParts pagination infinite loop (#3679)

## Motivation and Context
<!--- Why is this change required? What problem does it solve? -->
<!--- If it fixes an open issue, please link to the issue here -->
The paginator for the S3 `ListParts` operation could loop forever:
https://github.com/awslabs/aws-sdk-rust/issues/1143. This is because
most paginated operations use an empty next token as an indication that
pagination is exhausted. `ListParts` instead sets the final next token
as `0` causing the pagination to loop back to the first page and loop
forever. Instead of an empty next token `ListParts` uses `IsTruncated =
false` to indicate that pagination has been exhausted.

## Description
<!--- Describe your changes in detail -->
* Added a new trait `isTruncatedPaginatorTrait`
* Add that trait to the S3 `ListPartsOutput` shape
* Use the presence of that trait to vary the logic setting the
`is_empty` value in the paginator.
  * If the trait is absent it looks for an empty next token as always 
* if the trait is present the value is set based on the value of the
response's `is_truncated` field


## Testing
<!--- Please describe in detail how you tested your changes -->
<!--- Include details of your testing environment, and the tests you ran
to -->
<!--- see how your change affects other areas of the code, etc. -->
* Added integration test confirming that pagination terminates when a
response contains `is_truncated = false`

**Note:** I'm still working on turning this into a model test rather
than an S3 specific integration test, but I wanted to get some feedback
on the actual fix while I'm figuring that out)

## Checklist
<!--- If a checkbox below is not applicable, then please DELETE it
rather than leaving it unchecked -->
- [x] I have updated `CHANGELOG.next.toml` if I made changes to the
smithy-rs codegen or runtime crates
- [x] I have updated `CHANGELOG.next.toml` if I made changes to the AWS
SDK, generated SDK code, or SDK runtime crates

----

_By submitting this pull request, I confirm that you can use, modify,
copy, and redistribute this contribution, under the terms of your
choice._
parent 0cbeef35
Loading
Loading
Loading
Loading
+6 −0
Original line number Diff line number Diff line
@@ -29,6 +29,12 @@ references = ["smithy-rs#3675"]
meta = { "breaking" = false, "tada" = false, "bug" = true, "target" = "client" }
author = "dastrom"

[[aws-sdk-rust]]
message = "Fix S3 ListParts API paginator infinite loop."
references = ["aws-sdk-rust#1143"]
meta = { "breaking" = false, "tada" = false, "bug" = true }
author = "landonxjames"

[[aws-smithy-runtime-api]]
message = "Add conversions from smithy StatusCode to http StatusCode."
references = ["smithy-rs#3637"]
+8 −1
Original line number Diff line number Diff line
@@ -10,6 +10,7 @@ import software.amazon.smithy.rust.codegen.client.smithy.customizations.DocsRsMe
import software.amazon.smithy.rust.codegen.client.smithy.customize.ClientCodegenDecorator
import software.amazon.smithy.rust.codegen.client.smithy.customize.CombinedClientCodegenDecorator
import software.amazon.smithy.rustsdk.customize.DisabledAuthDecorator
import software.amazon.smithy.rustsdk.customize.IsTruncatedPaginatorDecorator
import software.amazon.smithy.rustsdk.customize.RemoveDefaultsDecorator
import software.amazon.smithy.rustsdk.customize.apigateway.ApiGatewayDecorator
import software.amazon.smithy.rustsdk.customize.applyDecorators
@@ -70,6 +71,7 @@ val DECORATORS: List<ClientCodegenDecorator> =
            S3Decorator(),
            S3ExpressDecorator(),
            S3ExtendedRequestIdDecorator(),
            IsTruncatedPaginatorDecorator(),
        ),
        S3ControlDecorator().onlyApplyTo("com.amazonaws.s3control#AWSS3ControlServiceV20180820"),
        STSDecorator().onlyApplyTo("com.amazonaws.sts#AWSSecurityTokenServiceV20110615"),
@@ -78,7 +80,12 @@ val DECORATORS: List<ClientCodegenDecorator> =
        TimestreamDecorator().onlyApplyTo("com.amazonaws.timestreamquery#Timestream_20181101"),
        // Only build docs-rs for linux to reduce load on docs.rs
        listOf(
            DocsRsMetadataDecorator(DocsRsMetadataSettings(targets = listOf("x86_64-unknown-linux-gnu"), allFeatures = true)),
            DocsRsMetadataDecorator(
                DocsRsMetadataSettings(
                    targets = listOf("x86_64-unknown-linux-gnu"),
                    allFeatures = true,
                ),
            ),
        ),
    ).flatten()

+44 −0
Original line number Diff line number Diff line
/*
 * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
 * SPDX-License-Identifier: Apache-2.0
 */

package software.amazon.smithy.rustsdk.customize

import software.amazon.smithy.model.Model
import software.amazon.smithy.model.shapes.ServiceShape
import software.amazon.smithy.model.shapes.Shape
import software.amazon.smithy.model.shapes.ShapeId
import software.amazon.smithy.model.shapes.StructureShape
import software.amazon.smithy.model.transform.ModelTransformer
import software.amazon.smithy.rust.codegen.client.smithy.ClientRustSettings
import software.amazon.smithy.rust.codegen.client.smithy.customize.ClientCodegenDecorator
import software.amazon.smithy.rust.codegen.client.smithy.traits.IsTruncatedPaginatorTrait
import software.amazon.smithy.rust.codegen.core.util.letIf
import java.util.logging.Logger

/**
 * Decorator for adding isTruncatedPaginator trait
 */
class IsTruncatedPaginatorDecorator : ClientCodegenDecorator {
    override val name: String = "IsTruncatedPaginatorDecorator"
    override val order: Byte = 0
    private val logger: Logger = Logger.getLogger(javaClass.name)
    private val operationsWithIsTruncatedPaginator = setOf(ShapeId.from("com.amazonaws.s3#ListPartsOutput"))

    override fun transformModel(
        service: ServiceShape,
        model: Model,
        settings: ClientRustSettings,
    ): Model =
        ModelTransformer.create().mapShapes(model) { shape ->
            shape.letIf(isInIsTruncatedList(shape)) {
                logger.info("Adding IsTruncatedPaginator trait to $it")
                (it as StructureShape).toBuilder().addTrait(IsTruncatedPaginatorTrait()).build()
            }
        }

    private fun isInIsTruncatedList(shape: Shape): Boolean {
        return shape.isStructureShape && operationsWithIsTruncatedPaginator.contains(shape.id)
    }
}
+1 −0
Original line number Diff line number Diff line
@@ -155,6 +155,7 @@ class S3Decorator : ClientCodegenDecorator {
                                    )
                                }
                            }

                            else -> {}
                        }
                    }
+159 −0
Original line number Diff line number Diff line
/*
 * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
 * SPDX-License-Identifier: Apache-2.0
 */

package software.amazon.smithy.rustsdk.customize

import org.junit.jupiter.api.Test
import software.amazon.smithy.model.shapes.ShapeId
import software.amazon.smithy.model.shapes.StructureShape
import software.amazon.smithy.model.transform.ModelTransformer
import software.amazon.smithy.rust.codegen.client.smithy.traits.IsTruncatedPaginatorTrait
import software.amazon.smithy.rust.codegen.core.rustlang.CargoDependency
import software.amazon.smithy.rust.codegen.core.rustlang.rustTemplate
import software.amazon.smithy.rust.codegen.core.smithy.RuntimeType.Companion.preludeScope
import software.amazon.smithy.rust.codegen.core.testutil.asSmithyModel
import software.amazon.smithy.rust.codegen.core.testutil.integrationTest
import software.amazon.smithy.rust.codegen.core.util.letIf
import software.amazon.smithy.rustsdk.AwsRuntimeType
import software.amazon.smithy.rustsdk.awsSdkIntegrationTest

class IsTruncatedPaginatorTest {
    private val model =
        """
        namespace test

        use aws.protocols#restXml
        use aws.api#service
        use smithy.rules#endpointRuleSet

        @restXml
        @service(sdkId: "fake")
        @endpointRuleSet({
            "version": "1.0",
            "rules": [{ "type": "endpoint", "conditions": [], "endpoint": { "url": "https://example.com" } }],
            "parameters": {
                "Region": { "required": false, "type": "String", "builtIn": "AWS::Region" },
            }
        })
        service TestService {
            operations: [PaginatedList]
        }

        @readonly
        @optionalAuth
        @http(uri: "/PaginatedList", method: "POST")
        @paginated(inputToken: "nextToken", outputToken: "nextToken",
                   pageSize: "maxResults", items: "items")
        operation PaginatedList {
            input: GetFoosInput,
            output: GetFoosOutput
        }

        structure GetFoosInput {
            maxResults: Integer,
            nextToken: String
        }

        structure GetFoosOutput {
            nextToken: String,
            items: StringList,
            isTruncated: Boolean,
        }

        list StringList {
            member: String
        }
        """.asSmithyModel()

    @Test
    fun `isTruncated paginators work`() {
        // Adding IsTruncated trait to the output shape
        val modifiedModel =
            ModelTransformer.create().mapShapes(model) { shape ->
                shape.letIf(shape.isStructureShape && shape.toShapeId() == ShapeId.from("test#GetFoosOutput")) {
                    (it as StructureShape).toBuilder().addTrait(IsTruncatedPaginatorTrait()).build()
                }
            }

        awsSdkIntegrationTest(modifiedModel) { context, rustCrate ->
            val rc = context.runtimeConfig
            val moduleName = context.moduleUseName()
            rustCrate.integrationTest("is_truncated_paginator") {
                rustTemplate(
                    """
                    ##![cfg(feature = "test-util")]

                    use $moduleName::Config;
                    use $moduleName::Client;
                    use #{Region};
                    use aws_smithy_runtime::client::http::test_util::{ReplayEvent, StaticReplayClient};
                    use aws_smithy_types::body::SdkBody;

                    fn mk_response(part_marker: u8) -> http::Response<SdkBody> {
                        let (part_num_marker, next_num_marker, is_truncated) = if part_marker < 3 {
                            (part_marker, part_marker + 1, true)
                        } else {
                            (part_marker, 0, false)
                        };
                        let body = format!(
                            "<?xml version=\"1.0\" encoding=\"UTF-8\"?>\n
                        <GetFoosOutput
                            xmlns=\"http://s3.amazonaws.com/doc/2006-03-01/\">
                            <token>{part_num_marker}</token>
                            <nextToken>{next_num_marker}</nextToken>
                            <isTruncated>{is_truncated}</isTruncated>
                        </GetFoosOutput>"
                        );
                        http::Response::builder().body(SdkBody::from(body)).unwrap()
                    }

                    fn mk_request() -> http::Request<SdkBody> {
                        http::Request::builder()
                            .uri("https://some-test-bucket.s3.us-east-1.amazonaws.com/test.txt?part-number-marker=PartNumberMarker&uploadId=UploadId")
                            .body(SdkBody::empty())
                            .unwrap()
                    }

                    ##[#{tokio}::test]
                    async fn is_truncated_pagination_does_not_loop() {
                        let http_client = StaticReplayClient::new(vec![
                            ReplayEvent::new(mk_request(), mk_response(0)),
                            ReplayEvent::new(mk_request(), mk_response(1)),
                            ReplayEvent::new(mk_request(), mk_response(2)),
                            ReplayEvent::new(mk_request(), mk_response(3)),
                            //The events below should never be called because the pagination should
                            //terminate with the event above
                            ReplayEvent::new(mk_request(), mk_response(0)),
                            ReplayEvent::new(mk_request(), mk_response(1)),
                        ]);

                        let config = Config::builder()
                            .region(Region::new("fake"))
                            .http_client(http_client.clone())
                            .with_test_defaults()
                            .build();
                        let client = Client::from_conf(config);

                        let list_parts_res = client
                            .paginated_list()
                            .max_results(1)
                            .into_paginator()
                            .send()
                            .collect::<Vec<_>>()
                            .await;

                        // Confirm that the pagination stopped calling the http client after the
                        // first page with is_truncated = false
                        assert_eq!(list_parts_res.len(), 4)
                    }
                    """,
                    *preludeScope,
                    "tokio" to CargoDependency.Tokio.toType(),
                    "Region" to AwsRuntimeType.awsTypes(rc).resolve("region::Region"),
                )
            }
        }
    }
}
Loading