1920843497

Committed 11 Jul 2025 11:00AM UTC coverage: 65.508% (+0.03%) from 65.48%

Build # 1920843497

Build Type

push

gitlab-ci

Committed by

web-flow

Commit Message

Merge pull request #782 from alfrunes/MEN-7743

MEN-7743: Distributed authentication rate limits (common parts)

Run Details

186 of 312 new or added lines in 6 files covered. (59.62%)

3 existing lines in 2 files now uncovered.

32347 of 49379 relevant lines covered (65.51%)

1.38 hits per line

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

87.14

/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

        // defaultLimiter is the default rate limiter if no group matches.
        defaultLimiter *templateEventLimiter
        // 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
}

// 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 a new HTTPLimiter using defaultLimiter as the
// default ratelimiter using defaultEventTemplate for grouping events.
// See HTTPLimiter.
func NewHTTPLimiter(
        defaultLimiter EventLimiter,
        defaultEventTemplate string,
) (*HTTPLimiter, error) {
        template, err := template.New("").Parse(defaultEventTemplate)
        if err != nil {
                return nil, fmt.Errorf("invalid eventTemplate: %w", err)
        }
        return &HTTPLimiter{
                rootTemplate:         template.New("").Option("missingkey=zero"),
                httpMux:              http.NewServeMux(),
                templateDataCallback: defaultTemplateData,
                defaultLimiter: &templateEventLimiter{
                        limiter:       defaultLimiter,
                        eventTemplate: template,
                },
                limiterGroups: make(map[string]*templateEventLimiter),
        }, nil
}

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) 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
}

