feat(sso): use SSOTokenProvider in SSOCredentialProvider (#4145)

* feat(sso): use SSOTokenProvider in SSOCredentialProvider

* feat(sso): use SSOTokenProvider in SSOCredentialProvider implementation

* feat(credential-provider-sso): use SSOTokenProvider when new config format is detected

* feat(credential-provider-sso): accept suggestion

Co-authored-by: Trivikram Kamat <[email protected]>

* feat(credential-provider-sso): accept suggestion

Co-authored-by: Trivikram Kamat <[email protected]>

* feat(credential-provider-sso): accept suggestion

Co-authored-by: Trivikram Kamat <[email protected]>

* feat(credential-provider-sso): update error message in sso token retrieval

Co-authored-by: Trivikram Kamat <[email protected]>

Author: kuhe

packages/credential-provider-ini/src/resolveSsoCredentials.spec.ts

@@ -24,12 +24,14 @@ describe(resolveSsoCredentials.name, () => {
     sso_role_name: "mock_sso_role_name",
   });
 
-  const getMockValidatedSsoProfile = () => ({
-    sso_start_url: "mock_validated_sso_start_url",
-    sso_account_id: "mock_validated_sso_account_id",
-    sso_region: "mock_validated_sso_region",
-    sso_role_name: "mock_validated_sso_role_name",
-  });
+  const getMockValidatedSsoProfile = <T>(add: T = {} as T) =>
+    ({
+      sso_start_url: "mock_validated_sso_start_url",
+      sso_account_id: "mock_validated_sso_account_id",
+      sso_region: "mock_validated_sso_region",
+      sso_role_name: "mock_validated_sso_role_name",
+      ...add,
+    });
 
   afterEach(() => {
     jest.clearAllMocks();
@@ -95,4 +97,28 @@ describe(resolveSsoCredentials.name, () => {
       ssoRoleName: mockValidatedProfile.sso_role_name,
     });
   });
+
+  it("calls fromSSO with optional sso session name", async () => {
+    const mockProfile = getMockOriginalSsoProfile();
+    const mockValidatedProfile = getMockValidatedSsoProfile({
+      sso_session: "test-session",
+    });
+
+    const mockCreds: Credentials = {
+      accessKeyId: "mockAccessKeyId",
+      secretAccessKey: "mockSecretAccessKey",
+    };
+
+    (validateSsoProfile as jest.Mock).mockReturnValue(mockValidatedProfile);
+    (fromSSO as jest.Mock).mockReturnValue(() => Promise.resolve(mockCreds));
+
+    await resolveSsoCredentials(mockProfile);
+    expect(fromSSO).toHaveBeenCalledWith({
+      ssoStartUrl: mockValidatedProfile.sso_start_url,
+      ssoAccountId: mockValidatedProfile.sso_account_id,
+      ssoRegion: mockValidatedProfile.sso_region,
+      ssoRoleName: mockValidatedProfile.sso_role_name,
+      ssoSession: mockValidatedProfile.sso_session,
+    });
+  });
 });

packages/credential-provider-ini/src/resolveSsoCredentials.ts

@@ -4,10 +4,11 @@ import { SsoProfile } from "@aws-sdk/credential-provider-sso";
 export { isSsoProfile } from "@aws-sdk/credential-provider-sso";
 
 export const resolveSsoCredentials = (data: Partial<SsoProfile>) => {
-  const { sso_start_url, sso_account_id, sso_region, sso_role_name } = validateSsoProfile(data);
+  const { sso_start_url, sso_account_id, sso_session, sso_region, sso_role_name } = validateSsoProfile(data);
   return fromSSO({
     ssoStartUrl: sso_start_url,
     ssoAccountId: sso_account_id,
+    ssoSession: sso_session,
     ssoRegion: sso_region,
     ssoRoleName: sso_role_name,
   })();

packages/credential-provider-sso/package.json

@@ -27,6 +27,7 @@
     "@aws-sdk/client-sso": "*",
     "@aws-sdk/property-provider": "*",
     "@aws-sdk/shared-ini-file-loader": "*",
+    "@aws-sdk/token-providers": "*",
     "@aws-sdk/types": "*",
     "tslib": "^2.3.1"
   },

packages/credential-provider-sso/src/fromSSO.spec.ts

