Support PKCE and refresh token rotation for connected app logins

cumulusci/cli/org.py

@@ -134,6 +134,7 @@ def setup_client(connected_app, login_url=None, sandbox=None) -> OAuth2Client:
         auth_uri=f"{base_uri}/services/oauth2/authorize",
         token_uri=f"{base_uri}/services/oauth2/token",
         scope="web full refresh_token",
+        pkce=True,
     )
     return OAuth2Client(sf_client_config)
 

cumulusci/core/config/org_config.py

@@ -1,7 +1,7 @@
 import os
 import re
 from collections import defaultdict, namedtuple
-from contextlib import contextmanager
+from contextlib import contextmanager, nullcontext
 from datetime import date, datetime
 from typing import Optional
 from urllib.parse import urlparse
@@ -14,9 +14,12 @@
 from cumulusci.core.exceptions import (
     CumulusCIException,
     DependencyResolutionError,
+    OrgCannotBeLoaded,
+    OrgNotFound,
     ServiceNotConfigured,
 )
 from cumulusci.oauth.client import OAuth2Client, OAuth2ClientConfig
+from cumulusci.oauth.exceptions import OAuth2Error
 from cumulusci.oauth.salesforce import SANDBOX_LOGIN_URL, jwt_session
 from cumulusci.utils import parse_api_datetime
 from cumulusci.utils.fileutils import open_fs_resource
@@ -26,6 +29,13 @@
 SKIP_REFRESH = os.environ.get("CUMULUSCI_DISABLE_REFRESH")
 SANDBOX_MYDOMAIN_RE = re.compile(r"\.cs\d+\.my\.(.*)salesforce\.com")
 MYDOMAIN_RE = re.compile(r"\.my\.(.*)salesforce\.com")
+# Checking whether an access token still works is safe to retry, so it gives up
+# instead of keeping other cci processes waiting for the org's lock.
+USERINFO_TIMEOUT = 20  # seconds
+# Refreshing isn't safe to retry with refresh token rotation: if Salesforce rotated
+# the token but its response never arrived, sending the old token again would
+# revoke the whole token family. So only connecting has a timeout.
+TOKEN_REFRESH_TIMEOUT = (10, None)  # (connect, read) seconds
 
 
 VersionInfo = namedtuple("VersionInfo", ["id", "number"])
@@ -60,6 +70,7 @@ class OrgConfig(BaseConfig):
     default: bool
     client_id: str
     refresh_token: str
+    refresh_token_rotation: bool
     client_secret: str
     connected_app: str
     serialization_format: str
@@ -80,6 +91,9 @@ def __init__(self, config: dict, name: str, keychain=None, global_org=False):
         self._installed_packages = None
         self._is_person_accounts_enabled = None
         self._multiple_currencies_is_enabled = False
+        # The refresh token as of when this org was loaded from the keychain or last
+        # saved, or None if it didn't come from the keychain
+        self._loaded_refresh_token = None
 
         super().__init__(config)
 
@@ -90,7 +104,10 @@ def refresh_oauth_token(self, keychain, connected_app=None, is_sandbox=False):
         this is done using the Oauth2 JWT flow.
 
         Otherwise it is done using the Oauth2 Refresh Token flow using the connected app
