|
1 | 1 | from __future__ import annotations |
2 | 2 |
|
| 3 | +import hashlib |
| 4 | +import html |
| 5 | +import ipaddress |
3 | 6 | import logging |
| 7 | +import socket |
| 8 | +import time |
| 9 | +from collections import OrderedDict |
| 10 | +from urllib.parse import urljoin, urlparse |
4 | 11 |
|
| 12 | +import aiohttp |
5 | 13 | import black |
6 | 14 | import markdown |
| 15 | +from bs4 import BeautifulSoup |
7 | 16 | from fastapi import APIRouter, Depends, HTTPException, Request, Response, status |
8 | 17 | from open_webui.config import DATA_DIR, ENABLE_ADMIN_EXPORT |
9 | 18 | from open_webui.constants import ERROR_MESSAGES |
| 19 | +from open_webui.env import REDIS_KEY_PREFIX |
10 | 20 | from open_webui.models.chats import ChatTitleMessagesForm |
11 | 21 | from open_webui.utils.auth import get_admin_user, get_verified_user |
12 | 22 | from open_webui.utils.code_interpreter import execute_code_jupyter |
@@ -121,3 +131,302 @@ async def download_db(user=Depends(get_admin_user)): |
121 | 131 | media_type='application/octet-stream', |
122 | 132 | filename='webui.db', |
123 | 133 | ) |
| 134 | + |
| 135 | + |
| 136 | +_FAVICON_CACHE: OrderedDict = OrderedDict() |
| 137 | +_FAVICON_CACHE_MAX = 1000 |
| 138 | +_FAVICON_CACHE_TTL = 7 * 24 * 3600 |
| 139 | +_FAVICON_TIMEOUT = aiohttp.ClientTimeout(total=5) |
| 140 | +_MAX_ICON_BYTES = 256 * 1024 # 256 KiB |
| 141 | +_MAX_HTML_BYTES = 250 * 1024 # 250 KiB |
| 142 | +_MAX_REDIRECTS = 3 |
| 143 | +_ALLOWED_ICON_TYPES = { |
| 144 | + 'image/x-icon', |
| 145 | + 'image/vnd.microsoft.icon', |
| 146 | + 'image/png', |
| 147 | + 'image/jpeg', |
| 148 | + 'image/gif', |
| 149 | + 'image/webp', |
| 150 | +} |
| 151 | +_BADGE_PALETTE = [ |
| 152 | + '#e74c3c', '#e67e22', '#2ecc71', '#3498db', |
| 153 | + '#9b59b6', '#1abc9c', '#e91e63', '#f39c12', |
| 154 | +] |
| 155 | + |
| 156 | + |
| 157 | +def _svg_badge(hostname: str) -> bytes: |
| 158 | + if hostname: |
| 159 | + letter = html.escape(hostname[0].upper()) |
| 160 | + idx = hashlib.sha256(hostname.encode()).digest()[0] % len(_BADGE_PALETTE) |
| 161 | + color = _BADGE_PALETTE[idx] |
| 162 | + else: |
| 163 | + letter = '?' |
| 164 | + color = '#95a5a6' |
| 165 | + svg = ( |
| 166 | + f'<svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 32 32">' |
| 167 | + f'<circle cx="16" cy="16" r="16" fill="{color}"/>' |
| 168 | + f'<text x="16" y="21" font-size="16" font-family="sans-serif" ' |
| 169 | + f'fill="white" text-anchor="middle">{letter}</text>' |
| 170 | + f'</svg>' |
| 171 | + ) |
| 172 | + return svg.encode() |
| 173 | + |
| 174 | + |
| 175 | +def _normalize_hostname(url: str) -> str | None: |
| 176 | + try: |
| 177 | + url = url.strip() |
| 178 | + if '://' not in url: |
| 179 | + url = 'https://' + url |
| 180 | + parsed = urlparse(url) |
| 181 | + hostname = (parsed.hostname or '').lower().rstrip('.') |
| 182 | + if not hostname: |
| 183 | + return None |
| 184 | + try: |
| 185 | + hostname = hostname.encode('idna').decode('ascii') |
| 186 | + except (UnicodeError, UnicodeDecodeError): |
| 187 | + pass |
| 188 | + return hostname or None |
| 189 | + except Exception: |
| 190 | + return None |
| 191 | + |
| 192 | + |
| 193 | +def _is_safe_url(url: str) -> bool: |
| 194 | + """Return True if url has an http/https scheme and a globally routable target.""" |
| 195 | + try: |
| 196 | + parsed = urlparse(url) |
| 197 | + if parsed.scheme not in ('http', 'https'): |
| 198 | + return False |
| 199 | + hostname = parsed.hostname |
| 200 | + if not hostname: |
| 201 | + return False |
| 202 | + try: |
| 203 | + if not ipaddress.ip_address(hostname).is_global: # raw IP |
| 204 | + return False |
| 205 | + except ValueError: |
| 206 | + pass # hostname; _SafeResolver handles it |
| 207 | + return True |
| 208 | + except Exception: |
| 209 | + return False |
| 210 | + |
| 211 | + |
| 212 | +class _SafeResolver(aiohttp.ThreadedResolver): |
| 213 | + async def resolve(self, hostname, port=0, family=socket.AF_UNSPEC): |
| 214 | + results = await super().resolve(hostname, port, family) |
| 215 | + for r in results: |
| 216 | + if not ipaddress.ip_address(r['host']).is_global: |
| 217 | + raise OSError(f'non-global IP: {r["host"]}') |
| 218 | + return results |
| 219 | + |
| 220 | + |
| 221 | +async def _fetch_bounded( |
| 222 | + session: aiohttp.ClientSession, |
| 223 | + url: str, |
| 224 | + max_bytes: int, |
| 225 | + allowed_types: set[str] | None, |
| 226 | +) -> tuple[bytes, str] | None: |
| 227 | + """GET with capped size and per-hop SSRF validation.""" |
| 228 | + current = url |
| 229 | + for _ in range(_MAX_REDIRECTS + 1): |
| 230 | + try: |
| 231 | + async with session.get( |
| 232 | + current, allow_redirects=False, timeout=_FAVICON_TIMEOUT |
| 233 | + ) as resp: |
| 234 | + if resp.status in (301, 302, 303, 307, 308): |
| 235 | + location = resp.headers.get('Location', '') |
| 236 | + if not location: |
| 237 | + return None |
| 238 | + next_url = urljoin(current, location) |
| 239 | + if not _is_safe_url(next_url): |
| 240 | + return None |
| 241 | + current = next_url |
| 242 | + continue |
| 243 | + |
| 244 | + if resp.status != 200: |
| 245 | + return None |
| 246 | + |
| 247 | + ct = resp.headers.get('Content-Type', '').split(';')[0].strip().lower() |
| 248 | + if allowed_types is not None and ct not in allowed_types: |
| 249 | + return None |
| 250 | + |
| 251 | + data = b'' |
| 252 | + async for chunk in resp.content.iter_chunked(8192): |
| 253 | + data += chunk |
| 254 | + if len(data) > max_bytes: |
| 255 | + return None |
| 256 | + return data, ct |
| 257 | + except Exception: |
| 258 | + return None |
| 259 | + return None |
| 260 | + |
| 261 | + |
| 262 | +async def _discover_favicon(hostname: str) -> tuple[bytes, str] | None: |
| 263 | + # aiohttp bypasses the resolver for raw numeric IPs (is_ip_address() short-circuits |
| 264 | + # _resolve_host). This is the sole guard for that path. |
| 265 | + try: |
| 266 | + if not ipaddress.ip_address(hostname).is_global: |
| 267 | + return None |
| 268 | + except ValueError: |
| 269 | + pass # hostname; _SafeResolver handles it |
| 270 | + |
| 271 | + connector = aiohttp.TCPConnector(resolver=_SafeResolver(), use_dns_cache=False) |
| 272 | + headers = {'User-Agent': 'Mozilla/5.0 (compatible; open-webui/favicon-proxy)'} |
| 273 | + async with aiohttp.ClientSession(connector=connector, headers=headers) as session: |
| 274 | + # try /favicon.ico |
| 275 | + result = await _fetch_bounded( |
| 276 | + session, |
| 277 | + f'https://{hostname}/favicon.ico', |
| 278 | + _MAX_ICON_BYTES, |
| 279 | + _ALLOWED_ICON_TYPES, |
| 280 | + ) |
| 281 | + if result: |
| 282 | + return result |
| 283 | + |
| 284 | + # fall back to HTML <link> discovery |
| 285 | + html_result = await _fetch_bounded( |
| 286 | + session, |
| 287 | + f'https://{hostname}/', |
| 288 | + _MAX_HTML_BYTES, |
| 289 | + {'text/html'}, |
| 290 | + ) |
| 291 | + if not html_result: |
| 292 | + return None |
| 293 | + |
| 294 | + html_bytes, _ = html_result |
| 295 | + try: |
| 296 | + soup = BeautifulSoup(html_bytes, 'html.parser') |
| 297 | + except Exception: |
| 298 | + return None |
| 299 | + |
| 300 | + candidates: list[tuple[int, int, str]] = [] |
| 301 | + for tag in soup.find_all('link'): |
| 302 | + rel = tag.get('rel') or [] |
| 303 | + if isinstance(rel, str): |
| 304 | + rel = [rel] |
| 305 | + rel_lower = [r.lower() for r in rel] |
| 306 | + |
| 307 | + href = tag.get('href') or '' |
| 308 | + if not href or href.startswith('data:'): |
| 309 | + continue |
| 310 | + tag_type = (tag.get('type') or '').lower() |
| 311 | + if 'svg' in tag_type: |
| 312 | + continue |
| 313 | + |
| 314 | + if 'apple-touch-icon' in rel_lower and 'precomposed' not in rel_lower: |
| 315 | + priority, size = 0, 180 |
| 316 | + elif 'icon' in rel_lower: |
| 317 | + priority = 1 |
| 318 | + size = 0 |
| 319 | + sizes_str = (tag.get('sizes') or '').strip() |
| 320 | + if sizes_str and sizes_str.lower() != 'any': |
| 321 | + try: |
| 322 | + size = max( |
| 323 | + int(p.lower().split('x')[0]) |
| 324 | + for p in sizes_str.split() |
| 325 | + if 'x' in p.lower() |
| 326 | + ) |
| 327 | + except Exception: |
| 328 | + size = 0 |
| 329 | + elif 'shortcut' in rel_lower: |
| 330 | + priority, size = 2, 0 |
| 331 | + elif 'apple-touch-icon' in rel_lower: # precomposed |
| 332 | + priority, size = 3, 0 |
| 333 | + else: |
| 334 | + continue |
| 335 | + |
| 336 | + full_href = urljoin(f'https://{hostname}/', href) |
| 337 | + candidates.append((priority, -size, full_href)) |
| 338 | + |
| 339 | + candidates.sort(key=lambda x: (x[0], x[1])) |
| 340 | + |
| 341 | + for _, _, href in candidates: |
| 342 | + if not _is_safe_url(href): |
| 343 | + continue |
| 344 | + icon_result = await _fetch_bounded( |
| 345 | + session, href, _MAX_ICON_BYTES, _ALLOWED_ICON_TYPES |
| 346 | + ) |
| 347 | + if icon_result: |
| 348 | + return icon_result |
| 349 | + |
| 350 | + return None |
| 351 | + |
| 352 | + |
| 353 | +def _cache_get(hostname: str) -> tuple[bytes, str] | None: |
| 354 | + entry = _FAVICON_CACHE.get(hostname) |
| 355 | + if entry is None: |
| 356 | + return None |
| 357 | + content, ct, expires_at = entry |
| 358 | + if time.monotonic() > expires_at: |
| 359 | + del _FAVICON_CACHE[hostname] |
| 360 | + return None |
| 361 | + _FAVICON_CACHE.move_to_end(hostname) |
| 362 | + return content, ct |
| 363 | + |
| 364 | + |
| 365 | +def _cache_set(hostname: str, content: bytes, ct: str) -> None: |
| 366 | + if hostname in _FAVICON_CACHE: |
| 367 | + _FAVICON_CACHE.move_to_end(hostname) |
| 368 | + _FAVICON_CACHE[hostname] = (content, ct, time.monotonic() + _FAVICON_CACHE_TTL) |
| 369 | + while len(_FAVICON_CACHE) > _FAVICON_CACHE_MAX: |
| 370 | + _FAVICON_CACHE.popitem(last=False) |
| 371 | + |
| 372 | + |
| 373 | +_FAVICON_RESPONSE_HEADERS = { |
| 374 | + 'Cache-Control': 'public, max-age=604800', |
| 375 | + 'X-Content-Type-Options': 'nosniff', |
| 376 | +} |
| 377 | + |
| 378 | + |
| 379 | +async def _get_cached_favicon(redis, hostname: str) -> tuple[bytes, str] | None: |
| 380 | + """Return cached (content, ct) from Redis or in-memory, whichever is configured.""" |
| 381 | + if redis is not None: |
| 382 | + try: |
| 383 | + raw = await redis.get(f'{REDIS_KEY_PREFIX}:favicon:{hostname}') |
| 384 | + if raw: |
| 385 | + ct_bytes, _, content = raw.partition(b'\n') |
| 386 | + return content, ct_bytes.decode() |
| 387 | + except Exception: |
| 388 | + pass |
| 389 | + return None |
| 390 | + return _cache_get(hostname) |
| 391 | + |
| 392 | + |
| 393 | +async def _set_cached_favicon(redis, hostname: str, content: bytes, ct: str) -> None: |
| 394 | + """Write (content, ct) to Redis (with TTL) or in-memory, whichever is configured.""" |
| 395 | + if redis is not None: |
| 396 | + try: |
| 397 | + await redis.set( |
| 398 | + f'{REDIS_KEY_PREFIX}:favicon:{hostname}', |
| 399 | + ct.encode() + b'\n' + content, |
| 400 | + ex=_FAVICON_CACHE_TTL, |
| 401 | + ) |
| 402 | + except Exception: |
| 403 | + pass |
| 404 | + else: |
| 405 | + _cache_set(hostname, content, ct) |
| 406 | + |
| 407 | + |
| 408 | +@router.get('/favicon') |
| 409 | +async def get_favicon(request: Request, url: str = '', user=Depends(get_verified_user)): |
| 410 | + """Favicon proxy; always returns a response.""" |
| 411 | + hostname = _normalize_hostname(url) |
| 412 | + redis = request.app.state.redis |
| 413 | + |
| 414 | + if hostname: |
| 415 | + cached = await _get_cached_favicon(redis, hostname) |
| 416 | + if cached: |
| 417 | + content, ct = cached |
| 418 | + return Response(content=content, media_type=ct, headers=_FAVICON_RESPONSE_HEADERS) |
| 419 | + |
| 420 | + icon: tuple[bytes, str] | None = None |
| 421 | + if hostname: |
| 422 | + try: |
| 423 | + icon = await _discover_favicon(hostname) |
| 424 | + except Exception: |
| 425 | + pass |
| 426 | + |
| 427 | + content, ct = icon if icon else (_svg_badge(hostname or ''), 'image/svg+xml') |
| 428 | + |
| 429 | + if hostname: |
| 430 | + await _set_cached_favicon(redis, hostname, content, ct) |
| 431 | + |
| 432 | + return Response(content=content, media_type=ct, headers=_FAVICON_RESPONSE_HEADERS) |
0 commit comments