package httpcache import ( "bytes" "context" "net/http" "sync" "time" "github.com/darkweak/souin/api" "github.com/darkweak/souin/cache/coalescing" "github.com/darkweak/souin/plugins" "github.com/darkweak/souin/rfc" kratos_http "github.com/go-kratos/kratos/v2/transport/http" ) const ( getterContextCtxKey key = "getter_context" ) type ( key string httpcacheKratosPlugin struct { plugins.SouinBasePlugin Configuration *plugins.BaseConfiguration bufPool *sync.Pool } getterContext struct { next http.HandlerFunc rw http.ResponseWriter req *http.Request } ) // NewHTTPCacheFilter, allows the user to set up an HTTP cache system, // RFC-7234 compliant and supports the tag based cache purge, // distributed and not-distributed storage, key generation tweaking. // Use it with // httpcache.NewHTTPCacheFilter(httpcache.ParseConfiguration(config)) func NewHTTPCacheFilter(c plugins.BaseConfiguration) kratos_http.FilterFunc { s := &httpcacheKratosPlugin{} s.Configuration = &c s.bufPool = &sync.Pool{ New: func() interface{} { return new(bytes.Buffer) }, } s.Retriever = plugins.DefaultSouinPluginInitializerFromConfiguration(&c) s.RequestCoalescing = coalescing.Initialize() s.MapHandler = api.GenerateHandlerMap(s.Configuration, s.Retriever.GetTransport()) return s.handle } func (s *httpcacheKratosPlugin) handle(next http.Handler) http.Handler { return http.HandlerFunc(func(rw http.ResponseWriter, r *http.Request) { req := s.Retriever.GetContext().Method.SetContext(r) if b, handler := s.HandleInternally(req); b { handler(rw, req) return } if !plugins.CanHandle(req, s.Retriever) { rw.Header().Add("Cache-Status", "Souin; fwd=uri-miss") next.ServeHTTP(rw, r) return } customWriter := &plugins.CustomWriter{ Response: &http.Response{}, Buf: s.bufPool.Get().(*bytes.Buffer), Rw: rw, } req = s.Retriever.GetContext().SetContext(req) getterCtx := getterContext{next.ServeHTTP, customWriter, req} ctx := context.WithValue(req.Context(), getterContextCtxKey, getterCtx) req = req.WithContext(ctx) if plugins.HasMutation(req, rw) { next.ServeHTTP(rw, r) return } req.Header.Set("Date", time.Now().UTC().Format(time.RFC1123)) combo := ctx.Value(getterContextCtxKey).(getterContext) _ = plugins.DefaultSouinPluginCallback(customWriter, req, s.Retriever, nil, func(_ http.ResponseWriter, _ *http.Request) error { var e error combo.next.ServeHTTP(customWriter, r) combo.req.Response = customWriter.Response // nolint if combo.req.Response, e = s.Retriever.GetTransport().(*rfc.VaryTransport).UpdateCacheEventually(combo.req); e != nil { return e } _, _ = customWriter.Send() return e }) }) }