[gerrit_util] Add linked account detection for SSO

Bug: b/348024314
Change-Id: Ic982de2892769870805407c6a00856943133dd62
Reviewed-on: https://chromium-review.googlesource.com/c/chromium/tools/depot_tools/+/5651293
Reviewed-by: Robbie Iannucci <iannucci@chromium.org>
Reviewed-by: Yiwei Zhang <yiwzhang@google.com>
Commit-Queue: Allen Li <ayatane@chromium.org>

gerrit_util.py

@@ -171,16 +171,33 @@ class SSOHelper(object):
 ssoHelper = SSOHelper()
 
 
-def ShouldUseSSO() -> bool:
+def ShouldUseSSO(host: str) -> bool:
     """Return True if we should use SSO for the current user."""
+    LOGGER.debug("Determining whether we should use SSO...")
     if not newauth.Enabled():
+        LOGGER.debug("SSO=False: not opted in")
+        return False
+    if newauth.SkipSSO():
+        LOGGER.debug("SSO=False: set skip SSO config")
         return False
     if not ssoHelper.find_cmd():
+        LOGGER.debug("SSO=False: no SSO command")
         return False
     cwd = os.getcwd()
     email = scm.GIT.GetConfig(cwd, 'user.email', default='')
-    # TODO(ayatane): enable logic not finished, for linked accounts
-    return email.endswith('@google.com')
+    if email.endswith('@google.com'):
+        LOGGER.debug("SSO=True: email is google.com")
+        return True
+    if not email.endswith('@chromium.org'):
+        LOGGER.debug("SSO=False: not chromium.org")
+        return False
+    authenticator = SSOAuthenticator()
+    records = GetAccountEmails(host, 'self', authenticator=authenticator)
+    if any(email == r['email'] for r in records):
+        LOGGER.debug("SSO=True: email is linked to google.com")
+        return True
+    LOGGER.debug("SSO=False: unlinked chromium.org")
+    return False
 
 
 class Authenticator(object):

git_cl.py

@@ -2202,6 +2202,11 @@ class Changelist(object):
             # Still raise exception so that stack trace is printed.
             raise
 
+    def GetGerritHostShortName(self) -> str:
+        """Return the short name for the CL's Gerrit host."""
+        host = self.GetGerritHost()
+        return host.split('.')[0]
+
     def GetGerritHost(self):
         # Populate self._gerrit_host
         self.GetCodereviewServer()
@@ -3616,19 +3621,26 @@ def ConfigureGitRepoAuth() -> None:
     logging.debug('Configuring current Git repo authentication...')
     cl = Changelist()
     cwd = os.getcwd()
-    gerrit_host = cl.GetGerritHost()
-    if gerrit_util.ShouldUseSSO():
+    gerrit_url = cl.GetCodereviewServer()
+    if gerrit_util.ShouldUseSSO(cl.GetGerritHost()):
         scm.GIT.SetConfig(cwd,
-                          f'credential.{gerrit_host}.helper',
+                          f'credential.{gerrit_url}.helper',
                           None,
                           modify_all=True)
+        scm.GIT.SetConfig(cwd, 'protocol.sso.allow', 'always')
+        scm.GIT.SetConfig(
+            cwd, f'url.sso://{cl.GetGerritHostShortName()}/.insteadOf',
+            gerrit_url)
     else:
+        scm.GIT.SetConfig(cwd, 'protocol.sso.allow', None)
+        scm.GIT.SetConfig(
+            cwd, f'url.sso://{cl.GetGerritHostShortName()}/.insteadOf', None)
         scm.GIT.SetConfig(cwd,
-                          f'credential.{gerrit_host}.helper',
+                          f'credential.{gerrit_url}.helper',
                           '',
                           modify_all=True)
         scm.GIT.SetConfig(cwd,
-                          f'credential.{gerrit_host}.helper',
+                          f'credential.{gerrit_url}.helper',
                           'luci',
                           append=True)
 

tests/gerrit_util_test.py

@@ -712,5 +712,55 @@ class SSOHelperTest(unittest.TestCase):
         self.assertEqual(which.called, 1)
 
 
+class ShouldUseSSOTest(unittest.TestCase):
+
+    def setUp(self) -> None:
+        self.newauth = mock.patch('newauth.Enabled', return_value=True)
+        self.newauth.start()
+        self.sso = mock.patch('gerrit_util.ssoHelper.find_cmd',
+                              return_value='/fake/git-remote-sso')
+        self.sso.start()
+        return super().setUp()
+
+    def tearDown(self) -> None:
+        super().tearDown()
+        self.sso.stop()
+        self.newauth.stop()
+
+    def testDisabled(self):
+        self.newauth.return_value = False
+        self.assertFalse(gerrit_util.ShouldUseSSO('fake-host'))
+
+    def testMissingCommand(self):
+        self.sso.return_value = 'fake-host'
+        self.assertFalse(gerrit_util.ShouldUseSSO('fake-host'))
+
+    @mock.patch('scm.GIT.GetConfig', return_value='firefly@google.com')
+    def testGoogle(self, _):
+        self.assertTrue(gerrit_util.ShouldUseSSO('fake-host'))
+
+    @mock.patch('scm.GIT.GetConfig', return_value='firefly@gmail.com')
+    def testGmail(self, _):
+        self.assertFalse(gerrit_util.ShouldUseSSO('fake-host'))
+
+    @mock.patch('gerrit_util.GetAccountEmails',
+                return_value=[{
+                    'email': 'firefly@chromium.org'
+                }])
+    @mock.patch('scm.GIT.GetConfig', return_value='firefly@chromium.org')
+    def testLinkedChromium(self, _cfg, email):
+        self.assertTrue(gerrit_util.ShouldUseSSO('fake-host'))
+        email.assert_called_with('fake-host', 'self', authenticator=mock.ANY)
+
+    @mock.patch('gerrit_util.GetAccountEmails',
+                return_value=[{
+                    'email': 'firefly@google.com'
+                }])
+    @mock.patch('scm.GIT.GetConfig', return_value='firefly@chromium.org')
+    def testUnlinkedChromium(self, _cfg, email):
+        self.assertFalse(gerrit_util.ShouldUseSSO('fake-host'))
+        email.assert_called_with('fake-host', 'self', authenticator=mock.ANY)
+
+
 if __name__ == '__main__':
     unittest.main()