1804977809

Committed 07 May 2025 09:43AM UTC coverage: 65.278% (-0.01%) from 65.288%

Build # 1804977809

Build Type

push

gitlab-ci

Committed by

web-flow

Commit Message

Merge pull request #597 from alfrunes/MEN-7744

MEN-7744: Rate limit authenticated devices API

Run Details

30 of 57 new or added lines in 3 files covered. (52.63%)

16 existing lines in 6 files now uncovered.

31813 of 48735 relevant lines covered (65.28%)

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,
        keyPrefix string,
        interval time.Duration,
        limitFromContext RatelimitParamsFunc,
) rate.Limiter {
        return &fixedWindowRatelimiter{
                client:     client,
                paramsFunc: limitFromContext,
                nowFunc:    time.Now,
                keyPrefix:  keyPrefix,
                interval:   interval,
        }
}

type RatelimitParamsFunc func(context.Context) (burst uint64, eventID string, err error)

func FixedRatelimitParams(burst uint64) RatelimitParamsFunc {
        return func(ctx context.Context) (uint64, string, error) {
                return burst, "", nil
        }
}

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

        interval time.Duration
}

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, eventID string, epoch int64) string {
        if prefix == "" {
                prefix = "ratelimit"
        }
        return fmt.Sprintf("%s:id:%s:e:%d:c", prefix, eventID, epoch)
}

func (rl *fixedWindowRatelimiter) Reserve(ctx context.Context) (rate.Reservation, error) {
        now := rl.nowFunc()
        burst, eventID, err := rl.paramsFunc(ctx)
        if err != nil {
                return nil, err
        }
        epoch := epoch(now, rl.interval)
        key := fixedWindowKey(rl.keyPrefix, eventID, epoch)
        count := uint64(1)

        err = rl.client.SetArgs(ctx, key, count, redis.SetArgs{
                TTL:  rl.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 <= burst {
                return &simpleReservation{
                        delay:  0,
                        ok:     true,
                        tokens: burst - count,
                }, nil
        }
        return &simpleReservation{
                delay: now.Sub(time.UnixMilli((epoch + 1) *
                        rl.interval.Milliseconds())),
                ok:     false,
                tokens: 0,
        }, nil
}

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

mendersoftware / mender-server / 1804977809

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