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