diff --git a/backend/config/config.go b/backend/config/config.go index 9022a05..d4d11f0 100644 --- a/backend/config/config.go +++ b/backend/config/config.go @@ -63,8 +63,9 @@ type Config struct { } `mapstructure:"init_model"` Extension struct { - Baseurl string `mapstructure:"baseurl"` - Limit int `mapstructure:"limit"` + Baseurl string `mapstructure:"baseurl"` + LimitSecond int `mapstructure:"limit_second"` + Limit int `mapstructure:"limit"` } `mapstructure:"extension"` } @@ -100,7 +101,8 @@ func Init() (*Config, error) { v.SetDefault("init_model.key", "") v.SetDefault("init_model.url", "https://model-square.app.baizhi.cloud/v1") v.SetDefault("extension.baseurl", "https://release.baizhi.cloud") - v.SetDefault("extension.limit", 10) + v.SetDefault("extension.limit", 1) + v.SetDefault("extension.limit_second", 10) c := Config{} if err := v.Unmarshal(&c); err != nil { diff --git a/backend/go.mod b/backend/go.mod index e11226b..118ffea 100644 --- a/backend/go.mod +++ b/backend/go.mod @@ -20,6 +20,7 @@ require ( golang.org/x/crypto v0.39.0 golang.org/x/oauth2 v0.18.0 golang.org/x/text v0.26.0 + golang.org/x/time v0.11.0 ) require ( @@ -97,7 +98,6 @@ require ( golang.org/x/net v0.41.0 // indirect golang.org/x/sync v0.15.0 // indirect golang.org/x/sys v0.33.0 // indirect - golang.org/x/time v0.11.0 // indirect golang.org/x/tools v0.33.0 // indirect google.golang.org/appengine v1.6.8 // indirect google.golang.org/protobuf v1.34.2 // indirect diff --git a/backend/internal/user/handler/v1/user.go b/backend/internal/user/handler/v1/user.go index 804a903..019c894 100644 --- a/backend/internal/user/handler/v1/user.go +++ b/backend/internal/user/handler/v1/user.go @@ -11,6 +11,7 @@ import ( "time" "github.com/GoYoko/web" + "golang.org/x/time/rate" "github.com/chaitin/MonkeyCode/backend/config" "github.com/chaitin/MonkeyCode/backend/consts" @@ -33,9 +34,9 @@ type UserHandler struct { session *session.Session logger *slog.Logger cfg *config.Config - limitCh chan struct{} - vsixCache map[string]*CacheEntry // 缓存处理后的vsix文件 - cacheMu sync.RWMutex // 缓存读写锁 + vsixCache map[string]*CacheEntry + cacheMu sync.RWMutex + limiter *rate.Limiter } func NewUserHandler( @@ -54,8 +55,8 @@ func NewUserHandler( logger: logger, cfg: cfg, euse: euse, - limitCh: make(chan struct{}, cfg.Extension.Limit), vsixCache: make(map[string]*CacheEntry), + limiter: rate.NewLimiter(rate.Every(time.Duration(cfg.Extension.LimitSecond)*time.Second), cfg.Extension.Limit), } w.GET("/api/v1/static/vsix/:version", web.BaseHandler(u.VSIXDownload)) @@ -130,27 +131,22 @@ func (h *UserHandler) cleanExpiredCache() { // @Produce octet-stream // @Router /api/v1/static/vsix [get] func (h *UserHandler) VSIXDownload(c *web.Context) error { - h.limitCh <- struct{}{} - defer func() { - <-h.limitCh - }() + if !h.limiter.Allow() { + return c.String(http.StatusTooManyRequests, "Too Many Requests") + } v, err := h.euse.GetByVersion(c.Request().Context(), c.Param("version")) if err != nil { return err } - // 生成缓存键 cacheKey := h.generateCacheKey(v.Version, h.cfg.BaseUrl) - // 先检查缓存 h.cacheMu.RLock() if entry, exists := h.vsixCache[cacheKey]; exists { - // 检查缓存是否过期(1小时) if time.Since(entry.createdAt) < time.Hour { h.cacheMu.RUnlock() - // 从缓存返回数据 disposition := fmt.Sprintf("attachment; filename=monkeycode-%s.vsix", v.Version) c.Response().Header().Set("Content-Type", "application/octet-stream") c.Response().Header().Set("Content-Disposition", disposition) @@ -162,15 +158,11 @@ func (h *UserHandler) VSIXDownload(c *web.Context) error { } h.cacheMu.RUnlock() - // 缓存未命中或已过期,需要重新生成 - - // 使用buffer来捕获生成的数据 var buf bytes.Buffer if err := vsix.ChangeVsixEndpoint(v.Path, "extension/package.json", h.cfg.BaseUrl, &buf); err != nil { return err } - // 将结果存入缓存 data := buf.Bytes() h.cacheMu.Lock() h.vsixCache[cacheKey] = &CacheEntry{ @@ -179,10 +171,8 @@ func (h *UserHandler) VSIXDownload(c *web.Context) error { } h.cacheMu.Unlock() - // 异步清理过期缓存 go h.cleanExpiredCache() - // 返回数据给客户端 disposition := fmt.Sprintf("attachment; filename=monkeycode-%s.vsix", v.Version) c.Response().Header().Set("Content-Type", "application/octet-stream") c.Response().Header().Set("Content-Disposition", disposition)