summaryrefslogtreecommitdiff
blob: 3aef43e5c17b40bc3fc656489ae5fc97af46e91d (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
This is a combination of upstream commits:
b368a675955707db4e940da29a1043871a3781b6
21ea03e0f874991086d2f1bcdc285216878bd566

Fixes https://bugs.gentoo.org/808791
Fixes https://bugs.gentoo.org/810317

diff --git a/go.mod b/go.mod
index 548c0590f..22a8833e2 100644
--- a/go.mod
+++ b/go.mod
@@ -150,7 +150,7 @@ require (
 	github.com/ryanuber/go-glob v1.0.0
 	github.com/samuel/go-zookeeper v0.0.0-20190923202752-2cc03de413da
 	github.com/sasha-s/go-deadlock v0.2.0
-	github.com/sethvargo/go-limiter v0.3.0
+	github.com/sethvargo/go-limiter v0.7.1
 	github.com/shirou/gopsutil v3.21.5+incompatible
 	github.com/stretchr/testify v1.7.0
 	github.com/tidwall/pretty v1.0.1 // indirect
diff --git a/go.sum b/go.sum
index c5b3b410d..98a5dd0a8 100644
--- a/go.sum
+++ b/go.sum
@@ -1120,8 +1120,8 @@ github.com/sean-/conswriter v0.0.0-20180208195008-f5ae3917a627/go.mod h1:7zjs06q
 github.com/sean-/pager v0.0.0-20180208200047-666be9bf53b5/go.mod h1:BeybITEsBEg6qbIiqJ6/Bqeq25bCLbL7YFmpaFfJDuM=
 github.com/sean-/seed v0.0.0-20170313163322-e2103e2c3529 h1:nn5Wsu0esKSJiIVhscUtVbo7ada43DJhG55ua/hjS5I=
 github.com/sean-/seed v0.0.0-20170313163322-e2103e2c3529/go.mod h1:DxrIzT+xaE7yg65j358z/aeFdxmN0P9QXhEzd20vsDc=
-github.com/sethvargo/go-limiter v0.3.0 h1:yRMc+Qs2yqw6YJp6UxrO2iUs6DOSq4zcnljbB7/rMns=
-github.com/sethvargo/go-limiter v0.3.0/go.mod h1:C0kbSFbiriE5k2FFOe18M1YZbAR2Fiwf72uGu0CXCcU=
+github.com/sethvargo/go-limiter v0.7.1 h1:wWNhTj0pxjyJ7wuJHpRJpYwJn+bUnjYfw2a85eu5w9U=
+github.com/sethvargo/go-limiter v0.7.1/go.mod h1:C0kbSFbiriE5k2FFOe18M1YZbAR2Fiwf72uGu0CXCcU=
 github.com/shirou/gopsutil v3.21.5+incompatible h1:OloQyEerMi7JUrXiNzy8wQ5XN+baemxSl12QgIzt0jc=
 github.com/shirou/gopsutil v3.21.5+incompatible/go.mod h1:5b4v6he4MtMOwMlS0TUMTu2PcXUg8+E1lC7eC3UO/RA=
 github.com/shopspring/decimal v0.0.0-20180709203117-cd690d0c9e24 h1:pntxY8Ary0t43dCZ5dqY4YTJCObLY1kIXl0uzMv+7DE=
diff --git a/http/util.go b/http/util.go
index 0550a93c7..cbb364843 100644
--- a/http/util.go
+++ b/http/util.go
@@ -48,7 +48,7 @@ func rateLimitQuotaWrapping(handler http.Handler, core *vault.Core) http.Handler
 			return
 		}
 
-		quotaResp, err := core.ApplyRateLimitQuota(&quotas.Request{
+		quotaResp, err := core.ApplyRateLimitQuota(r.Context(), &quotas.Request{
 			Type:          quotas.TypeRateLimit,
 			Path:          path,
 			MountPath:     strings.TrimPrefix(core.MatchingMount(r.Context(), path), ns.Path),
diff --git a/vault/core.go b/vault/core.go
index 3b6e461fd..27741e8c6 100644
--- a/vault/core.go
+++ b/vault/core.go
@@ -2744,7 +2744,7 @@ func (c *Core) setupQuotas(ctx context.Context, isPerfStandby bool) error {
 
 // ApplyRateLimitQuota checks the request against all the applicable quota rules.
 // If the given request's path is exempt, no rate limiting will be applied.
-func (c *Core) ApplyRateLimitQuota(req *quotas.Request) (quotas.Response, error) {
+func (c *Core) ApplyRateLimitQuota(ctx context.Context, req *quotas.Request) (quotas.Response, error) {
 	req.Type = quotas.TypeRateLimit
 
 	resp := quotas.Response{
@@ -2758,7 +2758,7 @@ func (c *Core) ApplyRateLimitQuota(req *quotas.Request) (quotas.Response, error)
 			return resp, nil
 		}
 
-		return c.quotaManager.ApplyQuota(req)
+		return c.quotaManager.ApplyQuota(ctx, req)
 	}
 
 	return resp, nil
diff --git a/vault/quotas/quotas.go b/vault/quotas/quotas.go
index 68cc72f9f..80ee59521 100644
--- a/vault/quotas/quotas.go
+++ b/vault/quotas/quotas.go
@@ -168,7 +168,7 @@ type Manager struct {
 // Quota represents the common properties of every quota type
 type Quota interface {
 	// allow checks the if the request is allowed by the quota type implementation.
-	allow(*Request) (Response, error)
+	allow(context.Context, *Request) (Response, error)
 
 	// quotaID is the identifier of the quota rule
 	quotaID() string
@@ -181,7 +181,7 @@ type Quota interface {
 
 	// close defines any cleanup behavior that needs to be executed when a quota
 	// rule is deleted.
-	close() error
+	close(context.Context) error
 
 	// handleRemount takes in the new mount path in the quota
 	handleRemount(string)
@@ -287,7 +287,7 @@ func (m *Manager) setQuotaLocked(ctx context.Context, qType string, quota Quota,
 	// If there already exists an entry in the db, remove that first.
 	if raw != nil {
 		quota := raw.(Quota)
-		if err := quota.close(); err != nil {
+		if err := quota.close(ctx); err != nil {
 			return err
 		}
 		err = txn.Delete(qType, raw)
@@ -518,7 +518,7 @@ func (m *Manager) DeleteQuota(ctx context.Context, qType string, name string) er
 	}
 
 	quota := raw.(Quota)
-	if err := quota.close(); err != nil {
+	if err := quota.close(ctx); err != nil {
 		return err
 	}
 
@@ -541,7 +541,7 @@ func (m *Manager) DeleteQuota(ctx context.Context, qType string, name string) er
 // ApplyQuota runs the request against any quota rule that is applicable to it. If
 // there are multiple quota rule that matches the request parameters, rule that
 // takes precedence will be used to allow/reject the request.
-func (m *Manager) ApplyQuota(req *Request) (Response, error) {
+func (m *Manager) ApplyQuota(ctx context.Context, req *Request) (Response, error) {
 	var resp Response
 
 	quota, err := m.QueryQuota(req)
@@ -562,7 +562,7 @@ func (m *Manager) ApplyQuota(req *Request) (Response, error) {
 		return resp, nil
 	}
 
-	return quota.allow(req)
+	return quota.allow(ctx, req)
 }
 
 // SetEnableRateLimitAuditLogging updates the operator preference regarding the
diff --git a/vault/quotas/quotas_rate_limit.go b/vault/quotas/quotas_rate_limit.go
index 64117b002..ad58b2af3 100644
--- a/vault/quotas/quotas_rate_limit.go
+++ b/vault/quotas/quotas_rate_limit.go
@@ -1,6 +1,7 @@
 package quotas
 
 import (
+	"context"
 	"encoding/hex"
 	"fmt"
 	"math"
@@ -264,7 +265,7 @@ func (rlq *RateLimitQuota) QuotaName() string {
 // returned if the request ID or address is empty. If the path is exempt, the
 // quota will not be evaluated. Otherwise, the client rate limiter is retrieved
 // by address and the rate limit quota is checked against that limiter.
-func (rlq *RateLimitQuota) allow(req *Request) (Response, error) {
+func (rlq *RateLimitQuota) allow(ctx context.Context, req *Request) (Response, error) {
 	resp := Response{
 		Headers: make(map[string]string),
 	}
@@ -300,7 +301,11 @@ func (rlq *RateLimitQuota) allow(req *Request) (Response, error) {
 		}
 	}
 
-	limit, remaining, reset, allow := rlq.store.Take(req.ClientAddress)
+	limit, remaining, reset, allow, err := rlq.store.Take(ctx, req.ClientAddress)
+	if err != nil {
+		return resp, err
+	}
+
 	resp.Allowed = allow
 	resp.Headers[httplimit.HeaderRateLimitLimit] = strconv.FormatUint(limit, 10)
 	resp.Headers[httplimit.HeaderRateLimitRemaining] = strconv.FormatUint(remaining, 10)
@@ -320,13 +325,13 @@ func (rlq *RateLimitQuota) allow(req *Request) (Response, error) {
 
 // close stops the current running client purge loop.
 // It should be called with the write lock held.
-func (rlq *RateLimitQuota) close() error {
+func (rlq *RateLimitQuota) close(ctx context.Context) error {
 	if rlq.purgeBlocked {
 		close(rlq.closePurgeBlockedCh)
 	}
 
 	if rlq.store != nil {
-		return rlq.store.Close()
+		return rlq.store.Close(ctx)
 	}
 
 	return nil
diff --git a/vault/quotas/quotas_rate_limit_test.go b/vault/quotas/quotas_rate_limit_test.go
index 27225e338..21f35dac3 100644
--- a/vault/quotas/quotas_rate_limit_test.go
+++ b/vault/quotas/quotas_rate_limit_test.go
@@ -37,7 +37,7 @@ func TestNewRateLimitQuota(t *testing.T) {
 			err := tc.rlq.initialize(logging.NewVaultLogger(log.Trace), metricsutil.BlackholeSink())
 			require.Equal(t, tc.expectErr, err != nil, err)
 			if err == nil {
-				require.Nil(t, tc.rlq.close())
+				require.Nil(t, tc.rlq.close(context.Background()))
 			}
 		})
 	}
@@ -46,7 +46,7 @@ func TestNewRateLimitQuota(t *testing.T) {
 func TestRateLimitQuota_Close(t *testing.T) {
 	rlq := NewRateLimitQuota("test-rate-limiter", "qa", "/foo/bar", 16.7, time.Second, time.Minute)
 	require.NoError(t, rlq.initialize(logging.NewVaultLogger(log.Trace), metricsutil.BlackholeSink()))
-	require.NoError(t, rlq.close())
+	require.NoError(t, rlq.close(context.Background()))
 
 	time.Sleep(time.Second) // allow enough time for purgeClientsLoop to receive on closeCh
 	require.False(t, rlq.getPurgeBlocked(), "expected blocked client purging to be disabled after explicit close")
@@ -66,14 +66,14 @@ func TestRateLimitQuota_Allow(t *testing.T) {
 	}
 
 	require.NoError(t, rlq.initialize(logging.NewVaultLogger(log.Trace), metricsutil.BlackholeSink()))
-	defer rlq.close()
+	defer rlq.close(context.Background())
 
 	var wg sync.WaitGroup
 
 	reqFunc := func(addr string, atomicNumAllow, atomicNumFail *atomic.Int32) {
 		defer wg.Done()
 
-		resp, err := rlq.allow(&Request{ClientAddress: addr})
+		resp, err := rlq.allow(context.Background(), &Request{ClientAddress: addr})
 		if err != nil {
 			return
 		}
@@ -141,7 +141,7 @@ func TestRateLimitQuota_Allow_WithBlock(t *testing.T) {
 	}
 
 	require.NoError(t, rlq.initialize(logging.NewVaultLogger(log.Trace), metricsutil.BlackholeSink()))
-	defer rlq.close()
+	defer rlq.close(context.Background())
 	require.True(t, rlq.getPurgeBlocked())
 
 	var wg sync.WaitGroup
@@ -149,7 +149,7 @@ func TestRateLimitQuota_Allow_WithBlock(t *testing.T) {
 	reqFunc := func(addr string, atomicNumAllow, atomicNumFail *atomic.Int32) {
 		defer wg.Done()
 
-		resp, err := rlq.allow(&Request{ClientAddress: addr})
+		resp, err := rlq.allow(context.Background(), &Request{ClientAddress: addr})
 		if err != nil {
 			return
 		}
@@ -221,5 +221,5 @@ func TestRateLimitQuota_Update(t *testing.T) {
 	require.NoError(t, qm.SetQuota(context.Background(), TypeRateLimit.String(), quota, true))
 	require.NoError(t, qm.SetQuota(context.Background(), TypeRateLimit.String(), quota, true))
 
-	require.Nil(t, quota.close())
+	require.Nil(t, quota.close(context.Background()))
 }
diff --git a/vault/quotas/quotas_util.go b/vault/quotas/quotas_util.go
index dc2fcdfac..7c0732f67 100644
--- a/vault/quotas/quotas_util.go
+++ b/vault/quotas/quotas_util.go
@@ -40,7 +40,7 @@ func (*entManager) Reset() error {
 
 type LeaseCountQuota struct{}
 
-func (l LeaseCountQuota) allow(request *Request) (Response, error) {
+func (l LeaseCountQuota) allow(_ context.Context, _ *Request) (Response, error) {
 	panic("implement me")
 }
 
@@ -56,7 +56,7 @@ func (l LeaseCountQuota) initialize(logger log.Logger, sink *metricsutil.Cluster
 	panic("implement me")
 }
 
-func (l LeaseCountQuota) close() error {
+func (l LeaseCountQuota) close(_ context.Context) error {
 	panic("implement me")
 }