-        configured in the keychain's connected_app service.
+        configured in the keychain's connected_app service. If the connected app uses
+        refresh token rotation, the stored access token is reused while it still works,
+        and a new refresh token is saved to the keychain right away, because the old one
+        no longer works.
 
         Also refreshes user and org info that is cached in the org config.
         """
@@ -107,13 +124,79 @@ def refresh_oauth_token(self, keychain, connected_app=None, is_sandbox=False):
                     auth_url=auth_url,
                     is_sandbox=is_sandbox,
                 )
+                if info != self.config:
+                    self.config.update(info)
             else:
-                info = self._refresh_token(keychain, connected_app)
-            if info != self.config:
-                self.config.update(info)
+                # With refresh token rotation, each refresh invalidates the refresh token
+                # it used and the access token issued before it, and reusing an old refresh
+                # token revokes the whole token family. So under a lock, start from the
+                # latest stored tokens, refresh only if the stored access token no longer
+                # works (which would end other cci processes' sessions), and save the new
+                # tokens at once.
+                with self._refresh_lock():
+                    self._load_latest_tokens()
+                    if not (self.refresh_token_rotation and self._access_token_works()):
+                        info = self._refresh_token(keychain, connected_app)
+                        new_refresh_token = info.get("refresh_token")
+                        rotated = (
+                            new_refresh_token
+                            and new_refresh_token != self.refresh_token
+                        )
+                        self.config.update(info)
+                        if rotated:
+                            self.config["refresh_token_rotation"] = True
+                            if self.keychain:
+                                self.save()
         self._load_userinfo()
         self._load_orginfo()
 
+    def _refresh_lock(self):
+        if self.keychain and self.name:
+            return self.keychain.org_lock(self.name, self.global_org)
+        return nullcontext()
+
+    def _load_latest_tokens(self):
+        """Use the tokens as currently stored in the keychain if another process
+        changed them after this org config was loaded from the keychain or saved.
+
+        Tokens in an org config that didn't come from the keychain, like one that
+        `cci org connect` just made, replace the stored ones."""
+        if not self.keychain or not self.name or self._loaded_refresh_token is None:
+            return
+        try:
+            stored_config = self.keychain.reload_org(self.name).config
+        except (OrgNotFound, OrgCannotBeLoaded):
+            # Nothing usable is stored, so this config's tokens replace it.
+            return
+        # If the org was connected again as another org or user, the stored tokens
+        # don't go with the rest of this config. Keep this config's own, as before.
+        # Only the org and user IDs at the end of the identity URL are compared,
+        # because its host can differ between logging in and refreshing.
+        stored_id, own_id = stored_config.get("id"), self.config.get("id")
+        if stored_id and own_id and stored_id.split("/")[-2:] != own_id.split("/")[-2:]:
+            return
+        stored_token = stored_config.get("refresh_token")
+        if stored_token and stored_token != self._loaded_refresh_token:
+            for key in ("refresh_token", "access_token", "refresh_token_rotation"):
+                if key in stored_config:
+                    self.config[key] = stored_config[key]
+            self._loaded_refresh_token = stored_token
+
+    def _access_token_works(self) -> bool:
+        if not self.access_token:
+            return False
+        response = requests.get(
+            self.instance_url + "/services/oauth2/userinfo",
+            headers={"Authorization": "Bearer " + self.access_token},
+            timeout=USERINFO_TIMEOUT,
+        )
+        if response.status_code in (401, 403):
+            return False
+        # Other errors, like a 503 during maintenance, don't mean the token is dead,
+        # and refreshing anyway would end other cci processes' sessions.
+        response.raise_for_status()
+        return True
+
     @contextmanager
     def save_if_changed(self):
         orig_config = self.config.copy()
@@ -150,7 +233,27 @@ def _refresh_token(self, keychain, connected_app):
             scope="web full refresh_token",
         )
         sf_oauth = self.OAuth2Client(sf_oauth_config)
-        return sf_oauth.refresh_token(self.refresh_token)
+        try:
+            return sf_oauth.refresh_token(
+                self.refresh_token, timeout=TOKEN_REFRESH_TIMEOUT
+            )
+        except OAuth2Error as e:
+            error = None
+            if e.response is not None:
+                try:
+                    error = e.response.json()
+                except ValueError:
+                    pass
+            if not isinstance(error, dict) or error.get("error") != "invalid_grant":
+                raise
+            # invalid_grant also covers problems that reconnecting won't fix,
+            # like an inactive user, so show what Salesforce said.
+            description = error.get("error_description") or "invalid_grant"
+            raise OAuth2Error(
+                f"Salesforce rejected the refresh token for org {self.name}: "
+                f"{description}. If the token expired from inactivity or was revoked, "
+                f"run `cci org connect {self.name}` to connect the org again."
+            ) from e
 
     @property
     def lightning_base_url(self):
@@ -370,7 +473,15 @@ def reset_installed_packages(self):
 
     def save(self):
         assert self.keychain, "Keychain was not set on OrgConfig"
-        self.keychain.set_org(self, self.global_org)
+        if self.refresh_token:
+            # Another cci process may have rotated the tokens since this org was
+            # loaded. Saving the old ones would bring back a dead refresh token.
+            with self._refresh_lock():
+                self._load_latest_tokens()
+                self.keychain.set_org(self, self.global_org)
+                self._loaded_refresh_token = self.refresh_token
+        else:
+            self.keychain.set_org(self, self.global_org)
 
     def get_domain(self):
         instance_url = self.config.get("instance_url", "")