Enhance `get_token` Function for Robust Credential Validation and Simplified Response Handling

[Umpire2018]

Technical Design

• Credentials Validation: Modify the endpoint to first check if the credentials.credentials is valid and contains a usable JWT. If invalid or missing, generate a new session ID using uuid.uuid4().
• Simplified Response Handling: Instead of using JSONResponse explicitly, return a simple Python dictionary. FastAPI's default response handler will automatically convert this to a JSON response, making the code cleaner and more Pythonic.

Original Code

def get_sid_from_token(token: str) -> str:
    """Gets the session id from a JWT token."""
    try:
        payload = jwt.decode(token, JWT_SECRET, algorithms=['HS256'])
        if payload is None:
            logger.error('Invalid token')
            return ''
        return payload['sid']
    except Exception as e:
        logger.exception('Error decoding token: %s', e)
        return ''


@app.get('/api/auth')
async def get_token(
        credentials: HTTPAuthorizationCredentials = Depends(security_scheme),
):
    """
    Get token for authentication when starts a websocket connection.
    """
    sid = get_sid_from_token(credentials.credentials) or str(uuid.uuid4())
    token = sign_token({'sid': sid})
    return JSONResponse(
        status_code=status.HTTP_200_OK,
        content={'token': token},
    )

Test Script

import aiohttp
import asyncio

async def get_jwt_token(valid_credentials=True):
    url = 'http://localhost:3000/api/auth'
    token = 'Token' if valid_credentials else ''
    headers = {
        'Authorization': f'Bearer {token}'
    }
    async with aiohttp.ClientSession() as session:
        async with session.get(url, headers=headers) as response:
            print(f"Testing with {'valid' if valid_credentials else 'invalid'} credentials...")
            if response.status == 200:
                data = await response.json()
                print(f"Success: Received JWT Token: {data.get('token')}")
            else:
                print(f"Fail: {response.status}, reason: {response.reason}")
                
async def main():
    await get_jwt_token(valid_credentials=True)  # Test with valid credentials
    await get_jwt_token(valid_credentials=False)  # Test with invalid credentials

asyncio.run(main())

Result

INFO:     Uvicorn running on http://127.0.0.1:3000 (Press CTRL+C to quit)
06:37:27 - opendevin:ERROR: auth.py:18 - Error decoding token: Not enough segments
Traceback (most recent call last):
  File "/home/test/.cache/pypoetry/virtualenvs/opendevin-P7zClqqX-py3.12/lib/python3.12/site-packages/jwt/api_jws.py", line 257, in _load
    signing_input, crypto_segment = jwt.rsplit(b".", 1)
    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
ValueError: not enough values to unpack (expected 2, got 1)

The above exception was the direct cause of the following exception:

Traceback (most recent call last):
  File "/home/test/arno/OpenDevin/opendevin/server/auth/auth.py", line 12, in get_sid_from_token
    payload = jwt.decode(token, JWT_SECRET, algorithms=['HS256'])
              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/test/.cache/pypoetry/virtualenvs/opendevin-P7zClqqX-py3.12/lib/python3.12/site-packages/jwt/api_jwt.py", line 210, in decode
    decoded = self.decode_complete(
              ^^^^^^^^^^^^^^^^^^^^^
  File "/home/test/.cache/pypoetry/virtualenvs/opendevin-P7zClqqX-py3.12/lib/python3.12/site-packages/jwt/api_jwt.py", line 151, in decode_complete
    decoded = api_jws.decode_complete(
              ^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/test/.cache/pypoetry/virtualenvs/opendevin-P7zClqqX-py3.12/lib/python3.12/site-packages/jwt/api_jws.py", line 198, in decode_complete
    payload, signing_input, header, signature = self._load(jwt)
                                                ^^^^^^^^^^^^^^^
  File "/home/test/.cache/pypoetry/virtualenvs/opendevin-P7zClqqX-py3.12/lib/python3.12/site-packages/jwt/api_jws.py", line 260, in _load
    raise DecodeError("Not enough segments") from err
jwt.exceptions.DecodeError: Not enough segments
INFO:     127.0.0.1:54796 - "GET /api/auth HTTP/1.1" 200 OK
INFO:     127.0.0.1:54808 - "GET /api/auth HTTP/1.1" 403 Forbidden

Explanation

1. When authentication is attempted using HTTPAuthorizationCredentials, the get_sid_from_token function tries to decode credentials.credentials, but since it's not the expected JWT, this triggers a jwt.exceptions.DecodeError. This specific error can be caught using the InvalidTokenError type, returning an empty string '', after which a new sid is regenerated in the get_token function.

2. When HTTPAuthorizationCredentials does not contain any token, it results in a 403 Forbidden error, which is the correct behavior.

Result with enhanced code

06:16:35 - opendevin:ERROR: auth.py:31 - Invalid token
06:16:35 - opendevin:INFO: listen.py:73 - Invalid or missing credentials, generating new session ID: 6f85cf7d-48d2-4fb0-9a9a-ffbcc77e20a4
INFO:     127.0.0.1:45208 - "GET /api/auth HTTP/1.1" 200 OK