ghproxy/proxy/error.go

package proxy

import (
	"bytes"
	"crypto/sha256"
	"encoding/gob"
	"encoding/hex"
	"sync"

	"fmt"
	"html/template"
	"io/fs"

	"github.com/WJQSERVER-STUDIO/logger"
	"github.com/cloudwego/hertz/pkg/app"
	lru "github.com/hashicorp/golang-lru/v2"
)

// 日志模块
var (
	logw       = logger.Logw
	logDump    = logger.LogDump
	logDebug   = logger.LogDebug
	logInfo    = logger.LogInfo
	logWarning = logger.LogWarning
	logError   = logger.LogError
)

func HandleError(c *app.RequestContext, message string) {
	ErrorPage(c, NewErrorWithStatusLookup(500, message))
	logError("Error handled: %s", message)
}

type GHProxyErrors struct {
	StatusCode   int
	StatusDesc   string
	StatusText   string
	HelpInfo     string
	ErrorMessage string
}

var (
	ErrInvalidURL = &GHProxyErrors{
		StatusCode: 400,
		StatusDesc: "Bad Request",
		StatusText: "无效请求",
		HelpInfo:   "请求的URL格式不正确，请检查后重试。",
	}
	ErrAuthHeaderUnavailable = &GHProxyErrors{
		StatusCode: 401,
		StatusDesc: "Unauthorized",
		StatusText: "认证失败",
		HelpInfo:   "缺少或无效的鉴权信息。",
	}
	ErrForbidden = &GHProxyErrors{
		StatusCode: 403,
		StatusDesc: "Forbidden",
		StatusText: "权限不足",
		HelpInfo:   "您没有权限访问此资源。",
	}
	ErrNotFound = &GHProxyErrors{
		StatusCode: 404,
		StatusDesc: "Not Found",
		StatusText: "页面未找到",
		HelpInfo:   "抱歉，您访问的页面不存在。",
	}
	ErrTooManyRequests = &GHProxyErrors{
		StatusCode: 429,
		StatusDesc: "Too Many Requests",
		StatusText: "请求过于频繁",
		HelpInfo:   "您的请求过于频繁，请稍后再试。",
	}
	ErrInternalServerError = &GHProxyErrors{
		StatusCode: 500,
		StatusDesc: "Internal Server Error",
		StatusText: "服务器内部错误",
		HelpInfo:   "服务器处理您的请求时发生错误，请稍后重试或联系管理员。",
	}
)

var statusErrorMap map[int]*GHProxyErrors

func init() {
	statusErrorMap = map[int]*GHProxyErrors{
		ErrInvalidURL.StatusCode:            ErrInvalidURL,
		ErrAuthHeaderUnavailable.StatusCode: ErrAuthHeaderUnavailable,
		ErrForbidden.StatusCode:             ErrForbidden,
		ErrNotFound.StatusCode:              ErrNotFound,
		ErrTooManyRequests.StatusCode:       ErrTooManyRequests,
		ErrInternalServerError.StatusCode:   ErrInternalServerError,
	}
}

func NewErrorWithStatusLookup(statusCode int, errMsg string) *GHProxyErrors {
	baseErr, found := statusErrorMap[statusCode]

	if found {
		return &GHProxyErrors{
			StatusCode:   baseErr.StatusCode,
			StatusDesc:   baseErr.StatusDesc,
			StatusText:   baseErr.StatusText,
			HelpInfo:     baseErr.HelpInfo,
			ErrorMessage: errMsg,
		}
	} else {
		return &GHProxyErrors{
			StatusCode:   statusCode,
			ErrorMessage: errMsg,
		}
	}
}

var errPagesFs fs.FS

func InitErrPagesFS(pages fs.FS) error {
	var err error
	errPagesFs, err = fs.Sub(pages, "pages/err")
	if err != nil {
		return err
	}
	return nil
}

type ErrorPageData struct {
	StatusCode   int
	StatusDesc   string
	StatusText   string
	HelpInfo     string
	ErrorMessage string
}

