1959198787

Committed 31 Jul 2025 12:04PM UTC coverage: 65.551% (+0.001%) from 65.55%

Build # 1959198787

Build Type

push

gitlab-ci

Committed by

web-flow

Commit Message

Merge pull request #833 from alfrunes/ratelimits-config-agian

Update ratelimits configuration interface

Run Details

30 of 46 new or added lines in 2 files covered. (65.22%)

4 existing lines in 2 files now uncovered.

32387 of 49407 relevant lines covered (65.55%)

1.39 hits per line

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

81.69

/backend/pkg/rate/limit_http.go

package rate

import (
        "bytes"
        "context"
        "encoding/json"
        "errors"
        "fmt"
        "math"
        "net/http"
        "strconv"
        "text/template"
        "time"

        "github.com/gin-gonic/gin"

        "github.com/mendersoftware/mender-server/pkg/identity"
        "github.com/mendersoftware/mender-server/pkg/requestid"
        "github.com/mendersoftware/mender-server/pkg/rest.utils"
)

// TooManyRequestsError is the error type returned when hitting the rate
// limits.
type TooManyRequestsError struct {
        Delay time.Duration
}

func (err *TooManyRequestsError) Error() string {
        return "too many requests"
}

type templateEventLimiter struct {
        limiter       EventLimiter
        eventTemplate *template.Template
}

// HTTPLimiter combines a set of EventLimiter (groups) and a http.ServeMux
// for routing API endpoints to a EventLimiter.
//
// Every EventLimiter comes with an eventTemplate for grouping events by
// a Go template expression. The data to the Go template can be customized by
// calling WithTemplateDataFunc and WithTemplateFuncs.
//
// Additional routes are added by calling MatchHTTPPattern which accepts a HTTP
// pattern and a target group which also accepts Go template expressions.
//
// The HTTPLimiter implements both the standard http.Handler and
// gin.HandlerFunc (see (*HTTPLimiter).MiddlewareGin) for use with either the
// standard library or as a middleware for the gin-gonic framework.
type HTTPLimiter struct {
        // rootTemplate stores the root Template that is inherited by all
        // template string parameters (event and group strings).
        rootTemplate *template.Template
        // templateDataCallback is a callback that is called before executing
        // template strings and should output the template data used.
        //
        // See defaultTemplateData.
        templateDataCallback func(r *http.Request) any

        // httpMux implements the request router to a ratelimiter match instance
        httpMux *http.ServeMux

        // limiterGroups maps the group name to rate limiters.
        limiterGroups map[string]*templateEventLimiter

        // rewriteRequests decides whether to autmatically rewrite the request
        // URL using X-Forwarded-* headers.
        rewriteRequests bool

        // fallbackReservation is the reservation returned when no rules
        // matches. It is either returning a constant OK Reservation (default)
        // or Reject all requests if WithRejectUnmatched is set.
        fallbackReservation Reservation
}

// defaultTemplateData is the default callback for executing template strings.
func defaultTemplateData(r *http.Request) any {
        id := identity.FromContext(r.Context())
        ctx := map[string]any{
                "Identity": id,
                // "Request":  r, // not needed, but could be handy
        }
        return ctx
}

// NewHTTPLimiter initializes an empty HTTPLimiter.
// The limiter is built by calling AddRateLimitGroup and AddMatchExpression.
func NewHTTPLimiter() *HTTPLimiter {
        return &HTTPLimiter{
                rootTemplate:         template.New("").Option("missingkey=zero"),
                httpMux:              http.NewServeMux(),
                templateDataCallback: defaultTemplateData,
                limiterGroups:        make(map[string]*templateEventLimiter),
                fallbackReservation:  okReservation{},
        }
}

func (h *HTTPLimiter) WithTemplateDataFunc(f func(*http.Request) any) *HTTPLimiter {
        h.templateDataCallback = f
        return h
}

