1 package cache
2
3 import (
4 "bytes"
5 "crypto/sha1"
6 "errors"
7 "github.com/gin-gonic/gin"
8 "io"
9 "net/http"
10 "net/url"
11 "time"
12 )
13
14 const (
15 DEFAULT = time.Duration(0)
16 FOREVER = time.Duration(-1)
17 CACHE_MIDDLEWARE_KEY = "gincontrib.cache"
18 )
19
20 var (
21 PageCachePrefix = "gincontrib.page.cache"
22 ErrCacheMiss = errors.New("cache: key not found.")
23 ErrNotStored = errors.New("cache: not stored.")
24 ErrNotSupport = errors.New("cache: not support.")
25 )
26
27 type CacheStore interface {
28 Get(key string, value interface{}) error
29 Set(key string, value interface{}, expire time.Duration) error
30 Add(key string, value interface{}, expire time.Duration) error
31 Replace(key string, data interface{}, expire time.Duration) error
32 Delete(key string) error
33 Increment(key string, data uint64) (uint64, error)
34 Decrement(key string, data uint64) (uint64, error)
35 Flush() error
36 }
37
38 type responseCache struct {
39 status int
40 header http.Header
41 data []byte
42 }
43
44 type cachedWriter struct {
45 gin.ResponseWriter
46 status int
47 written bool
48 store CacheStore
49 expire time.Duration
50 key string
51 }
52
53 func urlEscape(prefix string, u string) string {
54 key := url.QueryEscape(u)
55 if len(key) > 200 {
56 h := sha1.New()
57 io.WriteString(h, u)
58 key = string(h.Sum(nil))
59 }
60 var buffer bytes.Buffer
61 buffer.WriteString(prefix)
62 buffer.WriteString(":")
63 buffer.WriteString(key)
64 return buffer.String()
65 }
66
67 func newCachedWriter(store CacheStore, expire time.Duration, writer gin.ResponseWriter, key string) *cachedWriter {
68 return &cachedWriter{writer, 0, false, store, expire, key}
69 }
70
71 func (w *cachedWriter) WriteHeader(code int) {
72 w.status = code
73 w.written = true
74 w.ResponseWriter.WriteHeader(code)
75 }
76
77 func (w *cachedWriter) Status() int {
78 return w.status
79 }
80
81 func (w *cachedWriter) Written() bool {
82 return w.written
83 }
84
85 func (w *cachedWriter) Write(data []byte) (int, error) {
86 ret, err := w.ResponseWriter.Write(data)
87 if err == nil {
88
89 store := w.store
90 val := responseCache{
91 w.status,
92 w.Header(),
93 data,
94 }
95 err = store.Set(w.key, val, w.expire)
96 if err != nil {
97
98 }
99 }
100 return ret, err
101 }
102
103
104 func Cache(store *CacheStore) gin.HandlerFunc {
105 return func(c *gin.Context) {
106 c.Set(CACHE_MIDDLEWARE_KEY, store)
107 c.Next()
108 }
109 }
110
111 func SiteCache(store CacheStore, expire time.Duration) gin.HandlerFunc {
112
113 return func(c *gin.Context) {
114 var cache responseCache
115 url := c.Request.URL
116 key := urlEscape(PageCachePrefix, url.RequestURI())
117 if err := store.Get(key, &cache); err != nil {
118 c.Next()
119 } else {
120 c.Writer.WriteHeader(cache.status)
121 for k, vals := range cache.header {
122 for _, v := range vals {
123 c.Writer.Header().Add(k, v)
124 }
125 }
126 c.Writer.Write(cache.data)
127 }
128 }
129 }
130
131
132 func CachePage(store CacheStore, expire time.Duration, handle gin.HandlerFunc) gin.HandlerFunc {
133
134 return func(c *gin.Context) {
135 var cache responseCache
136 url := c.Request.URL
137 key := urlEscape(PageCachePrefix, url.RequestURI())
138 if err := store.Get(key, &cache); err != nil {
139
140 writer := newCachedWriter(store, expire, c.Writer, key)
141 c.Writer = writer
142 handle(c)
143 } else {
144 c.Writer.WriteHeader(cache.status)
145 for k, vals := range cache.header {
146 for _, v := range vals {
147 c.Writer.Header().Add(k, v)
148 }
149 }
150 c.Writer.Write(cache.data)
151 }
152 }
153 }
154
View as plain text