#219

Committed 01 Nov 2024 02:50PM UTC coverage: 95.359% (+1.5%) from 93.83%

Build # #219

Build Type

push

github

Committed by

web-flow

Commit Message

Merge pull request #577 from alovajs/changeset-release/main

ci: release

Run Details

1698 of 1787 branches covered (95.02%)

Branch coverage included in aggregate %.

5801 of 6077 relevant lines covered (95.46%)

223.07 hits per line

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

90.85

/packages/server/src/hooks/rateLimit.ts

import HookedMethod from '@/HookedMethod';
import { createServerHook } from '@/helper';
import { createAssert, getOptions, isFn, uuid } from '@alova/shared';
import { AlovaGenerics, AlovaGlobalCacheAdapter, Method } from 'alova';
import { IRateLimiterStoreOptions, RateLimiterRes } from 'rate-limiter-flexible';
import RateLimiterStoreAbstract from 'rate-limiter-flexible/lib/RateLimiterStoreAbstract.js';

type StoreResult = [points: number, expireTime: number];

interface LimitHandlerOptions<AG extends AlovaGenerics> {
  /** storage key */
  key?: string | ((method: Method<AG>) => string);
}

/**
 * Rate limit, there can only be a maximum of [points] requests within [duration] seconds
 *
 * Usage scenarios:
 * 1. Request restrictions. For example, when node acts as an intermediate layer to request downstream services, under an API with serious resource consumption, it is restricted through IP to avoid consuming downstream server resources.
 * 2. Prevent password brute force cracking. When the downstream server throws login errors multiple times in a row, restrict it by IP or user name.
 * 3. As a sending limit for sendCaptcha, it prevents users from frequently sending verification codes.
 */
export interface RateLimitOptions {
  /**
   * The maximum quantity that can be consumed within the duration
   * @default 4
   */
  points?: number;

  /**
   * Points reset time, unit ms
   * @default 4000
   */
  duration?: number;

  /**
   * Namespace, prevents conflicts when multiple limiters use the same storage medium
   */
  keyPrefix?: string;

  /**
   * The following two parameters are consumption interval control
   * */
  execEvenly?: boolean;
  execEvenlyMinDelayMs?: number;

  /**
   * After reaching the rate limit, [blockDuration]ms will be extended. For example, if the password is incorrect 5 times within 1 hour, it will be locked for 24 hours. This 24 hours is this parameter.
   */
  blockDuration?: number;

  /**
   * Custom storage adapter, defaults to methodObj.context.l2Cache if not set
   */
  storage?: AlovaGlobalCacheAdapter;
}

const assert = createAssert('RateLimit');

class RateLimiterStore extends RateLimiterStoreAbstract {
  constructor(
    protected storage: AlovaGlobalCacheAdapter,
    options: IRateLimiterStoreOptions
  ) {
    super(options);
  }

  /**
   * parses raw data from store to RateLimiterRes object.
   */
  _getRateLimiterRes(key: string | number, changedPoints: number, result: StoreResult) {
    const [consumed = 0, expireTime = 0] = result ?? [];
    const msBeforeNext = expireTime > 0 ? Math.max(expireTime - Date.now(), 0) : -1;
    const isFirstInDuration = !consumed || changedPoints === consumed;
    const currentConsumedPoints = isFirstInDuration ? changedPoints : consumed;

    const res = new RateLimiterRes(
      Math.max(0, this.points - currentConsumedPoints),
      msBeforeNext,
      isFirstInDuration ? changedPoints : consumed,
      isFirstInDuration
    );

    return res;
  }

  async _upsert(key: string | number, points: number, msDuration: number, forceExpire = false) {
    key = key.toString();
    const isNeverExpired = msDuration <= 0;
    const expireTime = isNeverExpired ? -1 : Date.now() + msDuration;
    const newRecord = [points, expireTime];
    if (!forceExpire) {
      const [oldPoints = 0, oldExpireTime = 0] = ((await this.storage.get(key)) ?? []) as StoreResult;

      // if haven't expired yet
      if (isNeverExpired || (!isNeverExpired && oldExpireTime > Date.now())) {
        newRecord[0] += oldPoints;
      }

      if (!isNeverExpired && oldExpireTime > Date.now()) {
        newRecord[1] = oldExpireTime;
      }
    }

    await this.storage.set(key.toString(), newRecord);

    // need to return the record after upsert
    return newRecord;
  }

