1769701362

Committed 15 Apr 2025 01:54PM UTC coverage: 65.264%. First build

Build # 1769701362

Build Type

Pull #202

gitlab-ci

Committed by

alfrunes

Commit Message

test: Fix deviceauth.cache tests after changing client initialization

Signed-off-by: Alf-Rune Siqveland <alf.rune@northern.tech>

Pull Request Pull Request #202: MEN-7733: Rate limits for devices APIs

Run Details

144 of 389 new or added lines in 8 files covered. (37.02%)

31803 of 48730 relevant lines covered (65.26%)

1.37 hits per line

Source File
Press 'n' to go to next uncovered line, 'b' for previous

0.0

/backend/pkg/redis/ratelimit.go

package redis

import (
        "context"
        "errors"
        "fmt"
        "time"

        "github.com/redis/go-redis/v9"

        "github.com/mendersoftware/mender-server/pkg/rate"
)

func NewFixedWindowRateLimiter(
        client Client,
        paramsFromContext RatelimitParamsFunc,
) rate.Limiter {
        return &fixedWindowRatelimiter{
                client:     client,
                paramsFunc: paramsFromContext,
                nowFunc:    time.Now,
        }
}

type RatelimitParams struct {
        Burst     uint64
        Interval  time.Duration
        KeyPrefix string
}

type RatelimitParamsFunc func(context.Context) (*RatelimitParams, error)

func FixedRatelimitParams(params RatelimitParams) RatelimitParamsFunc {
        return func(ctx context.Context) (*RatelimitParams, error) {
                return &params, nil
        }
}

type fixedWindowRatelimiter struct {
        client     Client
        paramsFunc RatelimitParamsFunc
        nowFunc    func() time.Time
}

type simpleReservation struct {
        ok     bool
        tokens uint64
        delay  time.Duration
}

func (r *simpleReservation) OK() bool {
        return r.ok
}

func (r *simpleReservation) Delay() time.Duration {
        return r.delay
}

func (r *simpleReservation) Tokens() uint64 {
        return r.tokens
}

func epoch(t time.Time, interval time.Duration) int64 {
        return t.UnixMilli() / interval.Milliseconds()
}

func fixedWindowKey(prefix string, epoch int64) string {
        if prefix == "" {
                prefix = "ratelimit"
        }
        return fmt.Sprintf("%s:e:%d:c", prefix, epoch)
}

func (rl *fixedWindowRatelimiter) Reserve(ctx context.Context) (rate.Reservation, error) {
        now := rl.nowFunc()
        params, err := rl.paramsFunc(ctx)
        if err != nil {
                return nil, err
        } else if params == nil {
                return &simpleReservation{
                        ok: true,
                }, nil
        }
        epoch := epoch(now, params.Interval)
        key := fixedWindowKey(params.KeyPrefix, epoch)
        count := uint64(1)

        err = rl.client.SetArgs(ctx, key, count, redis.SetArgs{
                TTL:  params.Interval,
                Mode: `NX`,
        }).Err()
        if errors.Is(err, redis.Nil) {
                count, err = rl.client.Incr(ctx, key).Uint64()
        }
        if err != nil {
                return nil, fmt.Errorf("redis: error computing rate limit: %w", err)
        }
        if count <= params.Burst {
                return &simpleReservation{
                        delay:  0,
                        ok:     true,
                        tokens: params.Burst - count,
                }, nil
        }
        return &simpleReservation{
                delay: now.Sub(time.UnixMilli((epoch + 1) *
                        params.Interval.Milliseconds())),
                ok:     false,
                tokens: 0,
        }, nil
}

func (rl *fixedWindowRatelimiter) Tokens(ctx context.Context) (uint64, error) {
        params, err := rl.paramsFunc(ctx)
        if err != nil {
                return 0, err
        }
        count, err := rl.client.Get(ctx,
                fixedWindowKey(params.KeyPrefix, epoch(rl.nowFunc(), params.Interval)),
        ).Uint64()
        if errors.Is(err, redis.Nil) {
                return params.Burst, nil
        } else if err != nil {
                return 0, fmt.Errorf("redis: error getting free tokens: %w", err)
        } else if count > params.Burst {
                return 0, nil
        }
        return params.Burst - count, nil
}

