diff --git a/packages/backend-app-api/src/services/implementations/httpRouter/createCredentialsBarrier.ts b/packages/backend-app-api/src/services/implementations/httpRouter/createCredentialsBarrier.ts index 6e5e06c59e..72b93220ee 100644 --- a/packages/backend-app-api/src/services/implementations/httpRouter/createCredentialsBarrier.ts +++ b/packages/backend-app-api/src/services/implementations/httpRouter/createCredentialsBarrier.ts @@ -15,6 +15,7 @@ */ import { + CacheService, HttpAuthService, HttpRouterServiceAuthPolicy, RootConfigService, @@ -24,6 +25,7 @@ import { RequestHandler } from 'express'; import { pathToRegexp } from 'path-to-regexp'; import { rateLimit } from 'express-rate-limit'; import { readDurationFromConfig } from '@backstage/config'; +import { RateLimitStore } from './rateLimitStore'; export function createPathPolicyPredicate(policyPath: string) { if (policyPath === '/' || policyPath === '*') { @@ -42,11 +44,12 @@ export function createPathPolicyPredicate(policyPath: string) { export function createCredentialsBarrier(options: { httpAuth: HttpAuthService; config: RootConfigService; + cache: CacheService; }): { middleware: RequestHandler; addAuthPolicy: (policy: HttpRouterServiceAuthPolicy) => void; } { - const { httpAuth, config } = options; + const { httpAuth, config, cache } = options; const disableDefaultAuthPolicy = config.getOptionalBoolean( 'backend.auth.dangerouslyDisableDefaultAuthPolicy', @@ -78,7 +81,8 @@ export function createCredentialsBarrier(options: { windowMs, max, standardHeaders: true, // Return rate limit info in the `RateLimit-*` headers - legacyHeaders: false, // Disable the `X-RateLimit-*` headers + legacyHeaders: false, // Disable the `X-RateLimit-*` headers, + store: RateLimitStore.fromOptions({ cache }), }); const middleware: RequestHandler = (req, res, next) => { diff --git a/packages/backend-app-api/src/services/implementations/httpRouter/httpRouterServiceFactory.ts b/packages/backend-app-api/src/services/implementations/httpRouter/httpRouterServiceFactory.ts index e72d091b96..50f9e5c97a 100644 --- a/packages/backend-app-api/src/services/implementations/httpRouter/httpRouterServiceFactory.ts +++ b/packages/backend-app-api/src/services/implementations/httpRouter/httpRouterServiceFactory.ts @@ -39,20 +39,32 @@ export const httpRouterServiceFactory = createServiceFactory( (options?: HttpRouterFactoryOptions) => ({ service: coreServices.httpRouter, deps: { + cache: coreServices.cache, plugin: coreServices.pluginMetadata, config: coreServices.rootConfig, lifecycle: coreServices.lifecycle, rootHttpRouter: coreServices.rootHttpRouter, httpAuth: coreServices.httpAuth, }, - async factory({ httpAuth, config, plugin, rootHttpRouter, lifecycle }) { + async factory({ + httpAuth, + config, + cache, + plugin, + rootHttpRouter, + lifecycle, + }) { const getPath = options?.getPath ?? (id => `/api/${id}`); const path = getPath(plugin.getId()); const router = PromiseRouter(); rootHttpRouter.use(path, router); - const credentialsBarrier = createCredentialsBarrier({ httpAuth, config }); + const credentialsBarrier = createCredentialsBarrier({ + httpAuth, + config, + cache, + }); router.use(createLifecycleMiddleware({ lifecycle })); router.use(credentialsBarrier.middleware); diff --git a/packages/backend-app-api/src/services/implementations/httpRouter/rateLimitStore.test.ts b/packages/backend-app-api/src/services/implementations/httpRouter/rateLimitStore.test.ts new file mode 100644 index 0000000000..843cb83ed1 --- /dev/null +++ b/packages/backend-app-api/src/services/implementations/httpRouter/rateLimitStore.test.ts @@ -0,0 +1,184 @@ +/* + * Copyright 2024 The Backstage Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +import { ServiceMock, mockServices } from '@backstage/backend-test-utils'; +import { RateLimitStore } from './rateLimitStore'; +import { CacheService } from '@backstage/backend-plugin-api'; + +jest.setSystemTime(Date.parse('2024-03-15T08:39:11.869Z')); + +describe('RateLimitStore', () => { + let cacheServiceMock: ServiceMock; + let rateLimitStore: RateLimitStore; + + beforeEach(() => { + jest.clearAllMocks(); + cacheServiceMock = mockServices.cache.mock(); + rateLimitStore = RateLimitStore.fromOptions({ + prefix: 'rl_', + cache: cacheServiceMock, + }); + }); + + it('should initialize with default options', () => { + expect(rateLimitStore.windowMs).toBe(15 * 60 * 1000); + }); + + it('should initialize with custom options', () => { + rateLimitStore.init({ windowMs: 10 * 60 * 1000 }); + expect(rateLimitStore.windowMs).toBe(10 * 60 * 1000); + }); + + it('should get rate limit info for existing key', async () => { + const key = 'existingKey'; + const existingValue = { + totalHits: 5, + resetTime: new Date(), + }; + + cacheServiceMock.get.mockResolvedValue({ + totalHits: existingValue.totalHits, + resetTime: existingValue.resetTime.getTime(), + }); + + const returnedValue = await rateLimitStore.get(key); + + expect(returnedValue).toEqual(existingValue); + expect(cacheServiceMock.get).toHaveBeenCalledWith( + rateLimitStore.prefixKey(key), + ); + }); + + it('should get default rate limit info for non-existing key', async () => { + const key = 'nonExistingKey'; + const defaultValue = { + totalHits: 0, + resetTime: new Date(Date.now() + rateLimitStore.windowMs), + }; + + cacheServiceMock.get.mockResolvedValue(undefined); + + const returnedValue = await rateLimitStore.get(key); + + expect(returnedValue).toEqual(defaultValue); + expect(cacheServiceMock.get).toHaveBeenCalledWith( + rateLimitStore.prefixKey(key), + ); + }); + + it('should increment rate limit for existing key', async () => { + const key = 'exisitingKey'; + const existingValue = { + totalHits: 5, + resetTime: new Date(), + }; + + cacheServiceMock.get.mockResolvedValue({ + totalHits: existingValue.totalHits, + resetTime: existingValue.resetTime.getTime(), + }); + + const incrementedValue = await rateLimitStore.increment(key); + + expect(incrementedValue.totalHits).toBe(existingValue.totalHits + 1); + expect(incrementedValue.resetTime).toEqual(existingValue.resetTime); + expect(cacheServiceMock.set).toHaveBeenCalledWith( + rateLimitStore.prefixKey(key), + { + totalHits: existingValue.totalHits + 1, + resetTime: existingValue.resetTime.getTime(), + }, + { ttl: rateLimitStore.windowMs }, + ); + }); + + it('should increment rate limit for non-existing key', async () => { + const key = 'nonExistingKey'; + + cacheServiceMock.get.mockResolvedValue(undefined); + + const incrementedValue = await rateLimitStore.increment(key); + + expect(incrementedValue.totalHits).toBe(1); + expect(incrementedValue.resetTime).toBeInstanceOf(Date); + expect(cacheServiceMock.set).toHaveBeenCalledWith( + rateLimitStore.prefixKey(key), + { + totalHits: incrementedValue.totalHits, + resetTime: incrementedValue.resetTime?.getTime(), + }, + { ttl: rateLimitStore.windowMs }, + ); + }); + + it('should decrement rate limit for existing key', async () => { + const key = 'exisitingKey'; + const existingValue = { + totalHits: 5, + resetTime: new Date(), + }; + + cacheServiceMock.get.mockResolvedValue({ + totalHits: existingValue.totalHits, + resetTime: existingValue.resetTime.getTime(), + }); + + await rateLimitStore.decrement(key); + + expect(cacheServiceMock.set).toHaveBeenCalledWith( + rateLimitStore.prefixKey(key), + { + totalHits: existingValue.totalHits - 1, + resetTime: existingValue.resetTime.getTime(), + }, + { ttl: rateLimitStore.windowMs }, + ); + }); + + it('should not decrement rate limit below zero', async () => { + const key = 'testKey'; + const existingValue = { + totalHits: 0, + resetTime: new Date(), + }; + + cacheServiceMock.get.mockResolvedValue({ + totalHits: existingValue.totalHits, + resetTime: existingValue.resetTime.getTime(), + }); + + await rateLimitStore.decrement(key); + + expect(cacheServiceMock.set).toHaveBeenCalledWith( + rateLimitStore.prefixKey(key), + { + totalHits: 0, + resetTime: existingValue.resetTime.getTime(), + }, + { ttl: rateLimitStore.windowMs }, + ); + }); + + it('should reset rate limit for existing key', async () => { + const key = 'exitingKey'; + + await rateLimitStore.resetKey(key); + + expect(cacheServiceMock.delete).toHaveBeenCalledWith( + rateLimitStore.prefixKey(key), + ); + }); +}); diff --git a/packages/backend-app-api/src/services/implementations/httpRouter/rateLimitStore.ts b/packages/backend-app-api/src/services/implementations/httpRouter/rateLimitStore.ts new file mode 100644 index 0000000000..c111fbc2d2 --- /dev/null +++ b/packages/backend-app-api/src/services/implementations/httpRouter/rateLimitStore.ts @@ -0,0 +1,119 @@ +/* + * Copyright 2024 The Backstage Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +import { CacheService } from '@backstage/backend-plugin-api'; +import type { + Store, + Options, + IncrementResponse, + ClientRateLimitInfo, +} from 'express-rate-limit'; + +type CacheStoreOptions = { + /** + * Optional field to differentiate hit countswhen multiple rate-limits are in use + */ + prefix?: string; + /** + * The cache service to use for storing the hit counts. + */ + cache: CacheService; +}; + +type CacheStoreValue = { + totalHits: number; + resetTime: number; +}; + +/** + * A `Store` that stores the hit count for each client. + */ +export class RateLimitStore implements Store { + /** + * The duration of time before which all hit counts are reset (in milliseconds). + * default: 15 minutes + */ + windowMs: number = 15 * 60 * 1000; + prefix: string; + #cache: CacheService; + + private constructor(options: CacheStoreOptions) { + this.prefix = options.prefix ?? 'rl_'; + this.#cache = options.cache; + } + + static fromOptions(options: CacheStoreOptions): RateLimitStore { + return new RateLimitStore(options); + } + + #getDefaultValue() { + return { totalHits: 0, resetTime: Date.now() + this.windowMs }; + } + + init(options: Partial): void { + if (options.windowMs) { + this.windowMs = options.windowMs; + } + } + + prefixKey(key: string): string { + return `${this.prefix}${key}`; + } + + async get(key: string): Promise { + const value = + (await this.#cache.get(this.prefixKey(key))) ?? + this.#getDefaultValue(); + return { + totalHits: value.totalHits, + resetTime: new Date(value.resetTime), + }; + } + + async increment(key: string): Promise { + const value = + (await this.#cache.get(this.prefixKey(key))) ?? + this.#getDefaultValue(); + const totalHits = value.totalHits + 1; + const resetTime = value.resetTime; + await this.#cache.set( + this.prefixKey(key), + { totalHits, resetTime }, + { ttl: this.windowMs }, + ); + return { + totalHits, + resetTime: new Date(resetTime), + }; + } + + async decrement(key: string): Promise { + const value = + (await this.#cache.get(this.prefixKey(key))) ?? + this.#getDefaultValue(); + const totalHits = value.totalHits > 0 ? value.totalHits - 1 : 0; + const resetTime = value.resetTime; + await this.#cache.set( + this.prefixKey(key), + { totalHits, resetTime }, + { ttl: this.windowMs }, + ); + } + + async resetKey(key: string): Promise { + await this.#cache.delete(this.prefixKey(key)); + } +}