From 5ecada8f5442e1fd3c3ac3e9626bed2f90ad57d1 Mon Sep 17 00:00:00 2001
From: Jan-Lukas Else <mail@jlelse.de>
Date: Wed, 1 Jan 2020 16:03:13 +0100
Subject: [PATCH] Improve extension handling for uploaded files

---
 mediaendpoint.go | 19 +++++++++++++------
 1 file changed, 13 insertions(+), 6 deletions(-)

diff --git a/mediaendpoint.go b/mediaendpoint.go
index 281ad4e..60297d9 100644
--- a/mediaendpoint.go
+++ b/mediaendpoint.go
@@ -4,7 +4,9 @@ import (
 	"crypto/sha256"
 	"fmt"
 	"io"
+	"mime"
 	"net/http"
+	"path/filepath"
 	"strings"
 )
 
@@ -45,13 +47,18 @@ func HandleMedia(w http.ResponseWriter, r *http.Request) {
 			return
 		}
 		fileName := fmt.Sprintf("%x", h.Sum(nil))
-		mimeType := header.Header.Get("content-type")
-		originalName := header.Filename
-		if strings.Contains(mimeType, "jpeg") || strings.Contains(originalName, ".jpeg") || strings.Contains(originalName, ".jpg") {
-			fileName += ".jpg"
-		} else if strings.Contains(mimeType, "png") || strings.Contains(originalName, ".png") {
-			fileName += ".png"
+		fileExtension := filepath.Ext(header.Filename)
+		if len(fileExtension) == 0 {
+			// Find correct file extension if original filename does not contain one
+			mimeType := header.Header.Get("content-type")
+			if len(mimeType) > 0 {
+				allExtensions, _ := mime.ExtensionsByType(mimeType)
+				if len(allExtensions) > 0 {
+					fileExtension = allExtensions[0]
+				}
+			}
 		}
+		fileName += fileExtension
 		location, err := SelectedMediaStorage.Upload(fileName, file)
 		if err != nil {
 			w.WriteHeader(http.StatusInternalServerError)