feat: optimized structure & added cache headers

server/controllers/avatar/get_avatar.go

@@ -1,15 +1,19 @@
 package avatar
 
 import (
+ "bytes"
  "context"
+ "net/http"
  "strings"
+ "time"
 
  "github.com/AH-dark/bytestring"
  "github.com/cloudwego/hertz/pkg/app"
  "github.com/samber/lo"
  "github.com/uptrace/opentelemetry-go-extra/otelzap"
  "go.uber.org/zap"
 
+ "github.com/AH-dark/gravatar-with-qq-avatar/pkg/cryptor"
  "github.com/AH-dark/gravatar-with-qq-avatar/services/avatar"
 )
 
@@ -23,7 +27,7 @@ type GetAvatarRequest struct {
 }
 
 func (h *handlers) GetAvatar(ctx context.Context, c *app.RequestContext) {
- ctx, span := tracer.Start(ctx, "server.controllers.avatar.GetAvatar")
+ ctx, span := tracer.Start(ctx, "server.controllers.avatarData.GetAvatar")
  defer span.End()
 
  var req GetAvatarRequest
@@ -53,23 +57,31 @@ func (h *handlers) GetAvatar(ctx context.Context, c *app.RequestContext) {
  c.Header("Content-Type", "image/png")
  }
 
- qqAvatar, err := h.AvatarService.GetQQAvatar(ctx, req.Hash, args)
+ avatarData, lastModified, err := h.AvatarService.GetAvatar(ctx, req.Hash, args)
  if err != nil {
- otelzap.L().Ctx(ctx).Warn("get qq avatar failed", zap.Error(err))
- } else {
- c.Status(200)
- c.SetBodyStream(qqAvatar, -1)
+ otelzap.L().Ctx(ctx).Error("get avatar data failed", zap.Error(err))
+ c.AbortWithStatus(http.StatusInternalServerError)
  return
  }
 
- gravatar, err := h.AvatarService.GetGravatar(ctx, req.Hash, args)
- if err != nil {
- otelzap.L().Ctx(ctx).Error("get gravatar failed", zap.Error(err))
- c.AbortWithStatus(500)
+ if avatarData == nil {
+ c.NotFound()
+ return
+ }
+
+ c.Data(http.StatusOK, lo.If(args.EnableWebp, "image/webp").Else("image/png"), avatarData)
+ if !lastModified.IsZero() {
+ c.Header("Last-Modified", lastModified.UTC().Format(time.RFC1123))
+ }
+
+ md5 := cryptor.Md5(avatarData)
+ c.Header("ETag", bytestring.BytesToString(md5))
+ if bytes.Equal(md5, c.GetHeader("If-None-Match")) {
+ c.Status(http.StatusNotModified)
  return
  }
- defer gravatar.Close()
 
- c.Status(200)
- c.SetBodyStream(gravatar, -1)
+ c.Header("Cache-Control", "public, max-age=86400, immutable")
+ c.Header("Expires", time.Now().Add(86400*time.Second).UTC().Format(time.RFC1123))
+ c.Header("X-Content-Type-Options", "nosniff")
 }

services/avatar/get_gravatar.go

@@ -3,15 +3,23 @@ package avatar
 import (
  "context"
  "fmt"
- "github.com/uptrace/opentelemetry-go-extra/otelzap"
- "go.uber.org/zap"
- "io"
+ "image"
+ "image/jpeg"
  "net/url"
  "strconv"
+ "time"
+
+ "github.com/uptrace/opentelemetry-go-extra/otelzap"
+ "go.uber.org/zap"
 )
 
-func (s *service) GetGravatar(ctx context.Context, hash string, args GetAvatarArgs) (io.ReadCloser, error) {
- ctx, span := tracer.Start(ctx, "service.AvatarService.GetGravatar")
+type GetGravatarResult struct {
+ Avatar image.Image
+ LastModified time.Time
+}
+
+func (s *service) getGravatar(ctx context.Context, hash string, args GetAvatarArgs) (GetGravatarResult, error) {
+ ctx, span := tracer.Start(ctx, "service.AvatarService.getGravatar")
  defer span.End()
 
  values := url.Values{}
@@ -28,9 +36,34 @@ func (s *service) GetGravatar(ctx context.Context, hash string, args GetAvatarAr
  u, err := url.Parse(fmt.Sprintf("https://www.gravatar.com/avatar/%s", hash))
  if err != nil {
  otelzap.L().Ctx(ctx).Error("parse gravatar url failed", zap.Error(err))
- return nil, err
+ return GetGravatarResult{}, err
  }
  u.RawQuery = values.Encode()
 
- return s.downloadAvatar(ctx, u.String())
+ stream, resp, err := s.downloadAvatar(ctx, u.String())
+ if err != nil {
+ otelzap.L().Ctx(ctx).Error("download gravatar failed", zap.Error(err))
+ return GetGravatarResult{}, err
+ }
+
+ img, err := jpeg.Decode(stream)
+ if err != nil {
+ otelzap.L().Ctx(ctx).Error("failed to decode gravatar", zap.Error(err))
+ return GetGravatarResult{}, err
+ }
+
+ if err := stream.Close(); err != nil {
+ otelzap.L().Ctx(ctx).Warn("close stream failed", zap.Error(err))
+ }
+
+ lastModified, err := time.Parse(time.RFC1123, resp.Header.Get("Last-Modified"))
+ if err != nil {
+ otelzap.L().Ctx(ctx).Warn("parse last modified failed", zap.Error(err))
+ lastModified = time.Time{}
+ }
+
+ return GetGravatarResult{
+ Avatar: img,
+ LastModified: lastModified,
+ }, nil
 }

services/avatar/get_qq_avatar.go

@@ -1,22 +1,18 @@
 package avatar
 
 import (
- "bytes"
  "context"
  "fmt"
- "github.com/kolesa-team/go-webp/webp"
+ "go.uber.org/zap"
+ "image"
  "image/jpeg"
- "image/png"
- "io"
 
- "github.com/cloudwego/hertz/pkg/common/bytebufferpool"
  "github.com/nfnt/resize"
  "github.com/uptrace/opentelemetry-go-extra/otelzap"
- "go.uber.org/zap"
 )
 
-func (s *service) GetQQAvatar(ctx context.Context, hash string, args GetAvatarArgs) (io.ReadCloser, error) {
- ctx, span := tracer.Start(ctx, "service.AvatarService.GetQQAvatar")
+func (s *service) getQQAvatar(ctx context.Context, hash string, args GetAvatarArgs) (image.Image, error) {
+ ctx, span := tracer.Start(ctx, "service.AvatarService.getQQAvatar")
  defer span.End()
 
  // get qq avatar
@@ -26,7 +22,7 @@ func (s *service) GetQQAvatar(ctx context.Context, hash string, args GetAvatarAr
  return nil, err
  }
 
- stream, err := s.downloadAvatar(ctx, fmt.Sprintf("https://q.qlogo.cn/headimg_dl?dst_uin=%d&spec=640&img_type=jpg", qqid))
+ stream, _, err := s.downloadAvatar(ctx, fmt.Sprintf("https://q.qlogo.cn/headimg_dl?dst_uin=%d&spec=640&img_type=jpg", qqid))
  if err != nil {
  otelzap.L().Ctx(ctx).Error("download qq avatar failed", zap.Error(err))
  return nil, err
@@ -44,20 +40,5 @@ func (s *service) GetQQAvatar(ctx context.Context, hash string, args GetAvatarAr
 
  img = resize.Resize(uint(args.Size), uint(args.Size), img, resize.Lanczos3)
 
- buf := bytebufferpool.Get()
- defer bytebufferpool.Put(buf)
-
- if args.EnableWebp {
- if err := webp.Encode(buf, img, nil); err != nil {
- otelzap.L().Ctx(ctx).Error("encode qq avatar failed", zap.Error(err))
- return nil, err
- }
- } else {
- if err := png.Encode(buf, img); err != nil {
- otelzap.L().Ctx(ctx).Error("encode qq avatar failed", zap.Error(err))
- return nil, err
- }
- }
-
- return io.NopCloser(bytes.NewReader(buf.Bytes())), nil
+ return img, nil
 }

services/avatar/service.go

@@ -2,13 +2,20 @@ package avatar
 
 import (
  "context"
+ "errors"
  "fmt"
+ "image/png"
+ "io"
+ "net/http"
+ "time"
+
+ "github.com/cloudwego/hertz/pkg/common/bytebufferpool"
  "github.com/imroc/req/v3"
+ "github.com/kolesa-team/go-webp/webp"
  "github.com/uptrace/opentelemetry-go-extra/otelzap"
  "go.opentelemetry.io/otel"
  "go.uber.org/fx"
  "go.uber.org/zap"
- "io"
 
  "github.com/AH-dark/gravatar-with-qq-avatar/database/dal"
 )
@@ -25,8 +32,7 @@ type GetAvatarArgs struct {
 }
 
 type Service interface {
- GetGravatar(ctx context.Context, hash string, args GetAvatarArgs) (io.ReadCloser, error)
- GetQQAvatar(ctx context.Context, hash string, args GetAvatarArgs) (io.ReadCloser, error)
+ GetAvatar(ctx context.Context, hash string, args GetAvatarArgs) ([]byte, time.Time, error)
 }
 
 type service struct {
@@ -38,24 +44,64 @@ func NewService(s service) Service {
  return &s
 }
 
-func (s *service) downloadAvatar(ctx context.Context, url string) (io.ReadCloser, error) {
+func (s *service) downloadAvatar(ctx context.Context, url string) (io.ReadCloser, *req.Response, error) {
  ctx, span := tracer.Start(ctx, "service.AvatarService.downloadAvatar")
  defer span.End()
 
  resp, err := req.NewRequest().SetContext(ctx).Get(url)
  if err != nil {
  otelzap.L().Ctx(ctx).Error("download avatar failed", zap.Error(err))
- return nil, err
+ return nil, resp, err
  }
 
  if resp.StatusCode != 200 || resp.IsErrorState() {
  otelzap.L().Ctx(ctx).Error("download avatar failed")
- return nil, fmt.Errorf("download avatar failed, status code: %d", resp.StatusCode)
+ return nil, resp, errors.New(http.StatusText(resp.StatusCode))
  }
 
  if resp.IsSuccessState() {
- return resp.Body, nil
+ return resp.Body, resp, nil
+ }
+
+ return nil, resp, fmt.Errorf("download avatar failed, status code: %d", resp.StatusCode)
+}
+
+func (s *service) GetAvatar(ctx context.Context, hash string, args GetAvatarArgs) ([]byte, time.Time, error) {
+ ctx, span := tracer.Start(ctx, "service.AvatarService.GetAvatar")
+ defer span.End()
+
+ var lastModified time.Time
+
+ // get qq avatar
+ img, err := s.getQQAvatar(ctx, hash, args)
+ if err != nil {
+ res, err := s.getGravatar(ctx, hash, args)
+ if errors.Is(err, errors.New("404 Not Found")) {
+ return nil, time.Time{}, nil
+ } else if err != nil {
+ otelzap.L().Ctx(ctx).Warn("get gravatar failed", zap.Error(err))
+ return nil, time.Time{}, err
+ }
+
+ img = res.Avatar
+ lastModified = res.LastModified
+ }
+
+ // encode image
+ b := bytebufferpool.Get()
+ defer bytebufferpool.Put(b)
+
+ if args.EnableWebp {
+ if err := webp.Encode(b, img, nil); err != nil {
+ otelzap.L().Ctx(ctx).Error("encode webp failed", zap.Error(err))
+ return nil, time.Time{}, err
+ }
+ } else {
+ if err := png.Encode(b, img); err != nil {
+ otelzap.L().Ctx(ctx).Error("encode png failed", zap.Error(err))
+ return nil, time.Time{}, err
+ }
  }
 
- return nil, fmt.Errorf("download avatar failed, status code: %d", resp.StatusCode)
+ return b.Bytes(), lastModified, nil
 }