2020-07-29 20:45:26 +00:00
|
|
|
package main
|
|
|
|
|
|
|
|
import (
|
2020-11-25 13:59:48 +00:00
|
|
|
"bytes"
|
|
|
|
"crypto/sha256"
|
2020-11-20 14:33:20 +00:00
|
|
|
"fmt"
|
2020-11-25 13:59:48 +00:00
|
|
|
"io"
|
2020-11-09 15:40:12 +00:00
|
|
|
"io/ioutil"
|
2020-07-29 20:45:26 +00:00
|
|
|
"net/http"
|
2020-07-30 14:43:04 +00:00
|
|
|
"net/http/httptest"
|
2020-11-20 14:33:20 +00:00
|
|
|
"strconv"
|
2020-07-29 20:45:26 +00:00
|
|
|
"time"
|
2020-10-19 21:02:57 +00:00
|
|
|
|
2020-11-26 08:44:44 +00:00
|
|
|
"github.com/araddon/dateparse"
|
2020-11-20 14:33:20 +00:00
|
|
|
lru "github.com/hashicorp/golang-lru"
|
2020-10-19 21:02:57 +00:00
|
|
|
"golang.org/x/sync/singleflight"
|
2020-07-29 20:45:26 +00:00
|
|
|
)
|
|
|
|
|
2020-11-20 14:33:20 +00:00
|
|
|
const (
|
|
|
|
cacheInternalExpirationHeader = "GoBlog-Expire"
|
|
|
|
)
|
|
|
|
|
2020-10-19 21:33:08 +00:00
|
|
|
var (
|
|
|
|
cacheGroup singleflight.Group
|
2020-11-20 14:33:20 +00:00
|
|
|
cacheLru *lru.Cache
|
2020-10-19 21:33:08 +00:00
|
|
|
)
|
2020-09-22 14:42:36 +00:00
|
|
|
|
2020-11-20 14:33:20 +00:00
|
|
|
func initCache() (err error) {
|
2020-11-20 15:24:18 +00:00
|
|
|
cacheLru, err = lru.New(500)
|
2020-11-20 14:33:20 +00:00
|
|
|
return
|
2020-10-19 21:33:08 +00:00
|
|
|
}
|
2020-09-22 14:42:36 +00:00
|
|
|
|
2020-07-30 19:18:13 +00:00
|
|
|
func cacheMiddleware(next http.Handler) http.Handler {
|
2020-07-29 20:45:26 +00:00
|
|
|
return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
2020-08-04 17:42:09 +00:00
|
|
|
if appConfig.Cache.Enable &&
|
2020-10-19 21:02:57 +00:00
|
|
|
// check method
|
2020-10-20 16:15:15 +00:00
|
|
|
(r.Method == http.MethodGet || r.Method == http.MethodHead) &&
|
|
|
|
// check bypass query
|
2020-11-20 14:33:20 +00:00
|
|
|
!(r.URL.Query().Get("cache") == "0" || r.URL.Query().Get("cache") == "false") {
|
2020-10-19 21:33:08 +00:00
|
|
|
key := cacheKey(r)
|
2020-10-19 21:02:57 +00:00
|
|
|
// Get cache or render it
|
2020-10-19 21:33:08 +00:00
|
|
|
cacheInterface, _, _ := cacheGroup.Do(key, func() (interface{}, error) {
|
|
|
|
return getCache(key, next, r), nil
|
2020-10-19 21:02:57 +00:00
|
|
|
})
|
|
|
|
cache := cacheInterface.(*cacheItem)
|
2020-07-30 19:18:13 +00:00
|
|
|
// copy cached headers
|
2020-10-19 21:02:57 +00:00
|
|
|
for k, v := range cache.header {
|
2020-07-30 19:18:13 +00:00
|
|
|
w.Header()[k] = v
|
|
|
|
}
|
2020-11-26 08:44:44 +00:00
|
|
|
setCacheHeaders(w, cache)
|
2020-11-25 13:59:48 +00:00
|
|
|
// check conditional request
|
|
|
|
if ifNoneMatchHeader := r.Header.Get("If-None-Match"); ifNoneMatchHeader != "" && ifNoneMatchHeader == cache.hash {
|
|
|
|
// send 304
|
|
|
|
w.WriteHeader(http.StatusNotModified)
|
|
|
|
return
|
|
|
|
}
|
2020-11-26 08:44:44 +00:00
|
|
|
if ifModifiedSinceHeader := r.Header.Get("If-Modified-Since"); ifModifiedSinceHeader != "" {
|
|
|
|
t, err := dateparse.ParseAny(ifModifiedSinceHeader)
|
|
|
|
if err == nil && t.Unix() >= cache.creationTime {
|
|
|
|
// send 304
|
|
|
|
w.WriteHeader(http.StatusNotModified)
|
|
|
|
return
|
|
|
|
}
|
|
|
|
}
|
2020-10-19 21:02:57 +00:00
|
|
|
// set status code
|
|
|
|
w.WriteHeader(cache.code)
|
2020-07-30 19:18:13 +00:00
|
|
|
// write cached body
|
2020-10-19 21:02:57 +00:00
|
|
|
_, _ = w.Write(cache.body)
|
2020-07-30 19:08:41 +00:00
|
|
|
return
|
2020-07-29 20:45:26 +00:00
|
|
|
}
|
2020-07-30 19:18:13 +00:00
|
|
|
next.ServeHTTP(w, r)
|
|
|
|
return
|
2020-07-29 20:45:26 +00:00
|
|
|
})
|
|
|
|
}
|
|
|
|
|
2020-10-19 21:33:08 +00:00
|
|
|
func cacheKey(r *http.Request) string {
|
2020-10-20 16:15:15 +00:00
|
|
|
return r.URL.String()
|
2020-10-19 21:33:08 +00:00
|
|
|
}
|
|
|
|
|
2020-11-26 08:44:44 +00:00
|
|
|
func setCacheHeaders(w http.ResponseWriter, cache *cacheItem) {
|
2020-11-20 14:33:20 +00:00
|
|
|
w.Header().Del(cacheInternalExpirationHeader)
|
2020-11-26 08:44:44 +00:00
|
|
|
w.Header().Set("ETag", cache.hash)
|
|
|
|
w.Header().Set("Last-Modified", time.Unix(cache.creationTime, 0).UTC().Format(http.TimeFormat))
|
2020-11-26 16:40:40 +00:00
|
|
|
if w.Header().Get("Cache-Control") == "" {
|
|
|
|
if cache.expiration != 0 {
|
|
|
|
expiresIn := cache.creationTime + int64(cache.expiration) - time.Now().Unix()
|
|
|
|
// Set expires time
|
|
|
|
w.Header().Set("Cache-Control", fmt.Sprintf("public,max-age=%d,stale-while-revalidate=%d", expiresIn, cache.expiration))
|
|
|
|
} else {
|
|
|
|
w.Header().Set("Cache-Control", fmt.Sprintf("public,max-age=%d,s-max-age=%d,stale-while-revalidate=%d", appConfig.Cache.Expiration, appConfig.Cache.Expiration/3, appConfig.Cache.Expiration))
|
|
|
|
}
|
2020-11-20 14:33:20 +00:00
|
|
|
}
|
2020-07-30 19:08:41 +00:00
|
|
|
}
|
|
|
|
|
2020-10-19 21:02:57 +00:00
|
|
|
type cacheItem struct {
|
2020-11-20 14:33:20 +00:00
|
|
|
expiration int
|
2020-11-25 13:59:48 +00:00
|
|
|
creationTime int64
|
|
|
|
hash string
|
2020-10-19 21:02:57 +00:00
|
|
|
code int
|
|
|
|
header http.Header
|
|
|
|
body []byte
|
2020-10-07 15:35:52 +00:00
|
|
|
}
|
|
|
|
|
2020-10-19 21:33:08 +00:00
|
|
|
func (c *cacheItem) expired() bool {
|
2020-11-20 14:33:20 +00:00
|
|
|
if c.expiration != 0 {
|
|
|
|
return c.creationTime < time.Now().Unix()-int64(c.expiration)
|
|
|
|
}
|
|
|
|
return false
|
2020-10-19 21:33:08 +00:00
|
|
|
}
|
|
|
|
|
2020-11-20 14:33:20 +00:00
|
|
|
func getCache(key string, next http.Handler, r *http.Request) (item *cacheItem) {
|
|
|
|
if lruItem, ok := cacheLru.Get(key); ok {
|
|
|
|
item = lruItem.(*cacheItem)
|
|
|
|
}
|
|
|
|
if item == nil || item.expired() {
|
2020-10-19 21:02:57 +00:00
|
|
|
// No cache available
|
2020-10-19 21:33:08 +00:00
|
|
|
// Record request
|
2020-10-19 21:02:57 +00:00
|
|
|
recorder := httptest.NewRecorder()
|
|
|
|
next.ServeHTTP(recorder, r)
|
2020-10-19 21:33:08 +00:00
|
|
|
// Cache values from recorder
|
2020-11-09 15:40:12 +00:00
|
|
|
result := recorder.Result()
|
|
|
|
body, _ := ioutil.ReadAll(result.Body)
|
2020-11-25 13:59:48 +00:00
|
|
|
_ = result.Body.Close()
|
|
|
|
h := sha256.New()
|
|
|
|
_, _ = io.Copy(h, bytes.NewReader(body))
|
|
|
|
hash := fmt.Sprintf("%x", h.Sum(nil))
|
2020-11-20 14:33:20 +00:00
|
|
|
exp, _ := strconv.Atoi(result.Header.Get(cacheInternalExpirationHeader))
|
2020-10-20 16:15:15 +00:00
|
|
|
item = &cacheItem{
|
2020-11-20 14:33:20 +00:00
|
|
|
expiration: exp,
|
2020-11-25 13:59:48 +00:00
|
|
|
creationTime: time.Now().Unix(),
|
|
|
|
hash: hash,
|
2020-11-09 15:40:12 +00:00
|
|
|
code: result.StatusCode,
|
|
|
|
header: result.Header,
|
|
|
|
body: body,
|
2020-10-20 16:15:15 +00:00
|
|
|
}
|
2020-10-19 21:02:57 +00:00
|
|
|
// Save cache
|
2020-11-20 14:33:20 +00:00
|
|
|
cacheLru.Add(key, item)
|
2020-10-19 21:02:57 +00:00
|
|
|
}
|
|
|
|
return item
|
2020-07-29 20:45:26 +00:00
|
|
|
}
|
2020-08-01 13:16:21 +00:00
|
|
|
|
2020-10-06 17:07:48 +00:00
|
|
|
func purgeCache() {
|
2020-11-20 14:33:20 +00:00
|
|
|
cacheLru.Purge()
|
|
|
|
}
|
|
|
|
|
|
|
|
func setInternalCacheExpirationHeader(w http.ResponseWriter, expiration int) {
|
|
|
|
w.Header().Set(cacheInternalExpirationHeader, strconv.Itoa(expiration))
|
2020-08-01 13:16:21 +00:00
|
|
|
}
|