From 77b1edcd0fef991f76c6e4db07d547484f64390c Mon Sep 17 00:00:00 2001 From: "Timothy J. Baek" Date: Mon, 1 Apr 2024 13:24:48 -0700 Subject: [PATCH] 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):