fix(s3): update endpoints for writeGetObjectResponse for object lambda (#3662)

AllanZhengYP · web-flow · commit bb9f18edb622 · 2022-06-10T16:08:16.000-07:00
* fix(s3): update endpoints for writeGetObjectResponse for object lambda

* fix(middleware-sdk-s3): set chunked encoding for writeGetObjectResponse

* test(client-s3): set chunked encoding for writeGetObjectResponse
diff --git a/clients/client-s3/src/commands/WriteGetObjectResponseCommand.ts b/clients/client-s3/src/commands/WriteGetObjectResponseCommand.ts
@@ -1,4 +1,5 @@
 // smithy-typescript generated code
+import { getWriteGetObjectResponseEndpointPlugin } from "@aws-sdk/middleware-sdk-s3";
 import { getSerdePlugin } from "@aws-sdk/middleware-serde";
 import { HttpRequest as __HttpRequest, HttpResponse as __HttpResponse } from "@aws-sdk/protocol-http";
 import { Command as $Command } from "@aws-sdk/smithy-client";
@@ -93,6 +94,7 @@ export class WriteGetObjectResponseCommand extends $Command<
     options?: __HttpHandlerOptions
   ): Handler<WriteGetObjectResponseCommandInput, WriteGetObjectResponseCommandOutput> {
     this.middlewareStack.use(getSerdePlugin(configuration, this.serialize, this.deserialize));
+    this.middlewareStack.use(getWriteGetObjectResponseEndpointPlugin(configuration));
 
     const stack = clientStack.concat(this.middlewareStack);
 
diff --git a/clients/client-s3/test/S3.spec.ts b/clients/client-s3/test/S3.spec.ts
@@ -3,7 +3,7 @@ import { HttpRequest } from "@aws-sdk/protocol-http";
 import { BuildMiddleware, FinalizeRequestMiddleware, SerializeMiddleware } from "@aws-sdk/types";
 import chai from "chai";
 import chaiAsPromised from "chai-as-promised";
-import { PassThrough } from "stream";
+import { PassThrough, Readable } from "stream";
 
 import { S3 } from "../src/S3";
 
@@ -38,7 +38,7 @@ describe("endpoint", () => {
   });
 });
 