@@ -27,6 +27,8 @@ describe(fromSSO.name, () => {
     secretAccessKey: "mockSecretAccessKey",
   };
 
+  const mockProfileName = "mockProfileName";
+
   beforeEach(() => {
     (resolveSSOCredentials as jest.Mock).mockResolvedValue(mockCreds);
   });
@@ -36,7 +38,6 @@ describe(fromSSO.name, () => {
   });
 
   describe("all sso* values are not set", () => {
-    const mockProfileName = "mockProfileName";
     const mockInit = { profile: mockProfileName };
     const mockProfiles = { [mockProfileName]: mockSsoProfile };
 
@@ -97,6 +98,8 @@ describe(fromSSO.name, () => {
         ssoAccountId: mockValidatedSsoProfile.sso_account_id,
         ssoRegion: mockValidatedSsoProfile.sso_region,
         ssoRoleName: mockValidatedSsoProfile.sso_role_name,
+        profile: mockProfileName,
+        ssoSession: undefined,
       });
     });
   });
@@ -117,7 +120,12 @@ describe(fromSSO.name, () => {
   });
 
   it("calls resolveSSOCredentials if all sso* values are set", async () => {
-    const mockOptions = { ...mockSsoProfile, ssoClient: mockSsoClient };
+    const mockOptions = {
+      ...mockSsoProfile,
+      ssoClient: mockSsoClient,
+      profile: mockProfileName,
+      ssoSession: "sso-session-name",
+    };
     const receivedCreds = await fromSSO(mockOptions)();
     expect(receivedCreds).toStrictEqual(mockCreds);
     expect(resolveSSOCredentials).toHaveBeenCalledWith(mockOptions);

packages/credential-provider-sso/src/fromSSO.ts

@@ -1,6 +1,11 @@
 import { SSOClient } from "@aws-sdk/client-sso";
 import { CredentialsProviderError } from "@aws-sdk/property-provider";
-import { getProfileName, parseKnownFiles, SourceProfileInit } from "@aws-sdk/shared-ini-file-loader";
+import {
+  getProfileName,
+  loadSsoSessionData,
+  parseKnownFiles,
+  SourceProfileInit,
+} from "@aws-sdk/shared-ini-file-loader";
 import { CredentialProvider } from "@aws-sdk/types";
 
 import { isSsoProfile } from "./isSsoProfile";
@@ -13,6 +18,12 @@ export interface SsoCredentialsParameters {
    */
   ssoStartUrl: string;
 
+  /**
+   * SSO session identifier.
+   * Presence implies usage of the SSOTokenProvider.
+   */
+  ssoSession?: string;
+
   /**
    * The ID of the AWS account to use for temporary credentials.
    */
@@ -36,35 +47,85 @@ export interface FromSSOInit extends SourceProfileInit {
 /**
  * Creates a credential provider that will read from a credential_process specified
  * in ini files.
+ *
+ * The SSO credential provider must support both
+ *
+ * 1. the legacy profile format,
+ * @example
+ * ```
+ * [profile sample-profile]
+ * sso_account_id = 012345678901
+ * sso_region = us-east-1
+ * sso_role_name = SampleRole
+ * sso_start_url = https://www.....com/start
+ * ```
+ *
+ * 2. and the profile format for SSO Token Providers.
+ * @example
+ * ```
+ * [profile sso-profile]
+ * sso_session = dev
+ * sso_account_id = 012345678901
+ * sso_role_name = SampleRole
+ *
+ * [sso-session dev]
+ * sso_region = us-east-1
+ * sso_start_url = https://www.....com/start
+ * ```
  */
 export const fromSSO =
   (init: FromSSOInit & Partial<SsoCredentialsParameters> = {}): CredentialProvider =>
   async () => {
-    const { ssoStartUrl, ssoAccountId, ssoRegion, ssoRoleName, ssoClient } = init;
-    if (!ssoStartUrl && !ssoAccountId && !ssoRegion && !ssoRoleName) {
+    const { ssoStartUrl, ssoAccountId, ssoRegion, ssoRoleName, ssoClient, ssoSession } = init;
+    const profileName = getProfileName(init);
+
+    if (!ssoStartUrl && !ssoAccountId && !ssoRegion && !ssoRoleName && !ssoSession) {
       // Load the SSO config from shared AWS config file.
       const profiles = await parseKnownFiles(init);
-      const profileName = getProfileName(init);
       const profile = profiles[profileName];
 
+      if (profile.sso_session) {
+        const ssoSessions = await loadSsoSessionData(init);
+        const session = ssoSessions[profile.sso_session];
+        const conflictMsg = ` configurations in profile ${profileName} and sso-session ${profile.sso_session}`;
+        if (ssoRegion && ssoRegion !== session.sso_region) {
+          throw new CredentialsProviderError(`Conflicting SSO region` + conflictMsg, false);
+        }
+        if (ssoStartUrl && ssoStartUrl !== session.sso_start_url) {
+          throw new CredentialsProviderError(`Conflicting SSO start_url` + conflictMsg, false);
+        }
+        profile.sso_region = session.sso_region;
+        profile.sso_start_url = session.sso_start_url;
+      }
+
       if (!isSsoProfile(profile)) {
         throw new CredentialsProviderError(`Profile ${profileName} is not configured with SSO credentials.`);
       }
 
-      const { sso_start_url, sso_account_id, sso_region, sso_role_name } = validateSsoProfile(profile);
+      const { sso_start_url, sso_account_id, sso_region, sso_role_name, sso_session } = validateSsoProfile(profile);
       return resolveSSOCredentials({
         ssoStartUrl: sso_start_url,
+        ssoSession: sso_session,
         ssoAccountId: sso_account_id,
         ssoRegion: sso_region,
         ssoRoleName: sso_role_name,
         ssoClient: ssoClient,
+        profile: profileName,
       });
     } else if (!ssoStartUrl || !ssoAccountId || !ssoRegion || !ssoRoleName) {
       throw new CredentialsProviderError(
-        'Incomplete configuration. The fromSSO() argument hash must include "ssoStartUrl",' +
-          ' "ssoAccountId", "ssoRegion", "ssoRoleName"'
+        "Incomplete configuration. The fromSSO() argument hash must include " +
+          '"ssoStartUrl", "ssoAccountId", "ssoRegion", "ssoRoleName"'
       );
     } else {
-      return resolveSSOCredentials({ ssoStartUrl, ssoAccountId, ssoRegion, ssoRoleName, ssoClient });
+      return resolveSSOCredentials({
+        ssoStartUrl,
+        ssoSession,
+        ssoAccountId,
+        ssoRegion,
+        ssoRoleName,
+        ssoClient,
+        profile: profileName,
+      });
     }
   };

packages/credential-provider-sso/src/isSsoProfile.spec.ts

@@ -5,7 +5,7 @@ describe(isSsoProfile.name, () => {
     expect(isSsoProfile({})).toEqual(false);
   });
 
-  it.each(["sso_start_url", "sso_account_id", "sso_region", "sso_role_name"])(
+  it.each(["sso_start_url", "sso_account_id", "sso_region", "sso_session", "sso_role_name"])(
     "returns true if value at '%s' is of type string",
     (key) => {
       expect(isSsoProfile({ [key]: "string" })).toEqual(true);

packages/credential-provider-sso/src/isSsoProfile.ts

@@ -9,5 +9,6 @@ export const isSsoProfile = (arg: Profile): arg is Partial<SsoProfile> =>
   arg &&
   (typeof arg.sso_start_url === "string" ||
     typeof arg.sso_account_id === "string" ||
+    typeof arg.sso_session === "string" ||
     typeof arg.sso_region === "string" ||
     typeof arg.sso_role_name === "string");

packages/credential-provider-sso/src/resolveSSOCredentials.spec.ts

@@ -1,11 +1,22 @@
 import { GetRoleCredentialsCommand, SSOClient } from "@aws-sdk/client-sso";
 import { CredentialsProviderError } from "@aws-sdk/property-provider";
 import { getSSOTokenFromFile } from "@aws-sdk/shared-ini-file-loader";
+import * as tokenProviders from "@aws-sdk/token-providers";
 
 import { resolveSSOCredentials } from "./resolveSSOCredentials";
 
 jest.mock("@aws-sdk/shared-ini-file-loader");
 jest.mock("@aws-sdk/client-sso");
+jest.mock("@aws-sdk/token-providers", () => {
+  return {
+    fromSso: jest.fn(() => async () => {
+      return {
+        token: "mockAccessToken",
+        expiration: new Date(Date.now() + 6_000_000),
+      };
+    }),
+  };
+});
 
 describe(resolveSSOCredentials.name, () => {
   const mockToken = {
@@ -57,6 +68,16 @@ describe(resolveSSOCredentials.name, () => {
     }
   });
 
+  it("uses the SSOTokenProvider if SSO Session name is present", async () => {
+    await resolveSSOCredentials({
+      ...mockOptions,
+      ssoSession: "test-sso-session",
+    });
+    expect(tokenProviders.fromSso).toHaveBeenCalledWith({
+      profile: undefined,
+    });
+  });
+
   describe("throws error on expiration", () => {
     afterEach(async () => {
       const expectedError = new CredentialsProviderError(
@@ -134,18 +155,15 @@ describe(resolveSSOCredentials.name, () => {
     });
 
     it("returns valid credentials from sso.getRoleCredentials", async () => {
-      const receivedCreds = await resolveSSOCredentials(mockOptions);
-      expect(receivedCreds).toStrictEqual(receivedCreds);
+      await resolveSSOCredentials(mockOptions);
       expect(mockSsoSend).toHaveBeenCalledTimes(1);
     });
 
     it("creates SSO client with provided region, if client is not passed", async () => {
       const mockCustomSsoSend = jest.fn().mockResolvedValue({ roleCredentials: mockCreds });
       (SSOClient as jest.Mock).mockReturnValue({ send: mockCustomSsoSend });
 
-      const receivedCreds = await resolveSSOCredentials({ ...mockOptions, ssoClient: undefined });
-      expect(receivedCreds).toStrictEqual(receivedCreds);
-
+      await resolveSSOCredentials({ ...mockOptions, ssoClient: undefined });
       expect(mockCustomSsoSend).toHaveBeenCalledTimes(1);
     });
   });

packages/credential-provider-sso/src/resolveSSOCredentials.ts

@@ -1,6 +1,7 @@
 import { GetRoleCredentialsCommand, GetRoleCredentialsCommandOutput, SSOClient } from "@aws-sdk/client-sso";
 import { CredentialsProviderError } from "@aws-sdk/property-provider";
 import { getSSOTokenFromFile, SSOToken } from "@aws-sdk/shared-ini-file-loader";
+import { fromSso as getSsoTokenProvider } from "@aws-sdk/token-providers";
 import { Credentials } from "@aws-sdk/types";
 
 import { FromSSOInit, SsoCredentialsParameters } from "./fromSSO";
@@ -9,28 +10,46 @@ import { FromSSOInit, SsoCredentialsParameters } from "./fromSSO";
  * The time window (15 mins) that SDK will treat the SSO token expires in before the defined expiration date in token.
  * This is needed because server side may have invalidated the token before the defined expiration date.
  *
- * @internal
+ * @private
  */
 const EXPIRE_WINDOW_MS = 15 * 60 * 1000;
 
 const SHOULD_FAIL_CREDENTIAL_CHAIN = false;
 
+/**
+ * @private
+ */
 export const resolveSSOCredentials = async ({
   ssoStartUrl,
+  ssoSession,
   ssoAccountId,
   ssoRegion,
   ssoRoleName,
   ssoClient,
+  profile,
 }: FromSSOInit & SsoCredentialsParameters): Promise<Credentials> => {
   let token: SSOToken;
   const refreshMessage = `To refresh this SSO session run aws sso login with the corresponding profile.`;
-  try {
-    token = await getSSOTokenFromFile(ssoStartUrl);
-  } catch (e) {
-    throw new CredentialsProviderError(
-      `The SSO session associated with this profile is invalid. ${refreshMessage}`,
-      SHOULD_FAIL_CREDENTIAL_CHAIN
-    );
+
+  if (ssoSession) {
+    try {
+      const _token = await getSsoTokenProvider({ profile })();
+      token = {
+        accessToken: _token.token,
+        expiresAt: new Date(_token.expiration!).toISOString(),
+      };
+    } catch (e) {
+      throw new CredentialsProviderError(e.message, SHOULD_FAIL_CREDENTIAL_CHAIN);
+    }
+  } else {
+    try {
+      token = await getSSOTokenFromFile(ssoStartUrl);
+    } catch (e) {
+      throw new CredentialsProviderError(
+        `The SSO session associated with this profile is invalid. ${refreshMessage}`,
+        SHOULD_FAIL_CREDENTIAL_CHAIN
+      );
+    }
   }
 
   if (new Date(token.expiresAt).getTime() - Date.now() <= EXPIRE_WINDOW_MS) {

packages/credential-provider-sso/src/types.ts

@@ -17,6 +17,7 @@ export interface SSOToken {
  */
 export interface SsoProfile extends Profile {
   sso_start_url: string;
+  sso_session?: string;
   sso_account_id: string;
   sso_region: string;
   sso_role_name: string;