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 } }