@@ -15,12 +15,12 @@ crate::pg_enum! {
1515}
1616
1717#[ derive( Debug , Clone , Copy ) ]
18- pub struct PublishRateLimit {
18+ pub struct RateLimiter {
1919 pub rate : Duration ,
2020 pub burst : i32 ,
2121}
2222
23- impl Default for PublishRateLimit {
23+ impl Default for RateLimiter {
2424 fn default ( ) -> Self {
2525 let minutes = dotenvy:: var ( "WEB_NEW_PKG_RATE_LIMIT_RATE_MINUTES" )
2626 . unwrap_or_default ( )
@@ -39,7 +39,7 @@ impl Default for PublishRateLimit {
3939 }
4040}
4141
42- impl PublishRateLimit {
42+ impl RateLimiter {
4343 pub fn check_rate_limit ( & self , uploader : i32 , conn : & mut PgConnection ) -> AppResult < ( ) > {
4444 let bucket = self . take_token ( uploader, Utc :: now ( ) . naive_utc ( ) , conn) ?;
4545 if bucket. tokens >= 1 {
@@ -67,8 +67,10 @@ impl PublishRateLimit {
6767 ) -> QueryResult < Bucket > {
6868 use self :: publish_limit_buckets:: dsl:: * ;
6969
70+ let performed_action = LimitedAction :: PublishNew ;
71+
7072 let burst: i32 = publish_rate_overrides:: table
71- . find ( ( uploader, LimitedAction :: PublishNew ) )
73+ . find ( ( uploader, performed_action ) )
7274 . filter (
7375 publish_rate_overrides:: expires_at
7476 . is_null ( )
@@ -88,8 +90,13 @@ impl PublishRateLimit {
8890 ) ;
8991
9092 diesel:: insert_into ( publish_limit_buckets)
91- . values ( ( user_id. eq ( uploader) , tokens. eq ( burst) , last_refill. eq ( now) ) )
92- . on_conflict ( user_id)
93+ . values ( (
94+ user_id. eq ( uploader) ,
95+ action. eq ( performed_action) ,
96+ tokens. eq ( burst) ,
97+ last_refill. eq ( now) ,
98+ ) )
99+ . on_conflict ( ( user_id, action) )
93100 . do_update ( )
94101 . set ( (
95102 tokens. eq ( least ( burst, greatest ( 0 , tokens - 1 ) + tokens_to_add) ) ,
@@ -126,7 +133,7 @@ mod tests {
126133 let conn = & mut pg_connection ( ) ;
127134 let now = now ( ) ;
128135
129- let rate = PublishRateLimit {
136+ let rate = RateLimiter {
130137 rate : Duration :: from_secs ( 1 ) ,
131138 burst : 10 ,
132139 } ;
@@ -139,7 +146,7 @@ mod tests {
139146 } ;
140147 assert_eq ! ( expected, bucket) ;
141148
142- let rate = PublishRateLimit {
149+ let rate = RateLimiter {
143150 rate : Duration :: from_millis ( 50 ) ,
144151 burst : 20 ,
145152 } ;
@@ -159,7 +166,7 @@ mod tests {
159166 let conn = & mut pg_connection ( ) ;
160167 let now = now ( ) ;
161168
162- let rate = PublishRateLimit {
169+ let rate = RateLimiter {
163170 rate : Duration :: from_secs ( 1 ) ,
164171 burst : 10 ,
165172 } ;
@@ -180,7 +187,7 @@ mod tests {
180187 let conn = & mut pg_connection ( ) ;
181188 let now = now ( ) ;
182189
183- let rate = PublishRateLimit {
190+ let rate = RateLimiter {
184191 rate : Duration :: from_secs ( 1 ) ,
185192 burst : 10 ,
186193 } ;
@@ -206,7 +213,7 @@ mod tests {
206213 NaiveDateTime :: parse_from_str ( "2019-03-19T21:11:24.620401" , "%Y-%m-%dT%H:%M:%S%.f" )
207214 . unwrap ( ) ;
208215
209- let rate = PublishRateLimit {
216+ let rate = RateLimiter {
210217 rate : Duration :: from_millis ( 100 ) ,
211218 burst : 10 ,
212219 } ;
@@ -228,7 +235,7 @@ mod tests {
228235 let conn = & mut pg_connection ( ) ;
229236 let now = now ( ) ;
230237
231- let rate = PublishRateLimit {
238+ let rate = RateLimiter {
232239 rate : Duration :: from_millis ( 100 ) ,
233240 burst : 10 ,
234241 } ;
@@ -250,7 +257,7 @@ mod tests {
250257 let conn = & mut pg_connection ( ) ;
251258 let now = now ( ) ;
252259
253- let rate = PublishRateLimit {
260+ let rate = RateLimiter {
254261 rate : Duration :: from_secs ( 1 ) ,
255262 burst : 10 ,
256263 } ;
@@ -274,7 +281,7 @@ mod tests {
274281 let conn = & mut pg_connection ( ) ;
275282 let now = now ( ) ;
276283
277- let rate = PublishRateLimit {
284+ let rate = RateLimiter {
278285 rate : Duration :: from_secs ( 1 ) ,
279286 burst : 10 ,
280287 } ;
@@ -297,7 +304,7 @@ mod tests {
297304 let conn = & mut pg_connection ( ) ;
298305 let now = now ( ) ;
299306
300- let rate = PublishRateLimit {
307+ let rate = RateLimiter {
301308 rate : Duration :: from_secs ( 1 ) ,
302309 burst : 10 ,
303310 } ;
@@ -320,7 +327,7 @@ mod tests {
320327 let conn = & mut pg_connection ( ) ;
321328 let now = now ( ) ;
322329
323- let rate = PublishRateLimit {
330+ let rate = RateLimiter {
324331 rate : Duration :: from_secs ( 1 ) ,
325332 burst : 10 ,
326333 } ;
@@ -347,7 +354,7 @@ mod tests {
347354 let conn = & mut pg_connection ( ) ;
348355 let now = now ( ) ;
349356
350- let rate = PublishRateLimit {
357+ let rate = RateLimiter {
351358 rate : Duration :: from_secs ( 1 ) ,
352359 burst : 10 ,
353360 } ;
0 commit comments