Skip to content

Commit

Permalink
feat(login): add rate-limit
Browse files Browse the repository at this point in the history
  • Loading branch information
adrienZ committed Oct 7, 2024
1 parent 17c5ff0 commit 0c285f8
Show file tree
Hide file tree
Showing 5 changed files with 139 additions and 177 deletions.
18 changes: 13 additions & 5 deletions src/runtime/core/core.ts
Original file line number Diff line number Diff line change
Expand Up @@ -10,11 +10,12 @@ import { EmailVerificationCodesRepository } from "./repositories/EmailVerificati
import { ResetPasswordTokensRepository } from "./repositories/ResetPasswordTokensRepository";
import type { SlipAuthPublicSession } from "../types";
import { defaultIdGenerationMethod, isValidEmail, defaultEmailVerificationCodeGenerationMethod, defaultHashPasswordMethod, defaultVerifyPasswordMethod, defaultResetPasswordTokenIdMethod, defaultResetPasswordTokenHashMethod } from "./email-and-password-utils";
import { InvalidEmailOrPasswordError, InvalidEmailToResetPasswordError, InvalidPasswordToResetError, InvalidUserIdToResetPasswordError, RateLimitLoginError, ResetPasswordTokenExpiredError, SlipAuthRateLimiterError, UnhandledError } from "./errors/SlipAuthError.js";
import { InvalidEmailOrPasswordError, InvalidEmailToResetPasswordError, InvalidPasswordToResetError, InvalidUserIdToResetPasswordError, RateLimitLoginError, ResetPasswordTokenExpiredError, UnhandledError } from "./errors/SlipAuthError.js";
import type { Database } from "db0";
import { createDate, isWithinExpirationDate, TimeSpan } from "oslo";
import type { H3Event } from "h3";
import { SlipAuthRateLimiters } from "./rate-limit/SlipAuthRateLimiters";
import type { Storage } from "unstorage";

