From 77b1edcd0fef991f76c6e4db07d547484f64390c Mon Sep 17 00:00:00 2001
From: "Timothy J. Baek" <timothyjrbeck@gmail.com>
Date: Mon, 1 Apr 2024 13:24:48 -0700
Subject: [PATCH 1/2] fix: allowed hosts

---
 backend/apps/ollama/main.py | 7 +++++++
 1 file changed, 7 insertions(+)

diff --git a/backend/apps/ollama/main.py b/backend/apps/ollama/main.py
index f9ebdb98f..73c94c3a9 100644
--- a/backend/apps/ollama/main.py
+++ b/backend/apps/ollama/main.py
@@ -970,6 +970,13 @@ def parse_huggingface_url(hf_url):
 async def download_file_stream(
     ollama_url, file_url, file_path, file_name, chunk_size=1024 * 1024
 ):
+    allowed_hosts = ["https://huggingface.co/", "https://github.com/"]
+
+    if not any(file_url.startswith(host) for host in allowed_hosts):
+        raise ValueError(
+            "Invalid file_url. Only URLs from allowed hosts are permitted."
+        )
+
     done = False
 
     if os.path.exists(file_path):

From d72653cdea760be3908dfe82fd1ec84e1dfa48ca Mon Sep 17 00:00:00 2001
From: "Timothy J. Baek" <timothyjrbeck@gmail.com>
Date: Mon, 1 Apr 2024 14:01:05 -0700
Subject: [PATCH 2/2] fix: download allowed hosts

---
 backend/apps/ollama/main.py                    | 16 +++++++++-------
 src/lib/components/chat/Settings/Models.svelte |  3 +++
 2 files changed, 12 insertions(+), 7 deletions(-)

diff --git a/backend/apps/ollama/main.py b/backend/apps/ollama/main.py
index 73c94c3a9..818b53d93 100644
--- a/backend/apps/ollama/main.py
+++ b/backend/apps/ollama/main.py
@@ -970,13 +970,6 @@ def parse_huggingface_url(hf_url):
 async def download_file_stream(
     ollama_url, file_url, file_path, file_name, chunk_size=1024 * 1024
 ):
-    allowed_hosts = ["https://huggingface.co/", "https://github.com/"]
-
-    if not any(file_url.startswith(host) for host in allowed_hosts):
-        raise ValueError(
-            "Invalid file_url. Only URLs from allowed hosts are permitted."
-        )
-
     done = False
 
     if os.path.exists(file_path):
@@ -1036,6 +1029,14 @@ async def download_model(
     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:
         url_idx = 0
     url = app.state.OLLAMA_BASE_URLS[url_idx]
@@ -1044,6 +1045,7 @@ async def download_model(
 
     if file_name:
         file_path = f"{UPLOAD_DIR}/{file_name}"
+
         return StreamingResponse(
             download_file_stream(url, form_data.url, file_path, file_name),
         )
diff --git a/src/lib/components/chat/Settings/Models.svelte b/src/lib/components/chat/Settings/Models.svelte
index 6de76a483..baa2f6c4e 100644
--- a/src/lib/components/chat/Settings/Models.svelte
+++ b/src/lib/components/chat/Settings/Models.svelte
@@ -258,6 +258,9 @@
 					console.log(error);
 				}
 			}
+		} else {
+			const error = await fileResponse?.json();
+			toast.error(error?.detail ?? error);
 		}
 
 		if (uploaded) {