diff --git a/main.go b/main.go index 484d59f..2236e2a 100644 --- a/main.go +++ b/main.go @@ -17,6 +17,7 @@ import ( "ghproxy/middleware/loggin" "ghproxy/proxy" "ghproxy/rate" + "ghproxy/weakcache" "github.com/WJQSERVER-STUDIO/logger" "github.com/hertz-contrib/http2/factory" @@ -50,6 +51,10 @@ var ( pagesFS embed.FS ) +var ( + wcache *weakcache.Cache[string] // docker token缓存 +) + var ( logw = logger.Logw logDump = logger.LogDump @@ -360,6 +365,9 @@ func init() { setMemLimit(cfg) loadlist(cfg) setupRateLimit(cfg) + if cfg.Docker.Enabled { + wcache = proxy.InitWeakCache() + } if cfg.Server.Debug { runMode = "dev" @@ -468,23 +476,7 @@ func main() { proxy.RoutingHandler(cfg, limiter, iplimiter)(ctx, c) }) - // for 3.4.0 - - /* - r.GET("/v2/", func(ctx context.Context, c *app.RequestContext) { - proxy.GhcrRouting(cfg)(ctx, c) - - /* - //proxy.GhcrRouting(cfg)(ctx, c) - // 返回200与空json - //c.JSON(200, map[string]interface{}{}) - emptyJSON := "{}" - //emptyJSON := `{"name":"disable-list-tags","tags":[]}` - c.Header("Content-Type", "application/json") - c.Header("Content-Length", fmt.Sprint(len(emptyJSON))) - c.String(200, emptyJSON) - */ - /* + r.GET("/v2/", func(ctx context.Context, c *app.RequestContext) { emptyJSON := "{}" c.Header("Content-Type", "application/json") c.Header("Content-Length", fmt.Sprint(len(emptyJSON))) @@ -493,30 +485,18 @@ func main() { c.Status(200) c.Write([]byte(emptyJSON)) - */ + }) + + r.Any("/v2/:target/:user/:repo/*filepath", func(ctx context.Context, c *app.RequestContext) { + proxy.GhcrWithImageRouting(cfg)(ctx, c) + }) /* - w := adaptor.GetCompatResponseWriter(&c.Response) - - const emptyJSON = "{}" - w.Header().Set("Content-Type", "application/json") - w.Header().Set("Content-Length", fmt.Sprint(len(emptyJSON))) - w.Header().Del("Server") - - fmt.Fprint(w, emptyJSON) - */ - /* - }) - r.Any("/v2/:target/*filepath", func(ctx context.Context, c *app.RequestContext) { proxy.GhcrRouting(cfg)(ctx, c) }) */ - r.Any("/v2/*filepath", func(ctx context.Context, c *app.RequestContext) { - proxy.GhcrRouting(cfg)(ctx, c) - }) - r.NoRoute(func(ctx context.Context, c *app.RequestContext) { proxy.NoRouteHandler(cfg, limiter, iplimiter)(ctx, c) }) @@ -531,7 +511,7 @@ func main() { }() } - r.Spin() + defer wcache.StopCleanup() defer logger.Close() defer func() { if hertZfile != nil { @@ -541,5 +521,8 @@ func main() { } } }() + + r.Spin() + fmt.Println("Program Exit") } diff --git a/proxy/authparse.go b/proxy/authparse.go new file mode 100644 index 0000000..f353db9 --- /dev/null +++ b/proxy/authparse.go @@ -0,0 +1,62 @@ +package proxy + +import ( + "fmt" + "strings" +) + +// BearerAuthParams 用于存放解析出的 Bearer 认证参数 +type BearerAuthParams struct { + Realm string + Service string + Scope string +} + +// parseBearerWWWAuthenticateHeader 解析 Bearer 方案的 Www-Authenticate Header。 +// 它期望格式为 'Bearer key1="value1",key2="value2",...' +// 并尝试将已知参数解析到 BearerAuthParams struct 中。 +func parseBearerWWWAuthenticateHeader(headerValue string) (*BearerAuthParams, error) { + if headerValue == "" { + return nil, fmt.Errorf("header value is empty") + } + + // 检查 Scheme 是否是 "Bearer" + parts := strings.SplitN(headerValue, " ", 2) + if len(parts) != 2 || strings.ToLower(parts[0]) != "bearer" { + return nil, fmt.Errorf("invalid or non-bearer header format: got '%s'", headerValue) + } + paramsStr := parts[1] + + paramPairs := strings.Split(paramsStr, ",") + tempMap := make(map[string]string) + + for _, pair := range paramPairs { + trimmedPair := strings.TrimSpace(pair) + keyValue := strings.SplitN(trimmedPair, "=", 2) + if len(keyValue) != 2 { + logWarning("Skipping malformed parameter '%s' in Www-Authenticate header: %s", pair, headerValue) + continue + } + key := strings.TrimSpace(keyValue[0]) + value := strings.TrimSpace(keyValue[1]) + if strings.HasPrefix(value, "\"") && strings.HasSuffix(value, "\"") { + value = value[1 : len(value)-1] + } + tempMap[key] = value + } + + //从 map 中提取值并填充到 struct + authParams := &BearerAuthParams{} + + if realm, ok := tempMap["realm"]; ok { + authParams.Realm = realm + } + if service, ok := tempMap["service"]; ok { + authParams.Service = service + } + if scope, ok := tempMap["scope"]; ok { + authParams.Scope = scope + } + + return authParams, nil +} diff --git a/proxy/docker.go b/proxy/docker.go index cdfc743..d4db75e 100644 --- a/proxy/docker.go +++ b/proxy/docker.go @@ -2,39 +2,132 @@ package proxy import ( "context" + "encoding/json" "fmt" + "ghproxy/config" + "ghproxy/weakcache" + "io" "net/http" "strconv" + "strings" "github.com/WJQSERVER-STUDIO/go-utils/limitreader" "github.com/cloudwego/hertz/pkg/app" ) +var ( + dockerhubTarget = "registry-1.docker.io" + ghcrTarget = "ghcr.io" +) + +var cache *weakcache.Cache[string] + +type imageInfo struct { + User string + Repo string + Image string +} + +func InitWeakCache() *weakcache.Cache[string] { + cache = weakcache.NewCache[string](weakcache.DefaultExpiration, 100) + return cache +} + +/* func GhcrRouting(cfg *config.Config) app.HandlerFunc { return func(ctx context.Context, c *app.RequestContext) { - if cfg.Docker.Enabled { + + charToFind := '.' + reqTarget := c.Param("target") + path := "" + target := "" + + if strings.ContainsRune(reqTarget, charToFind) { + + path = c.Param("filepath") + if reqTarget == "docker.io" { + target = dockerhubTarget + } else if reqTarget == "ghcr.io" { + target = ghcrTarget + } else { + target = reqTarget + } + } else { + path = string(c.Request.RequestURI()) + } + + GhcrToTarget(ctx, c, cfg, target, path, nil) + + } +} +*/ + +func GhcrWithImageRouting(cfg *config.Config) app.HandlerFunc { + return func(ctx context.Context, c *app.RequestContext) { + + charToFind := '.' + reqTarget := c.Param("target") + reqImageUser := c.Param("user") + reqImageName := c.Param("repo") + reqFilePath := c.Param("filepath") + + path := fmt.Sprintf("%s/%s/%s", reqImageUser, reqImageName, reqFilePath) + target := "" + + if strings.ContainsRune(reqTarget, charToFind) { + + if reqTarget == "docker.io" { + target = dockerhubTarget + } else if reqTarget == "ghcr.io" { + target = ghcrTarget + } else { + target = reqTarget + } + } else { + path = string(c.Request.RequestURI()) + reqImageUser = c.Param("target") + reqImageName = c.Param("user") + } + image := &imageInfo{ + User: reqImageUser, + Repo: reqImageName, + Image: fmt.Sprintf("%s/%s", reqImageUser, reqImageName), + } + + GhcrToTarget(ctx, c, cfg, target, path, image) + + } + +} + +func GhcrToTarget(ctx context.Context, c *app.RequestContext, cfg *config.Config, target string, path string, image *imageInfo) { + if cfg.Docker.Enabled { + if target != "" { + GhcrRequest(ctx, c, "https://"+target+"/v2/"+path+string(c.Request.QueryString()), image, cfg, target) + + } else { if cfg.Docker.Target == "ghcr" { - GhcrRequest(ctx, c, "https://ghcr.io"+string(c.Request.RequestURI()), cfg, "ghcr") + GhcrRequest(ctx, c, "https://"+ghcrTarget+string(c.Request.RequestURI()), image, cfg, ghcrTarget) } else if cfg.Docker.Target == "dockerhub" { - GhcrRequest(ctx, c, "https://registry-1.docker.io"+string(c.Request.RequestURI()), cfg, "dockerhub") + GhcrRequest(ctx, c, "https://"+dockerhubTarget+string(c.Request.RequestURI()), image, cfg, dockerhubTarget) } else if cfg.Docker.Target != "" { // 自定义taget - GhcrRequest(ctx, c, "https://"+cfg.Docker.Target+string(c.Request.RequestURI()), cfg, "custom") + GhcrRequest(ctx, c, "https://"+cfg.Docker.Target+string(c.Request.RequestURI()), image, cfg, cfg.Docker.Target) } else { // 配置为空 ErrorPage(c, NewErrorWithStatusLookup(403, "Docker Target is not set")) return } - - } else { - ErrorPage(c, NewErrorWithStatusLookup(403, "Docker is not Allowed")) - return } + + } else { + ErrorPage(c, NewErrorWithStatusLookup(403, "Docker is not Allowed")) + return } } -func GhcrRequest(ctx context.Context, c *app.RequestContext, u string, cfg *config.Config, matcher string) { +func GhcrRequest(ctx context.Context, c *app.RequestContext, u string, image *imageInfo, cfg *config.Config, target string) { var ( method []byte @@ -55,12 +148,11 @@ func GhcrRequest(ctx context.Context, c *app.RequestContext, u string, cfg *conf method = c.Request.Method() - rb := client.NewRequestBuilder(string(method), u) + rb := ghcrclient.NewRequestBuilder(string(method), u) rb.NoDefaultHeaders() rb.SetBody(c.Request.BodyStream()) rb.WithContext(ctx) - //req, err = client.NewRequest(string(method), u, c.Request.BodyStream()) req, err = rb.Build() if err != nil { HandleError(c, fmt.Sprintf("Failed to create request: %v", err)) @@ -73,14 +165,62 @@ func GhcrRequest(ctx context.Context, c *app.RequestContext, u string, cfg *conf req.Header.Add(headerKey, headerValue) }) - resp, err = client.Do(req) + req.Header.Set("Host", target) + token, exist := cache.Get(image.Image) + if exist { + logDebug("Use Cache Token: %s", token) + req.Header.Set("Authorization", "Bearer "+token) + } + + resp, err = ghcrclient.Do(req) if err != nil { HandleError(c, fmt.Sprintf("Failed to send request: %v", err)) return } - // 错误处理(404) - if resp.StatusCode == 404 { + // 处理状态码 + if resp.StatusCode == 401 { + // 请求target /v2/路径 + if string(c.Request.URI().Path()) != "/v2/" { + resp.Body.Close() + token := ChallengeReq(target, image, ctx, c) + + // 更新kv + if token != "" { + logDump("Update Cache Token: %s", token) + cache.Put(image.Image, token) + } + + rb := ghcrclient.NewRequestBuilder(string(method), u) + rb.NoDefaultHeaders() + rb.SetBody(c.Request.BodyStream()) + rb.WithContext(ctx) + + req, err = rb.Build() + if err != nil { + HandleError(c, fmt.Sprintf("Failed to create request: %v", err)) + return + } + + c.Request.Header.VisitAll(func(key, value []byte) { + headerKey := string(key) + headerValue := string(value) + req.Header.Add(headerKey, headerValue) + }) + + req.Header.Set("Host", target) + if token != "" { + req.Header.Set("Authorization", "Bearer "+token) + } + + resp, err = ghcrclient.Do(req) + if err != nil { + HandleError(c, fmt.Sprintf("Failed to send request: %v", err)) + return + } + } + + } else if resp.StatusCode == 404 { // 错误处理(404) ErrorPage(c, NewErrorWithStatusLookup(404, "Page Not Found (From Github)")) return } @@ -101,8 +241,7 @@ func GhcrRequest(ctx context.Context, c *app.RequestContext, u string, cfg *conf bodySize = -1 } if err == nil && bodySize > sizelimit { - var finalURL string - finalURL = resp.Request.URL.String() + finalURL := resp.Request.URL.String() err = resp.Body.Close() if err != nil { logError("Failed to close response body: %v", err) @@ -116,7 +255,6 @@ func GhcrRequest(ctx context.Context, c *app.RequestContext, u string, cfg *conf // 复制响应头，排除需要移除的 header for key, values := range resp.Header { for _, value := range values { - //c.Header(key, value) c.Response.Header.Add(key, value) } } @@ -136,3 +274,78 @@ func GhcrRequest(ctx context.Context, c *app.RequestContext, u string, cfg *conf c.SetBodyStream(bodyReader, -1) } + +type AuthToken struct { + Token string `json:"token"` +} + +func ChallengeReq(target string, image *imageInfo, ctx context.Context, c *app.RequestContext) (token string) { + var resp401 *http.Response + var req401 *http.Request + var err error + + rb401 := ghcrclient.NewRequestBuilder("GET", "https://"+target+"/v2/") + rb401.NoDefaultHeaders() + rb401.WithContext(ctx) + rb401.AddHeader("User-Agent", "docker/28.1.1 go/go1.23.8 git-commit/01f442b kernel/6.12.25-amd64 os/linux arch/amd64 UpstreamClient(Docker-Client/28.1.1 ") + req401, err = rb401.Build() + if err != nil { + HandleError(c, fmt.Sprintf("Failed to create request: %v", err)) + return + } + req401.Header.Set("Host", target) + + resp401, err = ghcrclient.Do(req401) + if err != nil { + HandleError(c, fmt.Sprintf("Failed to send request: %v", err)) + return + } + defer resp401.Body.Close() + bearer, err := parseBearerWWWAuthenticateHeader(resp401.Header.Get("Www-Authenticate")) + if err != nil { + logError("Failed to parse Www-Authenticate header: %v", err) + return + } + + scope := fmt.Sprintf("repository:%s:pull", image.Image) + + getAuthRB := ghcrclient.NewRequestBuilder("GET", bearer.Realm). + NoDefaultHeaders(). + WithContext(ctx). + AddHeader("User-Agent", "docker/28.1.1 go/go1.23.8 git-commit/01f442b kernel/6.12.25-amd64 os/linux arch/amd64 UpstreamClient(Docker-Client/28.1.1 "). + SetHeader("Host", bearer.Service). + AddQueryParam("service", bearer.Service). + AddQueryParam("scope", scope) + + getAuthReq, err := getAuthRB.Build() + if err != nil { + logError("Failed to create request: %v", err) + return + } + + authResp, err := ghcrclient.Do(getAuthReq) + if err != nil { + logError("Failed to send request: %v", err) + return + } + + defer authResp.Body.Close() + + bodyBytes, err := io.ReadAll(authResp.Body) + if err != nil { + logError("Failed to read auth response body: %v", err) + return + } + + // 解码json + var authToken AuthToken + err = json.Unmarshal(bodyBytes, &authToken) + if err != nil { + logError("Failed to decode auth response body: %v", err) + return + } + token = authToken.Token + + return token + +} diff --git a/proxy/httpc.go b/proxy/httpc.go index 83de29b..4d8bc4c 100644 --- a/proxy/httpc.go +++ b/proxy/httpc.go @@ -12,10 +12,12 @@ import ( var BufferSize int = 32 * 1024 // 32KB var ( - tr *http.Transport - gittr *http.Transport - client *httpc.Client - gitclient *httpc.Client + tr *http.Transport + gittr *http.Transport + client *httpc.Client + gitclient *httpc.Client + ghcrtr *http.Transport + ghcrclient *httpc.Client ) func InitReq(cfg *config.Config) error { @@ -23,11 +25,13 @@ func InitReq(cfg *config.Config) error { if cfg.GitClone.Mode == "cache" { initGitHTTPClient(cfg) } + initGhcrHTTPClient(cfg) err := SetGlobalRateLimit(cfg) if err != nil { return err } return nil + } func initHTTPClient(cfg *config.Config) { @@ -77,6 +81,7 @@ func initHTTPClient(cfg *config.Config) { httpc.WithTransport(tr), ) } + } func initGitHTTPClient(cfg *config.Config) { @@ -147,3 +152,51 @@ func initGitHTTPClient(cfg *config.Config) { ) } } + +func initGhcrHTTPClient(cfg *config.Config) { + var proTolcols = new(http.Protocols) + proTolcols.SetHTTP1(true) + proTolcols.SetHTTP2(true) + if cfg.Httpc.Mode == "auto" { + + ghcrtr = &http.Transport{ + IdleConnTimeout: 30 * time.Second, + WriteBufferSize: 32 * 1024, // 32KB + ReadBufferSize: 32 * 1024, // 32KB + Protocols: proTolcols, + } + } else if cfg.Httpc.Mode == "advanced" { + ghcrtr = &http.Transport{ + MaxIdleConns: cfg.Httpc.MaxIdleConns, + MaxConnsPerHost: cfg.Httpc.MaxConnsPerHost, + MaxIdleConnsPerHost: cfg.Httpc.MaxIdleConnsPerHost, + WriteBufferSize: 32 * 1024, // 32KB + ReadBufferSize: 32 * 1024, // 32KB + Protocols: proTolcols, + } + } else { + // 错误的模式 + logError("unknown httpc mode: %s", cfg.Httpc.Mode) + fmt.Println("unknown httpc mode: ", cfg.Httpc.Mode) + logWarning("use Auto to Run HTTP Client") + fmt.Println("use Auto to Run HTTP Client") + ghcrtr = &http.Transport{ + IdleConnTimeout: 30 * time.Second, + WriteBufferSize: 32 * 1024, // 32KB + ReadBufferSize: 32 * 1024, // 32KB + } + } + if cfg.Outbound.Enabled { + initTransport(cfg, ghcrtr) + } + if cfg.Server.Debug { + ghcrclient = httpc.New( + httpc.WithTransport(ghcrtr), + httpc.WithDumpLog(), + ) + } else { + ghcrclient = httpc.New( + httpc.WithTransport(ghcrtr), + ) + } +} diff --git a/proxy/match.go b/proxy/match.go index 9db0c23..cee4eaa 100644 --- a/proxy/match.go +++ b/proxy/match.go @@ -17,9 +17,12 @@ func Matcher(rawPath string, cfg *config.Config) (string, string, string, *GHPro // 匹配 "https://github.com"开头的链接 if strings.HasPrefix(rawPath, "https://github.com") { remainingPath := strings.TrimPrefix(rawPath, "https://github.com") - if strings.HasPrefix(remainingPath, "/") { - remainingPath = strings.TrimPrefix(remainingPath, "/") - } + /* + if strings.HasPrefix(remainingPath, "/") { + remainingPath = strings.TrimPrefix(remainingPath, "/") + } + */ + remainingPath = strings.TrimPrefix(remainingPath, "/") // 预期格式/user/repo/more... // 取出user和repo和最后部分 parts := strings.Split(remainingPath, "/") diff --git a/proxy/utils.go b/proxy/utils.go index 21e7886..029dc26 100644 --- a/proxy/utils.go +++ b/proxy/utils.go @@ -13,8 +13,7 @@ func listCheck(cfg *config.Config, c *app.RequestContext, user string, repo stri // 白名单检查 if cfg.Whitelist.Enabled { - var whitelist bool - whitelist = auth.CheckWhitelist(user, repo) + whitelist := auth.CheckWhitelist(user, repo) if !whitelist { ErrorPage(c, NewErrorWithStatusLookup(403, fmt.Sprintf("Whitelist Blocked repo: %s/%s", user, repo))) logInfo("%s %s %s %s %s Whitelist Blocked repo: %s/%s", c.ClientIP(), c.Method(), rawPath, c.Request.Header.UserAgent(), c.Request.Header.GetProtocol(), user, repo) @@ -24,8 +23,7 @@ func listCheck(cfg *config.Config, c *app.RequestContext, user string, repo stri // 黑名单检查 if cfg.Blacklist.Enabled { - var blacklist bool - blacklist = auth.CheckBlacklist(user, repo) + blacklist := auth.CheckBlacklist(user, repo) if blacklist { ErrorPage(c, NewErrorWithStatusLookup(403, fmt.Sprintf("Blacklist Blocked repo: %s/%s", user, repo))) logInfo("%s %s %s %s %s Blacklist Blocked repo: %s/%s", c.ClientIP(), c.Method(), rawPath, c.Request.Header.UserAgent(), c.Request.Header.GetProtocol(), user, repo) diff --git a/weakcache/weakcache.go b/weakcache/weakcache.go new file mode 100644 index 0000000..32d897c --- /dev/null +++ b/weakcache/weakcache.go @@ -0,0 +1,258 @@ +package weakcache + +import ( + "container/list" + "sync" + "time" + "weak" // Go 1.24 引入的 weak 包 +) + +// DefaultExpiration 默认过期时间，这里设置为 15 分钟。 +// 这是一个导出的常量，方便用户使用包时引用默认值。 +const DefaultExpiration = 5 * time.Minute + +// cleanupInterval 是后台清理 Go routine 的扫描间隔，这里设置为 5 分钟。 +// 这是一个内部常量，不导出。 +const cleanupInterval = 2 * time.Minute + +// cacheEntry 缓存项的内部结构。不导出。 +type cacheEntry[T any] struct { + Value T + Expiration time.Time + key string // 存储key，方便在list.Element中引用 +} + +// Cache 是一个基于 weak.Pointer, 带有过期和大小上限 (FIFO) 的泛型缓存。 +// 这是一个导出的类型。 +type Cache[T any] struct { + mu sync.RWMutex + + // 修正：缓存存储：key -> weak.Pointer 到 cacheEntry 结构体 (而不是指向结构体的指针) + // weak.Make(*cacheEntry[T]) 返回 weak.Pointer[cacheEntry[T]] + data map[string]weak.Pointer[cacheEntry[T]] + + // FIFO 链表：存储 key 的 list.Element + // 链表头部是最近放入的，尾部是最早放入的（最老的） + fifoList *list.List + // FIFO 元素的映射：key -> *list.Element + fifoMap map[string]*list.Element + + defaultExpiration time.Duration + maxSize int // 缓存最大容量，0 表示无限制 + + stopCleanup chan struct{} + wg sync.WaitGroup // 用于等待清理 Go routine 退出 +} + +// NewCache 创建一个新的缓存实例。 +// expiration: 新添加项的默认过期时间。如果为 0，则使用 DefaultExpiration。 +// maxSize: 缓存的最大容量，0 表示无限制。当达到上限时，采用 FIFO 策略淘汰。 +// 这是一个导出的构造函数。 +func NewCache[T any](expiration time.Duration, maxSize int) *Cache[T] { + if expiration <= 0 { + expiration = DefaultExpiration + } + + c := &Cache[T]{ + // 修正：初始化 map，值类型已修正 + data: make(map[string]weak.Pointer[cacheEntry[T]]), + fifoList: list.New(), + fifoMap: make(map[string]*list.Element), + defaultExpiration: expiration, + maxSize: maxSize, + stopCleanup: make(chan struct{}), + } + // 启动后台清理 Go routine + c.wg.Add(1) + go c.cleanupLoop() + return c +} + +// Put 将值放入缓存。如果 key 已存在，会更新其值和过期时间。 +// 这是导出的方法。 +func (c *Cache[T]) Put(key string, value T) { + c.mu.Lock() + defer c.mu.Unlock() + + now := time.Now() + expiration := now.Add(c.defaultExpiration) + + // 如果 key 已经存在，更新其值和过期时间。 + // 在 FIFO 策略中， Put 更新不改变其在链表中的位置，除非旧的 entry 已经被 GC。 + if elem, ok := c.fifoMap[key]; ok { + // 从 data map 中获取弱引用，wp 的类型现在是 weak.Pointer[cacheEntry[T]] + if wp, dataOk := c.data[key]; dataOk { + // wp.Value() 返回 *cacheEntry[T]， entry 的类型现在是 *cacheEntry[T] + entry := wp.Value() + if entry != nil { + // 旧的 cacheEntry 仍在内存中，直接更新 + entry.Value = value + entry.Expiration = expiration + // 在严格 FIFO 中，更新不移动位置 + return + } + // 如果 weak.Pointer.Value() 为 nil，说明之前的 cacheEntry 已经被 GC 了 + // 此时需要创建一个新的 entry，并将其从旧位置移除，再重新添加 + c.fifoList.Remove(elem) + delete(c.fifoMap, key) + } else { + c.fifoList.Remove(elem) + delete(c.fifoMap, key) + } + } + + // 新建缓存项 (注意这里是结构体值，而不是指针) + // weak.Make 接收的是指针 *T + entry := &cacheEntry[T]{ // 创建结构体指针 + Value: value, + Expiration: expiration, + key: key, // 存储 key + } + + // 将新的 *cacheEntry[T] 包装成 weak.Pointer[cacheEntry[T]] 存入 data map + // weak.Make(entry) 现在返回 weak.Pointer[cacheEntry[T]]，类型匹配 data map 的值类型 + c.data[key] = weak.Make(entry) + + // 添加到 FIFO 链表头部 (最近放入/更新的在头部) + // PushFront 返回新的 list.Element + c.fifoMap[key] = c.fifoList.PushFront(key) + + // 检查大小上限并进行淘汰 (淘汰尾部的最老项) + c.evictIfNeeded() +} + +// Get 从缓存中获取值。返回获取到的值和是否存在/是否有效。 +// 这是导出的方法。 +func (c *Cache[T]) Get(key string) (T, bool) { + c.mu.RLock() // 先读锁 + // 从 data map 中获取弱引用，wp 的类型现在是 weak.Pointer[cacheEntry[T]] + wp, ok := c.data[key] + c.mu.RUnlock() // 立即释放读锁，如果需要写操作（removeEntry）可以获得锁 + + var zero T // 零值 + + if !ok { + return zero, false + } + + // 尝试获取实际的 cacheEntry 指针 + // wp.Value() 返回 *cacheEntry[T]， entry 的类型现在是 *cacheEntry[T] + entry := wp.Value() + + if entry == nil { + // 对象已被GC回收，需要清理此弱引用 + c.removeEntry(key) // 内部会加写锁 + return zero, false + } + + // 检查过期时间 (通过 entry 指针访问字段) + if time.Now().After(entry.Expiration) { + // 逻辑上已过期 + c.removeEntry(key) // 内部会加写锁 + return zero, false + } + + // 在 FIFO 缓存中，Get 操作不改变项在链表中的位置 + return entry.Value, true // 通过 entry 指针访问值字段 +} + +// removeEntry 从缓存中移除项。 +// 这个方法是内部使用的，不导出。需要被调用者确保持有写锁，或者内部自己加锁。 +// 考虑到 Get 和 cleanupLoop 可能会调用，让其内部自己加锁更安全。 +func (c *Cache[T]) removeEntry(key string) { + c.mu.Lock() + defer c.mu.Unlock() + + // 从 data map 中删除 + delete(c.data, key) + + // 从 FIFO 链表和 fifoMap 中删除 + if elem, ok := c.fifoMap[key]; ok { + c.fifoList.Remove(elem) + delete(c.fifoMap, key) + } +} + +// evictIfNeeded 检查是否需要淘汰最老（FIFO 链表尾部）的项。 +// 这个方法是内部使用的，不导出。必须在持有写锁的情况下调用。 +func (c *Cache[T]) evictIfNeeded() { + if c.maxSize > 0 && c.fifoList.Len() > c.maxSize { + // 淘汰 FIFO 链表尾部的元素 (最老的) + oldest := c.fifoList.Back() + if oldest != nil { + keyToEvict := oldest.Value.(string) // 链表元素存储的是 key + c.fifoList.Remove(oldest) + delete(c.fifoMap, keyToEvict) + delete(c.data, keyToEvict) // 移除弱引用 + } + } +} + +// Size 返回当前缓存中的弱引用项数量。 +// 注意：这个数量可能包含已被 GC 回收但尚未清理的项。 +// 这是一个导出的方法。 +func (c *Cache[T]) Size() int { + c.mu.RLock() + defer c.mu.RUnlock() + return len(c.data) +} + +// cleanupLoop 后台清理 Go routine。不导出。 +func (c *Cache[T]) cleanupLoop() { + defer c.wg.Done() + // 使用内部常量 cleanupInterval + ticker := time.NewTicker(cleanupInterval) + defer ticker.Stop() + + for { + select { + case <-ticker.C: + c.cleanupExpiredAndGCed() + case <-c.stopCleanup: + return + } + } +} + +// cleanupExpiredAndGCed 扫描并清理已过期或已被 GC 回收的项。不导出。 +func (c *Cache[T]) cleanupExpiredAndGCed() { + c.mu.Lock() // 清理时需要写锁 + defer c.mu.Unlock() + + now := time.Now() + keysToRemove := make([]string, 0, len(c.data)) // 预估容量 + + // 遍历 data map 查找需要清理的键 + for key, wp := range c.data { + // wp 的类型是 weak.Pointer[cacheEntry[T]] + // wp.Value() 返回 *cacheEntry[T]， entry 的类型是 *cacheEntry[T] + entry := wp.Value() // 尝试获取强引用 + + if entry == nil { + // 已被 GC 回收 + keysToRemove = append(keysToRemove, key) + } else if now.After(entry.Expiration) { + // 逻辑过期 (通过 entry 指针访问字段) + keysToRemove = append(keysToRemove, key) + } + } + + // 执行删除操作 + for _, key := range keysToRemove { + // 从 data map 中删除 + delete(c.data, key) + // 从 FIFO 链表和 fifoMap 中删除 + // 需要再次检查 fifoMap，因为在持有锁期间，evictIfNeeded 可能已经移除了这个 key + if elem, ok := c.fifoMap[key]; ok { + c.fifoList.Remove(elem) + delete(c.fifoMap, key) + } + } +} + +// StopCleanup 停止后台清理 Go routine。 +// 这是一个导出的方法。 +func (c *Cache[T]) StopCleanup() { + close(c.stopCleanup) + c.wg.Wait() // 等待 Go routine 退出 +}