瀏覽代碼

[api] Store jwt access token in userprofile

- To communicate with external services as the user like Impala etc.
Harshg999 4 年之前
父節點
當前提交
6b6c9396c8

+ 14 - 8
desktop/core/src/desktop/auth/api_authentications.py

@@ -21,6 +21,7 @@ import jwt
 from rest_framework import authentication, exceptions
 
 from desktop.auth.backend import find_or_create_user, ensure_has_a_group, rewrite_user
+from desktop.conf import ENABLE_ORGANIZATIONS
 
 from useradmin.models import User
 
@@ -31,11 +32,17 @@ class JwtAuthentication(authentication.BaseAuthentication):
 
   def authenticate(self, request):
     authorization_header = request.META.get('HTTP_AUTHORIZATION')
+
     if not authorization_header:
       LOG.debug('JwtAuthentication: no authorization header')
       return None
 
-    access_token = authorization_header[len('Bearer '):]
+    header, access_token = authorization_header.split(' ')
+
+    if header != 'Bearer':
+      LOG.debug('JwtAuthentication: no Bearer header')
+      return None
+
     if not access_token:
       LOG.debug('JwtAuthentication: no Bearer value')
       return None
@@ -55,12 +62,11 @@ class JwtAuthentication(authentication.BaseAuthentication):
     ensure_has_a_group(user)
     user = rewrite_user(user)
 
-    # We should persist the token (to reuse for communicating with external services as the user, e.g. Impala)
-    # either via a new DB field (might just be cleaner, even if requires a DB migration) or in the json blob.
-    # if ENABLE_ORGANIZATIONS.get():
-    #   user.token = token
-    # else:
-    #   user.profile.update_data({'token': token})
-    #   user.profile.save()
+    # Persist the token (to reuse for communicating with external services as the user, e.g. Impala)
+    if ENABLE_ORGANIZATIONS.get():
+      user.token = access_token
+    else:
+      user.profile.update_data({'jwt_access_token': access_token})
+      user.profile.save()
 
     return (user, None)

+ 16 - 8
desktop/core/src/desktop/auth/api_authentications_tests.py

@@ -39,12 +39,12 @@ class TestJwtAuthentication():
     self.client = make_logged_in_client(username="test", groupname="default", recreate=True, is_superuser=False)
     self.user = rewrite_user(User.objects.get(username="test"))
 
-    HEADERS = {
-        "Authorization":
-        "Bearer eyJhbGciOiJSUzI1NiJ9.eyJhdWQiOlsid29ya2xvYWQtYXBwIiwicmFuZ2VyIl0sImV4cCI6MTYyNjI1Njg5MywiaWF0IjoxNjI2MjU2NTkzLCJpc3MiOiJDbG91ZGVyYTEiLCJqdGkiOiJpZDEiLCJzdWIiOiJ0ZXN0LXN1YmplY3QiLCJ1c2VyIjoidGVzdF91c2VyIn0.jvyVDxbWTAik0jbdUcIc9ZANNrJZUCWH-Pg7FloRhg0ZYAETd_AO3p5v_ppoMmVcPD2xBSrngA5J3_A_zPBvQ_hdDlpb0_-mCCJfGhC5tju4bI9EE9Akdn2FrrsqrvQQ8cPyGsIlvoIxrK1De4f74MmUaxfN7Hrrcue1PTY4u4IB9cWQqV9vIcX99Od5PUaNekLIee-I8gweqvfGEEsW7qWUM63nh59_TOB3LLq-YcEuaX1h_oiTATeCssjk_ee9RrJGLNyKmC0WJ4UrEWn8a_T3bwCy8CMe0zV5PSuuvPHy0FvnTo2il5SDjGimxKcbpgNiJdfblslu6i35DlfiWg"
+    self.sample_token = "eyJhbGciOiJSUzI1NiJ9.eyJhdWQiOlsid29ya2xvYWQtYXBwIiwicmFuZ2VyIl0sImV4cCI6MTYyNjI1Njg5MywiaWF0IjoxNjI2MjU2NTkzLCJpc3MiOiJDbG91ZGVyYTEiLCJqdGkiOiJpZDEiLCJzdWIiOiJ0ZXN0LXN1YmplY3QiLCJ1c2VyIjoidGVzdF91c2VyIn0.jvyVDxbWTAik0jbdUcIc9ZANNrJZUCWH-Pg7FloRhg0ZYAETd_AO3p5v_ppoMmVcPD2xBSrngA5J3_A_zPBvQ_hdDlpb0_-mCCJfGhC5tju4bI9EE9Akdn2FrrsqrvQQ8cPyGsIlvoIxrK1De4f74MmUaxfN7Hrrcue1PTY4u4IB9cWQqV9vIcX99Od5PUaNekLIee-I8gweqvfGEEsW7qWUM63nh59_TOB3LLq-YcEuaX1h_oiTATeCssjk_ee9RrJGLNyKmC0WJ4UrEWn8a_T3bwCy8CMe0zV5PSuuvPHy0FvnTo2il5SDjGimxKcbpgNiJdfblslu6i35DlfiWg"
+    self.request = MagicMock(
+      META={
+        "HTTP_AUTHORIZATION": "Bearer " + self.sample_token
       }
-
-    self.request = MagicMock(HEADERS=HEADERS)
+    )
 
 
   def test_authenticate_existing_user(self):
@@ -88,6 +88,14 @@ class TestJwtAuthentication():
       # Expired token
       jwt_decode.side_effect = exceptions.AuthenticationFailed('JwtAuthentication: Token expired')
       assert_raises(exceptions.AuthenticationFailed, JwtAuthentication().authenticate, self.request)
-  
-  # TODO:
-  # check user.token
+
+
+  def test_check_user_token_storage(self):
+    with patch('desktop.auth.api_authentications.jwt.decode') as jwt_decode:
+      jwt_decode.return_value = {
+        "userId": "test"
+      }
+      user, token = JwtAuthentication().authenticate(request=self.request)
+
+      assert_true('jwt_access_token' in user.profile.data)
+      assert_equal(user.profile.data['jwt_access_token'], self.sample_token)