1	package redis
2
3	import (
4	"context"
5	"errors"
6	"fmt"
7	"time"
8
9	"github.com/redis/go-redis/v9"
10
11	"github.com/mendersoftware/mender-server/pkg/rate"
12	)
13
14	func NewFixedWindowRateLimiter(
15	client Client,
16	paramsFromContext RatelimitParamsFunc,
NEW 17	) rate.Limiter {	×
NEW 18	return &fixedWindowRatelimiter{	×
NEW 19	client: client,	×
NEW 20	paramsFunc: paramsFromContext,	×
NEW 21	nowFunc: time.Now,	×
NEW 22	}	×
NEW 23	}	×
24
25	type RatelimitParams struct {
26	Burst uint64
27	Interval time.Duration
28	KeyPrefix string
29	}
30
31	type RatelimitParamsFunc func(context.Context) (*RatelimitParams, error)
32
NEW 33	func FixedRatelimitParams(params RatelimitParams) RatelimitParamsFunc {	×
NEW 34	return func(ctx context.Context) (*RatelimitParams, error) {	×
NEW 35	return &params, nil	×
NEW 36	}	×
37	}
38
39	type fixedWindowRatelimiter struct {
40	client Client
41	paramsFunc RatelimitParamsFunc
42	nowFunc func() time.Time
43	}
44
45	type simpleReservation struct {
46	ok bool
47	tokens uint64
48	delay time.Duration
49	}
50
NEW 51	func (r *simpleReservation) OK() bool {	×
NEW 52	return r.ok	×
NEW 53	}	×
54
NEW 55	func (r *simpleReservation) Delay() time.Duration {	×
NEW 56	return r.delay	×
NEW 57	}	×
58
NEW 59	func (r *simpleReservation) Tokens() uint64 {	×
NEW 60	return r.tokens	×
NEW 61	}	×
62
NEW 63	func epoch(t time.Time, interval time.Duration) int64 {	×
NEW 64	return t.UnixMilli() / interval.Milliseconds()	×
NEW 65	}	×
66
NEW 67	func fixedWindowKey(prefix string, epoch int64) string {	×
NEW 68	if prefix == "" {	×
NEW 69	prefix = "ratelimit"	×
NEW 70	}	×
NEW 71	return fmt.Sprintf("%s:e:%d:c", prefix, epoch)	×
72	}
73
NEW 74	func (rl *fixedWindowRatelimiter) Reserve(ctx context.Context) (rate.Reservation, error) {	×
NEW 75	now := rl.nowFunc()	×
NEW 76	params, err := rl.paramsFunc(ctx)	×
NEW 77	if err != nil {	×
NEW 78	return nil, err	×
NEW 79	} else if params == nil {	×
NEW 80	return &simpleReservation{	×
NEW 81	ok: true,	×
NEW 82	}, nil	×
NEW 83	}	×
NEW 84	epoch := epoch(now, params.Interval)	×
NEW 85	key := fixedWindowKey(params.KeyPrefix, epoch)	×
NEW 86	count := uint64(1)	×
NEW 87		×
NEW 88	err = rl.client.SetArgs(ctx, key, count, redis.SetArgs{	×
NEW 89	TTL: params.Interval,	×
NEW 90	Mode: `NX`,	×
NEW 91	}).Err()	×
NEW 92	if errors.Is(err, redis.Nil) {	×
NEW 93	count, err = rl.client.Incr(ctx, key).Uint64()	×
NEW 94	}	×
NEW 95	if err != nil {	×
NEW 96	return nil, fmt.Errorf("redis: error computing rate limit: %w", err)	×
NEW 97	}	×
NEW 98	if count <= params.Burst {	×
NEW 99	return &simpleReservation{	×
NEW 100	delay: 0,	×
NEW 101	ok: true,	×
NEW 102	tokens: params.Burst - count,	×
NEW 103	}, nil	×
NEW 104	}	×
NEW 105	return &simpleReservation{	×
NEW 106	delay: now.Sub(time.UnixMilli((epoch + 1) *	×
NEW 107	params.Interval.Milliseconds())),	×
NEW 108	ok: false,	×
NEW 109	tokens: 0,	×
NEW 110	}, nil	×
111	}
112
NEW 113	func (rl *fixedWindowRatelimiter) Tokens(ctx context.Context) (uint64, error) {	×
NEW 114	params, err := rl.paramsFunc(ctx)	×
NEW 115	if err != nil {	×
NEW 116	return 0, err	×
NEW 117	}	×
NEW 118	count, err := rl.client.Get(ctx,	×
NEW 119	fixedWindowKey(params.KeyPrefix, epoch(rl.nowFunc(), params.Interval)),	×
NEW 120	).Uint64()	×
NEW 121	if errors.Is(err, redis.Nil) {	×
NEW 122	return params.Burst, nil	×
NEW 123	} else if err != nil {	×
NEW 124	return 0, fmt.Errorf("redis: error getting free tokens: %w", err)	×
NEW 125	} else if count > params.Burst {	×
NEW 126	return 0, nil	×
NEW 127	}	×
NEW 128	return params.Burst - count, nil	×
129	}

mendersoftware / mender-server / 1769701362

Source File Press 'n' to go to next uncovered line, 'b' for previous

Source File
Press 'n' to go to next uncovered line, 'b' for previous