Change JWT and JTI to use a split token

Author: judgeaxl

driftbase/api/clients.py

@@ -12,22 +12,23 @@
 import http.client as http_client
 import json
 import logging
+
 import marshmallow as ma
+from drift.blueprint import Blueprint, abort
+from drift.core.extensions.jwt import current_user, issue_token
+from drift.core.extensions.jwt import split_token
+from drift.core.extensions.urlregistry import Endpoints
+from drift.utils import json_response
 from flask import request, url_for, g, current_app
 from flask.views import MethodView
-from drift.blueprint import Blueprint, abort
 from flask_marshmallow.fields import AbsoluteURLFor
 from marshmallow_sqlalchemy import SQLAlchemyAutoSchema
 
-from driftbase.richpresence import RichPresenceService
-
-from drift.core.extensions.jwt import current_user, issue_token
-from drift.core.extensions.urlregistry import Endpoints
-from drift.utils import json_response
 from driftbase.config import get_client_heartbeat_config
 from driftbase.models.db import (
     User, CorePlayer, Client, UserIdentity
 )
+from driftbase.richpresence import RichPresenceService
 from driftbase.utils import url_client
 
 log = logging.getLogger(__name__)
@@ -66,8 +67,8 @@ class Meta:
         exclude = ()
 
     client_url = AbsoluteURLFor('clients.entry',
-                     doc="Fully qualified URL of the client resource",
-                     client_id='<client_id>')
+                                doc="Fully qualified URL of the client resource",
+                                client_id='<client_id>')
 
 
 class ClientPostRequestSchema(ma.Schema):
@@ -206,14 +207,9 @@ def post(self, args):
         payload = dict(current_user)
         payload["client_id"] = client_id
         new_token = issue_token(payload)
-
-        jwt = new_token["token"]
-        jti = new_token['payload']["jti"]
+        _, token = split_token(new_token["token"])
 
         resource_url = url_client(client_id)
-        response_header = {
-            "Location": resource_url,
-        }
         log.info("Client %s for user %s / player %s has been registered",
                  client_id, user_id, player_id)
         heartbeat_period, heartbeat_timeout = get_client_heartbeat_config()
@@ -225,8 +221,8 @@ def post(self, args):
             "server_time": utcnow(),
             "next_heartbeat_seconds": heartbeat_period,
             "heartbeat_timeout": utcnow() + datetime.timedelta(seconds=heartbeat_timeout),
-            "jti": jti,
-            "jwt": jwt,
+            "jti": token,
+            "jwt": token,
         }
 
         message_data = {

driftbase/api/useridentities.py

@@ -10,7 +10,7 @@
 from drift.blueprint import Blueprint, abort
 from drift.core.extensions.urlregistry import Endpoints
 
-from drift.core.extensions.jwt import current_user, get_cached_token
+from drift.core.extensions.jwt import current_user, get_cached_token_payload
 
 from driftbase.models.db import User, CorePlayer, UserIdentity
 
@@ -138,7 +138,7 @@ def post(self, args):
             abort(http_client.NOT_FOUND, message="User %s is not active" % link_with_user_id)
 
         # Verify that link_with_user_id matches user_id in link_with_user_jti
-        link_with_user_jti_payload = get_cached_token(link_with_user_jti)
+        link_with_user_jti_payload = get_cached_token_payload(link_with_user_jti, g.conf)
         if link_with_user_jti_payload["user_id"] != link_with_user_id:
             log.warning("Request for a user identity switch with user_id %s which does not "
                         "match user_id %s from JWT",

tests/test_useridentities.py

@@ -87,7 +87,7 @@ def test_identities_add_gamecenter(self):
         self.auth(username="device_%s" % uuid_string())
         r = self.get("/").json()
         device_user_id = r["current_user"]["user_id"]
-        device_jti = r["current_user"]["jti"]
+        device_jti = self.token
 
         # switch to gamecenter user
         self.headers = headers_gamecenter