  /**
   * returns raw data by key or null if there is no key or expired.
   */
  async _get(key: string | number) {
    return Promise.resolve(this.storage.get(key.toString())).then(res => {
      if (!res) {
        return null;
      }

      const [, expireTime] = res as StoreResult;

      // if have expire time and it has expired
      if (expireTime > 0 && expireTime <= Date.now()) {
        return null;
      }

      return res;
    });
  }

  /**
   * returns true on deleted, false if key is not found.
   */
  async _delete(key: string | number) {
    try {
      await this.storage.remove(key.toString());
    } catch {
      return false;
    }

    return true;
  }
}

/**
 * The method instance modified by rateLimit, its extension method corresponds to the method of creating an instance in rate-limit-flexible, and the key is the key specified by calling rateLimit.
 * AlovaServerHook can currently only return unextended method types. It has not been changed to customizable returned extended method types.
 */
export class LimitedMethod<AG extends AlovaGenerics> extends HookedMethod<AG> {
  private keyGetter: () => string;

  constructor(
    method: Method<AG>,
    limiterKey: string | ((method: Method<AG>) => string),
    protected limiter: RateLimiterStore
  ) {
    super(method, force => this.consume().then(() => method.send(force)));
    this.keyGetter = isFn(limiterKey) ? () => limiterKey(method) : () => limiterKey;
  }

  private getLimiterKey() {
    return this.keyGetter();
  }

  /**
   * Get RateLimiterRes or null.
   */
  get(options?: { [key: string]: any }) {
    return this.limiter.get(this.getLimiterKey(), options);
  }

  /**
   * Set points by key.
   */
  set(points: number, msDuration: number) {
    return this.limiter.set(this.getLimiterKey(), points, msDuration / 1000);
  }

  /**
   * @param points default is 1
   */
  consume(points?: number) {
    return this.limiter.consume(this.getLimiterKey(), points);
  }

  /**
   * Increase number of consumed points in current duration.
   * @param points penalty points
   */
  penalty(points: number) {
    return this.limiter.penalty(this.getLimiterKey(), points);
  }

  /**
   * Decrease number of consumed points in current duration.
   * @param points reward points
   */
  reward(points: number) {
    return this.limiter.reward(this.getLimiterKey(), points);
  }

  /**
   * Block key for ms.
   */
  block(msDuration: number) {
    return this.limiter.block(this.getLimiterKey(), msDuration / 1000);
  }

  /**
   * Reset consumed points.
   */
  delete() {
    return this.limiter.delete(this.getLimiterKey());
  }
}

export function createRateLimiter(options: RateLimitOptions = {}) {
  const { points = 4, duration = 4 * 1000, keyPrefix, execEvenly, execEvenlyMinDelayMs, blockDuration } = options;

  const limitedMethodWrapper = createServerHook(
    <AG extends AlovaGenerics>(method: Method<AG>, handlerOptions: LimitHandlerOptions<AG> = {}) => {
      const { key = uuid() } = handlerOptions;
      const storage = options.storage ?? getOptions(method).l2Cache;

      assert(!!storage, 'storage is not define');
      const limiter = new RateLimiterStore(storage!, {
        points,
        duration: Math.floor(duration / 1000),
        keyPrefix,
        execEvenly,
        execEvenlyMinDelayMs,
        blockDuration: blockDuration ? Math.floor(blockDuration / 1000) : blockDuration,
        storeClient: {}
      });

      return new LimitedMethod<AG>(method, key, limiter);
    }
  );

  return limitedMethodWrapper;
}

export default createRateLimiter;

