From 3cfc41ea9e94ca6c865d4149dcd64f24412221ce Mon Sep 17 00:00:00 2001 From: Shihang Zhang Date: Sun, 12 Apr 2020 12:03:22 -0700 Subject: [PATCH] set expiration on token of incluster config and reload if expires --- config/incluster_config.py | 35 ++++++++++++++++++++++++--------- config/incluster_config_test.py | 30 ++++++++++++++++++++++++++++ 2 files changed, 56 insertions(+), 9 deletions(-) diff --git a/config/incluster_config.py b/config/incluster_config.py index 6f28a4ae..80853c28 100644 --- a/config/incluster_config.py +++ b/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=1) - 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,37 @@ 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 + def wrap(f): + in_cluster_config = self + def wrapped(self, identifier): + if identifier == 'authorization' and identifier in self.api_key and in_cluster_config.token_expires_at <= datetime.datetime.now(): + in_cluster_config._read_token_file() + 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() + self.token_expires_at = datetime.datetime.now() + self._token_refresh_period + 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) diff --git a/config/incluster_config_test.py b/config/incluster_config_test.py index 622b31b3..e5698021 100644 --- a/config/incluster_config_test.py +++ b/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()