Improve upload file type detection

Author: ego008

controller/user_file_upload.go

@@ -2,17 +2,17 @@ package controller
 
 import (
 	"bytes"
-	"github.com/ego008/goutils/json"
-	"github.com/ego008/sdb"
-	"github.com/valyala/fasthttp"
 	"goyoubbs/model"
 	"goyoubbs/util"
 	"image"
 	"image/jpeg"
 	"io"
 	"os"
-	"path"
 	"strconv"
+
+	"github.com/ego008/goutils/json"
+	"github.com/ego008/sdb"
+	"github.com/valyala/fasthttp"
 )
 
 const (
@@ -89,9 +89,19 @@ func (h *BaseHandler) FileUpload(ctx *fasthttp.RequestCtx) {
 			showPath = "/static/upload/" + saveName
 		}
 	} else {
-		fileSuffix = path.Ext(file.Filename) // source file suffix
-		saveName = imgKeyS + fileSuffix
-		showPath = "/static/upload/" + saveName
+		// is mp3 or mp4
+		mediaType := util.CheckMediaType(buff)
+		if len(mediaType) > 0 {
+			fileSuffix = "." + mediaType
+			saveName = imgKeyS + fileSuffix
+			showPath = "/static/upload/" + saveName
+		} else {
+			_, _ = ctx.WriteString(`{"Code":400,"Msg":"unknown image or media format"}`)
+			return
+		}
+		//fileSuffix = path.Ext(file.Filename) // source file suffix
+		//saveName = imgKeyS + fileSuffix
+		//showPath = "/static/upload/" + saveName
 	}
 
 	saveFullPath := h.App.Cf.Site.UploadDir + "/" + saveName

upload/11600925034256400078.jpg

@@ -3,20 +3,28 @@ package util
 import (
 	"bytes"
 	"errors"
-	"github.com/disintegration/imaging"
 	"image"
 	"image/gif"
 	"image/jpeg"
 	"image/png"
 	"net/http"
+
+	"github.com/disintegration/imaging"
 )
 
-var imgTable = map[string]string{
-	"image/jpeg": "jpg",
-	"image/jpg":  "jpg",
-	"image/gif":  "gif",
-	"image/png":  "png",
-}
+var (
+	imgTable = map[string]string{
+		"image/jpeg": "jpg",
+		"image/jpg":  "jpg",
+		"image/gif":  "gif",
+		"image/png":  "png",
+	}
+
+	mediaTable = map[string]string{
+		"audio/mpeg": "mp3",
+		"video/mp4":  "mp4",
+	}
+)
 
 func CheckImageType(buff []byte) string {
 	// why 512 bytes ? see http://golang.org/pkg/net/http/#DetectContentType
@@ -28,6 +36,64 @@ func CheckImageType(buff []byte) string {
 	return ""
 }
 
+func CheckMediaType(buff []byte) string {
+	// why 512 bytes ? see http://golang.org/pkg/net/http/#DetectContentType
+	//buff := make([]byte, 512)
+	fileType := http.DetectContentType(buff)
+	if v, ok := mediaTable[fileType]; ok {
+		return v
+	}
+	if isMP3(buff) {
+		return "mp3"
+	}
+	if isMP4(buff) {
+		return "mp4"
+	}
+	return ""
+}
+
+// 检测MP4文件的魔数
+func isMP4(buffer []byte) bool {
+	if len(buffer) < 8 {
+		return false
+	}
+
+	// MP4文件通常以ftyp开头
+	if string(buffer[4:8]) == "ftyp" {
+		return true
+	}
+
+	return false
+}
+
+// 检测MP3文件的魔数
+func isMP3(buffer []byte) bool {
+	if len(buffer) < 3 {
+		return false
+	}
+
+	// MP3文件帧头检测
+	// 检查ID3v2标签（以"ID3"开头）
+	if len(buffer) >= 3 && string(buffer[0:3]) == "ID3" {
+		return true
+	}
+
+	// 检查MPEG帧同步位
+	if len(buffer) >= 2 {
+		// 检查FF FB (MPEG版本1, 层3)
+		if buffer[0] == 0xFF && (buffer[1]&0xE0) == 0xE0 {
+			// 检查具体的MPEG音频帧
+			if (buffer[1]&0x18)>>3 == 0x03 { // MPEG版本1
+				if (buffer[1]&0x06)>>1 == 0x01 { // 层3
+					return true
+				}
+			}
+		}
+	}
+
+	return false
+}
+
 func GetImageObj(buff *bytes.Buffer) (image.Image, error) {
 	var img image.Image
 	var err error