feat(backend-app-api): create rate limit cache store

Signed-off-by: Camila Belo <camilaibs@gmail.com>

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) => {

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);

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<CacheService>;
+  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),
+    );
+  });
+});

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<Options>): void {
+    if (options.windowMs) {
+      this.windowMs = options.windowMs;
+    }
+  }
+
+  prefixKey(key: string): string {
+    return `${this.prefix}${key}`;
+  }
+
+  async get(key: string): Promise<ClientRateLimitInfo | undefined> {
+    const value =
+      (await this.#cache.get<CacheStoreValue>(this.prefixKey(key))) ??
+      this.#getDefaultValue();
+    return {
+      totalHits: value.totalHits,
+      resetTime: new Date(value.resetTime),
+    };
+  }
+
+  async increment(key: string): Promise<IncrementResponse> {
+    const value =
+      (await this.#cache.get<CacheStoreValue>(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<void> {
+    const value =
+      (await this.#cache.get<CacheStoreValue>(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<void> {
+    await this.#cache.delete(this.prefixKey(key));
+  }
+}