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"
"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,
})
}
})
}
}
|
