set expiration on token of incluster config and reload if expires

Author: zshihang

config/incluster_config.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 
 import os
+import datetime
 
 from kubernetes.client import Configuration
 
@@ -40,10 +41,11 @@ def __init__(self, token_filename,
         self._token_filename = token_filename
         self._cert_filename = cert_filename
         self._environ = environ
+        self._token_refresh_period = datetime.timedelta(minutes=5)
 
-    def load_and_set(self):
+    def load_and_set(self, refresh_token=True):
         self._load_config()
-        self._set_config()
+        self._set_config(refresh_token=refresh_token)
 
     def _load_config(self):
         if (SERVICE_HOST_ENV_NAME not in self._environ or
@@ -61,10 +63,7 @@ def _load_config(self):
         if not os.path.isfile(self._token_filename):
             raise ConfigException("Service token file does not exists.")
 
-        with open(self._token_filename) as f:
-            self.token = f.read()
-            if not self.token:
-                raise ConfigException("Token file exists but empty.")
+        self._read_token_file()
 
         if not os.path.isfile(self._cert_filename):
             raise ConfigException(
@@ -76,19 +75,38 @@ def _load_config(self):
 
         self.ssl_ca_cert = self._cert_filename
 
-    def _set_config(self):
+    def _set_config(self, refresh_token):
         configuration = Configuration()
         configuration.host = self.host
         configuration.ssl_ca_cert = self.ssl_ca_cert
         configuration.api_key['authorization'] = "bearer " + self.token
         Configuration.set_default(configuration)
+        if not refresh_token:
+            return
+        self.token_expires_at = datetime.datetime.now() + self._token_refresh_period
+        def wrap(f):
+            in_cluster_config = self
+            def wrapped(self, identifier):
+                if identifier == 'authorization' and in_cluster_config.token_expires_at <= datetime.datetime.now():
+                    in_cluster_config._read_token_file()
+                    in_cluster_config.token_expires_at = datetime.datetime.now() + in_cluster_config._token_refresh_period
+                    self.api_key[identifier] = "bearer " + in_cluster_config.token
+                return f(self, identifier)
+            return wrapped
+        Configuration.get_api_key_with_prefix = wrap(Configuration.get_api_key_with_prefix)
+
+    def _read_token_file(self):
+        with open(self._token_filename) as f:
+            self.token = f.read()
+            if not self.token:
+                raise ConfigException("Token file exists but empty.")
 
 
-def load_incluster_config():
+def load_incluster_config(refresh_token=True):
     """
     Use the service account kubernetes gives to pods to connect to kubernetes
     cluster. It's intended for clients that expect to be running inside a pod
     running on kubernetes. It will raise an exception if called from a process
     not running in a kubernetes environment."""
     InClusterConfigLoader(token_filename=SERVICE_TOKEN_FILENAME,
-                          cert_filename=SERVICE_CERT_FILENAME).load_and_set()
+                          cert_filename=SERVICE_CERT_FILENAME).load_and_set(refresh_token=refresh_token)

config/incluster_config_test.py

@@ -15,12 +15,17 @@
 import os
 import tempfile
 import unittest
+import datetime
+import time
+
+from kubernetes.client import Configuration
 
 from .config_exception import ConfigException
 from .incluster_config import (SERVICE_HOST_ENV_NAME, SERVICE_PORT_ENV_NAME,
                                InClusterConfigLoader, _join_host_port)
 
 _TEST_TOKEN = "temp_token"
+_TEST_NEW_TOKEN = "temp_new_token"
 _TEST_CERT = "temp_cert"
 _TEST_HOST = "127.0.0.1"
 _TEST_PORT = "80"
@@ -50,6 +55,12 @@ def _create_file_with_temp_content(self, content=""):
         os.close(handler)
         return name
 
+    def _overwrite_file_with_content(self, name, content=""):
+        handler = os.open(name, os.O_RDWR)
+        os.truncate(name, 0)
+        os.write(handler, str.encode(content))
+        os.close(handler)
+
     def get_test_loader(
             self,
             token_filename=None,
@@ -78,6 +89,25 @@ def test_load_config(self):
         self.assertEqual(cert_filename, loader.ssl_ca_cert)
         self.assertEqual(_TEST_TOKEN, loader.token)
 
+    def test_refresh_token(self):
+        loader = self.get_test_loader()
+        loader._token_refresh_period = datetime.timedelta(seconds=5)
+        loader.load_and_set()
+        config = Configuration()
+
+        self.assertEqual('bearer '+_TEST_TOKEN, config.get_api_key_with_prefix('authorization'))
+        self.assertEqual(_TEST_TOKEN, loader.token)
+        self.assertIsNotNone(loader.token_expires_at)
+
+        old_token = loader.token
+        old_token_expires_at = loader.token_expires_at
+        self._overwrite_file_with_content(loader._token_filename, _TEST_NEW_TOKEN)
+        time.sleep(5)
+
+        self.assertEqual('bearer '+_TEST_NEW_TOKEN, config.get_api_key_with_prefix('authorization'))
+        self.assertEqual(_TEST_NEW_TOKEN, loader.token)
+        self.assertGreater(loader.token_expires_at, old_token_expires_at)
+
     def _should_fail_load(self, config_loader, reason):
         try:
             config_loader.load_and_set()