func (h *HTTPLimiter) WithTemplateFuncs(funcs map[string]any) *HTTPLimiter {
        h.rootTemplate.Funcs(funcs)
        return h
}

func (h *HTTPLimiter) WithRewriteRequests(rewrite bool) *HTTPLimiter {
        h.rewriteRequests = rewrite
        return h
}

func (h *HTTPLimiter) WithRejectUnmatched() *HTTPLimiter {
        h.fallbackReservation = rejectReservation{}
        return h
}

func (h *HTTPLimiter) AddRateLimitGroup(limiter EventLimiter, group, eventTemplate string) error {
        t, err := h.rootTemplate.Clone()
        if err == nil {
                _, err = t.Parse(eventTemplate)
        }
        if err != nil {
                return fmt.Errorf("failed to compile event template: %w", err)
        }
        h.limiterGroups[group] = &templateEventLimiter{
                limiter:       limiter,
                eventTemplate: t,
        }
        return nil
}

// AddMatchExpression creates a new route using pattern to apply the
// rate limiter that is matched by groupTemplate Go Template.
func (h *HTTPLimiter) AddMatchExpression(
        pattern, groupTemplate string,
) error {
        var (
                t   *template.Template
                err error
        )
        // Compile eventTemplate:
        t, err = h.rootTemplate.Clone()
        if err == nil {
                _, err = t.Parse(groupTemplate)
        }
        if err != nil {
                return fmt.Errorf("error parsing group_template: %w", err)
        }
        limiterMatcher := matcher{
                HTTPLimiter:   h,
                groupTemplate: t,
        }
        h.httpMux.Handle(pattern, limiterMatcher)
        return nil
}

// matcher is the HTTPHandle
type matcher struct {
        *HTTPLimiter
        groupTemplate *template.Template
}

func (h *HTTPLimiter) handleRequest(r *http.Request) error {
        if h.rewriteRequests {
                r = rest.RewriteForwardedRequest(r)
        }
        res, err := h.Reserve(r)
        if err != nil {
                return err
        }
        if res == nil || res.OK() {
                return nil
        } else {
                return &TooManyRequestsError{
                        Delay: res.Delay(),
                }
        }
}

func handleError(ctx context.Context, w http.ResponseWriter, err error) {
        var tooManyRequests *TooManyRequestsError
        status := http.StatusInternalServerError
        hdr := w.Header()
        hdr.Set("Content-Type", "application/json")
        if errors.As(err, &tooManyRequests) {
                status = http.StatusTooManyRequests
                retryAfter := int64(math.Ceil(tooManyRequests.Delay.Abs().Seconds()))
                hdr.Set("Retry-After", strconv.FormatInt(retryAfter, 10))
        }
        w.WriteHeader(status)
        b, _ := json.Marshal(rest.Error{
                Err:       err.Error(),
                RequestID: requestid.FromContext(ctx),
        })
        _, _ = w.Write(b)
}

// ServeHTTP implements a basic http.Handler so that handler can be used
// as a handler for the mux. It will only write on errors and is expected
// to continue to the actual handler on success.
func (h *HTTPLimiter) ServeHTTP(w http.ResponseWriter, r *http.Request) {
        err := h.handleRequest(r)
        if err != nil {
                handleError(r.Context(), w, err)
        }
}

// MiddlewareGin implements rate limiting as a middleware for gin-gonic
// web framework.
func (h *HTTPLimiter) MiddlewareGin(c *gin.Context) {
        err := h.handleRequest(c.Request)
        if err != nil {
                _ = c.Error(err)
                handleError(c.Request.Context(), c.Writer, err)
                c.Abort()
        }
}

type okReservation struct{}

func (k okReservation) OK() bool             { return true }
func (k okReservation) Delay() time.Duration { return 0 }
func (k okReservation) Tokens() int64        { return math.MaxInt64 }

type rejectReservation struct{}

func (k rejectReservation) OK() bool             { return false }
func (k rejectReservation) Delay() time.Duration { return math.MaxInt64 }
func (k rejectReservation) Tokens() int64        { return -1 }

