refac

Co-Authored-By: Classic298 <27028174+Classic298@users.noreply.github.com>

Author: tjbck

backend/open_webui/models/users.py

@@ -279,6 +279,11 @@ async def insert_new_user(
         oauth: dict | None = None,
         db: AsyncSession | None = None,
     ) -> UserModel | None:
+        try:
+            profile_image_url = validate_profile_image_url(profile_image_url)
+        except ValueError:
+            profile_image_url = '/user.png'
+
         async with get_async_db_context(db) as session:
             user = UserModel(
                 **{
@@ -606,6 +611,11 @@ async def update_user_profile_image_url_by_id(
         profile_image_url: str,
         db: AsyncSession | None = None,
     ) -> UserModel | None:
+        try:
+            profile_image_url = validate_profile_image_url(profile_image_url)
+        except ValueError:
+            profile_image_url = '/user.png'
+
         async with get_async_db_context(db) as session:
             user = await session.get(User, id)
             if user is None:

backend/open_webui/utils/oauth.py

@@ -4,7 +4,6 @@
 import hashlib
 import json
 import logging
-import mimetypes
 import re
 import sys
 import time
@@ -85,6 +84,7 @@
 from open_webui.utils.auth import create_token, get_password_hash
 from open_webui.utils.groups import apply_default_group_assignment
 from open_webui.utils.misc import parse_duration
+from open_webui.utils.validate import validate_profile_image_url
 from starlette.responses import RedirectResponse
 
 
@@ -1648,12 +1648,17 @@ async def _process_picture_url(self, picture_url: str, access_token: str = None)
                     allow_redirects=AIOHTTP_CLIENT_ALLOW_REDIRECTS,
                 ) as resp:
                     if resp.ok:
+                        upstream_mime = (resp.headers.get('Content-Type', '') or '').split(';', 1)[0].strip().lower()
                         picture = await resp.read()
                         base64_encoded_picture = base64.b64encode(picture).decode('utf-8')
-                        guessed_mime_type = mimetypes.guess_type(picture_url)[0]
-                        if guessed_mime_type is None:
-                            guessed_mime_type = 'image/jpeg'
-                        return f'data:{guessed_mime_type};base64,{base64_encoded_picture}'
+                        try:
+                            return validate_profile_image_url(f'data:{upstream_mime};base64,{base64_encoded_picture}')
+                        except ValueError:
+                            log.warning(
+                                f'Rejected OAuth profile picture from {picture_url}: '
+                                f'MIME {upstream_mime!r} is not allowed'
+                            )
+                            return '/user.png'
                     else:
                         log.warning(f'Failed to fetch profile picture from {picture_url}')
                         return '/user.png'