feat(storage): enables location credentials provider

packages/storage/__tests__/providers/s3/apis/utils/resolveS3ConfigAndInput.test.ts

@@ -9,6 +9,11 @@ import {
 	StorageValidationErrorCode,
 	validationErrorMap,
 } from '../../../../../src/errors/types/validation';
+import {
+	CallbackPathStorageInput,
+	DeprecatedStorageInput,
+} from '../../../../../src/providers/s3/utils/resolveS3ConfigAndInput';
+import { INVALID_STORAGE_INPUT } from '../../../../../src/errors/constants';
 
 jest.mock('@aws-amplify/core', () => ({
 	ConsoleLogger: jest.fn(),
@@ -76,13 +81,11 @@ describe('resolveS3ConfigAndInput', () => {
 		}
 	});
 
-	it('should throw if identityId is not available', async () => {
+	it('should not throw if identityId is not available', async () => {
 		mockFetchAuthSession.mockResolvedValueOnce({
 			credentials,
 		});
-		await expect(resolveS3ConfigAndInput(Amplify, {})).rejects.toMatchObject(
-			validationErrorMap[StorageValidationErrorCode.NoIdentityId],
-		);
+		expect(async () => resolveS3ConfigAndInput(Amplify, {})).not.toThrow();
 	});
 
 	it('should resolve bucket from S3 config', async () => {
@@ -214,4 +217,74 @@ describe('resolveS3ConfigAndInput', () => {
 		});
 		expect(keyPrefix).toEqual('prefix');
 	});
+
+	describe('with locationCredentialsProvider', () => {
+		const mockLocationCredentialsProvider = jest
+			.fn()
+			.mockReturnValue({ credentials });
+		it('should resolve credentials without Amplify singleton', async () => {
+			mockGetConfig.mockReturnValue({
+				Storage: {
+					S3: {
+						bucket,
+						region,
+					},
+				},
+			});
+			const { s3Config } = await resolveS3ConfigAndInput(Amplify, {
+				locationCredentialsProvider: mockLocationCredentialsProvider,
+			});
+
+			if (typeof s3Config.credentials === 'function') {
+				const result = await s3Config.credentials();
+				expect(mockLocationCredentialsProvider).toHaveBeenCalled();
+				expect(result).toEqual(credentials);
+			} else {
+				throw new Error('Expect credentials to be a function');
+			}
+		});
+
+		describe('with deprecated or callback paths as inputs', () => {
+			const key = 'mock-value';
+			const prefix = 'mock-value';
+			const path = () => 'path';
+			const deprecatedInputs: DeprecatedStorageInput[] = [
+				{ prefix },
+				{ key },
+				{
+					source: { key },
+					destination: { key },
+				},
+			];
+			const callbackPathInputs: CallbackPathStorageInput[] = [
+				{ path },
+				{
+					destination: { path },
+					source: { path },
+				},
+			];
+
+			const testCases = [...deprecatedInputs, ...callbackPathInputs];
+
+			it.each(testCases)('should throw when input is %s', async input => {
+				const { s3Config } = await resolveS3ConfigAndInput(
+					Amplify,
+					{ locationCredentialsProvider: mockLocationCredentialsProvider },
+					input,
+				);
+				if (typeof s3Config.credentials === 'function') {
+					await expect(s3Config.credentials()).rejects.toThrow(
+						expect.objectContaining({
+							name: INVALID_STORAGE_INPUT,
+							message: 'The input needs to have a path as a string value.',
+							recoverySuggestion:
+								'Please provide a valid path as a string value for the input.',
+						}),
+					);
+				} else {
+					throw new Error('Expect credentials to be a function');
+				}
+			});
+		});
+	});
 });

packages/storage/src/errors/constants.ts

@@ -0,0 +1,4 @@
+// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// SPDX-License-Identifier: Apache-2.0
+
+export const INVALID_STORAGE_INPUT = 'InvalidStorageInput';

packages/storage/src/providers/s3/apis/downloadData.ts

@@ -115,7 +115,11 @@ const downloadDataJob =
 	> => {
 		const { options: downloadDataOptions } = downloadDataInput;
 		const { bucket, keyPrefix, s3Config, identityId } =
-			await resolveS3ConfigAndInput(Amplify, downloadDataOptions);
+			await resolveS3ConfigAndInput(
+				Amplify,
+				downloadDataOptions,
+				downloadDataInput,
+			);
 		const { inputType, objectKey } = validateStorageOperationInput(
 			downloadDataInput,
 			identityId,

packages/storage/src/providers/s3/apis/internal/copy.ts

@@ -40,8 +40,11 @@ const copyWithPath = async (
 	input: CopyWithPathInput,
 ): Promise<CopyWithPathOutput> => {
 	const { source, destination } = input;
-	const { s3Config, bucket, identityId } =
-		await resolveS3ConfigAndInput(amplify);
+	const { s3Config, bucket, identityId } = await resolveS3ConfigAndInput(
+		amplify,
+		{},
+		input,
+	);
 
 	assertValidationError(!!source.path, StorageValidationErrorCode.NoSourcePath);
 	assertValidationError(
@@ -92,10 +95,11 @@ export const copyWithKey = async (
 		s3Config,
 		bucket,
 		keyPrefix: sourceKeyPrefix,
-	} = await resolveS3ConfigAndInput(amplify, input.source);
+	} = await resolveS3ConfigAndInput(amplify, input.source, input);
 	const { keyPrefix: destinationKeyPrefix } = await resolveS3ConfigAndInput(
 		amplify,
 		input.destination,
+		input,
 	); // resolveS3ConfigAndInput does not make extra API calls or storage access if called repeatedly.
 
 	// TODO(ashwinkumar6) V6-logger: warn `You may copy files from another user if the source level is "protected", currently it's ${srcLevel}`

packages/storage/src/providers/s3/apis/internal/getProperties.ts

@@ -26,7 +26,7 @@ export const getProperties = async (
 ): Promise<GetPropertiesOutput | GetPropertiesWithPathOutput> => {
 	const { options: getPropertiesOptions } = input;
 	const { s3Config, bucket, keyPrefix, identityId } =
-		await resolveS3ConfigAndInput(amplify, getPropertiesOptions);
+		await resolveS3ConfigAndInput(amplify, getPropertiesOptions, input);
 	const { inputType, objectKey } = validateStorageOperationInput(
 		input,
 		identityId,

packages/storage/src/providers/s3/apis/internal/getUrl.ts

@@ -31,7 +31,7 @@ export const getUrl = async (
 ): Promise<GetUrlOutput | GetUrlWithPathOutput> => {
 	const { options: getUrlOptions } = input;
 	const { s3Config, keyPrefix, bucket, identityId } =
-		await resolveS3ConfigAndInput(amplify, getUrlOptions);
+		await resolveS3ConfigAndInput(amplify, getUrlOptions, input);
 	const { inputType, objectKey } = validateStorageOperationInput(
 		input,
 		identityId,

packages/storage/src/providers/s3/apis/internal/list.ts

@@ -58,7 +58,7 @@ export const list = async (
 		bucket,
 		keyPrefix: generatedPrefix,
 		identityId,
-	} = await resolveS3ConfigAndInput(amplify, options);
+	} = await resolveS3ConfigAndInput(amplify, options, input);
 
 	const { inputType, objectKey } = validateStorageOperationInputWithPrefix(
 		input,

packages/storage/src/providers/s3/apis/internal/remove.ts

@@ -25,7 +25,7 @@ export const remove = async (
 ): Promise<RemoveOutput | RemoveWithPathOutput> => {
 	const { options = {} } = input ?? {};
 	const { s3Config, keyPrefix, bucket, identityId } =
-		await resolveS3ConfigAndInput(amplify, options);
+		await resolveS3ConfigAndInput(amplify, options, input);
 
 	const { inputType, objectKey } = validateStorageOperationInput(
 		input,

packages/storage/src/providers/s3/apis/uploadData/multipart/uploadHandlers.ts

@@ -74,6 +74,7 @@ export const getMultipartUploadHandlers = (
 		const resolvedS3Options = await resolveS3ConfigAndInput(
 			Amplify,
 			uploadDataOptions,
+			uploadDataInput,
 		);
 
 		abortController = new AbortController();

packages/storage/src/providers/s3/apis/uploadData/putObjectJob.ts

@@ -29,7 +29,11 @@ export const putObjectJob =
 	async (): Promise<ItemWithKey | ItemWithPath> => {
 		const { options: uploadDataOptions, data } = uploadDataInput;
 		const { bucket, keyPrefix, s3Config, isObjectLockEnabled, identityId } =
-			await resolveS3ConfigAndInput(Amplify, uploadDataOptions);
+			await resolveS3ConfigAndInput(
+				Amplify,
+				uploadDataOptions,
+				uploadDataInput,
+			);
 		const { inputType, objectKey } = validateStorageOperationInput(
 			uploadDataInput,
 			identityId,

packages/storage/src/providers/s3/types/inputs.ts

@@ -35,19 +35,24 @@ import {
 	UploadDataOptionsWithPath,
 } from '../types';
 
+import { LocationCredentialsProvider } from './options';
+
 // TODO: support use accelerate endpoint option
 /**
  * @deprecated Use {@link CopyWithPathInput} instead.
  * Input type for S3 copy API.
  */
 export type CopyInput = StorageCopyInputWithKey<
 	CopySourceOptionsWithKey,
-	CopyDestinationOptionsWithKey
+	CopyDestinationOptionsWithKey,
+	{ locationCredentialsProvider?: LocationCredentialsProvider }
 >;
 /**
  * Input type with path for S3 copy API.
  */
-export type CopyWithPathInput = StorageCopyInputWithPath;
+export type CopyWithPathInput = StorageCopyInputWithPath<{
+	locationCredentialsProvider?: LocationCredentialsProvider;
+}>;
 
 /**
  * @deprecated Use {@link GetPropertiesWithPathInput} instead.

packages/storage/src/providers/s3/utils/resolveIdentityId.ts

@@ -0,0 +1,11 @@
+// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// SPDX-License-Identifier: Apache-2.0
+
+import { StorageValidationErrorCode } from '../../../errors/types/validation';
+import { assertValidationError } from '../../../errors/utils/assertValidationError';
+
+export const resolveIdentityId = (identityId?: string): string => {
+	assertValidationError(!!identityId, StorageValidationErrorCode.NoIdentityId);
+
+	return identityId;
+};

packages/storage/src/providers/s3/utils/resolveS3ConfigAndInput.ts

@@ -6,14 +6,27 @@ import { AmplifyClassV6, StorageAccessLevel } from '@aws-amplify/core';
 import { assertValidationError } from '../../../errors/utils/assertValidationError';
 import { StorageValidationErrorCode } from '../../../errors/types/validation';
 import { resolvePrefix as defaultPrefixResolver } from '../../../utils/resolvePrefix';
-import { ResolvedS3Config } from '../types/options';
+import {
+	LocationCredentialsProvider,
+	ResolvedS3Config,
+} from '../types/options';
+import {
+	StorageCopyInputWithPath,
+	StorageOperationInputWithKey,
+	StorageOperationInputWithPath,
+	StorageOperationInputWithPrefix,
+} from '../../../types/inputs';
+import { StorageError } from '../../../errors/StorageError';
+import { CopyInput } from '../types';
+import { INVALID_STORAGE_INPUT } from '../../../errors/constants';
 
 import { DEFAULT_ACCESS_LEVEL, LOCAL_TESTING_S3_ENDPOINT } from './constants';
 
 interface S3ApiOptions {
 	accessLevel?: StorageAccessLevel;
 	targetIdentityId?: string;
 	useAccelerateEndpoint?: boolean;
+	locationCredentialsProvider?: LocationCredentialsProvider;
 }
 
 interface ResolvedS3ConfigAndInput {
@@ -23,6 +36,14 @@ interface ResolvedS3ConfigAndInput {
 	isObjectLockEnabled?: boolean;
 	identityId?: string;
 }
+export type DeprecatedStorageInput =
+	| StorageOperationInputWithKey
+	| StorageOperationInputWithPrefix
+	| CopyInput;
+
+export type CallbackPathStorageInput =
+	| StorageOperationInputWithPath
+	| StorageCopyInputWithPath;
 
 /**
  * resolve the common input options for S3 API handlers from Amplify configuration and library options.
@@ -38,13 +59,13 @@ interface ResolvedS3ConfigAndInput {
 export const resolveS3ConfigAndInput = async (
 	amplify: AmplifyClassV6,
 	apiOptions?: S3ApiOptions,
+	input?: DeprecatedStorageInput | CallbackPathStorageInput,
 ): Promise<ResolvedS3ConfigAndInput> => {
 	/**
 	 * IdentityId is always cached in memory so we can safely make calls here. It
 	 * should be stable even for unauthenticated users, regardless of credentials.
 	 */
 	const { identityId } = await amplify.Auth.fetchAuthSession();
-	assertValidationError(!!identityId, StorageValidationErrorCode.NoIdentityId);
 
 	/**
 	 * A credentials provider function instead of a static credentials object is
@@ -53,7 +74,13 @@ export const resolveS3ConfigAndInput = async (
 	 * credentials if they are expired.
 	 */
 	const credentialsProvider = async () => {
-		const { credentials } = await amplify.Auth.fetchAuthSession();
+		if (isLocationCredentialsProvider(apiOptions)) {
+			assertStorageInput(input);
+		}
+
+		const { credentials } = isLocationCredentialsProvider(apiOptions)
+			? await apiOptions.locationCredentialsProvider()
+			: await amplify.Auth.fetchAuthSession();
 		assertValidationError(
 			!!credentials,
 			StorageValidationErrorCode.NoCredentials,
@@ -101,3 +128,72 @@ export const resolveS3ConfigAndInput = async (
 		isObjectLockEnabled,
 	};
 };
+
+const isLocationCredentialsProvider = (
+	options?: S3ApiOptions,
+): options is S3ApiOptions & {
+	locationCredentialsProvider: LocationCredentialsProvider;
+} => {
+	return !!options?.locationCredentialsProvider;
+};
+
+const isInputWithCallbackPath = (input?: CallbackPathStorageInput) => {
+	return (
+		((input as StorageOperationInputWithPath)?.path &&
+			typeof (input as StorageOperationInputWithPath).path === 'function') ||
+		((input as StorageCopyInputWithPath)?.destination?.path &&
+			typeof (input as StorageCopyInputWithPath).destination?.path ===
+				'function') ||
+		((input as StorageCopyInputWithPath)?.source?.path &&
+			typeof (input as StorageCopyInputWithPath).source?.path === 'function')
+	);
+};
+
+const isDeprecatedInput = (
+	input?: DeprecatedStorageInput | CallbackPathStorageInput,
+): input is DeprecatedStorageInput => {
+	return (
+		isInputWithKey(input) ||
+		isInputWithPrefix(input) ||
+		isInputWithCopySourceOrDestination(input)
+	);
+};
+const assertStorageInput = (
+	input?: DeprecatedStorageInput | CallbackPathStorageInput,
+) => {
+	if (isDeprecatedInput(input) || isInputWithCallbackPath(input)) {
+		throw new StorageError({
+			name: INVALID_STORAGE_INPUT,
+			message: 'The input needs to have a path as a string value.',
+			recoverySuggestion:
+				'Please provide a valid path as a string value for the input.',
+		});
+	}
+};
+
+const isInputWithKey = (
+	input?: DeprecatedStorageInput | CallbackPathStorageInput,
+): input is StorageOperationInputWithKey => {
+	return !!(
+		(input as StorageOperationInputWithKey)?.key &&
+		typeof (input as StorageOperationInputWithKey).key === 'string'
+	);
+};
+const isInputWithPrefix = (
+	input?: DeprecatedStorageInput | CallbackPathStorageInput,
+): input is StorageOperationInputWithPrefix => {
+	return !!(
+		(input as StorageOperationInputWithPrefix)?.prefix &&
+		typeof (input as StorageOperationInputWithPrefix).prefix === 'string'
+	);
+};
+const isInputWithCopySourceOrDestination = (
+	input?: CopyInput | CallbackPathStorageInput,
+): input is CopyInput => {
+	return !!(
+		((input as CopyInput)?.source?.key &&
+			typeof (input as CopyInput).source?.key === 'string') ||
+		((input as CopyInput)?.destination?.key &&
+			typeof (input as CopyInput).destination?.key === 'string')
+	);
+};