Wormhole/internal/forwardplugin/response_cache.go

package forwardplugin

import (
	"container/list"
	"encoding/json"
	"fmt"
	"net/http"
	"sync"
	"time"
)

func init() {
	Register("response_cache", newResponseCachePlugin)
}

type responseCacheConfig struct {
	MaxMB   int `json:"max_mb"`
	TTLSec  int `json:"ttl_sec"`
}

type responseCachePlugin struct {
	forwardID uint
	maxBytes  int64
	ttl       time.Duration
	store     *responseCacheStore
}

func newResponseCachePlugin(raw json.RawMessage) (Plugin, error) {
	var cfg responseCacheConfig
	if err := json.Unmarshal(raw, &cfg); err != nil {
		return nil, fmt.Errorf("response_cache 配置无效")
	}
	return &responseCachePlugin{
		maxBytes: int64(cfg.MaxMB) * 1024 * 1024,
		ttl:      time.Duration(cfg.TTLSec) * time.Second,
	}, nil
}

func (p *responseCachePlugin) bindForwardID(id uint) {
	p.forwardID = id
	if p.maxBytes <= 0 {
		return
	}
	p.store = globalResponseCaches.get(id, p.maxBytes, p.ttl)
}

func (p *responseCachePlugin) Type() string { return "response_cache" }

func (p *responseCachePlugin) ValidateConfig(raw json.RawMessage) error {
	var cfg responseCacheConfig
	if err := json.Unmarshal(raw, &cfg); err != nil {
		return fmt.Errorf("response_cache 配置无效")
	}
	if cfg.MaxMB < 0 {
		return fmt.Errorf("response_cache.max_mb 不能为负数")
	}
	if cfg.MaxMB > 0 && cfg.TTLSec <= 0 {
		return nil
	}
	if cfg.TTLSec < 0 {
		return fmt.Errorf("response_cache.ttl_sec 不能为负数")
	}
	return nil
}

func (p *responseCachePlugin) NormalizeConfig(raw json.RawMessage) (json.RawMessage, error) {
	var cfg responseCacheConfig
	if err := json.Unmarshal(raw, &cfg); err != nil {
		return nil, err
	}
	if cfg.MaxMB > 0 && cfg.TTLSec <= 0 {
		cfg.TTLSec = 300
	}
	return json.Marshal(cfg)
}

func (p *responseCachePlugin) NeedsBody() bool { return false }

func (p *responseCachePlugin) HandleRequest(ctx *Context, req *http.Request) (bool, error) {
	if p.store == nil || req.Method != http.MethodGet {
		return false, nil
	}
	key := cacheKey(req)
	if hit := p.store.get(key); hit != nil {
		ctx.CachedResp = hit.Clone()
		ctx.CachedResp.FromCache = true
		return true, nil
	}
	return false, nil
}

func (p *responseCachePlugin) HandleResponse(_ *Context, _ *BufferedResponse) error {
	return nil
}

func (p *responseCachePlugin) StoreRaw(ctx *Context, resp *BufferedResponse) {
	if p.store == nil || ctx.ClientReq.Method != http.MethodGet || resp.FromCache {
		return
	}
	if resp.StatusCode < 200 || resp.StatusCode >= 300 {
		return
	}
	if len(resp.Header.Values("Set-Cookie")) > 0 {
		return
	}
	p.store.put(cacheKey(ctx.ClientReq), resp.Clone())
}

func cacheKey(req *http.Request) string {
	if req.URL.RawQuery != "" {
		return req.URL.Path + "?" + req.URL.RawQuery
	}
	return req.URL.Path
}

type cacheEntry struct {
	key     string
	resp    *BufferedResponse
	expires time.Time
	size    int64
}

type responseCacheStore struct {
	mu        sync.Mutex
	maxBytes  int64
	ttl       time.Duration
	usedBytes int64
	items     map[string]*list.Element
	order     *list.List
}

func newResponseCacheStore(maxBytes int64, ttl time.Duration) *responseCacheStore {
	return &responseCacheStore{
		maxBytes: maxBytes,
		ttl:      ttl,
		items:    make(map[string]*list.Element),
		order:    list.New(),
	}
}

func (s *responseCacheStore) get(key string) *BufferedResponse {
	s.mu.Lock()
	defer s.mu.Unlock()
	el, ok := s.items[key]
	if !ok {
		return nil
	}
	ent := el.Value.(*cacheEntry)
	if time.Now().After(ent.expires) {
		s.removeElement(el)
		return nil
	}
	s.order.MoveToFront(el)
	return ent.resp
}

func (s *responseCacheStore) put(key string, resp *BufferedResponse) {
	size := int64(len(resp.Body))
	for k, v := range resp.Header {
		for _, vv := range v {
			size += int64(len(k) + len(vv))
		}
	}
	s.mu.Lock()
	defer s.mu.Unlock()
	if el, ok := s.items[key]; ok {
		old := el.Value.(*cacheEntry)
		s.usedBytes -= old.size
		s.order.Remove(el)
		delete(s.items, key)
	}
	for s.usedBytes+size > s.maxBytes && s.order.Len() > 0 {
		s.removeElement(s.order.Back())
	}
	ent := &cacheEntry{
		key:     key,
		resp:    resp,
		expires: time.Now().Add(s.ttl),
		size:    size,
	}
	el := s.order.PushFront(ent)
	s.items[key] = el
	s.usedBytes += size
}

func (s *responseCacheStore) removeElement(el *list.Element) {
	ent := el.Value.(*cacheEntry)
	delete(s.items, ent.key)
	s.order.Remove(el)
	s.usedBytes -= ent.size
}

type responseCacheManager struct {
	mu    sync.Mutex
	caches map[uint]*responseCacheStore
}

var globalResponseCaches = &responseCacheManager{caches: make(map[uint]*responseCacheStore)}

func (m *responseCacheManager) get(forwardID uint, maxBytes int64, ttl time.Duration) *responseCacheStore {
	m.mu.Lock()
	defer m.mu.Unlock()
	if c, ok := m.caches[forwardID]; ok {
		c.maxBytes = maxBytes
		c.ttl = ttl
		return c
	}
	c := newResponseCacheStore(maxBytes, ttl)
	m.caches[forwardID] = c
	return c
}

func BindForwardID(chain *Chain, forwardID uint) {
	if chain == nil {
		return
	}
	for _, p := range chain.plugins {
		if rc, ok := p.(*responseCachePlugin); ok {
			rc.bindForwardID(forwardID)
		}
	}
}

func InvalidateForwardCache(forwardID uint) {
	globalResponseCaches.mu.Lock()
	defer globalResponseCaches.mu.Unlock()
	delete(globalResponseCaches.caches, forwardID)
}