Improve extension handling for uploaded files

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)