// ToCacheKey 为 ErrorPageData 生成一个唯一的 SHA256 字符串键。
// 使用 gob 序列化来确保结构体内容到字节序列的顺序一致性，然后计算哈希。
func (d ErrorPageData) ToCacheKey() string {
	var buf bytes.Buffer
	enc := gob.NewEncoder(&buf)
	err := enc.Encode(d)
	if err != nil {
		logError("Failed to gob encode ErrorPageData for cache key: %v", err)
		return ""
	}

	hasher := sha256.New()
	hasher.Write(buf.Bytes())
	return hex.EncodeToString(hasher.Sum(nil))
}

func ErrPageUnwarper(errInfo *GHProxyErrors) ErrorPageData {
	return ErrorPageData{
		StatusCode:   errInfo.StatusCode,
		StatusDesc:   errInfo.StatusDesc,
		StatusText:   errInfo.StatusText,
		HelpInfo:     errInfo.HelpInfo,
		ErrorMessage: errInfo.ErrorMessage,
	}
}

// SizedLRUCache 实现了基于字节大小限制的 LRU 缓存。
// 它包装了 hashicorp/golang-lru/v2.Cache，并额外管理缓存的总字节大小。
type SizedLRUCache struct {
	cache        *lru.Cache[string, []byte]
	mu           sync.Mutex // 保护 currentBytes 字段
	maxBytes     int64      // 缓存的最大字节容量
	currentBytes int64      // 缓存当前占用的字节数
}

// NewSizedLRUCache 创建一个新的 SizedLRUCache 实例。
// 内部的 lru.Cache 的条目容量被设置为一个较大的值 (例如 10000)，
// 因为主要的逐出逻辑将由字节大小限制来控制。
func NewSizedLRUCache(maxBytes int64) (*SizedLRUCache, error) {
	if maxBytes <= 0 {
		return nil, fmt.Errorf("maxBytes must be positive")
	}

	c := &SizedLRUCache{
		maxBytes: maxBytes,
	}

	// 创建内部 LRU 缓存，并提供一个 OnEvictedFunc 回调函数。
	// 当内部 LRU 缓存因其自身的条目容量限制或 RemoveOldest 方法被调用而逐出条目时，
	// 此回调函数会被执行，从而更新 currentBytes。
	var err error
	c.cache, err = lru.NewWithEvict[string, []byte](10000, func(key string, value []byte) {
		c.mu.Lock()
		defer c.mu.Unlock()
		c.currentBytes -= int64(len(value))
		logDebug("LRU evicted key: %s, size: %d, current total: %d", key, len(value), c.currentBytes)
	})
	if err != nil {
		return nil, err
	}
	return c, nil
}

// Get 从缓存中检索值。
func (c *SizedLRUCache) Get(key string) ([]byte, bool) {
	return c.cache.Get(key)
}

// Add 向缓存中添加或更新一个键值对，并在必要时执行逐出以满足字节限制。
func (c *SizedLRUCache) Add(key string, value []byte) {
	c.mu.Lock() // 保护 currentBytes 和逐出逻辑
	defer c.mu.Unlock()

	itemSize := int64(len(value))

	// 如果待添加的条目本身就大于缓存的最大容量，则不进行缓存。
	if itemSize > c.maxBytes {
		logWarning("Item key %s (size %d) larger than cache max capacity %d. Not caching.", key, itemSize, c.maxBytes)
		return
	}

	// 如果键已存在，则首先从 currentBytes 中减去旧值的大小，并从内部 LRU 中移除旧条目。
	if oldVal, ok := c.cache.Get(key); ok {
		c.currentBytes -= int64(len(oldVal))
		c.cache.Remove(key)
		logDebug("Key %s exists, removed old size %d. Current total: %d", key, len(oldVal), c.currentBytes)
	}

	// 主动逐出最旧的条目，直到有足够的空间容纳新条目。
	for c.currentBytes+itemSize > c.maxBytes && c.cache.Len() > 0 {
		_, oldVal, existed := c.cache.RemoveOldest()
		if !existed {
			logWarning("Attempted to remove oldest, but item not found.")
			break
		}
		logDebug("Proactively evicted item (size %d) to free space. Current total: %d", len(oldVal), c.currentBytes)
	}

	// 添加新条目到内部 LRU 缓存。
	c.cache.Add(key, value)
	c.currentBytes += itemSize // 手动增加新条目的大小到 currentBytes。
	logDebug("Item added: key %s, size: %d, current total: %d", key, itemSize, c.currentBytes)
}

