/
rate_limit.go
284 lines (237 loc) · 7.41 KB
/
rate_limit.go
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
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
package rules
import (
"context"
"fmt"
"time"
"github.com/lightninglabs/lightning-terminal/firewalldb"
"github.com/lightninglabs/lightning-terminal/litrpc"
"google.golang.org/protobuf/proto"
"gopkg.in/macaroon-bakery.v2/bakery"
)
var (
// Compile-time checks to ensure that RateLimit, RateLimitMgr
// and RateLimitEnforcer implement the appropriate Manager, Enforcer
// and Values interface.
_ Manager = (*RateLimitMgr)(nil)
_ Enforcer = (*RateLimitEnforcer)(nil)
_ Values = (*RateLimit)(nil)
)
// RateLimitName is the string identifier of the RateLimitMgr values.
const RateLimitName = "rate-limit"
// RateLimitMgr represents the rate limit values.
type RateLimitMgr struct{}
// Stop cleans up the resources held by the manager.
//
// NOTE: This is part of the Manager interface.
func (r *RateLimitMgr) Stop() error {
return nil
}
// NewEnforcer constructs a new RateLimit rule enforcer using the passed values
// and config.
//
// NOTE: This is part of the Manager interface.
func (r *RateLimitMgr) NewEnforcer(cfg Config, values Values) (Enforcer,
error) {
limits, ok := values.(*RateLimit)
if !ok {
return nil, fmt.Errorf("values must be of type "+
"RateLimit, got %T", values)
}
return &RateLimitEnforcer{
rateLimitConfig: cfg,
RateLimit: limits,
}, nil
}
// NewValueFromProto converts the given proto value into a RateLimit Value
// object.
//
// NOTE: This is part of the Manager interface.
func (r *RateLimitMgr) NewValueFromProto(v *litrpc.RuleValue) (Values, error) {
rv, ok := v.Value.(*litrpc.RuleValue_RateLimit)
if !ok {
return nil, fmt.Errorf("incorrect RuleValue type")
}
budget := rv.RateLimit
readLim := budget.ReadLimit
writeLim := budget.WriteLimit
return &RateLimit{
ReadLimit: &Rate{
Iterations: readLim.Iterations,
NumHours: readLim.NumHours,
},
WriteLimit: &Rate{
Iterations: writeLim.Iterations,
NumHours: writeLim.NumHours,
},
}, nil
}
// EmptyValue returns a new RateLimit instance.
func (r *RateLimitMgr) EmptyValue() Values {
return &RateLimit{}
}
// rateLimitConfig is the config required by RateLimitMgr. It can be derived
// from the main rules Config struct.
type rateLimitConfig interface {
GetActionsDB() firewalldb.ActionsDB
GetMethodPerms() func(string) ([]bakery.Op, bool)
}
// RateLimitEnforcer enforces requests and responses against a RateLimit rule.
type RateLimitEnforcer struct {
rateLimitConfig
*RateLimit
}
// HandleResponse handles and possible alters a response. This is a noop for the
// RateLimitMgr values.
//
// NOTE: this is part of the Rule interface.
func (r *RateLimitEnforcer) HandleResponse(_ context.Context, _ string,
_ proto.Message) (proto.Message, error) {
return nil, nil
}
// HandleRequest checks the validity of a request. It checks if the request is a
// read or a write request. Then, using the past actions DB, it determines if
// letting this request through would violate the rate limit rules.
//
// NOTE: this is part of the Rule interface.
func (r *RateLimitEnforcer) HandleRequest(ctx context.Context, uri string,
_ proto.Message) (proto.Message, error) {
// First, we need to classify if this is a read or write call.
read := r.isRead(uri)
// Based on the above, we can extract the relevant rate limit values
// that apply for this call.
rateLim := r.WriteLimit
if read {
rateLim = r.ReadLimit
}
// Now we need to go and count all the previous read or write actions.
actions, err := r.GetActionsDB().ListActions(ctx)
if err != nil {
return nil, err
}
// Determine the start time of the actions window.
startTime := time.Now().Add(
-time.Duration(rateLim.NumHours) * time.Hour,
)
// Now count all relevant actions which have taken place after the
// start time.
var count uint32
for _, action := range actions {
if read != r.isRead(action.Method) {
continue
}
if action.PerformedAt.Before(startTime) {
continue
}
count++
}
if count >= rateLim.Iterations {
return nil, fmt.Errorf("too many requests received")
}
return nil, nil
}
// HandleErrorResponse handles and possible alters an error. This is a noop for
// the RateLimitEnforcer rule.
//
// NOTE: this is part of the Enforcer interface.
func (r *RateLimitEnforcer) HandleErrorResponse(_ context.Context, _ string,
_ error) (error, error) {
return nil, nil
}
// isRead is a helper that returns true if the given method/URI only requires
// read-permissions and false otherwise.
func (r *RateLimitEnforcer) isRead(method string) bool {
perms, ok := r.GetMethodPerms()(method)
if !ok {
return false
}
for _, p := range perms {
if p.Action != "read" {
return false
}
}
return true
}
// Rate describes a rate limit in iterations per number of hours.
type Rate struct {
Iterations uint32 `json:"iterations"`
NumHours uint32 `json:"num_hours"`
}
// RateLimit represents the rules values.
type RateLimit struct {
WriteLimit *Rate `json:"write_limit"`
ReadLimit *Rate `json:"read_limit"`
}
// VerifySane checks that the value of the values is ok given the min and max
// allowed values.
//
// NOTE: this is part of the Values interface.
func (r *RateLimit) VerifySane(minVal, maxVal Values) error {
minRL, ok := minVal.(*RateLimit)
if !ok {
return fmt.Errorf("min value is not of type RateLimit")
}
maxRL, ok := maxVal.(*RateLimit)
if !ok {
return fmt.Errorf("max value is not of type RateLimit")
}
// Check that our read limit is between the min and max.
if r.ReadLimit.lessThan(minRL.ReadLimit) ||
maxRL.ReadLimit.lessThan(r.ReadLimit) {
return fmt.Errorf("read limit is not between the min and max")
}
// Check that our write limit is between the min and max.
if r.WriteLimit.lessThan(minRL.WriteLimit) ||
maxRL.WriteLimit.lessThan(r.WriteLimit) {
return fmt.Errorf("write limit is not between the min and max")
}
return nil
}
// lessThan is a helper function that checks if the current rate is less than
// another rate.
func (r *Rate) lessThan(other *Rate) bool {
return float64(r.Iterations)/float64(r.NumHours) <
float64(other.Iterations)/float64(other.NumHours)
}
// RuleName returns the name of the rule that these values are to be used with.
//
// NOTE: this is part of the Values interface.
func (r *RateLimit) RuleName() string {
return RateLimitName
}
// ToProto converts the rule Values to the litrpc counterpart.
//
// NOTE: this is part of the Values interface.
func (r *RateLimit) ToProto() *litrpc.RuleValue {
return &litrpc.RuleValue{
Value: &litrpc.RuleValue_RateLimit{
RateLimit: &litrpc.RateLimit{
ReadLimit: &litrpc.Rate{
Iterations: r.ReadLimit.Iterations,
NumHours: r.ReadLimit.NumHours,
},
WriteLimit: &litrpc.Rate{
Iterations: r.WriteLimit.Iterations,
NumHours: r.WriteLimit.NumHours,
},
},
},
}
}
// PseudoToReal attempts to convert any appropriate pseudo fields in the rule
// Values to their corresponding real values. It uses the passed PrivacyMapDB to
// find the real values. This is a no-op for the RateLimit rule.
//
// NOTE: this is part of the Values interface.
func (r *RateLimit) PseudoToReal(_ firewalldb.PrivacyMapDB) (Values,
error) {
return r, nil
}
// RealToPseudo converts the rule Values to a new one that uses pseudo keys,
// channel IDs, channel points etc. It returns a map of real to pseudo strings
// that should be persisted. This is a no-op for the RateLimit rule.
//
// NOTE: this is part of the Values interface.
func (r *RateLimit) RealToPseudo(_ firewalldb.PrivacyMapReader) (Values,
map[string]string, error) {
return r, nil, nil
}