From aa0cd480604069519ebd9b52b0d629e33287341c Mon Sep 17 00:00:00 2001
From: Haitao Li <39936070+hligit@users.noreply.github.com>
Date: Fri, 6 May 2022 01:35:06 +1000
Subject: [PATCH] k8s: Support refreshing service account tokens (#2287)

Since Kubernetes v1.21, with projected service account token feature, service account tokens expire in 1 hour. Kubernetes clients are expected to reread the token file to refresh the token.

This patch re-reads the token file very minute for in-cluster config.

Fixes #2286

Signed-off-by: Haitao Li <hli@atlassian.com>
---
 patroni/dcs/kubernetes.py | 35 +++++++++++++++++++++++++----------
 tests/test_kubernetes.py  | 23 +++++++++++++++++++++++
 2 files changed, 48 insertions(+), 10 deletions(-)

diff --git a/patroni/dcs/kubernetes.py b/patroni/dcs/kubernetes.py
index 25969baf..06d90572 100644
--- a/patroni/dcs/kubernetes.py
+++ b/patroni/dcs/kubernetes.py
@@ -48,14 +48,29 @@ class K8sConfig(object):
 
     def __init__(self):
         self.pool_config = {'maxsize': 10, 'num_pools': 10}  # configuration for urllib3.PoolManager
+        self._token_expires_at = datetime.datetime.max
         self._make_headers()
 
+    def _set_token(self, token):
+        self._headers['authorization'] = 'Bearer ' + token
+
     def _make_headers(self, token=None, **kwargs):
         self._headers = urllib3.make_headers(user_agent=USER_AGENT, **kwargs)
         if token:
-            self._headers['authorization'] = 'Bearer ' + token
+            self._set_token(token)
 
-    def load_incluster_config(self, ca_certs=SERVICE_CERT_FILENAME):
+    def _read_token_file(self):
+        if not os.path.isfile(SERVICE_TOKEN_FILENAME):
+            raise self.ConfigException('Service token file does not exists.')
+        with open(SERVICE_TOKEN_FILENAME) as f:
+            token = f.read()
+            if not token:
+                raise self.ConfigException('Token file exists but empty.')
+            self._token_expires_at = datetime.datetime.now() + self._token_refresh_interval
+            return token
+
+    def load_incluster_config(self, ca_certs=SERVICE_CERT_FILENAME,
+                              token_refresh_interval=datetime.timedelta(minutes=1)):
         if SERVICE_HOST_ENV_NAME not in os.environ or SERVICE_PORT_ENV_NAME not in os.environ:
             raise self.ConfigException('Service host/port is not set.')
         if not os.environ[SERVICE_HOST_ENV_NAME] or not os.environ[SERVICE_PORT_ENV_NAME]:
@@ -67,14 +82,9 @@ class K8sConfig(object):
             if not f.read():
                 raise self.ConfigException('Cert file exists but empty.')
         self.pool_config['ca_certs'] = ca_certs
-
-        if not os.path.isfile(SERVICE_TOKEN_FILENAME):
-            raise self.ConfigException('Service token file does not exists.')
-        with open(SERVICE_TOKEN_FILENAME) as f:
-            token = f.read()
-            if not token:
-                raise self.ConfigException('Token file exists but empty.')
-            self._make_headers(token=token)
+        self._token_refresh_interval = token_refresh_interval
+        token = self._read_token_file()
+        self._make_headers(token=token)
         self._server = uri('https', (os.environ[SERVICE_HOST_ENV_NAME], os.environ[SERVICE_PORT_ENV_NAME]))
 
     @staticmethod
@@ -109,6 +119,11 @@ class K8sConfig(object):
 
     @property
     def headers(self):
+        if self._token_expires_at <= datetime.datetime.now():
+            try:
+                self._set_token(self._read_token_file())
+            except Exception as e:
+                logger.error('Failed to refresh service account token: %r', e)
         return self._headers.copy()
 
 
diff --git a/tests/test_kubernetes.py b/tests/test_kubernetes.py
index 92339b5a..982edb5d 100644
--- a/tests/test_kubernetes.py
+++ b/tests/test_kubernetes.py
@@ -1,3 +1,4 @@
+import datetime
 import json
 import socket
 import time
@@ -79,6 +80,28 @@ class TestK8sConfig(unittest.TestCase):
                 self.assertRaises(k8s_config.ConfigException, k8s_config.load_incluster_config)
             k8s_config.load_incluster_config()
             self.assertEqual(k8s_config.server, 'https://a:1')
+            self.assertEqual(k8s_config.headers.get('authorization'), 'Bearer a')
+
+    def test_refresh_token(self):
+        with patch('os.environ', {SERVICE_HOST_ENV_NAME: 'a', SERVICE_PORT_ENV_NAME: '1'}),\
+                patch('os.path.isfile', Mock(side_effect=[True, True, False, True, True, True])),\
+                patch.object(builtins, 'open', Mock(side_effect=[
+                    mock_open(read_data='cert')(), mock_open(read_data='a')(),
+                    mock_open()(), mock_open(read_data='b')(), mock_open(read_data='c')()])):
+            k8s_config.load_incluster_config(token_refresh_interval=datetime.timedelta(milliseconds=100))
+            self.assertEqual(k8s_config.headers.get('authorization'), 'Bearer a')
+            time.sleep(0.1)
+            # token file doesn't exist
+            self.assertEqual(k8s_config.headers.get('authorization'), 'Bearer a')
+            # token file is empty
+            self.assertEqual(k8s_config.headers.get('authorization'), 'Bearer a')
+            # token refreshed
+            self.assertEqual(k8s_config.headers.get('authorization'), 'Bearer b')
+            time.sleep(0.1)
+            # token refreshed
+            self.assertEqual(k8s_config.headers.get('authorization'), 'Bearer c')
+            # no need to refresh token
+            self.assertEqual(k8s_config.headers.get('authorization'), 'Bearer c')
 
     def test_load_kube_config(self):
         config = {