set expiration on token of incluster config and reload if expires

This commit is contained in:
Shihang Zhang 2020-04-12 12:03:22 -07:00
parent b87a5fe119
commit 3cfc41ea9e
2 changed files with 56 additions and 9 deletions

View File

@ -13,6 +13,7 @@
# limitations under the License.
import os
import datetime
from kubernetes.client import Configuration
@ -40,10 +41,11 @@ class InClusterConfigLoader(object):
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 @@ class InClusterConfigLoader(object):
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 @@ class InClusterConfigLoader(object):
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)

View File

@ -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 @@ class InClusterConfigTest(unittest.TestCase):
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 @@ class InClusterConfigTest(unittest.TestCase):
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()