From 4ebe57134446afe6aa5bbd33494d91831035a30f Mon Sep 17 00:00:00 2001
From: Jeremy Lewi <jeremy+github@lewi.us>
Date: Tue, 21 Jan 2020 12:50:34 -0800
Subject: [PATCH] Fix cleanup logic for IAM policy bindings (#566)

* Fix cleanup logic for IAM policy bindings

* We are not properly GC'ing policy bindings for deleted service accounts.

* The problem is that we only consider service accounts matching a certain
  regex and that regex isn't matching service accounts for our auto-deployed
  clusters.

* Using a regex should be unnecessary. If a service account doesn't exist
  that should be a sufficient criterion that the policy bindings should be
  deleted.

Related to: kubeflow/testing#543

* Fix typo.

* Fix syntax issue.
---
 py/kubeflow/testing/cleanup_ci.py      | 14 +++++++++++++-
 py/kubeflow/testing/cleanup_ci_test.py | 16 ++++++++++++++++
 2 files changed, 29 insertions(+), 1 deletion(-)

diff --git a/py/kubeflow/testing/cleanup_ci.py b/py/kubeflow/testing/cleanup_ci.py
index 4d4a317d7..613dafe8e 100644
--- a/py/kubeflow/testing/cleanup_ci.py
+++ b/py/kubeflow/testing/cleanup_ci.py
@@ -759,18 +759,23 @@ def cleanup_service_accounts(args):
 
 def trim_unused_bindings(iamPolicy, accounts):
   keepBindings = []
+  deleted_bindings = set()
+  kept_bindings = set()
   for binding in iamPolicy['bindings']:
     members_to_keep = []
     members_to_delete = []
     for member in binding['members']:
       if not member.startswith('serviceAccount:'):
         members_to_keep.append(member)
+        kept_bindings.add(member)
       else:
         accountEmail = member[15:]
-        if (not is_match(accountEmail)) or (accountEmail in accounts):
+        if accountEmail in accounts:
           members_to_keep.append(member)
+          kept_bindings.add(member)
         else:
           members_to_delete.append(member)
+          deleted_bindings.add(member)
     if members_to_keep:
       binding['members'] = members_to_keep
       keepBindings.append(binding)
@@ -779,6 +784,11 @@ def trim_unused_bindings(iamPolicy, accounts):
         members_to_delete))
   iamPolicy['bindings'] = keepBindings
 
+  logging.info("Removing bindings for following service accounts which "
+               "do not exist: %s", "\n".join(deleted_bindings))
+  logging.info("Keeping bindings for following service accounts which "
+               "still exist: %s", "\n".join(kept_bindings))
+
 def cleanup_service_account_bindings(args):
   logging.info("Cleanup service account bindings")
 
@@ -795,6 +805,8 @@ def cleanup_service_account_bindings(args):
       break
     next_page_token = service_accounts["nextPageToken"]
 
+  logging.info("The following service accounts exist so bindings will not be"
+               "deleted:\n%s", "\n".join(accounts))
   resourcemanager = discovery.build('cloudresourcemanager', 'v1', credentials=credentials)
   logging.info("Get IAM policy for project %s", args.project)
   iamPolicy = resourcemanager.projects().getIamPolicy(resource=args.project, body={}).execute()
diff --git a/py/kubeflow/testing/cleanup_ci_test.py b/py/kubeflow/testing/cleanup_ci_test.py
index de59c58ce..d7bee966b 100644
--- a/py/kubeflow/testing/cleanup_ci_test.py
+++ b/py/kubeflow/testing/cleanup_ci_test.py
@@ -1,4 +1,5 @@
 from kubeflow.testing import cleanup_ci
+import collections
 import logging
 import pytest
 
@@ -16,6 +17,21 @@ def test_match_disk():
   pvc = "gke-zresubmit-unittest-pvc-e3bf5be4-987b-11e9-8266-42010a8e00e9"
   assert cleanup_ci.is_match(pvc, patterns=cleanup_ci.E2E_PATTERNS)
 
+
+def test_match_service_accounts():
+  test_case = collections.namedtuple("test_case", ("input", "expected"))
+
+  cases = [
+    test_case("kf-vmaster-0121-b11-user@"
+              "kubeflow-ci-deployment.iam.gserviceaccount.com",
+              cleanup_ci.AUTO_INFRA)
+  ]
+
+  for c in cases:
+    actual = cleanup_ci.name_to_infra_type(c.input)
+
+    assert actual == c.expected
+
 if __name__ == "__main__":
   logging.basicConfig(
       level=logging.INFO,