1	import HookedMethod from '@/HookedMethod';	1✔
2	import { createServerHook } from '@/helper';
3	import { createAssert, getOptions, isFn, uuid } from '@alova/shared';
4	import { AlovaGenerics, AlovaGlobalCacheAdapter, Method } from 'alova';
5	import { IRateLimiterStoreOptions, RateLimiterRes } from 'rate-limiter-flexible';
6	import RateLimiterStoreAbstract from 'rate-limiter-flexible/lib/RateLimiterStoreAbstract.js';
7
8	type StoreResult = [points: number, expireTime: number];
9
10	interface LimitHandlerOptions<AG extends AlovaGenerics> {
11	/** storage key */
12	key?: string \| ((method: Method<AG>) => string);
13	}
14
15	/**
16	* Rate limit, there can only be a maximum of [points] requests within [duration] seconds
17	*
18	* Usage scenarios:
19	* 1. Request restrictions. For example, when node acts as an intermediate layer to request downstream services, under an API with serious resource consumption, it is restricted through IP to avoid consuming downstream server resources.
20	* 2. Prevent password brute force cracking. When the downstream server throws login errors multiple times in a row, restrict it by IP or user name.
21	* 3. As a sending limit for sendCaptcha, it prevents users from frequently sending verification codes.
22	*/
23	export interface RateLimitOptions {
24	/**
25	* The maximum quantity that can be consumed within the duration
26	* @default 4
27	*/
28	points?: number;
29
30	/**
31	* Points reset time, unit ms
32	* @default 4000
33	*/
34	duration?: number;
35
36	/**
37	* Namespace, prevents conflicts when multiple limiters use the same storage medium
38	*/
39	keyPrefix?: string;
40
41	/**
42	* The following two parameters are consumption interval control
43	* */
44	execEvenly?: boolean;
45	execEvenlyMinDelayMs?: number;
46
47	/**
48	* After reaching the rate limit, [blockDuration]ms will be extended. For example, if the password is incorrect 5 times within 1 hour, it will be locked for 24 hours. This 24 hours is this parameter.
49	*/
50	blockDuration?: number;
51
52	/**
53	* Custom storage adapter, defaults to methodObj.context.l2Cache if not set
54	*/
55	storage?: AlovaGlobalCacheAdapter;
56	}
57
58	const assert = createAssert('RateLimit');	1✔
59
60	class RateLimiterStore extends RateLimiterStoreAbstract {	1✔
61	constructor(	1✔
62	protected storage: AlovaGlobalCacheAdapter,	21✔
63	options: IRateLimiterStoreOptions	21✔
64	) {	21✔
65	super(options);	21✔
66	}	21✔
67
68	/**
69	* parses raw data from store to RateLimiterRes object.
70	*/
71	_getRateLimiterRes(key: string \| number, changedPoints: number, result: StoreResult) {	1✔
72	const [consumed = 0, expireTime = 0] = result ?? [];	40!
73	const msBeforeNext = expireTime > 0 ? Math.max(expireTime - Date.now(), 0) : -1;	40!
74	const isFirstInDuration = !consumed \|\| changedPoints === consumed;	40✔
75	const currentConsumedPoints = isFirstInDuration ? changedPoints : consumed;	40✔
76
77	const res = new RateLimiterRes(	40✔
78	Math.max(0, this.points - currentConsumedPoints),	40✔
79	msBeforeNext,	40✔
80	isFirstInDuration ? changedPoints : consumed,	40✔
81	isFirstInDuration	40✔
82	);	40✔
83
84	return res;	40✔
85	}	40✔
86
87	async _upsert(key: string \| number, points: number, msDuration: number, forceExpire = false) {	1✔
88	key = key.toString();	29✔
89	const isNeverExpired = msDuration <= 0;	29✔
90	const expireTime = isNeverExpired ? -1 : Date.now() + msDuration;	29!
91	const newRecord = [points, expireTime];	29✔
92	if (!forceExpire) {	29✔
93	const [oldPoints = 0, oldExpireTime = 0] = ((await this.storage.get(key)) ?? []) as StoreResult;	27✔
94
95	// if haven't expired yet
96	if (isNeverExpired \|\| (!isNeverExpired && oldExpireTime > Date.now())) {	27✔
97	newRecord[0] += oldPoints;	17✔
98	}	17✔
99
100	if (!isNeverExpired && oldExpireTime > Date.now()) {	27✔
101	newRecord[1] = oldExpireTime;	17✔
102	}	17✔
103	}	27✔
104
105	await this.storage.set(key.toString(), newRecord);	29✔
106
107	// need to return the record after upsert
108	return newRecord;	29✔
109	}	29✔
110
111	/**
112	* returns raw data by key or null if there is no key or expired.
113	*/
114	async _get(key: string \| number) {	1✔
115	return Promise.resolve(this.storage.get(key.toString())).then(res => {	20✔
116	if (!res) {	20✔
117	return null;	6✔
118	}	6✔
119
120	const [, expireTime] = res as StoreResult;	14✔
121
122	// if have expire time and it has expired
123	if (expireTime > 0 && expireTime <= Date.now()) {	20✔
124	return null;	1✔
125	}	1✔
126
127	return res;	13✔
128	});	20✔
129	}	20✔
130
131	/**
132	* returns true on deleted, false if key is not found.
133	*/
134	async _delete(key: string \| number) {	1✔
135	try {	2✔
136	await this.storage.remove(key.toString());	2✔
137	} catch {	2!
138	return false;	×
139	}	×
140
141	return true;	2✔
142	}	2✔
143	}	1✔
144
145	/**
146	* The method instance modified by rateLimit, its extension method corresponds to the method of creating an instance in rate-limit-flexible, and the key is the key specified by calling rateLimit.
147	* AlovaServerHook can currently only return unextended method types. It has not been changed to customizable returned extended method types.
148	*/
149	export class LimitedMethod<AG extends AlovaGenerics> extends HookedMethod<AG> {	1✔
150	private keyGetter: () => string;
151
152	constructor(	1✔
153	method: Method<AG>,	21✔
154	limiterKey: string \| ((method: Method<AG>) => string),	21✔
155	protected limiter: RateLimiterStore	21✔
156	) {	21✔
157	super(method, force => this.consume().then(() => method.send(force)));	21✔
158	this.keyGetter = isFn(limiterKey) ? () => limiterKey(method) : () => limiterKey;	21✔
159	}	21✔
160
161	private getLimiterKey() {	1✔
162	return this.keyGetter();	49✔
163	}	49✔
164
165	/**
166	* Get RateLimiterRes or null.
167	*/
168	get(options?: { [key: string]: any }) {	1✔
169	return this.limiter.get(this.getLimiterKey(), options);	20✔
170	}	20✔
171
172	/**
173	* Set points by key.
174	*/
175	set(points: number, msDuration: number) {	1✔
176	return this.limiter.set(this.getLimiterKey(), points, msDuration / 1000);	×
177	}	×
178
179	/**
180	* @param points default is 1
181	*/
182	consume(points?: number) {	1✔
183	return this.limiter.consume(this.getLimiterKey(), points);	27✔
184	}	27✔
185
186	/**
187	* Increase number of consumed points in current duration.
188	* @param points penalty points
189	*/
190	penalty(points: number) {	1✔
191	return this.limiter.penalty(this.getLimiterKey(), points);	×
192	}	×
193
194	/**
195	* Decrease number of consumed points in current duration.
196	* @param points reward points
197	*/
198	reward(points: number) {	1✔
199	return this.limiter.reward(this.getLimiterKey(), points);	×
200	}	×
201
202	/**
203	* Block key for ms.
204	*/
205	block(msDuration: number) {	1✔
206	return this.limiter.block(this.getLimiterKey(), msDuration / 1000);	×
207	}	×
208
209	/**
210	* Reset consumed points.
211	*/
212	delete() {	1✔
213	return this.limiter.delete(this.getLimiterKey());	2✔
214	}	2✔
215	}	1✔
216
217	export function createRateLimiter(options: RateLimitOptions = {}) {	1✔
218	const { points = 4, duration = 4 * 1000, keyPrefix, execEvenly, execEvenlyMinDelayMs, blockDuration } = options;	6✔
219
220	const limitedMethodWrapper = createServerHook(	6✔
221	<AG extends AlovaGenerics>(method: Method<AG>, handlerOptions: LimitHandlerOptions<AG> = {}) => {	6✔
222	const { key = uuid() } = handlerOptions;	21✔
223	const storage = options.storage ?? getOptions(method).l2Cache;	21✔
224
225	assert(!!storage, 'storage is not define');	21✔
226	const limiter = new RateLimiterStore(storage!, {	21✔
227	points,	21✔
228	duration: Math.floor(duration / 1000),	21✔
229	keyPrefix,	21✔
230	execEvenly,	21✔
231	execEvenlyMinDelayMs,	21✔
232	blockDuration: blockDuration ? Math.floor(blockDuration / 1000) : blockDuration,	21✔
233	storeClient: {}	21✔
234	});	21✔
235
236	return new LimitedMethod<AG>(method, key, limiter);	21✔
237	}	21✔
238	);	6✔
239
240	return limitedMethodWrapper;	6✔
241	}	6✔
242
243	export default createRateLimiter;	1✔

alovajs / alova / #219

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