1943211337

Committed 23 Jul 2025 08:34AM UTC coverage: 65.677% (+0.2%) from 65.465%

Build # 1943211337

Build Type

Pull #816

gitlab-ci

Committed by

alfrunes

Commit Message

refactor: move integration fixtures to conftest.py and remove unused imports

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

Pull Request Pull Request #816: Fixes and imporvements to a few integraiton tests

Run Details

29280 of 44582 relevant lines covered (65.68%)

1.44 hits per line

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

0.0

/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 {	×
29	return "too many requests"	×
30	}	×
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 {	×
75	id := identity.FromContext(r.Context())	×
76	ctx := map[string]any{	×
77	"Identity": id,	×
78	// "Request": r, // not needed, but could be handy	×
79	}	×
80	return ctx	×
81	}	×
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) {	×
90	template, err := template.New("").Parse(defaultEventTemplate)	×
91	if err != nil {	×
92	return nil, fmt.Errorf("invalid eventTemplate: %w", err)	×
93	}	×
94	return &HTTPLimiter{	×
95	rootTemplate: template.New("").Option("missingkey=zero"),	×
96	httpMux: http.NewServeMux(),	×
97	templateDataCallback: defaultTemplateData,	×
98	defaultLimiter: &templateEventLimiter{	×
99	limiter: defaultLimiter,	×
100	eventTemplate: template,	×
101	},	×
102	limiterGroups: make(map[string]*templateEventLimiter),	×
103	}, nil	×
104	}
105
106	func (h HTTPLimiter) WithTemplateDataFunc(f func(http.Request) any) *HTTPLimiter {	×
107	h.templateDataCallback = f	×
108	return h	×
109	}	×
110
111	func (h HTTPLimiter) WithTemplateFuncs(funcs map[string]any) HTTPLimiter {	×
112	h.rootTemplate.Funcs(funcs)	×
113	return h	×
114	}	×
115
116	func (h HTTPLimiter) WithRewriteRequests(rewrite bool) HTTPLimiter {	×
117	h.rewriteRequests = rewrite	×
118	return h	×
119	}	×
120
121	func (h *HTTPLimiter) AddRateLimitGroup(limiter EventLimiter, group, eventTemplate string) error {	×
122	t, err := h.rootTemplate.Clone()	×
123	if err == nil {	×
124	_, err = t.Parse(eventTemplate)	×
125	}	×
126	if err != nil {	×
127	return fmt.Errorf("failed to compile event template: %w", err)	×
128	}	×
129	h.limiterGroups[group] = &templateEventLimiter{	×
130	limiter: limiter,	×
131	eventTemplate: t,	×
132	}	×
133	return nil	×
134	}
135
136	func (h *HTTPLimiter) AddMatchExpression(
137	pattern, groupTemplate string,
138	) error {	×
139	var (	×
140	t *template.Template	×
141	err error	×
142	)	×
143	if groupTemplate != "" {	×
144	// Compile eventTemplate:	×
145	t, err = h.rootTemplate.Clone()	×
146	if err == nil {	×
147	_, err = t.Parse(groupTemplate)	×
148	}	×
149	if err != nil {	×
150	return fmt.Errorf("error parsing group_template: %w", err)	×
151	}	×
152	}
153	limiterMatcher := matcher{	×
154	HTTPLimiter: h,	×
155	groupTemplate: t,	×
156	}	×
157	h.httpMux.Handle(pattern, limiterMatcher)	×
158	return nil	×
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 {	×
168	if h.rewriteRequests {	×
169	r = rest.RewriteForwardedRequest(r)	×
170	}	×
171	res, err := h.Reserve(r)	×
172	if err != nil {	×
173	return err	×
174	}	×
175	if res == nil \|\| res.OK() {	×
176	return nil	×
177	} else {	×
178	return &TooManyRequestsError{	×
179	Delay: res.Delay(),	×
180	}	×
181	}	×
182	}
183
184	func handleError(ctx context.Context, w http.ResponseWriter, err error) {	×
185	var tooManyRequests *TooManyRequestsError	×
186	status := http.StatusInternalServerError	×
187	hdr := w.Header()	×
188	hdr.Set("Content-Type", "application/json")	×
189	if errors.As(err, &tooManyRequests) {	×
190	status = http.StatusTooManyRequests	×
191	retryAfter := int64(math.Ceil(tooManyRequests.Delay.Abs().Seconds()))	×
192	hdr.Set("Retry-After", strconv.FormatInt(retryAfter, 10))	×
193	}	×
194	w.WriteHeader(status)	×
195	b, _ := json.Marshal(rest.Error{	×
196	Err: err.Error(),	×
197	RequestID: requestid.FromContext(ctx),	×
198	})	×
199	_, _ = w.Write(b)	×
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) {	×
206	err := h.handleRequest(r)	×
207	if err != nil {	×
208	handleError(r.Context(), w, err)	×
209	}	×
210	}
211
212	// MiddlewareGin implements rate limiting as a middleware for gin-gonic
213	// web framework.
214	func (h HTTPLimiter) MiddlewareGin(c gin.Context) {	×
215	err := h.handleRequest(c.Request)	×
216	if err != nil {	×
217	_ = c.Error(err)	×
218	handleError(c.Request.Context(), c.Writer, err)	×
219	c.Abort()	×
220	}	×
221	}
222
223	type okReservation struct{}
224
225	func (k okReservation) OK() bool { return true }	×
226	func (k okReservation) Delay() time.Duration { return 0 }	×
227	func (k okReservation) Tokens() int64 { return math.MaxInt64 }	×
228
229	func (m HTTPLimiter) Reserve(r http.Request) (Reservation, error) {	×
230	var b bytes.Buffer	×
231	eventLimiter := m.defaultLimiter	×
232	templateData := m.templateDataCallback(r)	×
233	ctx := r.Context()	×
234	h, _ := m.httpMux.Handler(r)	×
235	hh, ok := h.(matcher)	×
236	if ok {	×
237	if hh.groupTemplate != nil {	×
238	err := hh.groupTemplate.Execute(&b, templateData)	×
239	if err != nil {	×
240	return nil, fmt.Errorf("error executing ratelimit group template: %w", err)	×
241	}	×
242	eventLimiter = m.limiterGroups[b.String()]	×
243	if eventLimiter == nil {	×
244	return okReservation{}, nil	×
245	}	×
246	b.Reset()	×
247	}
248	}
249	err := eventLimiter.eventTemplate.Execute(&b, templateData)	×
250	if err != nil {	×
251	return nil, fmt.Errorf("error executing template for event ID: %w", err)	×
252	}	×
253	return eventLimiter.limiter.ReserveEvent(ctx, b.String())	×
254	}

mendersoftware / mender-server / 1943211337

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