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,11 @@
 from cumulusci.core.exceptions import (
     CumulusCIException,
     DependencyResolutionError,
+    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
@@ -60,6 +62,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
@@ -90,7 +93,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 +113,67 @@ 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:
+                                # save() would take the lock again, which is already held
+                                self.keychain.set_org(self, self.global_org)
         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, in case
+        another process rotated them after this org config was loaded."""
+        if not self.keychain or not self.name:
+            return
+        try:
+            stored_config = self.keychain.reload_org(self.name).config
+        except OrgNotFound:
+            return
+        stored_token = stored_config.get("refresh_token")
+        if stored_token and stored_token != self.refresh_token:
+            for key in ("refresh_token", "access_token", "refresh_token_rotation"):
+                if key in stored_config:
+                    self.config[key] = stored_config[key]
+
+    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},
+        )
+        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 +210,16 @@ 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)
+        except OAuth2Error as e:
+            if "invalid_grant" not in str(e):
+                raise
+            raise OAuth2Error(
+                f"The refresh token for org {self.name} is no longer valid. "
+                "It may have expired from inactivity or been revoked. "
+                f"Run `cci org connect {self.name}` to connect the org again."
+            ) from e
 
     @property
     def lightning_base_url(self):
@@ -370,7 +439,14 @@ 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_rotation:
+            # 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)
+        else:
+            self.keychain.set_org(self, self.global_org)
 
     def get_domain(self):
         instance_url = self.config.get("instance_url", "")

cumulusci/core/config/tests/test_config.py

@@ -7,6 +7,7 @@
 from unittest import mock
 
 import pytest
+import requests
 import responses
 import yaml
 from github3.exceptions import NotFoundError
@@ -41,6 +42,7 @@
     BaseProjectKeychain,
 )
 from cumulusci.core.source import LocalFolderSource
+from cumulusci.oauth.exceptions import OAuth2Error
 from cumulusci.tests.util import CURRENT_SF_API_VERSION, DummyKeychain
 from cumulusci.utils import temporary_dir, touch
 from cumulusci.utils.version_strings import StrictVersion
@@ -914,6 +916,229 @@ def test_refresh_oauth_token__other_connected_app(self, OAuth2Client):
         assert client_config.client_id == "OTHER_ID"
         refresh_token.assert_called_once_with(mock.sentinel.refresh_token)
 
+    def _org_config_in_keychain(self, refresh_token="OLD_TOKEN"):
+        project_config = BaseProjectConfig(UniversalConfig())
+        keychain = BaseProjectKeychain(project_config, None)
+        config = OrgConfig(
+            {
+                "refresh_token": refresh_token,
+                "instance_url": "http://instance_url_111.com",
+            },
+            "test",
+            keychain=keychain,
+        )
+        keychain.set_org(config)
+        config._load_userinfo = mock.Mock()
+        config._load_orginfo = mock.Mock()
+        return config, keychain
+
+    @mock.patch("cumulusci.core.config.org_config.OrgConfig.OAuth2Client")
+    def test_refresh_oauth_token__rotated_token_is_saved(self, OAuth2Client):
+        config, keychain = self._org_config_in_keychain()
+        OAuth2Client.return_value.refresh_token.return_value = {
+            "access_token": "asdf",
+            "refresh_token": "NEW_TOKEN",
+        }
+
+        with mock.patch.object(keychain, "set_org") as save:
+            config.refresh_oauth_token(keychain)
+
+        assert config.refresh_token == "NEW_TOKEN"
+        assert config.refresh_token_rotation
+        save.assert_called_once()
+
+    @mock.patch("cumulusci.core.config.org_config.OrgConfig.OAuth2Client")
+    def test_refresh_oauth_token__rotated_token_without_keychain(self, OAuth2Client):
+        config = OrgConfig(
+            {
+                "refresh_token": "OLD_TOKEN",
+                "instance_url": "http://instance_url_111.com",
+            },
+            "test",
+        )
+        config._load_userinfo = mock.Mock()
+        config._load_orginfo = mock.Mock()
+        keychain = BaseProjectKeychain(BaseProjectConfig(UniversalConfig()), None)
+        OAuth2Client.return_value.refresh_token.return_value = {
+            "access_token": "asdf",
+            "refresh_token": "NEW_TOKEN",
+        }
+
+        with mock.patch.object(keychain, "set_org") as save:
+            config.refresh_oauth_token(keychain)
+
+        assert config.refresh_token == "NEW_TOKEN"
+        assert config.refresh_token_rotation
+        save.assert_not_called()
+
+    @mock.patch("cumulusci.core.config.org_config.OrgConfig.OAuth2Client")
+    def test_refresh_oauth_token__unrotated_token_is_not_saved(self, OAuth2Client):
+        config, keychain = self._org_config_in_keychain()
+        OAuth2Client.return_value.refresh_token.return_value = {
+            "access_token": "asdf",
+        }
+
+        with mock.patch.object(keychain, "set_org") as save:
+            config.refresh_oauth_token(keychain)
+
+        assert config.refresh_token == "OLD_TOKEN"
+        save.assert_not_called()
+
+    @mock.patch("cumulusci.core.config.org_config.OrgConfig.OAuth2Client")
+    def test_refresh_oauth_token__uses_latest_stored_tokens(self, OAuth2Client):
+        config, keychain = self._org_config_in_keychain()
+        stored_org = OrgConfig(
+            {"refresh_token": "STORED_TOKEN", "access_token": "STORED_ACCESS"}, "test"
+        )
+        refresh_token = OAuth2Client.return_value.refresh_token
+        refresh_token.return_value = {"access_token": "asdf"}
+
+        with mock.patch.object(keychain, "reload_org", return_value=stored_org):
+            with mock.patch.object(config, "_access_token_works") as works:
+                config.refresh_oauth_token(keychain)
+
+        refresh_token.assert_called_once_with("STORED_TOKEN")
+        works.assert_not_called()
+
+    @responses.activate
+    @mock.patch("cumulusci.core.config.org_config.OrgConfig.OAuth2Client")
+    def test_refresh_oauth_token__rotation_reuses_working_access_token(
+        self, OAuth2Client
+    ):
+        config, keychain = self._org_config_in_keychain()
+        config.config.update(
+            {"access_token": "WORKING", "refresh_token_rotation": True}
+        )
+        responses.add(
+            responses.GET, "http://instance_url_111.com/services/oauth2/userinfo"
+        )
+
+        with mock.patch.object(keychain, "set_org") as save:
+            config.refresh_oauth_token(keychain)
+
+        OAuth2Client.return_value.refresh_token.assert_not_called()
+        save.assert_not_called()
+        assert config.access_token == "WORKING"
+        assert config.refresh_token == "OLD_TOKEN"
+        assert responses.calls[0].request.headers["Authorization"] == "Bearer WORKING"
+
+    @pytest.mark.parametrize("userinfo_status", [401, 403])
+    @responses.activate
+    @mock.patch("cumulusci.core.config.org_config.OrgConfig.OAuth2Client")
+    def test_refresh_oauth_token__rotation_refreshes_expired_access_token(
+        self, OAuth2Client, userinfo_status
+    ):
+        config, keychain = self._org_config_in_keychain()
+        config.config.update(
+            {"access_token": "EXPIRED", "refresh_token_rotation": True}
+        )
+        responses.add(
+            responses.GET,
+            "http://instance_url_111.com/services/oauth2/userinfo",
+            status=userinfo_status,
+        )
+        refresh_token = OAuth2Client.return_value.refresh_token
+        refresh_token.return_value = {
+            "access_token": "NEW_ACCESS",
+            "refresh_token": "NEW_TOKEN",
+        }
+
+        with mock.patch.object(keychain, "set_org") as save:
+            config.refresh_oauth_token(keychain)
+
+        refresh_token.assert_called_once_with("OLD_TOKEN")
+        save.assert_called_once()
+        assert config.access_token == "NEW_ACCESS"
+        assert config.refresh_token == "NEW_TOKEN"
+
+    @responses.activate
+    @mock.patch("cumulusci.core.config.org_config.OrgConfig.OAuth2Client")
+    def test_refresh_oauth_token__rotation_userinfo_error(self, OAuth2Client):
+        config, keychain = self._org_config_in_keychain()
+        config.config.update(
+            {"access_token": "WORKING", "refresh_token_rotation": True}
+        )
+        responses.add(
+            responses.GET,
+            "http://instance_url_111.com/services/oauth2/userinfo",
+            status=503,
+        )
+
+        with pytest.raises(requests.exceptions.HTTPError):
+            config.refresh_oauth_token(keychain)
+
+        OAuth2Client.return_value.refresh_token.assert_not_called()
+
+    @mock.patch("cumulusci.core.config.org_config.OrgConfig.OAuth2Client")
+    def test_refresh_oauth_token__rotation_without_access_token(self, OAuth2Client):
+        config, keychain = self._org_config_in_keychain()
+        config.config["refresh_token_rotation"] = True
+        refresh_token = OAuth2Client.return_value.refresh_token
+        refresh_token.return_value = {"access_token": "asdf"}
+
+        config.refresh_oauth_token(keychain)
+
+        refresh_token.assert_called_once_with("OLD_TOKEN")
+
+    @mock.patch("cumulusci.core.config.org_config.OrgConfig.OAuth2Client")
+    def test_refresh_oauth_token__org_not_in_keychain(self, OAuth2Client):
+        project_config = BaseProjectConfig(UniversalConfig())
+        keychain = BaseProjectKeychain(project_config, None)
+        config = OrgConfig(
+            {
+                "refresh_token": "OLD_TOKEN",
+                "instance_url": "http://instance_url_111.com",
+            },
+            "test",
+            keychain=keychain,
+        )
+        config._load_userinfo = mock.Mock()
+        config._load_orginfo = mock.Mock()
+        refresh_token = OAuth2Client.return_value.refresh_token
+        refresh_token.return_value = {"access_token": "asdf"}
+
+        config.refresh_oauth_token(keychain)
+
+        refresh_token.assert_called_once_with("OLD_TOKEN")
+
+    @mock.patch("cumulusci.core.config.org_config.OrgConfig.OAuth2Client")
+    def test_refresh_oauth_token__holds_org_lock(self, OAuth2Client):
+        config, keychain = self._org_config_in_keychain()
+        lock = mock.MagicMock()
+
+        def refresh_token(token):
+            lock.__enter__.assert_called_once()
+            lock.__exit__.assert_not_called()
+            return {"access_token": "asdf"}
+
+        OAuth2Client.return_value.refresh_token.side_effect = refresh_token
+
+        with mock.patch.object(keychain, "org_lock", return_value=lock) as org_lock:
+            config.refresh_oauth_token(keychain)
+
+        org_lock.assert_called_once_with("test", False)
+        lock.__exit__.assert_called_once()
+
+    @mock.patch("cumulusci.core.config.org_config.OrgConfig.OAuth2Client")
+    def test_refresh_oauth_token__invalid_grant(self, OAuth2Client):
+        config, keychain = self._org_config_in_keychain()
+        OAuth2Client.return_value.refresh_token.side_effect = OAuth2Error(
+            'OAuth failed\nstatus_code: 400\ncontent: b\'{"error":"invalid_grant"}\''
+        )
+
+        with pytest.raises(OAuth2Error, match="Run `cci org connect test`"):
+            config.refresh_oauth_token(keychain)
+
+    @mock.patch("cumulusci.core.config.org_config.OrgConfig.OAuth2Client")
+    def test_refresh_oauth_token__other_oauth_error(self, OAuth2Client):
+        config, keychain = self._org_config_in_keychain()
+        OAuth2Client.return_value.refresh_token.side_effect = OAuth2Error(
+            "OAuth failed\nstatus_code: 500"
+        )
+
+        with pytest.raises(OAuth2Error, match="status_code: 500"):
+            config.refresh_oauth_token(keychain)
+
     @responses.activate
     def test_load_user_info__bad_json(self):
         config = OrgConfig(