* add improved logging and decrease ttl * fix test TTL time --------- Co-authored-by: Tushar Bhatt <95581504+TusharBhatt1@users.noreply.github.com>
336 lines
8.9 KiB
TypeScript
336 lines
8.9 KiB
TypeScript
import type { RatelimitResponse } from "@unkey/ratelimit";
|
|
import { vi, describe, it, expect, beforeEach, afterEach } from "vitest";
|
|
|
|
import { hashAPIKey } from "@calcom/features/ee/api-keys/lib/apiKeys";
|
|
import { RedisService } from "@calcom/features/redis/RedisService";
|
|
import prisma from "@calcom/prisma";
|
|
|
|
import { handleAutoLock } from "../autoLock";
|
|
|
|
// Mock the dependencies
|
|
vi.mock("@calcom/features/redis/RedisService");
|
|
vi.mock("@calcom/features/ee/api-keys/lib/apiKeys", () => ({
|
|
hashAPIKey: vi.fn((key) => `hashed_${key}`),
|
|
}));
|
|
vi.mock("@calcom/prisma", () => ({
|
|
default: {
|
|
user: {
|
|
update: vi.fn(),
|
|
},
|
|
apiKey: {
|
|
findUnique: vi.fn(),
|
|
},
|
|
},
|
|
}));
|
|
|
|
describe("autoLock", () => {
|
|
const mockRedis = {
|
|
get: vi.fn(),
|
|
set: vi.fn(),
|
|
del: vi.fn(),
|
|
expire: vi.fn(),
|
|
};
|
|
|
|
beforeEach(() => {
|
|
// Reset all mocks before each test
|
|
vi.clearAllMocks();
|
|
vi.mocked(RedisService).mockImplementation(() => mockRedis as any);
|
|
|
|
// Mock environment variables
|
|
process.env.UPSTASH_REDIS_REST_TOKEN = "test-token";
|
|
process.env.UPSTASH_REDIS_REST_URL = "test-url";
|
|
});
|
|
|
|
afterEach(() => {
|
|
delete process.env.UPSTASH_REDIS_REST_TOKEN;
|
|
delete process.env.UPSTASH_REDIS_REST_URL;
|
|
});
|
|
|
|
describe("handleAutoLock", () => {
|
|
it("should return early if Upstash env variables are not set", async () => {
|
|
delete process.env.UPSTASH_REDIS_REST_TOKEN;
|
|
|
|
const rateLimitResponse: RatelimitResponse = {
|
|
success: false,
|
|
remaining: 0,
|
|
limit: 5,
|
|
reset: 0,
|
|
};
|
|
|
|
await handleAutoLock({
|
|
identifier: "test@example.com",
|
|
identifierType: "email",
|
|
rateLimitResponse,
|
|
});
|
|
expect(mockRedis.get).not.toHaveBeenCalled();
|
|
});
|
|
|
|
it("should handle identifier with keyword correctly", async () => {
|
|
const rateLimitResponse: RatelimitResponse = {
|
|
success: false,
|
|
remaining: 0,
|
|
limit: 5,
|
|
reset: 0,
|
|
};
|
|
|
|
mockRedis.get.mockResolvedValue("0");
|
|
|
|
await handleAutoLock({
|
|
identifier: "addSecondaryEmail.test@example.com",
|
|
identifierType: "email",
|
|
rateLimitResponse,
|
|
identifierKeyword: "addSecondaryEmail",
|
|
});
|
|
|
|
expect(mockRedis.get).toHaveBeenCalledWith("autolock:email:addSecondaryEmail:test@example.com.count");
|
|
});
|
|
|
|
it("should increment counter when below threshold", async () => {
|
|
const rateLimitResponse: RatelimitResponse = {
|
|
success: false,
|
|
remaining: 0,
|
|
limit: 5,
|
|
reset: 0,
|
|
};
|
|
|
|
mockRedis.get.mockResolvedValue("2");
|
|
|
|
await handleAutoLock({
|
|
identifier: "test@example.com",
|
|
identifierType: "email",
|
|
rateLimitResponse,
|
|
});
|
|
|
|
expect(mockRedis.set).toHaveBeenCalledWith("autolock:email:test@example.com.count", "3");
|
|
expect(mockRedis.expire).toHaveBeenCalledWith("autolock:email:test@example.com.count", 1800);
|
|
expect(prisma.user.update).not.toHaveBeenCalled();
|
|
});
|
|
|
|
it("should lock user when threshold is reached", async () => {
|
|
const rateLimitResponse: RatelimitResponse = {
|
|
success: false,
|
|
remaining: 0,
|
|
limit: 5,
|
|
reset: 0,
|
|
};
|
|
|
|
mockRedis.get.mockResolvedValue("4");
|
|
|
|
await handleAutoLock({
|
|
identifier: "test@example.com",
|
|
identifierType: "email",
|
|
rateLimitResponse,
|
|
});
|
|
|
|
expect(prisma.user.update).toHaveBeenCalledWith({
|
|
where: { email: "test@example.com" },
|
|
data: { locked: true },
|
|
select: {
|
|
id: true,
|
|
email: true,
|
|
username: true,
|
|
},
|
|
});
|
|
expect(mockRedis.del).toHaveBeenCalledWith("autolock:email:test@example.com.count");
|
|
});
|
|
|
|
it("should respect custom threshold and duration", async () => {
|
|
const rateLimitResponse: RatelimitResponse = {
|
|
success: false,
|
|
remaining: 0,
|
|
limit: 5,
|
|
reset: 0,
|
|
};
|
|
|
|
mockRedis.get.mockResolvedValue("1");
|
|
|
|
await handleAutoLock({
|
|
identifier: "test@example.com",
|
|
identifierType: "email",
|
|
rateLimitResponse,
|
|
autolockThreshold: 3,
|
|
autolockDuration: 30 * 60 * 1000, // 30 minutes
|
|
});
|
|
|
|
expect(mockRedis.set).toHaveBeenCalledWith("autolock:email:test@example.com.count", "2");
|
|
expect(mockRedis.expire).toHaveBeenCalledWith("autolock:email:test@example.com.count", 1800);
|
|
expect(prisma.user.update).not.toHaveBeenCalled();
|
|
});
|
|
|
|
it("should lock API key and associated user", async () => {
|
|
const rateLimitResponse: RatelimitResponse = {
|
|
success: false,
|
|
remaining: 0,
|
|
limit: 5,
|
|
reset: 0,
|
|
};
|
|
|
|
mockRedis.get.mockResolvedValue("4");
|
|
const testApiKey = "test_api_key_123";
|
|
|
|
// Mock API key lookup
|
|
vi.mocked(prisma.apiKey.findUnique).mockResolvedValue({
|
|
hashedKey: `hashed_${testApiKey}`,
|
|
user: { id: 456 },
|
|
} as any);
|
|
|
|
await handleAutoLock({
|
|
identifier: testApiKey,
|
|
identifierType: "apiKey",
|
|
rateLimitResponse,
|
|
});
|
|
|
|
// Verify the API key was hashed
|
|
expect(hashAPIKey).toHaveBeenCalledWith(testApiKey);
|
|
|
|
// Verify the associated user was locked
|
|
expect(prisma.user.update).toHaveBeenCalledWith({
|
|
where: { id: 456 },
|
|
data: { locked: true },
|
|
select: {
|
|
id: true,
|
|
email: true,
|
|
username: true,
|
|
},
|
|
});
|
|
});
|
|
|
|
it("should throw error when API key has no associated user", async () => {
|
|
const rateLimitResponse: RatelimitResponse = {
|
|
success: false,
|
|
remaining: 0,
|
|
limit: 5,
|
|
reset: 0,
|
|
};
|
|
|
|
mockRedis.get.mockResolvedValue("4"); // Over threshold to trigger lock
|
|
const testApiKey = "test_api_key_123";
|
|
|
|
// Mock API key lookup with no user
|
|
vi.mocked(prisma.apiKey.findUnique).mockResolvedValue(null);
|
|
|
|
// Expect handleAutoLock to throw the error from lockUser
|
|
await expect(async () => {
|
|
await handleAutoLock({
|
|
identifier: testApiKey,
|
|
identifierType: "apiKey",
|
|
rateLimitResponse,
|
|
});
|
|
}).rejects.toThrow("No user found for this API key.");
|
|
});
|
|
|
|
it("should not increment counter when rate limit is successful", async () => {
|
|
const rateLimitResponse: RatelimitResponse = {
|
|
success: true,
|
|
remaining: 5,
|
|
limit: 5,
|
|
reset: 0,
|
|
};
|
|
|
|
await handleAutoLock({
|
|
identifier: "test@example.com",
|
|
identifierType: "email",
|
|
rateLimitResponse,
|
|
});
|
|
|
|
expect(mockRedis.get).not.toHaveBeenCalled();
|
|
expect(mockRedis.set).not.toHaveBeenCalled();
|
|
});
|
|
|
|
it("should initialize counter when it doesn't exist", async () => {
|
|
const rateLimitResponse: RatelimitResponse = {
|
|
success: false,
|
|
remaining: 0,
|
|
limit: 5,
|
|
reset: 0,
|
|
};
|
|
|
|
mockRedis.get.mockResolvedValue(null);
|
|
|
|
await handleAutoLock({
|
|
identifier: "test@example.com",
|
|
identifierType: "email",
|
|
rateLimitResponse,
|
|
});
|
|
|
|
expect(mockRedis.set).toHaveBeenCalledWith("autolock:email:test@example.com.count", "1");
|
|
expect(mockRedis.expire).toHaveBeenCalledWith("autolock:email:test@example.com.count", 1800);
|
|
});
|
|
|
|
it("should handle Redis errors gracefully", async () => {
|
|
const rateLimitResponse: RatelimitResponse = {
|
|
success: false,
|
|
remaining: 0,
|
|
limit: 5,
|
|
reset: 0,
|
|
};
|
|
|
|
mockRedis.get.mockRejectedValue(new Error("Redis connection error"));
|
|
|
|
const result = await handleAutoLock({
|
|
identifier: "test@example.com",
|
|
identifierType: "email",
|
|
rateLimitResponse,
|
|
});
|
|
|
|
expect(result).toBe(false);
|
|
});
|
|
});
|
|
|
|
describe("lockUser", () => {
|
|
it("should lock user by email", async () => {
|
|
const rateLimitResponse: RatelimitResponse = {
|
|
success: false,
|
|
remaining: 0,
|
|
limit: 5,
|
|
reset: 0,
|
|
};
|
|
|
|
mockRedis.get.mockResolvedValue("4");
|
|
|
|
await handleAutoLock({
|
|
identifier: "test@example.com",
|
|
identifierType: "email",
|
|
rateLimitResponse,
|
|
});
|
|
|
|
expect(prisma.user.update).toHaveBeenCalledWith({
|
|
where: { email: "test@example.com" },
|
|
data: { locked: true },
|
|
select: {
|
|
id: true,
|
|
email: true,
|
|
username: true,
|
|
},
|
|
});
|
|
});
|
|
|
|
it("should lock user by userId", async () => {
|
|
const rateLimitResponse: RatelimitResponse = {
|
|
success: false,
|
|
remaining: 0,
|
|
limit: 5,
|
|
reset: 0,
|
|
};
|
|
|
|
mockRedis.get.mockResolvedValue("4");
|
|
|
|
await handleAutoLock({
|
|
identifier: "123",
|
|
identifierType: "userId",
|
|
rateLimitResponse,
|
|
});
|
|
|
|
expect(prisma.user.update).toHaveBeenCalledWith({
|
|
where: { id: 123 },
|
|
data: { locked: true },
|
|
select: {
|
|
id: true,
|
|
email: true,
|
|
username: true,
|
|
},
|
|
});
|
|
});
|
|
});
|
|
});
|