func (m *HTTPLimiter) Reserve(r *http.Request) (Reservation, error) {
        var b bytes.Buffer
        var eventLimiter *templateEventLimiter
        templateData := m.templateDataCallback(r)
        ctx := r.Context()
        h, _ := m.httpMux.Handler(r)
        hh, ok := h.(matcher)
        if ok && hh.groupTemplate != nil {
                err := hh.groupTemplate.Execute(&b, templateData)
                if err != nil {
                        return nil, fmt.Errorf("error executing ratelimit group template: %w", err)
                }
                eventLimiter = m.limiterGroups[b.String()]
                if eventLimiter == nil {
                        return m.fallbackReservation, nil
                }
                b.Reset()
        } else {
                return m.fallbackReservation, nil
        }
        err := eventLimiter.eventTemplate.Execute(&b, templateData)
        if err != nil {
                return nil, fmt.Errorf("error executing template for event ID: %w", err)
        }
        return eventLimiter.limiter.ReserveEvent(ctx, b.String())
}

1	package rate
2
3	import (
4	"bytes"
5	"context"
6	"encoding/json"
7	"errors"
8	"fmt"
9	"math"
10	"net/http"
11	"strconv"
12	"text/template"
13	"time"
14
15	"github.com/gin-gonic/gin"
16
17	"github.com/mendersoftware/mender-server/pkg/identity"
18	"github.com/mendersoftware/mender-server/pkg/requestid"
19	"github.com/mendersoftware/mender-server/pkg/rest.utils"
20	)
21
22	// TooManyRequestsError is the error type returned when hitting the rate
23	// limits.
24	type TooManyRequestsError struct {
25	Delay time.Duration
26	}
27
28	func (err *TooManyRequestsError) Error() string {	1✔
29	return "too many requests"	1✔
30	}	1✔
31
32	type templateEventLimiter struct {
33	limiter EventLimiter
34	eventTemplate *template.Template
35	}
36
37	// HTTPLimiter combines a set of EventLimiter (groups) and a http.ServeMux
38	// for routing API endpoints to a EventLimiter.
39	//
40	// Every EventLimiter comes with an eventTemplate for grouping events by
41	// a Go template expression. The data to the Go template can be customized by
42	// calling WithTemplateDataFunc and WithTemplateFuncs.
43	//
44	// Additional routes are added by calling MatchHTTPPattern which accepts a HTTP
45	// pattern and a target group which also accepts Go template expressions.
46	//
47	// The HTTPLimiter implements both the standard http.Handler and
48	// gin.HandlerFunc (see (*HTTPLimiter).MiddlewareGin) for use with either the
49	// standard library or as a middleware for the gin-gonic framework.
50	type HTTPLimiter struct {
51	// rootTemplate stores the root Template that is inherited by all
52	// template string parameters (event and group strings).
53	rootTemplate *template.Template
54	// templateDataCallback is a callback that is called before executing
55	// template strings and should output the template data used.
56	//
57	// See defaultTemplateData.
58	templateDataCallback func(r *http.Request) any
59
60	// httpMux implements the request router to a ratelimiter match instance
61	httpMux *http.ServeMux
62
63	// limiterGroups maps the group name to rate limiters.
64	limiterGroups map[string]*templateEventLimiter
65
66	// rewriteRequests decides whether to autmatically rewrite the request
67	// URL using X-Forwarded-* headers.
68	rewriteRequests bool
69
70	// fallbackReservation is the reservation returned when no rules
71	// matches. It is either returning a constant OK Reservation (default)
72	// or Reject all requests if WithRejectUnmatched is set.
73	fallbackReservation Reservation
74	}
75
76	// defaultTemplateData is the default callback for executing template strings.
77	func defaultTemplateData(r *http.Request) any {	1✔
78	id := identity.FromContext(r.Context())	1✔
79	ctx := map[string]any{	1✔
80	"Identity": id,	1✔
81	// "Request": r, // not needed, but could be handy	1✔
82	}	1✔
83	return ctx	1✔
84	}	1✔
85
86	// NewHTTPLimiter initializes an empty HTTPLimiter.
87	// The limiter is built by calling AddRateLimitGroup and AddMatchExpression.
88	func NewHTTPLimiter() *HTTPLimiter {	1✔
89	return &HTTPLimiter{	1✔
90	rootTemplate: template.New("").Option("missingkey=zero"),	1✔
91	httpMux: http.NewServeMux(),	1✔
92	templateDataCallback: defaultTemplateData,	1✔
93	limiterGroups: make(map[string]*templateEventLimiter),	1✔
94	fallbackReservation: okReservation{},	1✔
95	}	1✔
96	}	1✔
97
98	func (h HTTPLimiter) WithTemplateDataFunc(f func(http.Request) any) *HTTPLimiter {	1✔
99	h.templateDataCallback = f	1✔
100	return h	1✔
101	}	1✔
102
103	func (h HTTPLimiter) WithTemplateFuncs(funcs map[string]any) HTTPLimiter {	×
104	h.rootTemplate.Funcs(funcs)	×
105	return h	×
106	}	×
107
108	func (h HTTPLimiter) WithRewriteRequests(rewrite bool) HTTPLimiter {	×
109	h.rewriteRequests = rewrite	×
110	return h	×
111	}	×
112
NEW 113	func (h HTTPLimiter) WithRejectUnmatched() HTTPLimiter {	×
NEW 114	h.fallbackReservation = rejectReservation{}	×
NEW 115	return h	×
NEW 116	}	×
117
118	func (h *HTTPLimiter) AddRateLimitGroup(limiter EventLimiter, group, eventTemplate string) error {	1✔
119	t, err := h.rootTemplate.Clone()	1✔
120	if err == nil {	2✔
121	_, err = t.Parse(eventTemplate)	1✔
122	}	1✔
123	if err != nil {	1✔
124	return fmt.Errorf("failed to compile event template: %w", err)	×
125	}	×
126	h.limiterGroups[group] = &templateEventLimiter{	1✔
127	limiter: limiter,	1✔
128	eventTemplate: t,	1✔
129	}	1✔
130	return nil	1✔
131	}
132
133	// AddMatchExpression creates a new route using pattern to apply the
134	// rate limiter that is matched by groupTemplate Go Template.
135	func (h *HTTPLimiter) AddMatchExpression(
136	pattern, groupTemplate string,
137	) error {	1✔
138	var (	1✔
139	t *template.Template	1✔
140	err error	1✔
141	)	1✔
142	// Compile eventTemplate:	1✔
143	t, err = h.rootTemplate.Clone()	1✔
144	if err == nil {	2✔
145	_, err = t.Parse(groupTemplate)	1✔
146	}	1✔
147	if err != nil {	1✔
NEW 148	return fmt.Errorf("error parsing group_template: %w", err)	×
UNCOV 149	}	×
150	limiterMatcher := matcher{	1✔
151	HTTPLimiter: h,	1✔
152	groupTemplate: t,	1✔
153	}	1✔
154	h.httpMux.Handle(pattern, limiterMatcher)	1✔
155	return nil	1✔
156	}
157
158	// matcher is the HTTPHandle
159	type matcher struct {
160	*HTTPLimiter
161	groupTemplate *template.Template
162	}
163
164	func (h HTTPLimiter) handleRequest(r http.Request) error {	1✔
165	if h.rewriteRequests {	1✔
166	r = rest.RewriteForwardedRequest(r)	×
167	}	×
168	res, err := h.Reserve(r)	1✔
169	if err != nil {	2✔
170	return err	1✔
171	}	1✔
172	if res == nil \|\| res.OK() {	2✔
173	return nil	1✔
174	} else {	2✔
175	return &TooManyRequestsError{	1✔
176	Delay: res.Delay(),	1✔
177	}	1✔
178	}	1✔
179	}
180
181	func handleError(ctx context.Context, w http.ResponseWriter, err error) {	1✔
182	var tooManyRequests *TooManyRequestsError	1✔
183	status := http.StatusInternalServerError	1✔
184	hdr := w.Header()	1✔
185	hdr.Set("Content-Type", "application/json")	1✔
186	if errors.As(err, &tooManyRequests) {	2✔
187	status = http.StatusTooManyRequests	1✔
188	retryAfter := int64(math.Ceil(tooManyRequests.Delay.Abs().Seconds()))	1✔
189	hdr.Set("Retry-After", strconv.FormatInt(retryAfter, 10))	1✔
190	}	1✔
191	w.WriteHeader(status)	1✔
192	b, _ := json.Marshal(rest.Error{	1✔
193	Err: err.Error(),	1✔
194	RequestID: requestid.FromContext(ctx),	1✔
195	})	1✔
196	_, _ = w.Write(b)	1✔
197	}
198
199	// ServeHTTP implements a basic http.Handler so that handler can be used
200	// as a handler for the mux. It will only write on errors and is expected
201	// to continue to the actual handler on success.
202	func (h HTTPLimiter) ServeHTTP(w http.ResponseWriter, r http.Request) {	1✔
203	err := h.handleRequest(r)	1✔
204	if err != nil {	2✔
205	handleError(r.Context(), w, err)	1✔
206	}	1✔
207	}
208
209	// MiddlewareGin implements rate limiting as a middleware for gin-gonic
210	// web framework.
211	func (h HTTPLimiter) MiddlewareGin(c gin.Context) {	1✔
212	err := h.handleRequest(c.Request)	1✔
213	if err != nil {	2✔
214	_ = c.Error(err)	1✔
215	handleError(c.Request.Context(), c.Writer, err)	1✔
216	c.Abort()	1✔
217	}	1✔
218	}
219
220	type okReservation struct{}
221
222	func (k okReservation) OK() bool { return true }	1✔
223	func (k okReservation) Delay() time.Duration { return 0 }	×
224	func (k okReservation) Tokens() int64 { return math.MaxInt64 }	×
225
226	type rejectReservation struct{}
227
NEW 228	func (k rejectReservation) OK() bool { return false }	×
NEW 229	func (k rejectReservation) Delay() time.Duration { return math.MaxInt64 }	×
NEW 230	func (k rejectReservation) Tokens() int64 { return -1 }	×
231
232	func (m HTTPLimiter) Reserve(r http.Request) (Reservation, error) {	1✔
233	var b bytes.Buffer	1✔
234	var eventLimiter *templateEventLimiter	1✔
235	templateData := m.templateDataCallback(r)	1✔
236	ctx := r.Context()	1✔
237	h, _ := m.httpMux.Handler(r)	1✔
238	hh, ok := h.(matcher)	1✔
239	if ok && hh.groupTemplate != nil {	2✔
240	err := hh.groupTemplate.Execute(&b, templateData)	1✔
241	if err != nil {	2✔
242	return nil, fmt.Errorf("error executing ratelimit group template: %w", err)	1✔
243	}	1✔
244	eventLimiter = m.limiterGroups[b.String()]	1✔
245	if eventLimiter == nil {	2✔
246	return m.fallbackReservation, nil	1✔
247	}	1✔
248	b.Reset()	1✔
NEW 249	} else {	×
NEW 250	return m.fallbackReservation, nil	×
UNCOV 251	}	×
252	err := eventLimiter.eventTemplate.Execute(&b, templateData)	1✔
253	if err != nil {	2✔
254	return nil, fmt.Errorf("error executing template for event ID: %w", err)	1✔
255	}	1✔
256	return eventLimiter.limiter.ReserveEvent(ctx, b.String())	1✔
257	}

mendersoftware / mender-server / 1959198787

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