Implement cache and logger middlewares for gemini

2 files changed, 114 insertions, 4 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,
+				})
+			}
+		})
+	}
+}
diff --git a/src/gmi/gmi.go b/src/gmi/gmi.go
index 6835ea0..c2a32aa 100644
--- a/src/gmi/gmi.go
+++ b/src/gmi/gmi.go
@@ -13,6 +13,7 @@ import (
 
 	"git.sr.ht/~adnano/go-gemini"
 	"git.sr.ht/~adnano/go-gemini/certificate"
+	"github.com/mediocregopher/blog.mediocregopher.com/srv/cache"
 	"github.com/mediocregopher/blog.mediocregopher.com/srv/cfg"
 	"github.com/mediocregopher/blog.mediocregopher.com/srv/post"
 	"github.com/mediocregopher/mediocre-go-lib/v2/mctx"
@@ -23,6 +24,7 @@ import (
 // unless otherwise noted.
 type Params struct {
 	Logger *mlog.Logger
+	Cache  cache.Cache
 
 	PostStore      post.Store
 	PostAssetStore post.AssetStore
@@ -108,10 +110,11 @@ func New(params Params) (API, error) {
 
 	go func() {
 
-		ctx := mctx.WithAnnotator(context.Background(), &a.params)
+		err := a.srv.ListenAndServe(context.Background())
 
-		err := a.srv.ListenAndServe(ctx)
 		if err != nil && !errors.Is(err, context.Canceled) {
+
+			ctx := mctx.WithAnnotator(context.Background(), &a.params)
 			a.params.Logger.Fatal(ctx, "serving gemini server", err)
 		}
 	}()
@@ -123,6 +126,26 @@ func (a *api) Shutdown(ctx context.Context) error {
 	return a.srv.Shutdown(ctx)
 }
 
+func (a *api) logReqMiddleware(h gemini.Handler) gemini.Handler {
+
+	type logCtxKey string
+
+	return gemini.HandlerFunc(func(
+		ctx context.Context,
+		rw gemini.ResponseWriter,
+		r *gemini.Request,
+	) {
+
+		ctx = mctx.Annotate(ctx,
+			logCtxKey("url"), r.URL.String(),
+		)
+
+		h.ServeGemini(ctx, rw, r)
+
+		a.params.Logger.Info(ctx, "handled gemini request")
+	})
+}
+
 func indexMiddleware(h gemini.Handler) gemini.Handler {
 
 	return gemini.HandlerFunc(func(
@@ -212,8 +235,8 @@ func (a *api) handler() (gemini.Handler, error) {
 	h = mux
 	h = indexMiddleware(h)
 
-	// TODO logging
-	// TODO caching
+	h = a.logReqMiddleware(h)
+	h = cacheMiddleware(a.params.Cache)(h)
 
 	return h, nil
 }