Coverage for open_webui/utils/files.py: 17%

1import asyncio 

2import base64 

3import io 

4import mimetypes 

5import re 

6from pathlib import Path 

7from typing import Optional 

8 

9import aiofiles 

10from fastapi import ( 

11 APIRouter, 

12 Depends, 

13 HTTPException, 

14 Request, 

15 UploadFile, 

16) 

17from open_webui.env import ( 

18 AIOHTTP_CLIENT_ALLOW_REDIRECTS, 

19 AIOHTTP_CLIENT_SESSION_SSL, 

20 ENABLE_IMAGE_CONTENT_TYPE_EXTENSION_FALLBACK, 

21) 

22from open_webui.models.chats import Chats 

23from open_webui.models.files import Files 

24from open_webui.retrieval.web.utils import get_ssrf_safe_session, validate_url 

25from open_webui.routers.files import upload_file_handler 

26from open_webui.utils.access_control.files import has_access_to_file 

27from open_webui.routers.images import ( 

28 get_image_data, 

29 upload_image, 

30) 

31from open_webui.storage.provider import Storage 

32 

33BASE64_IMAGE_URL_PREFIX = re.compile(r'data:image/\w+;base64,', re.IGNORECASE) 

34MARKDOWN_IMAGE_URL_PATTERN = re.compile(r'!\[(.*?)\]\((.+?)\)', re.IGNORECASE) 

35FILE_CONTENT_URL_PATTERN = re.compile(r'^/api/v1/files/([^/?#]+)/content(?:[?#]|$)') 

36 

37# Extension-based MIME fallback, only used when ENABLE_IMAGE_CONTENT_TYPE_EXTENSION_FALLBACK is True. 

38_IMAGE_MIME_FALLBACK = { 

39 '.webp': 'image/webp', 

40 '.png': 'image/png', 

41 '.jpg': 'image/jpeg', 

42 '.jpeg': 'image/jpeg', 

43 '.gif': 'image/gif', 

44 '.svg': 'image/svg+xml', 

45 '.bmp': 'image/bmp', 

46 '.tiff': 'image/tiff', 

47 '.tif': 'image/tiff', 

48 '.ico': 'image/x-icon', 

49 '.heic': 'image/heic', 

50 '.heif': 'image/heif', 

51 '.avif': 'image/avif', 

52} 

53 

54 

55async def get_image_base64_from_url(url: str, user=None) -> Optional[str]: 

56 try: 

57 if url.startswith('http'): 

58 from open_webui.models.config import Config 

59 

60 max_bytes = None 

61 try: 

62 max_size_mb = int(await Config.get('rag.file.max_size') or 0) 

63 except (TypeError, ValueError): 

64 max_size_mb = 0 

65 if max_size_mb > 0: 

66 max_bytes = max_size_mb * 1024 * 1024 

67 

68 # Validate URL to prevent SSRF attacks against local/private networks. 

69 # allow_redirects=False prevents redirect-based SSRF: validate_url() is 

70 # called only on the originally-submitted URL; following 3xx redirects 

71 # without re-validation would let an attacker reach private IPs via a 

72 # public host that redirects internally (e.g. cloud-metadata exfil). 

73 await asyncio.to_thread(validate_url, url) 

74 # Fetch through an SSRF-safe session that re-checks the connect-time IP, so a 

75 # rebinding DNS answer that passed validate_url cannot reach an internal address. 

76 async with get_ssrf_safe_session() as session: 

77 async with session.get( 

78 url, 

79 ssl=AIOHTTP_CLIENT_SESSION_SSL, 

80 allow_redirects=AIOHTTP_CLIENT_ALLOW_REDIRECTS, 

81 headers={'Accept-Encoding': 'identity'}, 

82 ) as response: 

83 response.raise_for_status() 

84 # Accept-Encoding is only a request; the sender can still compress and pick our decompressed size. 

85 encodings = response.headers.getall('Content-Encoding', ()) 

86 if any(encoding.lower() not in ('', 'identity') for encoding in encodings): 

87 return None 

88 image_data = bytearray() 

89 total = 0 

90 async for chunk in response.content.iter_chunked(64 * 1024): 

91 total += len(chunk) 

92 if max_bytes is not None and total > max_bytes: 

93 return None 

94 image_data.extend(chunk) 

95 encoded_string = base64.b64encode(image_data).decode('utf-8') 

96 content_type = response.headers.get('Content-Type', 'image/png') 

97 return f'data:{content_type};base64,{encoded_string}' 

98 else: 

99 # Non-URL string — treat as file_id. Delegate to the canonical 

100 # file-ID resolver which enforces ownership/access checks. 

101 file_id_match = FILE_CONTENT_URL_PATTERN.match(url) 

102 return await get_image_base64_from_file_id(file_id_match.group(1) if file_id_match else url, user=user) 

103 

104 except Exception: 

105 return None 

106 

107 

108async def get_image_url_from_base64(request, base64_image_string, metadata, user): 

109 if BASE64_IMAGE_URL_PREFIX.match(base64_image_string): 