-describe("Accesspoint ARN", async () => {
+describe("Endpoints from ARN", () => {
   const endpointValidator: BuildMiddleware<any, any> = (next, context) => (args) => {
     // middleware intercept the request and return it early
     const request = args.request as HttpRequest;
@@ -52,51 +52,93 @@ describe("Accesspoint ARN", async () => {
     });
   };
 
-  it("should succeed with access point ARN", async () => {
-    const client = new S3({ region: "us-west-2" });
-    client.middlewareStack.add(endpointValidator, { step: "build", priority: "low" });
-    const result: any = await client.putObject({
-      Bucket: "arn:aws:s3:us-west-2:123456789012:accesspoint:myendpoint",
-      Key: "key",
-      Body: "body",
+  describe("Accesspoint ARN", async () => {
+    it("should succeed with access point ARN", async () => {
+      const client = new S3({ region: "us-west-2" });
+      client.middlewareStack.add(endpointValidator, { step: "build", priority: "low" });
+      const result: any = await client.putObject({
+        Bucket: "arn:aws:s3:us-west-2:123456789012:accesspoint:myendpoint",
+        Key: "key",
+        Body: "body",
+      });
+      expect(result.request.hostname).to.eql("myendpoint-123456789012.s3-accesspoint.us-west-2.amazonaws.com");
     });
-    expect(result.request.hostname).to.eql("myendpoint-123456789012.s3-accesspoint.us-west-2.amazonaws.com");
-  });
 
-  it("should sign request with region from ARN is useArnRegion is set", async () => {
-    const client = new S3({
-      region: "us-east-1",
-      useArnRegion: true,
-      credentials: { accessKeyId: "key", secretAccessKey: "secret" },
+    it("should sign request with region from ARN is useArnRegion is set", async () => {
+      const client = new S3({
+        region: "us-east-1",
+        useArnRegion: true,
+        credentials: { accessKeyId: "key", secretAccessKey: "secret" },
+      });
+      client.middlewareStack.add(endpointValidator, { step: "finalizeRequest", priority: "low" });
+      const result: any = await client.putObject({
+        Bucket: "arn:aws:s3:us-west-2:123456789012:accesspoint:myendpoint",
+        Key: "key",
+        Body: "body",
+      });
+      expect(result.request.hostname).to.eql("myendpoint-123456789012.s3-accesspoint.us-west-2.amazonaws.com");
+      // Sign request with us-west-2 region from bucket access point ARN
+      expect(result.request.headers.authorization).to.contain("/us-west-2/s3/aws4_request, SignedHeaders=");
     });
-    client.middlewareStack.add(endpointValidator, { step: "finalizeRequest", priority: "low" });
-    const result: any = await client.putObject({
-      Bucket: "arn:aws:s3:us-west-2:123456789012:accesspoint:myendpoint",
-      Key: "key",
-      Body: "body",
+  });
+
+  describe("Outposts ARN", () => {
+    it("should succeed with outposts ARN", async () => {
+      const OutpostId = "op-01234567890123456";
+      const AccountId = "123456789012";
+      const region = "us-west-2";
+      const credentials = { accessKeyId: "key", secretAccessKey: "secret" };
+      const client = new S3({ region: "us-east-1", credentials, useArnRegion: true });
+      client.middlewareStack.add(endpointValidator, { step: "finalizeRequest", priority: "low" });
+      const result: any = await client.putObject({
+        Bucket: `arn:aws:s3-outposts:${region}:${AccountId}:outpost/${OutpostId}/accesspoint/abc-111`,
+        Key: "key",
+        Body: "body",
+      });
+      expect(result.request.hostname).to.eql(`abc-111-${AccountId}.${OutpostId}.s3-outposts.us-west-2.amazonaws.com`);
+      const date = new Date().toISOString().slice(0, 10).replace(/-/g, ""); //20201029
+      expect(result.request.headers["authorization"]).contains(
+        `Credential=${credentials.accessKeyId}/${date}/${region}/s3-outposts/aws4_request`
+      );
     });
-    expect(result.request.hostname).to.eql("myendpoint-123456789012.s3-accesspoint.us-west-2.amazonaws.com");
-    // Sign request with us-west-2 region from bucket access point ARN
-    expect(result.request.headers.authorization).to.contain("/us-west-2/s3/aws4_request, SignedHeaders=");
   });
 
-  it("should succeed with outposts ARN", async () => {
-    const OutpostId = "op-01234567890123456";
-    const AccountId = "123456789012";
-    const region = "us-west-2";
-    const credentials = { accessKeyId: "key", secretAccessKey: "secret" };
-    const client = new S3({ region: "us-east-1", credentials, useArnRegion: true });
-    client.middlewareStack.add(endpointValidator, { step: "finalizeRequest", priority: "low" });
-    const result: any = await client.putObject({
-      Bucket: `arn:aws:s3-outposts:${region}:${AccountId}:outpost/${OutpostId}/accesspoint/abc-111`,
-      Key: "key",
-      Body: "body",
+  describe("Object Lambda ARN", () => {
+    it("should succeed with Object Lambda ARN", async () => {
+      const region = "us-east-1";
+      const credentials = { accessKeyId: "key", secretAccessKey: "secret" };
+      const client = new S3({ region, credentials });
+      client.middlewareStack.add(endpointValidator, { step: "finalizeRequest", priority: "low" });
+      const result: any = await client.putObject({
+        Bucket: "arn:aws:s3-object-lambda:us-east-1:123456789012:accesspoint/mybanner",
+        Key: "key",
+        Body: "body",
+      });
+      const date = new Date().toISOString().slice(0, 10).replace(/-/g, ""); //20201029
+      expect(result.request.hostname).to.eql("mybanner-123456789012.s3-object-lambda.us-east-1.amazonaws.com");
+      expect(result.request.headers["authorization"]).contains(
+        `Credential=${credentials.accessKeyId}/${date}/${region}/s3-object-lambda/aws4_request`
+      );
+    });
+
+    it("should update endpoint of WriteGetObjectResponse command", async () => {
+      const region = "us-east-1";
+      const credentials = { accessKeyId: "key", secretAccessKey: "secret" };
+      const client = new S3({ region, credentials });
+      client.middlewareStack.add(endpointValidator, { step: "finalizeRequest", priority: "low" });
+      const requestRoute = "route123";
+      const result: any = await client.writeGetObjectResponse({
+        RequestRoute: requestRoute,
+        RequestToken: "token",
+        Body: new Readable(),
+      });
+      const date = new Date().toISOString().slice(0, 10).replace(/-/g, ""); //20201029
+      expect(result.request.hostname).to.eql(`${requestRoute}.s3-object-lambda.us-east-1.amazonaws.com`);
+      expect(result.request.headers["authorization"]).contains(
+        `Credential=${credentials.accessKeyId}/${date}/${region}/s3-object-lambda/aws4_request`
+      );
+      expect(result.request.headers["transfer-encoding"]).to.equal("chunked");
     });
-    expect(result.request.hostname).to.eql(`abc-111-${AccountId}.${OutpostId}.s3-outposts.us-west-2.amazonaws.com`);
-    const date = new Date().toISOString().slice(0, 10).replace(/-/g, ""); //20201029
-    expect(result.request.headers["authorization"]).contains(
-      `Credential=${credentials.accessKeyId}/${date}/${region}/s3-outposts/aws4_request`
-    );
   });
 });
 
diff --git a/codegen/smithy-aws-typescript-codegen/src/main/java/software/amazon/smithy/aws/typescript/codegen/AddS3Config.java b/codegen/smithy-aws-typescript-codegen/src/main/java/software/amazon/smithy/aws/typescript/codegen/AddS3Config.java
@@ -172,6 +172,12 @@ public List<RuntimeClientPlugin> getClientPlugins() {
                                         (m, s, o) -> EXCEPTIONS_OF_200_OPERATIONS.contains(o.getId().getName(s))
                                                 && testServiceId(s))
                                 .build(),
+                RuntimeClientPlugin.builder()
+                                .withConventions(AwsDependency.S3_MIDDLEWARE.dependency,
+                                        "WriteGetObjectResponseEndpoint", HAS_MIDDLEWARE)
+                                .operationPredicate((m, s, o) -> testServiceId(s)
+                                                    && o.getId().getName(s).equals("WriteGetObjectResponse"))
+                                .build(),
                 RuntimeClientPlugin.builder()
                         .withConventions(AwsDependency.ADD_EXPECT_CONTINUE.dependency, "AddExpectContinue",
                                         HAS_MIDDLEWARE)
diff --git a/packages/middleware-sdk-s3/package.json b/packages/middleware-sdk-s3/package.json
@@ -19,6 +19,7 @@
   },
   "license": "Apache-2.0",
   "dependencies": {
+    "@aws-sdk/middleware-bucket-endpoint": "*",
     "@aws-sdk/protocol-http": "*",
     "@aws-sdk/types": "*",
     "@aws-sdk/util-arn-parser": "*",
diff --git a/packages/middleware-sdk-s3/src/index.ts b/packages/middleware-sdk-s3/src/index.ts
@@ -1,3 +1,4 @@
 export * from "./check-content-length-header";
 export * from "./throw-200-exceptions";
 export * from "./validate-bucket-name";
+export * from "./write-get-object-response-endpoint";
diff --git a/packages/middleware-sdk-s3/src/write-get-object-response-endpoint.spec.ts b/packages/middleware-sdk-s3/src/write-get-object-response-endpoint.spec.ts
@@ -0,0 +1,127 @@
+import { HttpRequest } from "@aws-sdk/protocol-http";
+
+import { writeGetObjectResponseEndpointMiddleware } from "./write-get-object-response-endpoint";
+
+describe("writeGetObjectResponseEndpointMiddlewareOptions", () => {
+  const mockNextHandler = jest.fn();
+  const mockRegionProvider = jest.fn();
+  const mockConfig = {
+    region: mockRegionProvider,
+    isCustomEndpoint: false,
+    disableHostPrefix: false,
+    runtime: "node",
+  };
+
+  beforeEach(() => {
+    jest.clearAllMocks();
+  });
+
+  it.each(["s3.amazonaws.com", "s3-external-1.amazonaws.com"])(
+    "should not update endpoint for global endpoints %s",
+    async (hostname) => {
+      const context = {} as any;
+      const handler = writeGetObjectResponseEndpointMiddleware(mockConfig)(mockNextHandler, context);
+      await handler({
+        request: new HttpRequest({ hostname }),
+        input: {
+          RequestRoute: "route",
+        },
+      });
+      expect(mockNextHandler.mock.calls.length).toBe(1);
+      expect(mockNextHandler.mock.calls[0][0].request.hostname).toEqual(hostname);
+      expect(context).toEqual({});
+      mockNextHandler.mockClear();
+    }
+  );
+
+  it.each([
+    ["us-west-2", "s3.us-west-2.amazonaws.com", "route1", "route1.s3-object-lambda.us-west-2.amazonaws.com"],
+    ["us-east-1", "s3.us-east-1.amazonaws.com", "route2", "route2.s3-object-lambda.us-east-1.amazonaws.com"],
+    [
+      "cn-northeast-1",
+      "s3.cn-northeast-1.amazonaws.com.cn",
+      "route",
+      "route.s3-object-lambda.cn-northeast-1.amazonaws.com.cn",
+    ],
+  ])("should update endpoint for regional endpoints", async (region, originalEndpoint, requestRoute, expected) => {
+    const context = {} as any;
+    mockRegionProvider.mockResolvedValueOnce(region);
+    const handler = writeGetObjectResponseEndpointMiddleware(mockConfig)(mockNextHandler, context);
+    await handler({
+      request: new HttpRequest({ hostname: originalEndpoint }),
+      input: {
+        RequestRoute: requestRoute,
+      },
+    });
+    expect(mockNextHandler.mock.calls.length).toBe(1);
+    expect(mockNextHandler.mock.calls[0][0].request.hostname).toBe(expected);
+    expect(context).toMatchObject({ signing_service: "s3-object-lambda" });
+    expect(mockNextHandler.mock.calls[0][0].request.headers["transfer-encoding"]).toBe("chunked");
+    mockNextHandler.mockClear();
+  });
+
+  it("should update endpoint for custom endpoints", async () => {
+    const context = {} as any;
+    const handler = writeGetObjectResponseEndpointMiddleware({ ...mockConfig, isCustomEndpoint: true })(
+      mockNextHandler,
+      context
+    );
+    await handler({
+      request: new HttpRequest({ hostname: "my-endpoint.com" }),
+      input: {
+        RequestRoute: "route",
+      },
+    });
+    expect(mockNextHandler.mock.calls.length).toBe(1);
+    expect(mockNextHandler.mock.calls[0][0].request.hostname).toBe("route.my-endpoint.com");
+    expect(mockNextHandler.mock.calls[0][0].request.headers["transfer-encoding"]).toBe("chunked");
+    expect(context).toMatchObject({ signing_service: "s3-object-lambda" });
+    mockNextHandler.mockClear();
+  });
+
+  it("should not prepend requestRoute parameter if disableHostPrefix is set", async () => {
+    mockRegionProvider.mockResolvedValueOnce("us-west-2");
+    const handler = writeGetObjectResponseEndpointMiddleware({ ...mockConfig, disableHostPrefix: true })(
+      mockNextHandler,
+      {} as any
+    );
+    await handler({
+      request: new HttpRequest({ hostname: "s3.us-west-2.amazonaws.com" }),
+      input: {
+        RequestRoute: "route",
+      },
+    });
+    expect(mockNextHandler.mock.calls.length).toBe(1);
+    expect(mockNextHandler.mock.calls[0][0].request.hostname).toBe("s3-object-lambda.us-west-2.amazonaws.com");
+  });
+
+  it("should not set chunked encoding if content-length is already set", async () => {
+    const context = {} as any;
+    const handler = writeGetObjectResponseEndpointMiddleware({ ...mockConfig })(mockNextHandler, context);
+    const headers = { "content-length": "123" };
+    await handler({
+      request: new HttpRequest({ hostname: "s3.us-west-2.amazonaws.com", headers }),
+      input: {
+        RequestRoute: "route",
+      },
+    });
+    expect(mockNextHandler.mock.calls.length).toBe(1);
+    expect(mockNextHandler.mock.calls[0][0].request.headers).toEqual(headers);
+  });
+
+  it("should not set the chunked encoding for non Node.js runtime", async () => {
+    const context = {} as any;
+    const handler = writeGetObjectResponseEndpointMiddleware({ ...mockConfig, runtime: "browser" })(
+      mockNextHandler,
+      context
+    );
+    await handler({
+      request: new HttpRequest({ hostname: "s3.us-west-2.amazonaws.com" }),
+      input: {
+        RequestRoute: "route",
+      },
+    });
+    expect(mockNextHandler.mock.calls.length).toBe(1);
+    expect(mockNextHandler.mock.calls[0][0].request.headers["transfer-encoding"]).not.toBeDefined();
+  });
+});
diff --git a/packages/middleware-sdk-s3/src/write-get-object-response-endpoint.ts b/packages/middleware-sdk-s3/src/write-get-object-response-endpoint.ts
