Browse Source

Merge pull request #1382 from open-webui/main

dev
Timothy Jaeryang Baek 1 year ago
parent
commit
6c3f2f8d38
2 changed files with 12 additions and 0 deletions
  1. 9 0
      backend/apps/ollama/main.py
  2. 3 0
      src/lib/components/chat/Settings/Models.svelte

+ 9 - 0
backend/apps/ollama/main.py

@@ -1029,6 +1029,14 @@ async def download_model(
     url_idx: Optional[int] = None,
     url_idx: Optional[int] = None,
 ):
 ):
 
 
+    allowed_hosts = ["https://huggingface.co/", "https://github.com/"]
+
+    if not any(form_data.url.startswith(host) for host in allowed_hosts):
+        raise HTTPException(
+            status_code=400,
+            detail="Invalid file_url. Only URLs from allowed hosts are permitted.",
+        )
+
     if url_idx == None:
     if url_idx == None:
         url_idx = 0
         url_idx = 0
     url = app.state.OLLAMA_BASE_URLS[url_idx]
     url = app.state.OLLAMA_BASE_URLS[url_idx]
@@ -1037,6 +1045,7 @@ async def download_model(
 
 
     if file_name:
     if file_name:
         file_path = f"{UPLOAD_DIR}/{file_name}"
         file_path = f"{UPLOAD_DIR}/{file_name}"
+
         return StreamingResponse(
         return StreamingResponse(
             download_file_stream(url, form_data.url, file_path, file_name),
             download_file_stream(url, form_data.url, file_path, file_name),
         )
         )

+ 3 - 0
src/lib/components/chat/Settings/Models.svelte

@@ -258,6 +258,9 @@
 					console.log(error);
 					console.log(error);
 				}
 				}
 			}
 			}
+		} else {
+			const error = await fileResponse?.json();
+			toast.error(error?.detail ?? error);
 		}
 		}
 
 
 		if (uploaded) {
 		if (uploaded) {