1 files changed, 87 insertions, 0 deletions
diff --git a/src/gmi/cache.go b/src/gmi/cache.go
new file mode 100644
index 0000000..0f2e975
--- /dev/null
+++ b/src/gmi/cache.go
@@ -0,0 +1,87 @@
+package gmi
+
+import (
+	"bytes"
+	"context"
+	"io"
+	"sync"
+
+	"git.sr.ht/~adnano/go-gemini"
+	"github.com/mediocregopher/blog.mediocregopher.com/srv/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,
+				})
+			}
+		})
+	}
+}