From 4850855448a77a4a0e71f9100aff053a8a779a5a Mon Sep 17 00:00:00 2001
From: AH-dark <ahdark@outlook.com>
Date: Sun, 7 Jan 2024 15:52:34 +0800
Subject: [PATCH] fix: handle for different image format

---
 services/avatar/get_gravatar.go  |  5 ++---
 services/avatar/get_qq_avatar.go |  9 ++++-----
 services/avatar/utils.go         | 27 +++++++++++++++++++++++++++
 3 files changed, 33 insertions(+), 8 deletions(-)
 create mode 100644 services/avatar/utils.go

diff --git a/services/avatar/get_gravatar.go b/services/avatar/get_gravatar.go
index 42cd26f..942069c 100644
--- a/services/avatar/get_gravatar.go
+++ b/services/avatar/get_gravatar.go
@@ -6,7 +6,6 @@ import (
 	"github.com/imroc/req/v3"
 	"github.com/samber/lo"
 	"image"
-	"image/jpeg"
 	"strconv"
 	"time"
 
@@ -71,9 +70,9 @@ func (s *service) getGravatar(ctx context.Context, hash string, args GetAvatarAr
 
 	defer resp.Body.Close()
 
-	img, err := jpeg.Decode(resp.Body)
+	img, err := parseImage(resp.GetHeader("Content-Type"), resp.Body)
 	if err != nil {
-		otelzap.L().Ctx(ctx).Error("failed to decode gravatar", zap.Error(err))
+		otelzap.L().Ctx(ctx).Error("parse image failed", zap.Error(err))
 		return GetGravatarResult{}, err
 	}
 
diff --git a/services/avatar/get_qq_avatar.go b/services/avatar/get_qq_avatar.go
index 318a475..61090da 100644
--- a/services/avatar/get_qq_avatar.go
+++ b/services/avatar/get_qq_avatar.go
@@ -3,14 +3,13 @@ package avatar
 import (
 	"context"
 	"fmt"
-	"github.com/imroc/req/v3"
-	"go.uber.org/zap"
 	"image"
-	"image/jpeg"
 	"strconv"
 
+	"github.com/imroc/req/v3"
 	"github.com/nfnt/resize"
 	"github.com/uptrace/opentelemetry-go-extra/otelzap"
+	"go.uber.org/zap"
 )
 
 func initQQClient() *req.Client {
@@ -67,9 +66,9 @@ func (s *service) getQQAvatar(ctx context.Context, hash string, args GetAvatarAr
 
 	defer resp.Body.Close()
 
-	img, err := jpeg.Decode(resp.Body)
+	img, err := parseImage(resp.GetHeader("Content-Type"), resp.Body)
 	if err != nil {
-		otelzap.L().Ctx(ctx).Error("decode qq avatar failed", zap.Error(err))
+		otelzap.L().Ctx(ctx).Error("parse image failed", zap.Error(err))
 		return nil, err
 	}
 
diff --git a/services/avatar/utils.go b/services/avatar/utils.go
new file mode 100644
index 0000000..9231b14
--- /dev/null
+++ b/services/avatar/utils.go
@@ -0,0 +1,27 @@
+package avatar
+
+import (
+	"github.com/kolesa-team/go-webp/webp"
+	"image"
+	"image/gif"
+	"image/jpeg"
+	"image/png"
+	"io"
+)
+
+func parseImage(contentType string, reader io.Reader) (img image.Image, err error) {
+	switch contentType {
+	case "image/jpeg":
+		img, err = jpeg.Decode(reader)
+	case "image/png":
+		img, err = png.Decode(reader)
+	case "image/webp":
+		img, err = webp.Decode(reader, nil)
+	case "image/gif":
+		img, err = gif.Decode(reader)
+	default:
+		img, err = jpeg.Decode(reader)
+	}
+
+	return
+}