From 3b95e89fa174cbeca6cde65dff592fc244f4c0d2 Mon Sep 17 00:00:00 2001
From: orangix <uleo8b8g@anonaddy.me>
Date: Fri, 30 Jan 2026 04:04:53 +0100
Subject: [PATCH] filter extensions for HandleMedia

---
 main.go        | 9 ++++++++-
 pages/media.go | 7 +++----
 2 files changed, 11 insertions(+), 5 deletions(-)

diff --git a/main.go b/main.go
index 0b99d4a..eaa130c 100644
--- a/main.go
+++ b/main.go
@@ -129,7 +129,14 @@ func main() {
 			r.SetPathValue("postID", component)
 			return pages.HandleGifv(w, r)
 		case strings.Contains(component, "."):
-			return pages.HandleMedia(w, r)
+			baseName, extension := utils.SplitNameExt(r.PathValue("component"))
+			r.SetPathValue("baseName", baseName)
+			r.SetPathValue("extension", extension)
+			switch extension {
+			case "png", "gif", "jpg", "jpeg", "webp":
+				return pages.HandleMedia(w, r)
+			}
+			fallthrough
 		default:
 			r.SetPathValue("postID", component)
 			return pages.HandlePost(w, r)
diff --git a/pages/media.go b/pages/media.go
index e327811..8ea24e5 100644
--- a/pages/media.go
+++ b/pages/media.go
@@ -12,12 +12,11 @@ import (
 func HandleMedia(w http.ResponseWriter, r *http.Request) error {
 	w.Header().Set("Cache-Control", "public,max-age=31557600")
 	w.Header().Set("Content-Security-Policy", "default-src 'none'; style-src 'self'; img-src 'self'")
-	splitName := strings.SplitN(r.PathValue("component"), ".", 2)
-	baseName, extension := splitName[0], splitName[1]
+	baseName, extension := r.PathValue("baseName"), r.PathValue("extension")
 	if strings.HasPrefix(r.URL.Path, "/stack") {
-		return handleMedia(w, r, "https://i.stack.imgur.com/"+strings.ReplaceAll(baseName, "stack/", "")+"."+extension)
+		return handleMedia(w, r, "https://i.stack.imgur.com/"+baseName[5:]+extension)
 	} else {
-		return handleMedia(w, r, "https://i.imgur.com/"+baseName+"."+extension)
+		return handleMedia(w, r, "https://i.imgur.com/"+baseName+extension)
 	}
 }