110 image_url = '' 

111 # Extract base64 image data from the line 

112 image_data, content_type = await get_image_data(base64_image_string) 

113 if image_data is not None: 

114 _, image_file = await upload_image( 

115 request, 

116 image_data, 

117 content_type, 

118 metadata, 

119 user, 

120 ) 

121 image_url = image_file['url'] 

122 

123 return image_url 

124 return None 

125 

126 

127async def convert_markdown_base64_images(request, content: str, metadata, user): 

128 MIN_REPLACEMENT_URL_LENGTH = 1024 

129 result_parts = [] 

130 last_end = 0 

131 

132 for match in MARKDOWN_IMAGE_URL_PATTERN.finditer(content): 

133 result_parts.append(content[last_end : match.start()]) 

134 base64_string = match.group(2) 

135 if len(base64_string) > MIN_REPLACEMENT_URL_LENGTH: 

136 url = await get_image_url_from_base64(request, base64_string, metadata, user) 

137 if url: 

138 result_parts.append(f'') 

139 else: 

140 result_parts.append(match.group(0)) 

141 else: 

142 result_parts.append(match.group(0)) 

143 last_end = match.end() 

144 

145 result_parts.append(content[last_end:]) 

146 return ''.join(result_parts) 

147 

148 

149def load_b64_audio_data(b64_str): 

150 try: 

151 if ',' in b64_str: 

152 header, b64_data = b64_str.split(',', 1) 

153 else: 

154 b64_data = b64_str 

155 header = 'data:audio/wav;base64' 

156 audio_data = base64.b64decode(b64_data) 

157 content_type = header.split(';')[0].split(':')[1] if ';' in header else 'audio/wav' 

158 return audio_data, content_type 

159 except Exception as e: 

160 print(f'Error decoding base64 audio data: {e}') 

161 return None, None 

162 

163 

164async def upload_audio(request, audio_data, content_type, metadata, user): 

165 audio_format = mimetypes.guess_extension(content_type) 

166 file = UploadFile( 

167 file=io.BytesIO(audio_data), 

168 filename=f'generated-{audio_format}', # will be converted to a unique ID on upload_file 

169 headers={ 

170 'content-type': content_type, 

171 }, 

172 ) 

173 file_item = await upload_file_handler( 

174 request, 

175 file=file, 

176 metadata=metadata, 

177 process=False, 

178 user=user, 

179 ) 

180 url = request.app.url_path_for('get_file_content_by_id', id=file_item.id) 

181 return url 

182 

183 

184async def get_audio_url_from_base64(request, base64_audio_string, metadata, user): 

185 if 'data:audio/wav;base64' in base64_audio_string: 

186 audio_url = '' 

187 # Extract base64 audio data from the line 

188 audio_data, content_type = load_b64_audio_data(base64_audio_string) 

189 if audio_data is not None: 

190 audio_url = await upload_audio( 

191 request, 

192 audio_data, 

193 content_type, 

194 metadata, 

195 user, 

196 ) 

197 return audio_url 

198 return None 

199 

200 

201async def get_file_url_from_base64(request, base64_file_string, metadata, user): 

202 if BASE64_IMAGE_URL_PREFIX.match(base64_file_string): 

203 return await get_image_url_from_base64(request, base64_file_string, metadata, user) 

204 elif 'data:audio/wav;base64' in base64_file_string: 

205 return await get_audio_url_from_base64(request, base64_file_string, metadata, user) 

206 return None 

207 

208 

209async def get_image_base64_from_file_id(id: str, user=None) -> Optional[str]: 

210 file = await Files.get_file_by_id(id) 

211 if not file: 

212 return None 

213 

214 # Gate file-by-id resolution by ownership to prevent exfiltration. 

215 # A caller could place another user's file_id in an image_url field; 

216 # without this check the server reads the file from disk, inlines it 

217 # base64 into the LLM request, and the content leaks via OCR/describe. 

218 # Owner, admin, and explicit read-grant holders are allowed. 

219 if user is None: 

220 return None 

221 if file.user_id != user.id and user.role != 'admin' and not await has_access_to_file(file.id, 'read', user): 

222 return None 

223 

224 try: 

225 file_path = await asyncio.to_thread(Storage.get_file, file.path) 

226 file_path = Path(file_path) 

227 

228 # Check if the file already exists in the cache 

229 if file_path.is_file(): 

230 async with aiofiles.open(file_path, 'rb') as image_file: 

231 encoded_string = base64.b64encode(await image_file.read()).decode('utf-8') 

232 content_type = mimetypes.guess_type(file_path.name)[0] or (file.meta or {}).get('content_type') 

233 if not content_type and ENABLE_IMAGE_CONTENT_TYPE_EXTENSION_FALLBACK: 

234 content_type = _IMAGE_MIME_FALLBACK.get(file_path.suffix.lower()) 

235 if not content_type: 

236 return None 

237 return f'data:{content_type};base64,{encoded_string}' 

238 else: 

239 return None 

240 except Exception: 

241 return None