diff --git a/tests/app/services/auth/test_token_service.py b/tests/app/services/auth/test_token_service.py index f5e47af5b41..55e7cdbe3bd 100644 --- a/tests/app/services/auth/test_token_service.py +++ b/tests/app/services/auth/test_token_service.py @@ -5,6 +5,14 @@ from invokeai.app.services.auth.token_service import TokenData, create_access_token, verify_token +# Minimum token length to safely modify middle characters for testing +# JWT tokens have format header.payload.signature and are typically >180 characters +MIN_TOKEN_LENGTH_FOR_MODIFICATION = 50 + +# Minimum signature length to safely modify middle characters for testing +# JWT signatures are typically 43 characters (base64-encoded HMAC-SHA256) +MIN_SIGNATURE_LENGTH_FOR_MODIFICATION = 10 + class TestTokenCreation: """Tests for JWT token creation.""" @@ -151,10 +159,13 @@ def test_verify_token_with_modified_payload(self): token = create_access_token(token_data) - # Try to modify the token by changing a character + # Try to modify the token by changing a character in the middle # JWT tokens are base64 encoded, so changing any character should invalidate the signature - if len(token) > 10: - modified_token = token[:-1] + ("X" if token[-1] != "X" else "Y") + # Note: We change a character in the middle to avoid Base64 padding issues where + # the last character might not affect the decoded value + if len(token) > MIN_TOKEN_LENGTH_FOR_MODIFICATION: + mid = len(token) // 2 + modified_token = token[:mid] + ("X" if token[mid] != "X" else "Y") + token[mid + 1 :] verified_data = verify_token(modified_token) assert verified_data is None @@ -288,11 +299,14 @@ def test_token_signature_verification(self): assert verify_token(token) is not None # Modified token should fail verification - if len(token) > 50: + if len(token) > MIN_TOKEN_LENGTH_FOR_MODIFICATION: # Change a character in the signature part (last part of JWT) parts = token.split(".") - if len(parts) == 3: - modified_signature = parts[2][:-1] + ("X" if parts[2][-1] != "X" else "Y") + if len(parts) == 3 and len(parts[2]) > MIN_SIGNATURE_LENGTH_FOR_MODIFICATION: + # Modify a character in the middle of the signature to avoid Base64 padding issues + # where the last few characters might not affect the decoded value + mid = len(parts[2]) // 2 + modified_signature = parts[2][:mid] + ("X" if parts[2][mid] != "X" else "Y") + parts[2][mid + 1 :] modified_token = f"{parts[0]}.{parts[1]}.{modified_signature}" assert verify_token(modified_token) is None