feat: add middleware.cache

internal/middleware/cache/cache.go

@@ -0,0 +1,129 @@
+package cache
+
+import (
+	"bytes"
+	"encoding/json"
+	"errors"
+	"github.com/loveuer/nf"
+	"github.com/loveuer/nf/nft/log"
+	"io"
+	"net/http"
+	"strconv"
+	"time"
+	"ultone/internal/database/cache"
+	"ultone/internal/model"
+)
+
+var (
+	defaultKeyFn = func(c *nf.Ctx) string {
+		return c.Request.URL.String()
+	}
+	defaultTimeout = 3600
+	defaultPrefix  = "midd:cache"
+)
+
+type Config struct {
+	// if return "" (won't cache)
+	KeyFn func(c *nf.Ctx) string
+
+	// cache timeout(seconds)
+	Timeout int
+
+	Prefix  string
+	Refresh bool
+}
+
+type store struct {
+	Body   []byte      `json:"body"`
+	Header http.Header `json:"header"`
+	When   int64       `json:"when"`
+}
+
+func New(cfgs ...Config) nf.HandlerFunc {
+	if cache.Client == nil {
+		log.Panic("[middleware.cache] database cache client is nil")
+	}
+
+	var cfg Config
+	if len(cfgs) > 0 {
+		cfg = cfgs[0]
+	}
+
+	if cfg.KeyFn == nil {
+		cfg.KeyFn = defaultKeyFn
+	}
+
+	if cfg.Timeout <= 0 {
+		cfg.Timeout = defaultTimeout
+	}
+
+	if cfg.Prefix == "" {
+		cfg.Prefix = defaultPrefix
+	}
+
+	return func(c *nf.Ctx) error {
+		var (
+			key string
+			err error
+			bs  []byte
+			res = new(store)
+		)
+
+		if key = cfg.KeyFn(c); key == "" {
+			return c.Next()
+		}
+
+		key = cfg.Prefix + ":" + key
+		duration := time.Duration(cfg.Timeout) * time.Second
+
+		if cfg.Refresh {
+			if bs, err = cache.Client.GetEx(c.Context(), key, duration); err != nil {
+				if !errors.Is(err, cache.ErrorKeyNotFound) {
+					log.Warn("[middleware.cache] cache get err: %s", err.Error())
+				}
+				goto FromNext
+			}
+		} else {
+			if bs, err = cache.Client.Get(c.Context(), key); err != nil {
+				if !errors.Is(err, cache.ErrorKeyNotFound) {
+					log.Warn("[middleware.cache] cache get err: %s", err.Error())
+				}
+				goto FromNext
+			}
+		}
+
+		if err = json.Unmarshal(bs, res); err != nil {
+			log.Warn("[middleware.cache] cache data unamrshal err: %s", err.Error())
+			goto FromNext
+		}
+
+		for key := range res.Header {
+			for idx := range res.Header[key] {
+				c.SetHeader(key, res.Header[key][idx])
+			}
+		}
+
+		c.SetHeader("X-Nf-Cache-At", strconv.Itoa(int(res.When)))
+
+		_, err = c.Write(res.Body)
+		return err
+
+	FromNext:
+
+		blw := &model.CustomResponseWriter{Body: bytes.NewBuffer(make([]byte, 0, 256)), ResponseWriter: c.Writer}
+		c.Writer = blw
+
+		rerr := c.Next()
+
+		resp, _ := io.ReadAll(blw.Body)
+
+		data := &store{Body: resp, Header: blw.Header().Clone(), When: time.Now().UnixMilli()}
+		cbs, _ := json.Marshal(data)
+
+		if err = cache.Client.SetEx(c.Context(), key, cbs, duration); err != nil {
+			log.Warn("[middleware.cache] cache client setex err: %s", err.Error())
+		}
+
+		return rerr
+	}
+}