package gmi import ( "bytes" "context" "io" "sync" "git.sr.ht/~adnano/go-gemini" "code.betamike.com/mediocregopher/mediocre-blog/src/cache" ) type cacheRW struct { gemini.ResponseWriter status gemini.Status mediaType string buf []byte } func (c *cacheRW) SetMediaType(mediaType string) { c.mediaType = mediaType c.ResponseWriter.SetMediaType(mediaType) } func (c *cacheRW) WriteHeader(status gemini.Status, meta string) { c.status = status c.ResponseWriter.WriteHeader(status, meta) } func (c *cacheRW) Write(b []byte) (int, error) { c.buf = append(c.buf, b...) return c.ResponseWriter.Write(b) } func cacheMiddleware(cache cache.Cache) func(h gemini.Handler) gemini.Handler { type entry struct { mediaType string body []byte } pool := sync.Pool{ New: func() interface{} { return new(bytes.Reader) }, } return func(h gemini.Handler) gemini.Handler { return gemini.HandlerFunc(func( ctx context.Context, rw gemini.ResponseWriter, r *gemini.Request, ) { id := r.URL.String() if value := cache.Get(id); value != nil { entry := value.(entry) if entry.mediaType != "" { rw.SetMediaType(entry.mediaType) } reader := pool.Get().(*bytes.Reader) defer pool.Put(reader) reader.Reset(entry.body) io.Copy(rw, reader) return } cacheRW := &cacheRW{ ResponseWriter: rw, status: gemini.StatusSuccess, } h.ServeGemini(ctx, cacheRW, r) if cacheRW.status == gemini.StatusSuccess { cache.Set(id, entry{ mediaType: cacheRW.mediaType, body: cacheRW.buf, }) } }) } }