From cbab94e901aec4650ab9e0eb19a4515e44d4ba62 Mon Sep 17 00:00:00 2001
From: Sam Chung <samc@seek.com.au>
Date: Sat, 7 May 2022 00:55:55 +1000
Subject: [PATCH] fix(lib-storage): add missing return keys (#2700)

* fix(lib-storage): add missing return keys

* feat: assert putResponse is defined

* feat: add encodeURIComponent to Bucket and Key

* refactor: move extra key insertion

* fix: correct wrong line being deleted

* chore: change to ternary operator

* chore: remove done from tests
---
 lib/lib-storage/src/Upload.spec.ts | 44 +++++++++++++++++++++++++++++-
 lib/lib-storage/src/Upload.ts      | 28 ++++++++++++++++---
 2 files changed, 67 insertions(+), 5 deletions(-)

diff --git a/lib/lib-storage/src/Upload.spec.ts b/lib/lib-storage/src/Upload.spec.ts
index bc36138ccbf5..91d224bdb4d8 100644
--- a/lib/lib-storage/src/Upload.spec.ts
+++ b/lib/lib-storage/src/Upload.spec.ts
@@ -21,13 +21,27 @@ const putObjectTaggingMock = jest.fn().mockResolvedValue({
   Success: "Tags have been applied!",
 });
 
+const endpointMock = jest.fn().mockResolvedValue({
+  hostname: "s3.region.amazonaws.com",
+  port: undefined,
+  protocol: "https:",
+  path: "/",
+  query: undefined,
+});
+
 jest.mock("@aws-sdk/client-s3", () => ({
   ...(jest.requireActual("@aws-sdk/client-s3") as {}),
   S3: jest.fn().mockReturnValue({
     send: sendMock,
+    config: {
+      endpoint: endpointMock,
+    },
   }),
   S3Client: jest.fn().mockReturnValue({
     send: sendMock,
+    config: {
+      endpoint: endpointMock,
+    },
   }),
   CreateMultipartUploadCommand: createMultipartMock,
   UploadPartCommand: uploadPartMock,
@@ -36,7 +50,7 @@ jest.mock("@aws-sdk/client-s3", () => ({
   PutObjectCommand: putObjectMock,
 }));
 
-import { S3 } from "@aws-sdk/client-s3";
+import { CompleteMultipartUploadCommandOutput, S3 } from "@aws-sdk/client-s3";
 import { Readable } from "stream";
 
 import { Progress, Upload } from "./index";
@@ -184,6 +198,34 @@ describe(Upload.name, () => {
     expect(putObjectTaggingMock).toHaveBeenCalledTimes(0);
   });
 
+  it("should return a Bucket, Key and Location fields when upload uses a PUT", async () => {
+    const buffer = Buffer.from("");
+    const actionParams = { ...params, Body: buffer };
+    const upload = new Upload({
+      params: actionParams,
+      client: new S3({}),
+    });
+
+    const result = (await upload.done()) as CompleteMultipartUploadCommandOutput;
+    expect(result.Key).toEqual("example-key");
+    expect(result.Bucket).toEqual("example-bucket");
+    expect(result.Location).toEqual("https://example-bucket.s3.region.amazonaws.com/example-key");
+  });
+
+  it("should return a Location field formatted in path style when forcePathStyle is true", async () => {
+    const buffer = Buffer.from("");
+    const actionParams = { ...params, Body: buffer };
+    const s3Client = new S3({});
+    s3Client.config.forcePathStyle = true;
+    const upload = new Upload({
+      params: actionParams,
+      client: s3Client,
+    });
+
+    const result = (await upload.done()) as CompleteMultipartUploadCommandOutput;
+    expect(result.Location).toEqual("https://s3.region.amazonaws.com/example-bucket/example-key");
+  });
+
   it("should upload using multi-part when parts are larger than part size", async () => {
     // create a string that's larger than 5MB.
     const partSize = 1024 * 1024 * 5;
diff --git a/lib/lib-storage/src/Upload.ts b/lib/lib-storage/src/Upload.ts
index 35b501fbc2f8..8625ba3ac675 100644
--- a/lib/lib-storage/src/Upload.ts
+++ b/lib/lib-storage/src/Upload.ts
@@ -2,16 +2,17 @@ import { AbortController, AbortSignal } from "@aws-sdk/abort-controller";
 import {
   CompletedPart,
   CompleteMultipartUploadCommand,
+  CompleteMultipartUploadCommandOutput,
   CreateMultipartUploadCommand,
   CreateMultipartUploadCommandOutput,
   PutObjectCommand,
   PutObjectCommandInput,
-  PutObjectCommandOutput,
   PutObjectTaggingCommand,
   ServiceOutputTypes,
   Tag,
   UploadPartCommand,
 } from "@aws-sdk/client-s3";
+import { extendedEncodeURIComponent } from "@aws-sdk/smithy-client";
 import { EventEmitter } from "events";
 
 import { byteLength } from "./bytelength";
@@ -55,7 +56,7 @@ export class Upload extends EventEmitter {
   uploadEvent?: string;
 
   private isMultiPart = true;
-  private putResponse?: PutObjectCommandOutput;
+  private putResponse?: CompleteMultipartUploadCommandOutput;
 
   constructor(options: Options) {
     super();
@@ -97,8 +98,27 @@ export class Upload extends EventEmitter {
   async __uploadUsingPut(dataPart: RawDataPart) {
     this.isMultiPart = false;
     const params = { ...this.params, Body: dataPart.data };
-    const putResult = await this.client.send(new PutObjectCommand(params));
-    this.putResponse = putResult;
+    const [putResult, endpoint] = await Promise.all([
+      this.client.send(new PutObjectCommand(params)),
+      this.client.config.endpoint(),
+    ]);
+
+    const locationKey = this.params
+      .Key!.split("/")
+      .map((segment) => extendedEncodeURIComponent(segment))
+      .join("/");
+    const locationBucket = extendedEncodeURIComponent(this.params.Bucket!);
+
+    const Location: string = this.client.config.forcePathStyle
+      ? `${endpoint.protocol}//${endpoint.hostname}/${locationBucket}/${locationKey}`
+      : `${endpoint.protocol}//${locationBucket}.${endpoint.hostname}/${locationKey}`;
+
+    this.putResponse = {
+      ...putResult,
+      Bucket: this.params.Bucket,
+      Key: this.params.Key,
+      Location,
+    };
     const totalSize = byteLength(dataPart.data);
     this.__notifyProgress({
       loaded: totalSize,