Accept both base64 and base64url tokens in the bidstream

This behaviour matches that of the .NET SDK.

Author: mcollins-ttd

tests/test_bidstream_client.py

@@ -280,5 +280,19 @@ def test_refresh_keys(self, mock_refresh_bidstream_keys):
                                                             client_secret_bytes)
 
 
+    def test_decrypt_v4_token_encoded_as_base64(self):
+        for scope in IdentityScope:
+            with self.subTest(scope=scope):
+                self.refresh(key_bidstream_response_json_default_keys(identity_scope=scope))
+
+                while True:
+                    token = generate_uid_token(scope, AdvertisingTokenVersion.ADVERTISING_TOKEN_V4)
+                    token = base64.b64encode(Uid2Base64UrlCoder.decode(token)).decode('utf-8')
+                    if ("=" in token) and ("/" in token) and ("+" in token):
+                        break
+            
+                self._decrypt_and_assert_success(token, AdvertisingTokenVersion.ADVERTISING_TOKEN_V4, scope)
+
+
 if __name__ == '__main__':
     unittest.main()

tests/test_utils.py

@@ -39,19 +39,18 @@
 phone_uid = "BEOGxroPLdcY7LrSiwjY52+X05V0ryELpJmoWAyXiwbZ"
 
 test_cases_all_scopes_all_versions = [
-    [IdentityScope.UID2, AdvertisingTokenVersion.ADVERTISING_TOKEN_V2],
-    [IdentityScope.UID2, AdvertisingTokenVersion.ADVERTISING_TOKEN_V3],
-    [IdentityScope.UID2, AdvertisingTokenVersion.ADVERTISING_TOKEN_V4],
-    [IdentityScope.EUID, AdvertisingTokenVersion.ADVERTISING_TOKEN_V2],
-    [IdentityScope.EUID, AdvertisingTokenVersion.ADVERTISING_TOKEN_V3],
-    [IdentityScope.EUID, AdvertisingTokenVersion.ADVERTISING_TOKEN_V4]
+    [scope, version]
+    for scope in IdentityScope
+    for version in AdvertisingTokenVersion
 ]
 
 test_cases_all_scopes_v3_v4_versions = [
-    [IdentityScope.UID2, AdvertisingTokenVersion.ADVERTISING_TOKEN_V3],
-    [IdentityScope.UID2, AdvertisingTokenVersion.ADVERTISING_TOKEN_V4],
-    [IdentityScope.EUID, AdvertisingTokenVersion.ADVERTISING_TOKEN_V3],
-    [IdentityScope.EUID, AdvertisingTokenVersion.ADVERTISING_TOKEN_V4]
+    [scope, version]
+    for scope in IdentityScope
+    for version in [
+        AdvertisingTokenVersion.ADVERTISING_TOKEN_V3,
+        AdvertisingTokenVersion.ADVERTISING_TOKEN_V4,
+    ]
 ]
 
 YESTERDAY = now + dt.timedelta(days=-1)

uid2_client/encryption.py

@@ -105,12 +105,25 @@ def _decrypt_token(token, keys, domain_name, client_type, now):
     elif token_bytes[1] == AdvertisingTokenVersion.ADVERTISING_TOKEN_V3.value:
         return _decrypt_token_v3(base64.b64decode(token), keys, domain_name, client_type, now, AdvertisingTokenVersion.ADVERTISING_TOKEN_V3)
     elif token_bytes[1] == AdvertisingTokenVersion.ADVERTISING_TOKEN_V4.value:
-        # same as V3 but use Base64URL encoding
-        return _decrypt_token_v3(Uid2Base64UrlCoder.decode(token), keys, domain_name, client_type, now, AdvertisingTokenVersion.ADVERTISING_TOKEN_V4)
+        # Accept either base64 or base64url encoding.
+        return _decrypt_token_v3(base64.b64decode(_base64url_to_base64(token)), keys, domain_name, client_type, now, AdvertisingTokenVersion.ADVERTISING_TOKEN_V4)
     else:
         return DecryptedToken.make_error(DecryptionStatus.VERSION_NOT_SUPPORTED)
 
 
+def _base64url_to_base64(value):
+    value = value.replace('-', '+').replace('_', '/')
+    input_size_mod4 = len(value) % 4
+    if input_size_mod4 == 0:
+        return value
+    elif input_size_mod4 == 2:
+        return value + '=='
+    elif input_size_mod4 == 3:
+        return value + '='
+    else:
+        raise EncryptionError('invalid payload')
+
+
 def _token_has_valid_lifetime(keys, client_type, generated_or_now, expires, now):
     #  generated_or_now allows "now" for token v2, since v2 does not contain a "token generated" field.
     #  v2 therefore checks against remaining lifetime rather than total lifetime
@@ -294,7 +307,6 @@ def encrypt(uid2, identity_scope, keys, keyset_id=None, **kwargs):
         return EncryptionDataResponse.make_error(EncryptionStatus.ENCRYPTION_FAILURE)
 
 
-
 # DEPRECATED, DO NOT CALL
 def encrypt_data(data, identity_scope, **kwargs):
     """Encrypt arbitrary binary data.