Selaa lähdekoodia

[jwt] Support JKU HA for public APIs JWT auth (#3360)

- When in HA mode, take healthy JKU to fetch the public key-set.
- Key server urls are comma separated configs in Hue in HA mode.

- Tested by adding new unit tests and also in HA setup
Harsh Gupta 2 vuotta sitten
vanhempi
commit
e77ec167cd

+ 27 - 2
desktop/core/src/desktop/auth/api_authentications.py

@@ -111,9 +111,11 @@ class JwtAuthentication(authentication.BaseAuthentication):
   def _handle_public_key(self, access_token):
     token_metadata = jwt.get_unverified_header(access_token)
     kid = token_metadata.get('kid')
+    key_server_url = self._handle_jku_ha()
 
-    if AUTH.JWT.KEY_SERVER_URL.get():
-      response = requests.get(AUTH.JWT.KEY_SERVER_URL.get(), verify=False)
+    if key_server_url:
+      LOG.debug('Fetching JWKS from URL: %s' % key_server_url)
+      response = requests.get(key_server_url, verify=False)
       jwk = json.loads(response.content)
 
       if jwk.get('keys'):
@@ -125,6 +127,29 @@ class JwtAuthentication(authentication.BaseAuthentication):
 
             return public_key_pem
 
+  def _handle_jku_ha(self):
+    res = None
+
+    key_server_urls = AUTH.JWT.KEY_SERVER_URL.get()
+    if not key_server_urls:
+      return None
+
+    if "," in key_server_urls:
+      key_server_urls_list = key_server_urls.split(',')
+
+      for jku in key_server_urls_list:
+        try:
+          res = requests.get(jku.rstrip('/'), verify=False)
+        except Exception as e:
+          if 'Failed to establish a new connection' in str(e):
+            LOG.warning('JKU %s is not available.' % jku)
+
+        # Check response for None and if response code is successful (200) or authentication needed (401), use that host URL.
+        if (res is not None) and (res.status_code in (200, 401)):
+          return jku
+    else:
+      # For non-HA, it's normal url string.
+      return key_server_urls
 
 class DummyCustomAuthentication(authentication.BaseAuthentication):
   """

+ 33 - 0
desktop/core/src/desktop/auth/api_authentications_tests.py

@@ -225,3 +225,36 @@ class TestJwtAuthentication():
         finally:
           for reset in resets:
             reset()
+
+
+  def test_handle_jku_ha(self):
+    with patch('desktop.auth.api_authentications.requests.get') as requests_get:
+      requests_get.return_value = Mock(status_code=200)
+
+      # Non-HA mode
+      reset = AUTH.JWT.KEY_SERVER_URL.set_for_testing('https://ext-authz:8000/api/v1/jwks.json')
+
+      try:
+        jku = JwtAuthentication()._handle_jku_ha()
+        assert_equal(jku, 'https://ext-authz:8000/api/v1/jwks.json')
+      finally:
+        reset()
+
+      # HA mode - where first URL sends 200 status code
+      reset = AUTH.JWT.KEY_SERVER_URL.set_for_testing('https://ext-authz:8000/api/v1/jwks.json, https://ext1-authz:8000/api/v1/jwks.json')
+
+      try:
+        jku = JwtAuthentication()._handle_jku_ha()
+        assert_equal(jku, 'https://ext-authz:8000/api/v1/jwks.json')
+      finally:
+        reset()
+
+      # When no JKU URL is healthy
+      requests_get.return_value = Mock(status_code=404)
+      reset = AUTH.JWT.KEY_SERVER_URL.set_for_testing('https://ext-authz:8000/api/v1/jwks.json, https://ext1-authz:8000/api/v1/jwks.json')
+
+      try:
+        jku = JwtAuthentication()._handle_jku_ha()
+        assert_equal(jku, None)
+      finally:
+        reset()

+ 1 - 1
desktop/core/src/desktop/conf.py

@@ -1231,7 +1231,7 @@ AUTH = ConfigSection(
         KEY_SERVER_URL=Config(
             key="key_server_url",
             default=None,
-            type=str,
+            type=coerce_string,
             help=_("Endpoint to fetch the public key from verification server.")
         ),
         ISSUER=Config(