116 lines
2.5 KiB
Go
116 lines
2.5 KiB
Go
package router
|
|
|
|
import (
|
|
"bytes"
|
|
"net/http"
|
|
"strings"
|
|
"sync"
|
|
"time"
|
|
|
|
"github.com/gin-gonic/gin"
|
|
)
|
|
|
|
const (
|
|
responseCacheTTL = 5 * time.Minute
|
|
responseCacheMaxBody = 4 << 20
|
|
responseCacheMaxKeys = 512
|
|
)
|
|
|
|
type responseCacheEntry struct {
|
|
status int
|
|
contentType string
|
|
body []byte
|
|
expiresAt time.Time
|
|
}
|
|
|
|
type responseCacheStore struct {
|
|
mu sync.RWMutex
|
|
items map[string]responseCacheEntry
|
|
}
|
|
|
|
func newResponseCacheStore() *responseCacheStore {
|
|
return &responseCacheStore{items: make(map[string]responseCacheEntry)}
|
|
}
|
|
|
|
func (s *responseCacheStore) get(key string) (responseCacheEntry, bool) {
|
|
now := time.Now()
|
|
|
|
s.mu.RLock()
|
|
entry, ok := s.items[key]
|
|
s.mu.RUnlock()
|
|
if !ok {
|
|
return responseCacheEntry{}, false
|
|
}
|
|
if now.After(entry.expiresAt) {
|
|
s.mu.Lock()
|
|
if current, exists := s.items[key]; exists && now.After(current.expiresAt) {
|
|
delete(s.items, key)
|
|
}
|
|
s.mu.Unlock()
|
|
return responseCacheEntry{}, false
|
|
}
|
|
return entry, true
|
|
}
|
|
|
|
func (s *responseCacheStore) set(key string, entry responseCacheEntry) {
|
|
s.mu.Lock()
|
|
if len(s.items) >= responseCacheMaxKeys {
|
|
now := time.Now()
|
|
for cacheKey, current := range s.items {
|
|
if now.After(current.expiresAt) || len(s.items) >= responseCacheMaxKeys {
|
|
delete(s.items, cacheKey)
|
|
}
|
|
}
|
|
}
|
|
s.items[key] = entry
|
|
s.mu.Unlock()
|
|
}
|
|
|
|
type cachedResponseWriter struct {
|
|
gin.ResponseWriter
|
|
body *bytes.Buffer
|
|
}
|
|
|
|
func (w cachedResponseWriter) Write(data []byte) (int, error) {
|
|
w.body.Write(data)
|
|
return w.ResponseWriter.Write(data)
|
|
}
|
|
|
|
func cacheableGETResponseMiddleware(store *responseCacheStore) gin.HandlerFunc {
|
|
return func(c *gin.Context) {
|
|
if c.Request.Method != http.MethodGet || strings.HasPrefix(c.Request.URL.Path, "/api/v1/download/") {
|
|
c.Next()
|
|
return
|
|
}
|
|
|
|
key := c.Request.URL.RequestURI()
|
|
if entry, ok := store.get(key); ok {
|
|
c.Header("Content-Type", entry.contentType)
|
|
c.Data(entry.status, entry.contentType, entry.body)
|
|
return
|
|
}
|
|
|
|
body := bytes.NewBuffer(nil)
|
|
writer := cachedResponseWriter{ResponseWriter: c.Writer, body: body}
|
|
c.Writer = writer
|
|
c.Next()
|
|
|
|
status := c.Writer.Status()
|
|
if status != http.StatusOK || body.Len() == 0 || body.Len() > responseCacheMaxBody {
|
|
return
|
|
}
|
|
|
|
contentType := c.Writer.Header().Get("Content-Type")
|
|
if !strings.Contains(contentType, "application/json") {
|
|
return
|
|
}
|
|
|
|
store.set(key, responseCacheEntry{
|
|
status: status,
|
|
contentType: contentType,
|
|
body: append([]byte(nil), body.Bytes()...),
|
|
expiresAt: time.Now().Add(responseCacheTTL),
|
|
})
|
|
}
|
|
}
|