Adding ProjectInfo to AccessQuota (#137)

Author: klaidliadon

cache.go

@@ -54,17 +54,13 @@ var (
 	_ UsageCache = (*RedisCache)(nil)
 )
 
-type Service interface {
-	GetService() proto.Service
-}
-
-func NewLimitCounter(svc Service, cfg RedisConfig, logger *slog.Logger) httprate.LimitCounter {
+func NewLimitCounter(svc proto.Service, cfg RedisConfig, logger *slog.Logger) httprate.LimitCounter {
 	if !cfg.Enabled {
 		return nil
 	}
 
 	prefix := redisRLPrefix
-	if s := svc.GetService().String(); s != "" {
+	if s := svc.String(); s != "" {
 		prefix = fmt.Sprintf("%s%s:", redisRLPrefix, s)
 	}
 
@@ -91,7 +87,7 @@ func NewLimitCounter(svc Service, cfg RedisConfig, logger *slog.Logger) httprate
 const (
 	defaultExpRedis = time.Hour
 	defaultExpLRU   = time.Minute
-	cacheVersion    = "v1"
+	cacheVersion    = "v2"
 )
 
 // usageKey returns the redis key for storing usage amount.

client.go

@@ -133,7 +133,7 @@ func (c *Client) FetchProjectQuota(ctx context.Context, projectID uint64, chainI
 			logger.Warn("failed to cache project quota", slog.Any("error", err))
 		}
 	}
-	if err := quota.AccessKey.ValidateChains(chainIDs); err != nil {
+	if err := quota.Info.ValidateChains(chainIDs); err != nil {
 		return quota, proto.ErrInvalidChain.WithCause(err)
 	}
 	return quota, nil
@@ -163,8 +163,11 @@ func (c *Client) FetchKeyQuota(ctx context.Context, accessKey, origin string, ch
 			logger.Warn("failed to cache access quota", slog.Any("error", err))
 		}
 	}
+	if err := quota.Info.ValidateChains(chainIDs); err != nil {
+		return quota, proto.ErrInvalidChain.WithCause(err)
+	}
 	// validate access key
-	if err := c.validateAccessKey(quota.AccessKey, origin, chainIDs); err != nil {
+	if err := c.validateAccessKey(quota.AccessKey, origin); err != nil {
 		return quota, err
 	}
 	return quota, nil
@@ -325,7 +328,7 @@ func (c *Client) ClearQuotaCacheByAccessKey(ctx context.Context, accessKey strin
 	return c.cache.QuotaCache.DeleteAccessQuota(ctx, accessKey)
 }
 
-func (c *Client) validateAccessKey(access *proto.AccessKey, origin string, chainIDs []uint64) (err error) {
+func (c *Client) validateAccessKey(access *proto.AccessKey, origin string) (err error) {
 	if !access.Active {
 		return proto.ErrAccessKeyNotFound
 	}
@@ -335,9 +338,6 @@ func (c *Client) validateAccessKey(access *proto.AccessKey, origin string, chain
 	if !access.ValidateService(c.service) {
 		return proto.ErrInvalidService
 	}
-	if err := access.ValidateChains(chainIDs); err != nil {
-		return proto.ErrInvalidChain.WithCause(err)
-	}
 	return nil
 }
 

middleware/common.go

@@ -85,7 +85,7 @@ func VerifyChains(ctx context.Context, chainIDs ...uint64) error {
 	if !ok {
 		return nil
 	}
-	if err := quota.AccessKey.ValidateChains(chainIDs); err != nil {
+	if err := quota.Info.ValidateChains(chainIDs); err != nil {
 		return proto.ErrInvalidChain.WithCause(err)
 	}
 	return nil

tests/mock/mem.go mock/mem.gotests/mock/mem.go renamed to mock/mem.go

@@ -6,16 +6,15 @@ import (
 	"time"
 
 	"github.com/0xsequence/authcontrol"
-	"github.com/0xsequence/quotacontrol/internal/store"
 	"github.com/0xsequence/quotacontrol/internal/usage"
 	"github.com/0xsequence/quotacontrol/proto"
 )
 
 // NewMemoryStore returns a new in-memory store.
 func NewMemoryStore() *MemoryStore {
 	ms := MemoryStore{
+		infos:       map[uint64]proto.ProjectInfo{},
 		limits:      map[uint64]proto.Limit{},
-		cycles:      map[uint64]proto.Cycle{},
 		accessKeys:  map[string]proto.AccessKey{},
 		usage:       map[proto.Service]usage.Record{},
 		users:       map[string]bool{},
@@ -38,16 +37,30 @@ type userPermission struct {
 type MemoryStore struct {
 	sync.Mutex
 	limits      map[uint64]proto.Limit
-	cycles      map[uint64]proto.Cycle
+	infos       map[uint64]proto.ProjectInfo
 	accessKeys  map[string]proto.AccessKey
 	usage       map[proto.Service]usage.Record
 	users       map[string]bool
 	projects    map[uint64]*authcontrol.Auth
 	permissions map[uint64]map[string]userPermission
 }
 
+func (m *MemoryStore) SetProjectInfo(ctx context.Context, projectID uint64, info *proto.ProjectInfo) error {
+	m.Lock()
+	if info != nil {
+		m.infos[projectID] = *info
+	} else {
+		delete(m.infos, projectID)
+	}
+	m.Unlock()
+	return nil
+}
+
 func (m *MemoryStore) SetAccessLimit(ctx context.Context, projectID uint64, config *proto.Limit) error {
 	m.Lock()
+	if _, ok := m.infos[projectID]; !ok {
+		m.infos[projectID] = proto.ProjectInfo{ID: projectID}
+	}
 	m.limits[projectID] = *config
 	m.Unlock()
 	return nil
@@ -63,25 +76,11 @@ func (m *MemoryStore) GetAccessLimit(ctx context.Context, projectID uint64, cycl
 	return &limit, nil
 }
 
-func (m *MemoryStore) SetAccessCycle(ctx context.Context, projectID uint64, cycle *proto.Cycle) error {
+func (m *MemoryStore) GetProjectInfo(ctx context.Context, projectID uint64, now time.Time) (*proto.ProjectInfo, error) {
 	m.Lock()
-	if cycle != nil {
-		m.cycles[projectID] = *cycle
-	} else {
-		delete(m.cycles, projectID)
-	}
-	m.Unlock()
-	return nil
-}
-
-func (m *MemoryStore) GetAccessCycle(ctx context.Context, projectID uint64, now time.Time) (*proto.Cycle, error) {
-	m.Lock()
-	cycle := m.cycles[projectID]
+	info := m.infos[projectID]
 	m.Unlock()
-	if cycle.Start.IsZero() && cycle.End.IsZero() {
-		return store.Cycle{}.GetAccessCycle(ctx, projectID, now)
-	}
-	return &cycle, nil
+	return &info, nil
 }
 
 func (m *MemoryStore) InsertAccessKey(ctx context.Context, access *proto.AccessKey) error {

tests/mock/server.go mock/server.gotests/mock/server.go renamed to mock/server.go

@@ -50,11 +50,11 @@ func NewServer(cfg *quotacontrol.Config) (server *Server, cleanup func()) {
 		PermissionCache: quotacontrol.NewRedisCache(client, time.Minute),
 	}
 	qcStore := quotacontrol.Store{
-		LimitStore:      store,
-		AccessKeyStore:  store,
-		UsageStore:      store,
-		CycleStore:      store,
-		PermissionStore: store,
+		ProjectInfoStore: store,
+		LimitStore:       store,
+		AccessKeyStore:   store,
+		UsageStore:       store,
+		PermissionStore:  store,
 	}
 
 	logger := qc.logger.With(slog.String("server", "server"))

proto/proto.go

@@ -3,42 +3,47 @@
 package proto
 
 import (
+	"cmp"
 	"fmt"
 	"slices"
 	"time"
 )
 
+// Ptr is an utility function to return a pointer to the value
 func Ptr[T any](v T) *T {
 	return &v
 }
 
-func (t *AccessKey) ValidateOrigin(rawOrigin string) bool {
+// ValidateOrigin checks if the given origin is allowed by the access key.
+func (a *AccessKey) ValidateOrigin(rawOrigin string) bool {
 	if rawOrigin == "" {
-		return !t.RequireOrigin
+		return !a.RequireOrigin
 	}
-	return t.AllowedOrigins.MatchAny(rawOrigin)
+	return a.AllowedOrigins.MatchAny(rawOrigin)
 }
 
-func (t *AccessKey) ValidateService(service Service) bool {
-	if len(t.AllowedServices) == 0 {
+// ValidateService checks if the given service is allowed by the access key.
+func (a *AccessKey) ValidateService(service Service) bool {
+	if len(a.AllowedServices) == 0 {
 		return true
 	}
-	for _, s := range t.AllowedServices {
+	for _, s := range a.AllowedServices {
 		if service == s {
 			return true
 		}
 	}
 	return false
 }
 
-func (t *AccessKey) ValidateChains(chainIDs []uint64) error {
-	if len(t.ChainIDs) == 0 {
+// ValidateChains checks if the given chain IDs are allowed by the project.
+func (i *ProjectInfo) ValidateChains(chainIDs []uint64) error {
+	if len(i.ChainIDs) == 0 {
 		return nil
 	}
 
 	invalid := make([]uint64, 0, len(chainIDs))
 	for _, id := range chainIDs {
-		if !slices.Contains(t.ChainIDs, id) {
+		if !slices.Contains(i.ChainIDs, id) {
 			invalid = append(invalid, id)
 		}
 	}
@@ -49,6 +54,7 @@ func (t *AccessKey) ValidateChains(chainIDs []uint64) error {
 	return nil
 }
 
+// Validate checks if the limit configuration is valid.
 func (l Limit) Validate() error {
 	for name, cfg := range l.ServiceLimit {
 		svc, ok := ParseService(name)
@@ -62,18 +68,21 @@ func (l Limit) Validate() error {
 	return nil
 }
 
+// GetSettings returns the service limit settings for the given service.
 func (l Limit) GetSettings(svc Service) (ServiceLimit, bool) {
 	settings, ok := l.ServiceLimit[svc.String()]
 	return settings, ok
 }
 
+// SetSetting sets the service limit settings for the given service.
 func (l *Limit) SetSetting(svc Service, limits ServiceLimit) {
 	if l.ServiceLimit == nil {
 		l.ServiceLimit = make(map[string]ServiceLimit)
 	}
 	l.ServiceLimit[svc.String()] = limits
 }
 
+// Validate checks if the service limit configuration is valid.
 func (l ServiceLimit) Validate() error {
 	if l.RateLimit < 1 {
 		return fmt.Errorf("rateLimit must be > 0")
@@ -104,6 +113,7 @@ func getOverThreshold(v, total, threshold int64) (int64, bool) {
 	return max(0, total-threshold), true
 }
 
+// GetSpendResult calculates the spend result and event type based on the service limit and usage
 func (l *ServiceLimit) GetSpendResult(v, total int64) (int64, *EventType) {
 	// valid usage
 	if total < l.FreeMax {
@@ -176,14 +186,23 @@ func (c *Cycle) GetEnd(now time.Time) time.Time {
 	return c.GetStart(now).AddDate(0, 1, -1)
 }
 
-func (c *Cycle) GetDuration(now time.Time) time.Duration {
-	return c.GetEnd(now).Sub(c.GetStart(now))
-}
+func (c *Cycle) SetInterval(from, to *time.Time, now time.Time) {
+	from = cmp.Or(from, &time.Time{})
+	to = cmp.Or(to, &time.Time{})
+
+	if !from.IsZero() && !to.IsZero() {
+		return
+	}
 
-func (c *Cycle) Advance(now time.Time) {
-	for c.End.Before(now) {
-		c.Start = c.Start.AddDate(0, 1, 0)
-		c.End = c.End.AddDate(0, 1, 0)
+	duration := c.GetEnd(now).Sub(c.GetStart(now))
+	switch {
+	case !to.IsZero():
+		*from = to.Add(-duration)
+	case !from.IsZero():
+		*to = from.Add(duration)
+	default:
+		*from = c.Start
+		*to = c.End
 	}
 }
 
@@ -194,14 +213,6 @@ func (u *UserPermission) CanAccess(perm UserPermission) bool {
 	return *u >= perm
 }
 
-func (e WebRPCError) WithMessage(message string) WebRPCError {
-	err := e
-	if message != "" {
-		err.Message = message
-	}
-	return err
-}
-
 func ParseService(v string) (Service, bool) {
 	raw, ok := Service_value[v]
 	if !ok {
@@ -233,7 +244,3 @@ func (x Service) GetName() string {
 	}
 	return ""
 }
-
-func (x Service) GetService() Service {
-	return x
-}