export class SlipAuthCore {
readonly #db: Database;
Expand Down Expand Up @@ -114,17 +115,20 @@ export class SlipAuthCore {
throw new InvalidEmailOrPasswordError("no password oauth user");
}

const rateLimitRes = await this.#rateLimiters.login.consumeIncremental(existingUser.id);
if (rateLimitRes instanceof SlipAuthRateLimiterError) {
throw new RateLimitLoginError(rateLimitRes.data);
const [isNotRateLimited, rateLimitResult] = await this.#rateLimiters.login.check(existingUser.id);
if (!isNotRateLimited) {
throw new RateLimitLoginError({
msBeforeNext: (rateLimitResult.updatedAt + rateLimitResult.timeout * 1000) - Date.now(),
});
}

const validPassword = await this.#passwordHashingMethods.verify(existingUser.password, password);
if (!validPassword) {
await this.#rateLimiters.login.increment(existingUser.id);
throw new InvalidEmailOrPasswordError("login invalid password");
}

await this.#rateLimiters.login.delete(existingUser.id);
await this.#rateLimiters.login.reset(existingUser.id);
const sessionToLoginId = this.#createRandomSessionId();
const sessionToLogin = await this.#repos.sessions.insert({
sessionId: sessionToLoginId,
Expand Down Expand Up @@ -373,6 +377,10 @@ export class SlipAuthCore {
const methods = fn();
this.#passwordHashingMethods = methods;
},

setLoginRateLimiter: (fn: () => Storage) => {
this.#rateLimiters.login.storage = fn();
},
};

public getUser({ userId }: { userId: string }) {
Expand Down
7 changes: 2 additions & 5 deletions src/runtime/core/rate-limit/SlipAuthRateLimiters.ts
Original file line number Diff line number Diff line change
Expand Up @@ -5,11 +5,8 @@ export class SlipAuthRateLimiters {

constructor() {
this.login = new Throttler({
points: 5,
duration: 0,
},
{
storage: createThrottlerStorage("./data/ratelimiter/login"),
timeoutSeconds: [0, 1, 2, 4, 8, 16, 30, 60, 180, 300],
storage: createThrottlerStorage(),
});
}
}
102 changes: 59 additions & 43 deletions src/runtime/core/rate-limit/Throttler.ts
Original file line number Diff line number Diff line change
@@ -1,53 +1,69 @@
import { RateLimiterMemory, RateLimiterRes, type IRateLimiterOptions } from "rate-limiter-flexible";
import { createStorage, type Storage } from "unstorage";
import fsDriver from "unstorage/drivers/fs";
import { SlipAuthRateLimiterError } from "../errors/SlipAuthError";
import memoryDriver from "unstorage/drivers/memory";

export function createThrottlerStorage(base: string = "./.data/cache/ratelimit"): Storage<number> {
return createStorage<number>({
driver: fsDriver({ base }),
export function createThrottlerStorage(): Storage<ThrottlingCounter> {
return createStorage<ThrottlingCounter>({
driver: memoryDriver(),
});
}

export class Throttler extends RateLimiterMemory {
storage: Storage<number>;
initialBlockDurationSeconds = 5;
#incrementFactor = 2;
/**
* @see https://github.com/pilcrowOnPaper/astro-email-password-2fa/blob/main/src/lib/server/rate-limit.ts
*
* added suport for unstorage
*/
export class Throttler {
public timeoutSeconds: number[];

constructor(depOptions: IRateLimiterOptions, options: { storage: Storage<number>, initialBlockDurationSeconds?: number }) {
super(depOptions);
this.storage = options.storage;
this.initialBlockDurationSeconds = options.initialBlockDurationSeconds ?? this.initialBlockDurationSeconds;
public storage: Storage<ThrottlingCounter>;

constructor({ timeoutSeconds, storage }: { timeoutSeconds: number[], storage: Storage }) {
this.timeoutSeconds = timeoutSeconds;
this.storage = storage;
}

public async check(key: string): Promise<
[true] | [false, ThrottlingCounter]
> {
const counter = await this.storage.getItem(key) ?? null;
const now = Date.now();
if (counter === null) {
return [true];
}

const valid = now - counter.updatedAt >= this.timeoutSeconds[counter.timeout] * 1000;

if (valid) {
return [true];
}
else {
return [false, counter];
}
}

// eslint-disable-next-line @typescript-eslint/no-explicit-any
async consumeIncremental(key: string | number, pointsToConsume?: number, options?: { [key: string]: any }): Promise<RateLimiterRes | SlipAuthRateLimiterError> {
const strKey = key.toString();
const cachedBlockDuration = await this.storage.getItem(strKey) ?? (this.initialBlockDurationSeconds / this.#incrementFactor);

return super.consume(strKey, pointsToConsume, options)
.then((res) => {
this.storage.setItem(strKey, cachedBlockDuration);

if (res.remainingPoints <= 0) {
this.block(strKey, cachedBlockDuration);
}

return res;
})
.catch((error) => {
let msBeforeNext;

if (error instanceof RateLimiterRes && error.remainingPoints === 0) {
// Increase block duration and ensure it stays within the 32-bit signed integer range
const newBlockDuration = Math.min(cachedBlockDuration * 2, Number.MAX_SAFE_INTEGER);
this.storage.setItem(strKey, newBlockDuration);
msBeforeNext = newBlockDuration * 1000;
}

return new SlipAuthRateLimiterError({
msBeforeNext: msBeforeNext,
});
});
public async increment(key: string): Promise<void> {
let counter = await this.storage.getItem(key) ?? null;

const now = Date.now();
if (counter === null) {
counter = {
timeout: 0,
updatedAt: now,
};
await this.storage.setItem(key, counter);
return;
}
counter.updatedAt = now;
counter.timeout = Math.min(counter.timeout + 1, this.timeoutSeconds.length - 1);
await this.storage.setItem(key, counter);
}

public async reset(key: string): Promise<void> {
await this.storage.removeItem(key);
}
}

interface ThrottlingCounter {
timeout: number
updatedAt: number
}
91 changes: 65 additions & 26 deletions tests/rate-limit.test.ts
Original file line number Diff line number Diff line change
@@ -1,9 +1,12 @@
import { describe, it, expect, vi, beforeEach } from "vitest";
import { describe, it, expect, vi, beforeEach, afterAll } from "vitest";
import sqlite from "db0/connectors/better-sqlite3";
import { createDatabase } from "db0";
import { SlipAuthCore } from "../src/runtime/core/core";
import { autoSetupTestsDatabase, createH3Event, testTablesNames } from "./test-helpers";
import { RateLimitLoginError, SlipAuthRateLimiterError } from "../src/runtime/core/errors/SlipAuthError";
import { InvalidEmailOrPasswordError, RateLimitLoginError } from "../src/runtime/core/errors/SlipAuthError";
import { createThrottlerStorage } from "../src/runtime/core/rate-limit/Throttler";

const testStorage = createThrottlerStorage();

const db = createDatabase(sqlite({
name: "rate-limit.test",
Expand All @@ -28,10 +31,21 @@ const mocks = vi.hoisted(() => {
};
});

describe("rate limit", () => {
beforeEach(() => {
function wait(ms: number) {
return new Promise(resolve => setTimeout(resolve, ms));
}

describe.sequential("rate limit", () => {
afterAll(
async () => {
// await testStorage.clear();
await testStorage.dispose();
});

beforeEach(async () => {
mocks.userCreatedCount = 0;
mocks.sessionCreatedCount = 0;
await testStorage.clear();

auth = new SlipAuthCore(
db,
Expand All @@ -51,6 +65,8 @@ describe("rate limit", () => {
return `session-id-${mocks.sessionCreatedCount}`;
});

auth.setters.setLoginRateLimiter(() => testStorage);

function sanitizePassword(str: string) {
return str.replaceAll("$", "") + "$";
}
Expand All @@ -68,44 +84,67 @@ describe("rate limit", () => {
});

describe("login", () => {
it.only("should allow 5 failed tries", async () => {
it("should allow 2 failed tries", async () => {
await auth.register(createH3Event(), defaultInsert);
const doAttempt = () => auth.login(createH3Event(), {
email: defaultInsert.email,
password: defaultInsert.password + "123",
});

const results = await Promise.all([
doAttempt(),
doAttempt(),
doAttempt(),
doAttempt(),
doAttempt(),
]);
const t1 = doAttempt();
await expect(t1).rejects.toBeInstanceOf(InvalidEmailOrPasswordError);
// will not rate-limit because timeout is 0
const t2 = doAttempt();
await expect(t2).rejects.toBeInstanceOf(InvalidEmailOrPasswordError);
});

expect(results.every(res => res instanceof SlipAuthRateLimiterError)).toBe(false);
it("should rate-limit 3 failed tries", async () => {
await auth.register(createH3Event(), defaultInsert);
const doAttempt = () => auth.login(createH3Event(), {
email: defaultInsert.email,
password: defaultInsert.password + "123",
});

const t1 = doAttempt();
await expect(t1).rejects.toBeInstanceOf(InvalidEmailOrPasswordError);
// will not rate-limit because timeout is 0
const t2 = doAttempt();
await expect(t2).rejects.toBeInstanceOf(InvalidEmailOrPasswordError);

const t3 = doAttempt();
await expect(t3).rejects.toBeInstanceOf(RateLimitLoginError);
});

it("should rate-limit 6 failed tries", async () => {
it("should increment timeout", async () => {
await auth.register(createH3Event(), defaultInsert);
const doAttempt = () => auth.login(createH3Event(), {
email: defaultInsert.email,
password: defaultInsert.password + "123",
}).catch(e => e);
});
vi.useFakeTimers();

const t1 = doAttempt();
await expect(t1).rejects.toBeInstanceOf(InvalidEmailOrPasswordError);
// will not rate-limit because timeout is 0
const t2 = doAttempt();
await expect(t2).rejects.toBeInstanceOf(InvalidEmailOrPasswordError);

const t3 = await doAttempt().catch(e => JSON.parse(JSON.stringify(e)));
expect(t3).toMatchObject({ data: { msBeforeNext: 1000 } });

vi.advanceTimersByTime(1000);

const results = await Promise.all([
doAttempt(),
doAttempt(),
doAttempt(),
doAttempt(),
doAttempt(),
doAttempt(),
]);
// not rate-limited because after timeout of 1000
const t4 = doAttempt();
await expect(t4).rejects.toBeInstanceOf(InvalidEmailOrPasswordError);

console.log(results);

const t5 = await doAttempt().catch(e => JSON.parse(JSON.stringify(e)));
expect(t5).toMatchObject({ data: { msBeforeNext: 2000 } });

expect(results.every(res => res instanceof RateLimitLoginError)).toBe(false);
// use previous timeout
vi.advanceTimersByTime(500);
const t6 = await doAttempt().catch(e => JSON.parse(JSON.stringify(e)));
expect(t6).toMatchObject({ data: { msBeforeNext: 1500 } });
});
});
});
Loading

0 comments on commit 0c285f8

Please sign in to comment.