package cache import ( "bytes" "crypto/sha1" "errors" "github.com/gin-gonic/gin" "io" "net/http" "net/url" "time" ) const ( DEFAULT = time.Duration(0) FOREVER = time.Duration(-1) CACHE_MIDDLEWARE_KEY = "gincontrib.cache" ) var ( PageCachePrefix = "gincontrib.page.cache" ErrCacheMiss = errors.New("cache: key not found.") ErrNotStored = errors.New("cache: not stored.") ErrNotSupport = errors.New("cache: not support.") ) type CacheStore interface { Get(key string, value interface{}) error Set(key string, value interface{}, expire time.Duration) error Add(key string, value interface{}, expire time.Duration) error Replace(key string, data interface{}, expire time.Duration) error Delete(key string) error Increment(key string, data uint64) (uint64, error) Decrement(key string, data uint64) (uint64, error) Flush() error } type responseCache struct { status int header http.Header data []byte } type cachedWriter struct { gin.ResponseWriter status int written bool store CacheStore expire time.Duration key string } func urlEscape(prefix string, u string) string { key := url.QueryEscape(u) if len(key) > 200 { h := sha1.New() io.WriteString(h, u) key = string(h.Sum(nil)) } var buffer bytes.Buffer buffer.WriteString(prefix) buffer.WriteString(":") buffer.WriteString(key) return buffer.String() } func newCachedWriter(store CacheStore, expire time.Duration, writer gin.ResponseWriter, key string) *cachedWriter { return &cachedWriter{writer, 0, false, store, expire, key} } func (w *cachedWriter) WriteHeader(code int) { w.status = code w.written = true w.ResponseWriter.WriteHeader(code) } func (w *cachedWriter) Status() int { return w.status } func (w *cachedWriter) Written() bool { return w.written } func (w *cachedWriter) Write(data []byte) (int, error) { ret, err := w.ResponseWriter.Write(data) if err == nil { //cache response store := w.store val := responseCache{ w.status, w.Header(), data, } err = store.Set(w.key, val, w.expire) if err != nil { // need logger } } return ret, err } // Cache Middleware func Cache(store *CacheStore) gin.HandlerFunc { return func(c *gin.Context) { c.Set(CACHE_MIDDLEWARE_KEY, store) c.Next() } } func SiteCache(store CacheStore, expire time.Duration) gin.HandlerFunc { return func(c *gin.Context) { var cache responseCache url := c.Request.URL key := urlEscape(PageCachePrefix, url.RequestURI()) if err := store.Get(key, &cache); err != nil { c.Next() } else { c.Writer.WriteHeader(cache.status) for k, vals := range cache.header { for _, v := range vals { c.Writer.Header().Add(k, v) } } c.Writer.Write(cache.data) } } } // Cache Decorator func CachePage(store CacheStore, expire time.Duration, handle gin.HandlerFunc) gin.HandlerFunc { return func(c *gin.Context) { var cache responseCache url := c.Request.URL key := urlEscape(PageCachePrefix, url.RequestURI()) if err := store.Get(key, &cache); err != nil { // replace writer writer := newCachedWriter(store, expire, c.Writer, key) c.Writer = writer handle(c) } else { c.Writer.WriteHeader(cache.status) for k, vals := range cache.header { for _, v := range vals { c.Writer.Header().Add(k, v) } } c.Writer.Write(cache.data) } } }