Bug 1920626 - Only update saved OAuth tokens if the scopes match. r=mkmelin

Differential Revision: https://phabricator.services.mozilla.com/D223252

--HG--
extra : rebase_source : bc89b0062adb5a8cc08d301f608c0e900cb99248

mailnews/base/src/OAuth2Module.sys.mjs

@@ -181,39 +181,41 @@ OAuth2Module.prototype = {
   },
   async setRefreshToken(token) {
     const scope = this._oauth.scope ?? this._scope;
+    const grantedScopes = scopeSet(scope);
 
-    // Check if we already have a login with this username, and modify the
-    // password on that, if we do.
+    // Update any existing logins matching this origin, username, and scope.
     const logins = Services.logins.findLogins(this._loginOrigin, null, "");
+    let didChangePassword = false;
     for (const login of logins) {
       if (login.username != this._username) {
         continue;
       }
 
-      if (token) {
-        log.debug(`Found existing login for ${this._loginOrigin}...`);
-        const propBag = Cc["@mozilla.org/hash-property-bag;1"].createInstance(
-          Ci.nsIWritablePropertyBag
-        );
-        if (token != login.password) {
-          log.debug("... changing password");
-          propBag.setProperty("password", token);
-        }
-        if (scope != login.httpRealm) {
+      const loginScopes = scopeSet(login.httpRealm);
+      if (grantedScopes.isSupersetOf(loginScopes)) {
+        if (grantedScopes.size == loginScopes.size) {
+          // The scope matches, just update the token.
           log.debug(
-            `... changing httpRealm from "${login.httpRealm}" to "${scope}"`
+            `Updating existing token for ${this._loginOrigin} with scope "${scope}"`
           );
-          propBag.setProperty("httpRealm", scope);
+          const propBag = Cc["@mozilla.org/hash-property-bag;1"].createInstance(
+            Ci.nsIWritablePropertyBag
+          );
+          propBag.setProperty("password", token);
+          Services.logins.modifyLogin(login, propBag);
+          didChangePassword = true;
+        } else {
+          // We've got a new token for this scope, remove the existing one.
+          log.debug(
+            `Removing superceded token for ${this._loginOrigin} with scope "${login.httpRealm}"`
+          );
+          Services.logins.removeLogin(login);
         }
-        Services.logins.modifyLogin(login, propBag);
-      } else {
-        Services.logins.removeLogin(login);
       }
-      return;
     }
 
-    // Unless the token is null, we need to create and fill in a new login
-    if (token) {
+    // Unless the token is null, we need to create and fill in a new login.
+    if (!didChangePassword && token) {
       log.debug(
         `Creating new login for ${this._loginOrigin} with httpRealm "${scope}"`
       );
@@ -260,8 +262,11 @@ OAuth2Module.prototype = {
 
         this._oauth.connect(shouldPrompt, false).then(
           async () => {
-            if (this._oauth.refreshToken != oldRefreshToken) {
-              // Refresh token changed; save it.
+            if (
+              this._oauth.refreshToken != oldRefreshToken ||
+              this._oauth.scope != this._scope
+            ) {
+              // Refresh token and/or scope changed; save it.
               await this.setRefreshToken(this._oauth.refreshToken);
             }
             if (this._prefRoot) {

mailnews/base/test/unit/test_oAuth2Module.js

@@ -321,6 +321,90 @@ add_task(async function testSetRefreshTokenWithNewScope() {
   Services.logins.removeAllLogins();
   OAuth2TestUtils.forgetObjects();
   OAuth2TestUtils.stopServer();
+  delete oAuth2Server.grantedScope;
+});
+
+add_task(async function testSetRefreshTokenPreservesOthers() {
+  // Create a server that makes a new token every time we use the current token.
+  const oAuth2Server = await OAuth2TestUtils.startServer();
+
+  // Tell the server to grant us a new scope. We won't be asking for it, but
+  // servers are weird.
+  oAuth2Server.grantedScope = "test_mail test_calendar";
+
+  await storeLogins([
+    ["https://test.test", "unknown_scope", "oscar@bar.invalid", "WRONG"],
+    ["oauth://test.test", "test_mail", "oscar@bar.invalid", "oscar-mail"],
+    [
+      "oauth://test.test",
+      "test_addressbook",
+      "oscar@bar.invalid",
+      "oscar-addressbook",
+    ],
+    [
+      "oauth://test.test",
+      "test_calendar",
+      "oscar@bar.invalid",
+      "oscar-calendar",
+    ],
+  ]);
+
+  // Connect.
+  const mod = new OAuth2Module();
+  mod.initFromHostname("test.test", "oscar@bar.invalid");
+  Assert.equal(
+    mod._oauth.refreshToken,
+    "",
+    "there should be no refresh token in memory"
+  );
+  mod._oauth.refreshToken = "refresh_token";
+
+  const deferred = Promise.withResolvers();
+  mod.connect(false, {
+    onSuccess: deferred.resolve,
+    onFailure: deferred.reject,
+  });
+  await deferred.promise;
+
+  Assert.equal(
+    mod._oauth.refreshToken,
+    "refresh_token",
+    "refresh token in memory should have been updated"
+  );
+  Assert.equal(
+    mod._oauth.scope,
+    "test_mail test_calendar",
+    "scope in memory should have been updated"
+  );
+
+  // Check that the new token was added and tokens it replaces are removed.
+  // This assumes that `getAllLogins` returns the logins in the order they
+  // were added. If this changes the test will need updating.
+  const logins = await Services.logins.getAllLogins();
+  Assert.equal(logins.length, 3, "there should be 3 remaining logins");
+
+  // Login with different origin is unchanged.
+  Assert.equal(logins[0].hostname, "https://test.test");
+  Assert.equal(logins[0].httpRealm, "unknown_scope");
+  Assert.equal(logins[0].username, "oscar@bar.invalid");
+  Assert.equal(logins[0].password, "WRONG");
+
+  // Token with a scope not granted in this test is unchanged.
+  Assert.equal(logins[1].hostname, "oauth://test.test");
+  Assert.equal(logins[1].httpRealm, "test_addressbook");
+  Assert.equal(logins[1].username, "oscar@bar.invalid");
+  Assert.equal(logins[1].password, "oscar-addressbook");
+
+  // Token with the new scopes replaces the individual tokens for those scopes.
+  Assert.equal(logins[2].hostname, "oauth://test.test");
+  Assert.equal(logins[2].httpRealm, "test_mail test_calendar");
+  Assert.equal(logins[2].username, "oscar@bar.invalid");
+  Assert.equal(logins[2].password, "refresh_token");
+
+  Services.logins.removeAllLogins();
+  OAuth2TestUtils.forgetObjects();
+  OAuth2TestUtils.stopServer();
+  delete oAuth2Server.grantedScope;
 });
 
 async function storeLogins(logins) {