goproxy/internal/cache/cache.go

package cache

import (
	"bytes"
	"crypto/md5"
	"encoding/hex"
	"io"
	"net/http"
	"strings"
	"sync"
	"time"
)

// Cache 缓存接口
type Cache interface {
	// Get 获取缓存
	Get(key string) (*http.Response, bool)
	// Set 设置缓存
	Set(key string, resp *http.Response)
	// Delete 删除缓存
	Delete(key string)
	// Clear 清空缓存
	Clear()
}

// MemoryCache 内存缓存实现
type MemoryCache struct {
	// 缓存内容
	items sync.Map
	// 过期时间
	ttl time.Duration
	// 清理间隔
	cleanupInterval time.Duration
	// 最大条目数
	maxEntries int
	// 当前条目数
	size int32
	// 互斥锁
	mu sync.Mutex
}

// CacheItem 缓存项
type CacheItem struct {
	response     *http.Response
	responseBody []byte
	expiry       time.Time
}

// NewMemoryCache 创建内存缓存
func NewMemoryCache(ttl, cleanupInterval time.Duration, maxEntries int) *MemoryCache {
	cache := &MemoryCache{
		ttl:             ttl,
		cleanupInterval: cleanupInterval,
		maxEntries:      maxEntries,
	}

	// 启动过期清理
	if cleanupInterval > 0 {
		go cache.startCleanup()
	}

	return cache
}

// Get 获取缓存
func (c *MemoryCache) Get(key string) (*http.Response, bool) {
	value, ok := c.items.Load(key)
	if !ok {
		return nil, false
	}

	item := value.(*CacheItem)
	if time.Now().After(item.expiry) {
		c.Delete(key)
		return nil, false
	}

	// 克隆响应，避免修改原始数据
	resp := cloneResponse(item.response, item.responseBody)
	return resp, true
}

// Set 设置缓存
func (c *MemoryCache) Set(key string, resp *http.Response) {
	// 检查缓存是否已满
	c.mu.Lock()
	if c.maxEntries > 0 && c.size >= int32(c.maxEntries) {
		c.mu.Unlock()
		return
	}
	c.size++
	c.mu.Unlock()

	// 读取并保存响应体
	var bodyBytes []byte
	if resp.Body != nil {
		bodyBytes, _ = io.ReadAll(resp.Body)
		resp.Body.Close()
		resp.Body = io.NopCloser(bytes.NewBuffer(bodyBytes))
	}

	item := &CacheItem{
		response:     resp,
		responseBody: bodyBytes,
		expiry:       time.Now().Add(c.ttl),
	}

	c.items.Store(key, item)
}

// Delete 删除缓存
func (c *MemoryCache) Delete(key string) {
	c.items.Delete(key)
	c.mu.Lock()
	c.size--
	if c.size < 0 {
		c.size = 0
	}
	c.mu.Unlock()
}

// Clear 清空缓存
func (c *MemoryCache) Clear() {
	c.items = sync.Map{}
	c.mu.Lock()
	c.size = 0
	c.mu.Unlock()
}

// startCleanup 启动过期清理
func (c *MemoryCache) startCleanup() {
	ticker := time.NewTicker(c.cleanupInterval)
	defer ticker.Stop()

	for range ticker.C {
		now := time.Now()
		c.items.Range(func(key, value interface{}) bool {
			item := value.(*CacheItem)
			if now.After(item.expiry) {
				c.Delete(key.(string))
			}
			return true
		})
	}
}

// GenerateCacheKey 生成缓存键
func GenerateCacheKey(req *http.Request) string {
	// 忽略一些可变的头部
	ignoredHeaders := map[string]bool{
		"Connection":          true,
		"Keep-Alive":          true,
		"Proxy-Authenticate":  true,
		"Proxy-Authorization": true,
		"TE":                  true,
		"Trailers":            true,
		"Transfer-Encoding":   true,
		"Upgrade":             true,
	}

	// 提取缓存键组件
	components := []string{
		req.Method,
		req.URL.String(),
	}

	// 添加选择性头部
	for key, values := range req.Header {
		if !ignoredHeaders[key] {
			for _, value := range values {
				components = append(components, key+":"+value)
			}
		}
	}

	// 连接并计算哈希
	data := strings.Join(components, "|")
	hash := md5.New()
	hash.Write([]byte(data))
	return hex.EncodeToString(hash.Sum(nil))
}

// cloneResponse 克隆HTTP响应
func cloneResponse(resp *http.Response, body []byte) *http.Response {
	clone := *resp
	clone.Body = io.NopCloser(bytes.NewBuffer(body))
	clone.Header = make(http.Header)

	for k, v := range resp.Header {
		clone.Header[k] = v
	}

	return &clone
}

// ShouldCache 判断请求是否应该缓存
func ShouldCache(req *http.Request, resp *http.Response) bool {
	// 只缓存GET请求
	if req.Method != http.MethodGet {
		return false
	}

	// 检查响应状态码
	if resp.StatusCode != http.StatusOK &&
		resp.StatusCode != http.StatusNotModified &&
		resp.StatusCode != http.StatusMovedPermanently &&
		resp.StatusCode != http.StatusPermanentRedirect {
		return false
	}

	// 检查Cache-Control头
	cacheControl := resp.Header.Get("Cache-Control")
	if strings.Contains(cacheControl, "no-store") ||
		strings.Contains(cacheControl, "no-cache") ||
		strings.Contains(cacheControl, "private") {
		return false
	}

	return true
}