summaryrefslogtreecommitdiff
path: root/src/gmi/cache.go
blob: 4701673630e3ee527dfe33ae3aaf388fb4bbf352 (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
package gmi

import (
	"bytes"
	"context"
	"io"
	"sync"

	"git.sr.ht/~adnano/go-gemini"
	"dev.mediocregopher.com/mediocre-blog.git/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,
				})
			}
		})
	}
}