func (h *HTTPLimiter) AddMatchExpression(
        pattern, groupTemplate string,
) error {
        var (
                t   *template.Template
                err error
        )
        if groupTemplate != "" {
                // 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 }

func (m *HTTPLimiter) Reserve(r *http.Request) (Reservation, error) {
        var b bytes.Buffer
        eventLimiter := m.defaultLimiter
        templateData := m.templateDataCallback(r)
        ctx := r.Context()
        h, _ := m.httpMux.Handler(r)
        hh, ok := h.(matcher)
        if ok {
                if 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 okReservation{}, nil
                        }
                        b.Reset()
                }
        }
        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	// defaultLimiter is the default rate limiter if no group matches.
64	defaultLimiter *templateEventLimiter
65	// limiterGroups maps the group name to rate limiters.
66	limiterGroups map[string]*templateEventLimiter
67
68	// rewriteRequests decides whether to autmatically rewrite the request
69	// URL using X-Forwarded-* headers.
70	rewriteRequests bool
71	}
72
73	// defaultTemplateData is the default callback for executing template strings.
74	func defaultTemplateData(r *http.Request) any {	1✔
75	id := identity.FromContext(r.Context())	1✔
76	ctx := map[string]any{	1✔
77	"Identity": id,	1✔
78	// "Request": r, // not needed, but could be handy	1✔
79	}	1✔
80	return ctx	1✔
81	}	1✔
82
83	// NewHTTPLimiter initializes a new HTTPLimiter using defaultLimiter as the
84	// default ratelimiter using defaultEventTemplate for grouping events.
85	// See HTTPLimiter.
86	func NewHTTPLimiter(
87	defaultLimiter EventLimiter,
88	defaultEventTemplate string,
89	) (*HTTPLimiter, error) {	1✔
90	template, err := template.New("").Parse(defaultEventTemplate)	1✔
91	if err != nil {	1✔
NEW 92	return nil, fmt.Errorf("invalid eventTemplate: %w", err)	×
NEW 93	}	×
94	return &HTTPLimiter{	1✔
95	rootTemplate: template.New("").Option("missingkey=zero"),	1✔
96	httpMux: http.NewServeMux(),	1✔
97	templateDataCallback: defaultTemplateData,	1✔
98	defaultLimiter: &templateEventLimiter{	1✔
99	limiter: defaultLimiter,	1✔
100	eventTemplate: template,	1✔
101	},	1✔
102	limiterGroups: make(map[string]*templateEventLimiter),	1✔
103	}, nil	1✔
104	}
105
106	func (h HTTPLimiter) WithTemplateDataFunc(f func(http.Request) any) *HTTPLimiter {	1✔
107	h.templateDataCallback = f	1✔
108	return h	1✔
109	}	1✔
110
NEW 111	func (h HTTPLimiter) WithTemplateFuncs(funcs map[string]any) HTTPLimiter {	×
NEW 112	h.rootTemplate.Funcs(funcs)	×
NEW 113	return h	×
NEW 114	}	×
115
NEW 116	func (h HTTPLimiter) WithRewriteRequests(rewrite bool) HTTPLimiter {	×
NEW 117	h.rewriteRequests = rewrite	×
NEW 118	return h	×
NEW 119	}	×
120
121	func (h *HTTPLimiter) AddRateLimitGroup(limiter EventLimiter, group, eventTemplate string) error {	1✔
122	t, err := h.rootTemplate.Clone()	1✔
123	if err == nil {	2✔
124	_, err = t.Parse(eventTemplate)	1✔
125	}	1✔
126	if err != nil {	1✔
NEW 127	return fmt.Errorf("failed to compile event template: %w", err)	×
NEW 128	}	×
129	h.limiterGroups[group] = &templateEventLimiter{	1✔
130	limiter: limiter,	1✔
131	eventTemplate: t,	1✔
132	}	1✔
133	return nil	1✔
134	}
135
136	func (h *HTTPLimiter) AddMatchExpression(
137	pattern, groupTemplate string,
138	) error {	1✔
139	var (	1✔
140	t *template.Template	1✔
141	err error	1✔
142	)	1✔
143	if groupTemplate != "" {	2✔
144	// Compile eventTemplate:	1✔
145	t, err = h.rootTemplate.Clone()	1✔
146	if err == nil {	2✔
147	_, err = t.Parse(groupTemplate)	1✔
148	}	1✔
149	if err != nil {	1✔
NEW 150	return fmt.Errorf("error parsing group_template: %w", err)	×
NEW 151	}	×
152	}
153	limiterMatcher := matcher{	1✔
154	HTTPLimiter: h,	1✔
155	groupTemplate: t,	1✔
156	}	1✔
157	h.httpMux.Handle(pattern, limiterMatcher)	1✔
158	return nil	1✔
159	}
160
161	// matcher is the HTTPHandle
162	type matcher struct {
163	*HTTPLimiter
164	groupTemplate *template.Template
165	}
166
167	func (h HTTPLimiter) handleRequest(r http.Request) error {	1✔
168	if h.rewriteRequests {	1✔
NEW 169	r = rest.RewriteForwardedRequest(r)	×
NEW 170	}	×
171	res, err := h.Reserve(r)	1✔
172	if err != nil {	2✔
173	return err	1✔
174	}	1✔
175	if res == nil \|\| res.OK() {	2✔
176	return nil	1✔
177	} else {	2✔
178	return &TooManyRequestsError{	1✔
179	Delay: res.Delay(),	1✔
180	}	1✔
181	}	1✔
182	}
183
184	func handleError(ctx context.Context, w http.ResponseWriter, err error) {	1✔
185	var tooManyRequests *TooManyRequestsError	1✔
186	status := http.StatusInternalServerError	1✔
187	hdr := w.Header()	1✔
188	hdr.Set("Content-Type", "application/json")	1✔
189	if errors.As(err, &tooManyRequests) {	2✔
190	status = http.StatusTooManyRequests	1✔
191	retryAfter := int64(math.Ceil(tooManyRequests.Delay.Abs().Seconds()))	1✔
192	hdr.Set("Retry-After", strconv.FormatInt(retryAfter, 10))	1✔
193	}	1✔
194	w.WriteHeader(status)	1✔
195	b, _ := json.Marshal(rest.Error{	1✔
196	Err: err.Error(),	1✔
197	RequestID: requestid.FromContext(ctx),	1✔
198	})	1✔
199	_, _ = w.Write(b)	1✔
200	}
201
202	// ServeHTTP implements a basic http.Handler so that handler can be used
203	// as a handler for the mux. It will only write on errors and is expected
204	// to continue to the actual handler on success.
205	func (h HTTPLimiter) ServeHTTP(w http.ResponseWriter, r http.Request) {	1✔
206	err := h.handleRequest(r)	1✔
207	if err != nil {	2✔
208	handleError(r.Context(), w, err)	1✔
209	}	1✔
210	}
211
212	// MiddlewareGin implements rate limiting as a middleware for gin-gonic
213	// web framework.
214	func (h HTTPLimiter) MiddlewareGin(c gin.Context) {	1✔
215	err := h.handleRequest(c.Request)	1✔
216	if err != nil {	2✔
217	_ = c.Error(err)	1✔
218	handleError(c.Request.Context(), c.Writer, err)	1✔
219	c.Abort()	1✔
220	}	1✔
221	}
222
223	type okReservation struct{}
224
225	func (k okReservation) OK() bool { return true }	1✔
NEW 226	func (k okReservation) Delay() time.Duration { return 0 }	×
NEW 227	func (k okReservation) Tokens() int64 { return math.MaxInt64 }	×
228
229	func (m HTTPLimiter) Reserve(r http.Request) (Reservation, error) {	1✔
230	var b bytes.Buffer	1✔
231	eventLimiter := m.defaultLimiter	1✔
232	templateData := m.templateDataCallback(r)	1✔
233	ctx := r.Context()	1✔
234	h, _ := m.httpMux.Handler(r)	1✔
235	hh, ok := h.(matcher)	1✔
236	if ok {	2✔
237	if hh.groupTemplate != nil {	2✔
238	err := hh.groupTemplate.Execute(&b, templateData)	1✔
239	if err != nil {	2✔
240	return nil, fmt.Errorf("error executing ratelimit group template: %w", err)	1✔
241	}	1✔
242	eventLimiter = m.limiterGroups[b.String()]	1✔
243	if eventLimiter == nil {	2✔
244	return okReservation{}, nil	1✔
245	}	1✔
246	b.Reset()	1✔
247	}
248	}
249	err := eventLimiter.eventTemplate.Execute(&b, templateData)	1✔
250	if err != nil {	2✔
251	return nil, fmt.Errorf("error executing template for event ID: %w", err)	1✔
252	}	1✔
253	return eventLimiter.limiter.ReserveEvent(ctx, b.String())	1✔
254	}

mendersoftware / mender-server / 1920843497

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