Shard long JWT token on Windows

Author: hovaesco

setup.py

@@ -46,6 +46,7 @@
     "pre-commit",
     "black",
     "isort",
+    "keyring"
 ]
 
 setup(

tests/unit/test_client.py

@@ -16,11 +16,12 @@
 import uuid
 from contextlib import nullcontext as does_not_raise
 from typing import Any, Dict, Optional
-from unittest import mock
+from unittest import TestCase, mock
 from urllib.parse import urlparse
 
 import gssapi
 import httpretty
+import keyring
 import pytest
 import requests
 from httpretty import httprettified
@@ -42,7 +43,12 @@
     _post_statement_requests,
 )
 from trino import __version__, constants
-from trino.auth import GSSAPIAuthentication, KerberosAuthentication, _OAuth2TokenBearer
+from trino.auth import (
+    GSSAPIAuthentication,
+    KerberosAuthentication,
+    _OAuth2KeyRingTokenCache,
+    _OAuth2TokenBearer,
+)
 from trino.client import (
     ClientSession,
     TrinoQuery,
@@ -1343,3 +1349,76 @@ def test_request_with_invalid_timezone(mock_get_and_post):
             ),
         )
     assert str(zinfo_error.value).startswith("'No time zone found with key")
+
+
+class TestShardedPassword(TestCase):
+    def test_store_short_password(self):
+        # set the keyring to mock class
+        keyring.set_keyring(MockKeyring())
+
+        host = "trino.com"
+        short_password = "x" * 10
+
+        cache = _OAuth2KeyRingTokenCache()
+        cache.store_token_to_cache(host, short_password)
+
+        retrieved_password = cache.get_token_from_cache(host)
+        self.assertEqual(short_password, retrieved_password)
+
+    def test_store_long_password(self):
+        # set the keyring to mock class
+        keyring.set_keyring(MockKeyring())
+
+        host = "trino.com"
+        long_password = "x" * 3000
+
+        cache = _OAuth2KeyRingTokenCache()
+        cache.store_token_to_cache(host, long_password)
+
+        retrieved_password = cache.get_token_from_cache(host)
+        self.assertEqual(long_password, retrieved_password)
+
+
+class MockKeyring(keyring.backend.KeyringBackend):
+    def __init__(self):
+        self.file_location = self._generate_test_root_dir()
+
+    @staticmethod
+    def _generate_test_root_dir():
+        import tempfile
+
+        return tempfile.mkdtemp(prefix="trino-python-client-unit-test-")
+
+    def file_path(self, servicename, username):
+        from os.path import join
+
+        file_location = self.file_location
+        file_name = f"{servicename}_{username}.txt"
+        return join(file_location, file_name)
+
+    def set_password(self, servicename, username, password):
+        file_path = self.file_path(servicename, username)
+
+        with open(file_path, "w") as file:
+            file.write(password)
+
+    def get_password(self, servicename, username):
+        import os
+
+        file_path = self.file_path(servicename, username)
+        if not os.path.exists(file_path):
+            return None
+
+        with open(file_path, "r") as file:
+            password = file.read()
+
+        return password
+
+    def delete_password(self, servicename, username):
+        import os
+
+        file_path = self.file_path(servicename, username)
+        if not os.path.exists(file_path):
+            return None
+
+        os.remove(file_path)

tests/unit/test_dbapi.py

@@ -119,7 +119,7 @@ def test_token_retrieved_once_per_auth_instance(sample_post_response_data, sampl
         conn2.cursor().execute("SELECT 2")
         conn2.cursor().execute("SELECT 3")
 
-    assert len(_get_token_requests(challenge_id)) == 2
+    assert len(_get_token_requests(challenge_id)) == 1
 
 
 @httprettified

trino/auth.py

@@ -25,7 +25,7 @@
 
 import trino.logging
 from trino.client import exceptions
-from trino.constants import HEADER_USER
+from trino.constants import HEADER_USER, MAX_NT_PASSWORD_SIZE
 
 logger = trino.logging.get_logger(__name__)
 
@@ -347,17 +347,49 @@ def is_keyring_available(self) -> bool:
             and not isinstance(self._keyring.get_keyring(), self._keyring.backends.fail.Keyring)
 
     def get_token_from_cache(self, key: Optional[str]) -> Optional[str]:
+        password = self._keyring.get_password(key, "token")
+
         try:
-            return self._keyring.get_password(key, "token")
+            password_as_dict = json.loads(str(password))
+            if password_as_dict.get("sharded_password"):
+                # if password was stored shared, reconstruct it
+                shard_count = int(password_as_dict.get("shard_count"))
+
+                password = ""
+                for i in range(shard_count):
+                    password += str(self._keyring.get_password(key, f"token__{i}"))
+
         except self._keyring.errors.NoKeyringError as e:
             raise trino.exceptions.NotSupportedError("Although keyring module is installed no backend has been "
                                                      "detected, check https://pypi.org/project/keyring/ for more "
                                                      "information.") from e
+        except ValueError:
+            pass
+
+        return password
 
     def store_token_to_cache(self, key: Optional[str], token: str) -> None:
+        # keyring is installed, so we can store the token for reuse within multiple threads
         try:
-            # keyring is installed, so we can store the token for reuse within multiple threads
-            self._keyring.set_password(key, "token", token)
+            # if not Windows or "small" password, stick to the default
+            if os.name != "nt" or len(token) < MAX_NT_PASSWORD_SIZE:
+                self._keyring.set_password(key, "token", token)
+            else:
+                logger.debug(f"password is {len(token)} characters, sharding it.")
+
+                password_shards = [
+                    token[i: i + MAX_NT_PASSWORD_SIZE] for i in range(0, len(token), MAX_NT_PASSWORD_SIZE)
+                ]
+                shard_info = {
+                    "sharded_password": True,
+                    "shard_count": len(password_shards),
+                }
+
+                # store the "shard info" as the "base" password
+                self._keyring.set_password(key, "token", json.dumps(shard_info))
+                # then store all shards with the shard number as postfix
+                for i, s in enumerate(password_shards):
+                    self._keyring.set_password(key, f"token__{i}", s)
         except self._keyring.errors.NoKeyringError as e:
             raise trino.exceptions.NotSupportedError("Although keyring module is installed no backend has been "
                                                      "detected, check https://pypi.org/project/keyring/ for more "

trino/constants.py

@@ -20,6 +20,7 @@
 DEFAULT_AUTH: Optional[Any] = None
 DEFAULT_MAX_ATTEMPTS = 3
 DEFAULT_REQUEST_TIMEOUT: float = 30.0
+MAX_NT_PASSWORD_SIZE: int = 1280
 
 HTTP = "http"
 HTTPS = "https"