const maxErrorPageCacheBytes = 512 * 1024 // 错误页面缓存的最大容量：512KB

var errorPageCache *SizedLRUCache

func init() {
	// 初始化 SizedLRUCache。
	var err error
	errorPageCache, err = NewSizedLRUCache(maxErrorPageCacheBytes)
	if err != nil {
		logError("Failed to initialize error page LRU cache: %v", err)
		panic(err)
	}
}

// parsedTemplateOnce 用于确保 HTML 模板只被解析一次。
var (
	parsedTemplateOnce sync.Once
	parsedTemplate     *template.Template
	parsedTemplateErr  error
)

// getParsedTemplate 用于获取缓存的解析后的 HTML 模板。
func getParsedTemplate() (*template.Template, error) {
	parsedTemplateOnce.Do(func() {
		tmplPath := "page.tmpl"
		// 确保 errPagesFs 已初始化。这要求在任何 ErrorPage 调用之前调用 InitErrPagesFS。
		if errPagesFs == nil {
			parsedTemplateErr = fmt.Errorf("errPagesFs not initialized. Call InitErrPagesFS first")
			return
		}
		parsedTemplate, parsedTemplateErr = template.ParseFS(errPagesFs, tmplPath)
		if parsedTemplateErr != nil {
			parsedTemplate = nil
		}
	})
	return parsedTemplate, parsedTemplateErr
}

// htmlTemplateRender 修改为使用缓存的模板。
func htmlTemplateRender(data interface{}) ([]byte, error) {
	tmpl, err := getParsedTemplate()
	if err != nil {
		return nil, fmt.Errorf("failed to get parsed template: %w", err)
	}
	if tmpl == nil {
		return nil, fmt.Errorf("template is nil after parsing")
	}

	// 创建一个 bytes.Buffer 用于存储渲染结果
	var buf bytes.Buffer

	err = tmpl.Execute(&buf, data)
	if err != nil {
		return nil, fmt.Errorf("failed to execute template: %w", err)
	}

	// 返回 buffer 的内容作为 []byte
	return buf.Bytes(), nil
}

func ErrorPage(c *app.RequestContext, errInfo *GHProxyErrors) {
	// 将 errInfo 转换为 ErrorPageData 结构体
	pageDataStruct := ErrPageUnwarper(errInfo)
	// 使用 ErrorPageData 生成一个唯一的 SHA256 缓存键
	cacheKey := pageDataStruct.ToCacheKey()
	if cacheKey == "" {
		c.JSON(errInfo.StatusCode, map[string]string{"error": errInfo.ErrorMessage})
		logWarning("Failed to generate cache key for error page: %v", errInfo)
		return
	}

	var pageData []byte
	var err error

	// 尝试从缓存中获取页面数据
	if cachedPage, found := errorPageCache.Get(cacheKey); found {
		pageData = cachedPage
		logDebug("Serving error page from cache (Key: %s)", cacheKey)
	} else {
		// 如果不在缓存中，则渲染页面
		pageData, err = htmlTemplateRender(pageDataStruct)
		if err != nil {
			c.JSON(errInfo.StatusCode, map[string]string{"error": errInfo.ErrorMessage})
			logWarning("Failed to render error page for status %d (Key: %s): %v", errInfo.StatusCode, cacheKey, err)
			return
		}

		// 将渲染结果存入缓存
		errorPageCache.Add(cacheKey, pageData)
		logDebug("Cached error page (Key: %s, Size: %d bytes)", cacheKey, len(pageData))
	}

	c.Data(errInfo.StatusCode, "text/html; charset=utf-8", pageData)
}