From c878bb8ca433f5041d026da48e190924d1dd821d Mon Sep 17 00:00:00 2001
From: Unlock Music Dev <dev@unlock-music.dev>
Date: Tue, 22 Nov 2022 06:45:18 +0800
Subject: [PATCH] feat: add image sniffer

---
 cmd/um/main.go          | 20 ++++++++++++++++++++
 internal/sniff/audio.go |  2 +-
 internal/sniff/image.go | 30 ++++++++++++++++++++++++++++++
 3 files changed, 51 insertions(+), 1 deletion(-)
 create mode 100644 internal/sniff/image.go

diff --git a/cmd/um/main.go b/cmd/um/main.go
index 62eb144..9673a83 100644
--- a/cmd/um/main.go
+++ b/cmd/um/main.go
@@ -2,6 +2,7 @@ package main
 
 import (
 	"bytes"
+	"context"
 	"errors"
 	"fmt"
 	"io"
@@ -208,6 +209,25 @@ func tryDecFile(inputFile string, outputDir string, allDec []common.NewDecoderFu
 		return err
 	}
 
+	if coverGetter, ok := dec.(common.CoverImageGetter); ok {
+		ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
+		defer cancel()
+
+		cover, err := coverGetter.GetCoverImage(ctx)
+		if err != nil {
+			logger.Warn("get cover image failed", zap.Error(err))
+		} else if imgExt, ok := sniff.ImageExtension(cover); !ok {
+			logger.Warn("sniff cover image type failed", zap.Error(err))
+		} else {
+			coverPath := filepath.Join(outputDir, inFilename+imgExt)
+			err = os.WriteFile(coverPath, cover, 0644)
+			if err != nil {
+				logger.Warn("write cover image failed", zap.Error(err))
+			}
+		}
+
+	}
+
 	// if source file need to be removed
 	if removeSource {
 		err := os.RemoveAll(inputFile)
diff --git a/internal/sniff/audio.go b/internal/sniff/audio.go
index 2f36276..12454dd 100644
--- a/internal/sniff/audio.go
+++ b/internal/sniff/audio.go
@@ -13,7 +13,7 @@ type Sniffer interface {
 
 var audioExtensions = map[string]Sniffer{
 	// ref: https://mimesniff.spec.whatwg.org
-	".mp3": prefixSniffer("ID3"),
+	".mp3": prefixSniffer("ID3"), // todo: check mp3 without ID3v2 tag
 	".ogg": prefixSniffer("OggS"),
 	".wav": prefixSniffer("RIFF"),
 
diff --git a/internal/sniff/image.go b/internal/sniff/image.go
new file mode 100644
index 0000000..ed567a0
--- /dev/null
+++ b/internal/sniff/image.go
@@ -0,0 +1,30 @@
+package sniff
+
+// ref: https://mimesniff.spec.whatwg.org
+var imageMIMEs = map[string]Sniffer{
+	"image/jpeg": prefixSniffer{0xFF, 0xD8, 0xFF},
+	"image/png":  prefixSniffer{'P', 'N', 'G', '\r', '\n', 0x1A, '\n'},
+	"image/bmp":  prefixSniffer("BM"),
+	"image/webp": prefixSniffer("RIFF"),
+	"image/gif":  prefixSniffer("GIF8"),
+}
+
+// ImageMIME sniffs the well-known image types, and returns its MIME.
+func ImageMIME(header []byte) (string, bool) {
+	for ext, sniffer := range imageMIMEs {
+		if sniffer.Sniff(header) {
+			return ext, true
+		}
+	}
+	return "", false
+}
+
+// ImageExtension is equivalent to ImageMIME, but returns file extension
+func ImageExtension(header []byte) (string, bool) {
+	ext, ok := ImageMIME(header)
+	if !ok {
+		return "", false
+	}
+	// todo: use mime.ExtensionsByType
+	return "." + ext[6:], true // "image/" is 6 bytes
+}