🐛 Workaround 429 rate limit error when loading profiles

AWS recently introduced rate liming for call to its SSO portal frontend
API. This can cause failure to load all profiles or can prevent
subsequent profiles loading to work as 429 response is returned.

To avoid triggering the error, we now load the profiles sequentially
from the SSO portal page.

.eslintrc.yml

@@ -14,6 +14,7 @@ globals:
   loadAwsProfiles: true
   saveAwsProfile: true
   saveAwsProfiles: true
+  removeAwsProfilesForPortalDomain: true
 parserOptions:
   ecmaVersion: latest
 rules: {}

src/content_scripts/profiles_info_loader.js

@@ -54,26 +54,32 @@ function extractAwsProfilesFromAccountSection(section) {
   return awsProfiles;
 }
 
-/*******************************************************************************
- * Main functions
- *******************************************************************************/
-
-async function findAndExpandAwsAccountSectionFolds(page) {
+async function* findAndExpandAwsAccountSectionFolds(page) {
   const accountsParentSection = findAwsAccountsParentSection(page);
   await expandSection(accountsParentSection);
   const accountSections = findAwsAccountSections(accountsParentSection);
-  await Promise.all(accountSections.map(expandSection));
+  // We open section sequentially with a small delay to avoid triggering
+  // the 429 too many requests error from AWS SSO Portal
+  for (const section of accountSections) {
+    yield expandSection(section);
+    await waitFor(50);
+  }
   return accountSections;
 }
 
-function extractAwsProfilesFromAllAccountSections(awsAccountSections) {
-  return awsAccountSections.map(extractAwsProfilesFromAccountSection).flat();
-}
-
 /*******************************************************************************
  * Main code
  *******************************************************************************/
 
-findAndExpandAwsAccountSectionFolds(document)
-  .then(extractAwsProfilesFromAllAccountSections)
-  .then(saveAwsProfiles);
+(async () => {
+  const profilesRemoved = await removeAwsProfilesForPortalDomain(findAwsPortalDomain());
+  const mergeWithPreviousProfileToKeepSettings = (profile) =>
+    Object.assign({}, profilesRemoved[profile.id] || {}, profile);
+
+  for await (let section of findAndExpandAwsAccountSectionFolds(document)) {
+    const profiles = extractAwsProfilesFromAccountSection(section).map(
+      mergeWithPreviousProfileToKeepSettings
+    );
+    await saveAwsProfiles(profiles);
+  }
+})();

src/lib/common.js

@@ -13,27 +13,38 @@ const getNextProfileColor = (index) => {
   return availableProfileColors[index % availableProfileColors.length];
 };
 
-function updateStoredProfile(storedProfiles, newProfile) {
-  const existingProfile = storedProfiles[newProfile.id] || {};
-  // We assign a default color if one doesn't exist yet
-  const color = getNextProfileColor(Object.keys(storedProfiles).length);
-  storedProfiles[newProfile.id] = Object.assign({ color }, existingProfile, newProfile);
-  return storedProfiles[newProfile.id];
+async function saveAwsProfile(awsProfile) {
+  const awsProfiles = await saveAwsProfiles([awsProfile]);
+  return awsProfiles[awsProfile.id];
 }
 
-async function saveAwsProfile(awsProfile) {
-  const storageContent = await browser.storage.local.get({ awsProfiles: {} });
-  const profile = updateStoredProfile(storageContent.awsProfiles, awsProfile);
-  await browser.storage.local.set(storageContent);
-  return profile;
+async function saveAwsProfiles(newAwsProfiles) {
+  const { awsProfiles = {} } = await browser.storage.local.get({ awsProfiles: {} });
+
+  newAwsProfiles.sort().forEach((profile) => {
+    const color = getNextProfileColor(Object.keys(awsProfiles).length);
+    awsProfiles[profile.id] = Object.assign({ color }, profile);
+  });
+
+  await browser.storage.local.set({ awsProfiles });
+  return awsProfiles;
 }
 
-async function saveAwsProfiles(awsProfiles) {
-  const currentStorage = await browser.storage.local.get({ awsProfiles: {} });
-  for (const profile of Object.values(awsProfiles).sort()) {
-    await updateStoredProfile(currentStorage.awsProfiles, profile);
-  }
-  await browser.storage.local.set(currentStorage);
+async function removeAwsProfilesForPortalDomain(portalDomain) {
+  const { awsProfiles = {} } = await browser.storage.local.get({ awsProfiles: {} });
+
+  const profilesToRemove = Object.values(awsProfiles).filter(
+    (profile) => profile.portalDomain === portalDomain
+  );
+
+  const removedProfiles = {};
+  profilesToRemove.forEach((profile) => {
+    removedProfiles[profile.id] = profile;
+    delete awsProfiles[profile.id];
+  });
+
+  await browser.storage.local.set({ awsProfiles });
+  return removedProfiles;
 }
 
 async function legacyLoadAwsProfiles() {
@@ -72,6 +83,7 @@ const common = {
   loadAwsProfiles,
   saveAwsProfile,
   saveAwsProfiles,
+  removeAwsProfilesForPortalDomain,
 };
 
 exports = common;

src/popup/popup.js

@@ -85,7 +85,6 @@ function createTableEntryFromAwsProfile(awsProfile) {
 }
 
 async function loadProfilesFromPortalPage() {
-  browser.storage.onChanged.addListener(refreshPopupDisplay);
   await browser.tabs.executeScript({ file: "/lib/common.js" });
   await browser.tabs.executeScript({ file: "/content_scripts/profiles_info_loader.js" });
 }
@@ -124,8 +123,6 @@ async function refreshPopupDisplay() {
   setPopupSectionVisibility("aws-access-portal", isOnAwsPortal);
   setPopupSectionVisibility("no-profile-info", !hasProfiles && !isOnAwsPortal);
   setPopupSectionVisibility("profiles", hasProfiles);
-
-  browser.storage.onChanged.removeListener(refreshPopupDisplay);
 }
 
 function installEventHandlers() {
@@ -148,3 +145,5 @@ refreshPopupDisplay().then(() => {
   installEventHandlers();
   document.querySelector("input#searchbox").focus();
 });
+
+browser.storage.onChanged.addListener(refreshPopupDisplay);

tests/background.test.js

@@ -69,6 +69,7 @@ test("Auto-populate with an AWS profile used by the user", async () => {
     configuration: { autoPopulateUsedProfiles: true },
   });
   // When
+  await browser.webRequest.onBeforeRequest.waitForListener();
   await browser.webRequest.onBeforeRequest.triggerListener(TEST_PROFILE_LOGIN_REQUEST);
   // Then
   const storageContent = await browser.storage.local.get();

tests/helper.js

@@ -103,6 +103,7 @@ const createFakeBrowser = (browserStorage, tabUrl) => {
       local: browserStorage || mockBrowserStorage(),
       onChanged: {
         removeListener: () => {},
+        addListener: () => {},
       },
     },
     tabs: {
@@ -212,7 +213,7 @@ const createFakePage = async (
     return { matches: true };
   };
 
-  dom.injectScripts = async (scripts, { waitCondition, waitTimeout } = {}) => {
+  dom.injectScripts = async (scripts) => {
     const scriptLoaders = scripts.map(
       (script) =>
         new Promise((resolve) => {
@@ -223,7 +224,6 @@ const createFakePage = async (
         })
     );
     await Promise.all(scriptLoaders);
-    await waitForCondition(waitCondition, waitTimeout);
   };
 
   // Workaround the issue that jsdom doesn't define these 2 ones

tests/profiles_info_loader.test.js

@@ -1,6 +1,11 @@
 const path = require("path");
 
-const { mockBrowserStorage, buildStorageContentForDomains, createFakePage } = require("./helper");
+const {
+  mockBrowserStorage,
+  buildStorageContentForDomains,
+  createFakePage,
+  waitForCondition,
+} = require("./helper");
 
 /*******************************************************************************
  * Constants
@@ -26,13 +31,19 @@ const createFakeAwsPortalPage = async (htmlFile, { browserStorage } = {}) => {
     browserStorage,
   });
 
-  page.injectScriptsAndWaitForStorageUpdate = async (scripts) => {
-    const beforeValue = await page.window.browser.storage.local.get({ awsProfiles: {} });
-    const browserStorageChanged = async () => {
-      const afterValue = await page.window.browser.storage.local.get({ awsProfiles: {} });
-      return JSON.stringify(beforeValue) != JSON.stringify(afterValue);
+  page.injectScriptsAndWaitForStorageUpdate = async (scripts, expectedProfilesCount) => {
+    await page.injectScripts(scripts);
+    await page.waitForBrowserStorageUpdate(expectedProfilesCount);
+  };
+
+  page.waitForBrowserStorageUpdate = async (expectedProfilesCount) => {
+    const browserStorageHasExpectedProfilesCount = async () => {
+      const storage = await page.window.browser.storage.local.get({ awsProfiles: {} });
+      return (
+        storage.awsProfiles && Object.keys(storage.awsProfiles).length === expectedProfilesCount
+      );
     };
-    await page.injectScripts(scripts, { waitCondition: browserStorageChanged });
+    await waitForCondition(browserStorageHasExpectedProfilesCount);
   };
 
   return page;
@@ -74,15 +85,20 @@ const testFiles = [
 test.each(testFiles)("Parse correctly AWS Portal page with $case", async ({ htmlFile }) => {
   // Given
   const awsPortalPage = await createFakeAwsPortalPage(htmlFile);
+  const expectedStorage = buildStorageContentForDomains("mysso");
+  const expectedProfilesCount = Object.keys(expectedStorage.awsProfiles).length;
   // We have to manually simulate section expansion click logic
   awsPortalPage.window.document
     .querySelectorAll("div.instance-section, div.logo")
     .forEach((elt) => elt.addEventListener("click", simulateExpandSection));
   // When
-  await awsPortalPage.injectScriptsAndWaitForStorageUpdate(PROFILES_INFO_LOADER_SCRIPTS);
+  await awsPortalPage.injectScriptsAndWaitForStorageUpdate(
+    PROFILES_INFO_LOADER_SCRIPTS,
+    expectedProfilesCount
+  );
   // Then
   const browserStorageContent = await awsPortalPage.window.browser.storage.local.get();
-  expect(browserStorageContent).toEqual(buildStorageContentForDomains("mysso"));
+  expect(browserStorageContent).toEqual(expectedStorage);
 });
 
 test("Merge profiles from two different AWS Portal pages in storage", async () => {
@@ -93,13 +109,17 @@ test("Merge profiles from two different AWS Portal pages in storage", async () =
       createFakeAwsPortalPage(`aws_portal.${domain}.html`, { browserStorage })
     )
   );
+  const expectedContent = buildStorageContentForDomains("mysso", "anothersso");
+  const expectedProfilesCount = Object.keys(expectedContent.awsProfiles).length;
   // When
   for (const page of awsPortalPages) {
-    await page.injectScriptsAndWaitForStorageUpdate(PROFILES_INFO_LOADER_SCRIPTS);
+    await page.injectScriptsAndWaitForStorageUpdate(
+      PROFILES_INFO_LOADER_SCRIPTS,
+      expectedProfilesCount
+    );
   }
   // Then
   const browserStorageContent = await browserStorage.get();
-  const expectedContent = buildStorageContentForDomains("mysso", "anothersso");
   expect(browserStorageContent).toEqual(expectedContent);
 });
 
@@ -111,16 +131,19 @@ test("Update profiles fom an existing AWS Portal pages in storage", async () =>
   const favoriteProfile = Object.values(storageContent.awsProfiles)[1];
   favoriteProfile.favorite = true;
   const expectedContent = structuredClone(storageContent);
+  const expectedProfilesCount = Object.keys(expectedContent.awsProfiles).length;
 
   // We add a new account not present in the portal and expect to be removed on update
-  awsProfiles.ToBeRemoved = awsProfiles.Production;
-  delete awsProfiles.Production;
+  awsProfiles.ToBeRemoved = Object.assign({}, Object.values(awsProfiles)[0], { id: "ToBeRemoved" });
 
   const browserStorage = mockBrowserStorage(storageContent);
   const awsPortalPage = await createFakeAwsPortalPage("aws_portal.mysso.html", { browserStorage });
 
   // When
-  await awsPortalPage.injectScriptsAndWaitForStorageUpdate(PROFILES_INFO_LOADER_SCRIPTS);
+  await awsPortalPage.injectScriptsAndWaitForStorageUpdate(
+    PROFILES_INFO_LOADER_SCRIPTS,
+    expectedProfilesCount
+  );
   // Then
   const browserStorageContent = await awsPortalPage.window.browser.storage.local.get();
   expect(browserStorageContent).toEqual(expectedContent);