[public-api] add rate limiting in server (#18953)

Author: akosyakov

components/server/src/api/rate-limited.ts

@@ -0,0 +1,23 @@
+/**
+ * Copyright (c) 2023 Gitpod GmbH. All rights reserved.
+ * Licensed under the GNU Affero General Public License (AGPL).
+ * See License.AGPL.txt in the project root for license information.
+ */
+
+import { IRateLimiterOptions } from "rate-limiter-flexible";
+
+const RATE_LIMIT_METADATA_KEY = Symbol("RateLimited");
+
+export function RateLimited(options: IRateLimiterOptions) {
+    return Reflect.metadata(RATE_LIMIT_METADATA_KEY, options);
+}
+
+export namespace RateLimited {
+    const defaultOptions: IRateLimiterOptions = {
+        points: 200,
+        duration: 60,
+    };
+    export function getOptions(target: Object, properyKey: string | symbol): IRateLimiterOptions {
+        return Reflect.getMetadata(RATE_LIMIT_METADATA_KEY, target, properyKey) || defaultOptions;
+    }
+}

components/server/src/api/server.ts

@@ -30,6 +30,13 @@ import { APIStatsService as StatsServiceAPI } from "./stats";
 import { APITeamsService as TeamsServiceAPI } from "./teams";
 import { APIUserService as UserServiceAPI } from "./user";
 import { WorkspaceServiceAPI } from "./workspace-service-api";
+import { IRateLimiterOptions, RateLimiterMemory, RateLimiterRedis, RateLimiterRes } from "rate-limiter-flexible";
+import { Redis } from "ioredis";
+import { RateLimited } from "./rate-limited";
+import { Config } from "../config";
+import { ApplicationError, ErrorCodes } from "@gitpod/gitpod-protocol/lib/messaging/error";
+import { UserService } from "../user/user-service";
+import { User } from "@gitpod/gitpod-protocol";
 
 decorate(injectable(), PublicAPIConverter);
 
@@ -46,6 +53,9 @@ export class API {
     @inject(HelloServiceAPI) private readonly helloServiceApi: HelloServiceAPI;
     @inject(SessionHandler) private readonly sessionHandler: SessionHandler;
     @inject(PublicAPIConverter) private readonly apiConverter: PublicAPIConverter;
+    @inject(Redis) private readonly redis: Redis;
+    @inject(Config) private readonly config: Config;
+    @inject(UserService) private readonly userService: UserService;
 
     listenPrivate(): http.Server {
         const app = express();
@@ -174,9 +184,30 @@ export class API {
 
                     const context = args[1] as HandlerContext;
 
+                    const rateLimit = async (subjectId: string) => {
+                        const key = `${grpc_service}/${grpc_method}`;
+                        const options = self.config.rateLimits?.[key] || RateLimited.getOptions(target, prop);
+                        try {
+                            await self.getRateLimitter(options).consume(`${subjectId}_${key}`);
+                        } catch (e) {
+                            if (e instanceof RateLimiterRes) {
+                                throw new ConnectError("rate limit exceeded", Code.ResourceExhausted, {
+                                    // http compatibility, can be respected by gRPC clients as well
+                                    // instead of doing an ad-hoc retry, the client can wait for the given amount of seconds
+                                    "Retry-After": e.msBeforeNext / 1000,
+                                    "X-RateLimit-Limit": options.points,
+                                    "X-RateLimit-Remaining": e.remainingPoints,
+                                    "X-RateLimit-Reset": new Date(Date.now() + e.msBeforeNext),
+                                });
+                            }
+                            throw e;
+                        }
+                    };
+
                     const apply = async <T>(): Promise<T> => {
-                        const user = await self.verify(context);
-                        context.user = user;
+                        const subjectId = await self.verify(context);
+                        await rateLimit(subjectId);
+                        context.user = await self.ensureFgaMigration(subjectId);
 
                         return Reflect.apply(target[prop as any], target, args);
                     };
@@ -211,16 +242,49 @@ export class API {
         };
     }
 
-    private async verify(context: HandlerContext) {
-        const user = await this.sessionHandler.verify(context.requestHeader.get("cookie") || "");
-        if (!user) {
+    private async verify(context: HandlerContext): Promise<string> {
+        const claims = await this.sessionHandler.verifyJWTCookie(context.requestHeader.get("cookie") || "");
+        const subjectId = claims?.sub;
+        if (!subjectId) {
             throw new ConnectError("unauthenticated", Code.Unauthenticated);
         }
-        const fgaChecksEnabled = await isFgaChecksEnabled(user.id);
+        return subjectId;
+    }
+
+    private async ensureFgaMigration(userId: string): Promise<User> {
+        const fgaChecksEnabled = await isFgaChecksEnabled(userId);
         if (!fgaChecksEnabled) {
             throw new ConnectError("unauthorized", Code.PermissionDenied);
         }
-        return user;
+        try {
+            return await this.userService.findUserById(userId, userId);
+        } catch (e) {
+            if (e instanceof ApplicationError && e.code === ErrorCodes.NOT_FOUND) {
+                throw new ConnectError("unauthorized", Code.PermissionDenied);
+            }
+            throw e;
+        }
+    }
+
+    private readonly rateLimiters = new Map<string, RateLimiterRedis>();
+    private getRateLimitter(options: IRateLimiterOptions): RateLimiterRedis {
+        const sortedKeys = Object.keys(options).sort();
+        const sortedObject: { [key: string]: any } = {};
+        for (const key of sortedKeys) {
+            sortedObject[key] = options[key as keyof IRateLimiterOptions];
+        }
+        const key = JSON.stringify(sortedObject);
+
+        let rateLimiter = this.rateLimiters.get(key);
+        if (!rateLimiter) {
+            rateLimiter = new RateLimiterRedis({
+                storeClient: this.redis,
+                ...options,
+                insuranceLimiter: new RateLimiterMemory(options),
+            });
+            this.rateLimiters.set(key, rateLimiter);
+        }
+        return rateLimiter;
     }
 
     static bindAPI(bind: interfaces.Bind): void {

components/server/src/api/teams.spec.db.ts

@@ -21,6 +21,9 @@ import { UserAuthentication } from "../user/user-authentication";
 import { WorkspaceService } from "../workspace/workspace-service";
 import { API } from "./server";
 import { SessionHandler } from "../session-handler";
+import { Redis } from "ioredis";
+import { UserService } from "../user/user-service";
+import { Config } from "../config";
 
 const expect = chai.expect;
 
@@ -39,6 +42,9 @@ export class APITeamsServiceSpec {
         this.container.bind(WorkspaceService).toConstantValue({} as WorkspaceService);
         this.container.bind(UserAuthentication).toConstantValue({} as UserAuthentication);
         this.container.bind(SessionHandler).toConstantValue({} as SessionHandler);
+        this.container.bind(Config).toConstantValue({} as Config);
+        this.container.bind(Redis).toConstantValue({} as Redis);
+        this.container.bind(UserService).toConstantValue({} as UserService);
 
         // Clean-up database
         const typeorm = testContainer.get<TypeORM>(TypeORM);

components/server/src/config.ts

@@ -17,6 +17,7 @@ import { log } from "@gitpod/gitpod-protocol/lib/util/logging";
 import { filePathTelepresenceAware } from "@gitpod/gitpod-protocol/lib/env";
 import { WorkspaceClassesConfig } from "./workspace/workspace-classes";
 import { PrebuildRateLimiters } from "./workspace/prebuild-rate-limiter";
+import { IRateLimiterOptions } from "rate-limiter-flexible";
 
 export const Config = Symbol("Config");
 export type Config = Omit<
@@ -172,9 +173,20 @@ export interface ConfigSerialized {
 
     /**
      * The configuration for the rate limiter we (mainly) use for the websocket API
+     * @deprecated used for JSON-RPC API, for gRPC use rateLimits
      */
     rateLimiter: RateLimiterConfig;
 
+    /**
+     * The configuration for the rate limiter we use for the gRPC API.
+     * As a primary means use RateLimited decorator.
+     * Only use this if you need to adjst in production, make sure to apply changes to the decorator as well.
+     * Key is of the form `<grpc_service>/<grpc_method>`
+     */
+    rateLimits?: {
+        [key: string]: IRateLimiterOptions;
+    };
+
     /**
      * The address content service clients connect to
      * Example: content-service:8080

components/server/src/session-handler.ts

@@ -15,6 +15,7 @@ import { Config } from "./config";
 import { WsNextFunction, WsRequestHandler } from "./express/ws-handler";
 import { reportJWTCookieIssued } from "./prometheus-metrics";
 import { UserService } from "./user/user-service";
+import { JwtPayload } from "jsonwebtoken";
 
 @injectable()
 export class SessionHandler {
@@ -103,29 +104,41 @@ export class SessionHandler {
     }
 
     async verify(cookie: string): Promise<User | undefined> {
-        const cookies = parseCookieHeader(cookie);
-        const jwtToken = cookies[this.getJWTCookieName(this.config)];
-        if (!jwtToken) {
-            log.debug("No JWT session present on request");
-            return undefined;
-        }
         try {
-            const claims = await this.authJWT.verify(jwtToken);
-            log.debug("JWT Session token verified", { claims });
-
+            const claims = await this.verifyJWTCookie(cookie);
+            if (!claims) {
+                return undefined;
+            }
             const subject = claims.sub;
             if (!subject) {
                 throw new Error("Subject is missing from JWT session claims");
             }
-
-            return await this.userService.findUserById(subject, subject);
+            return await this.userService.findUserById(claims.subject, claims.subject);
         } catch (err) {
             log.warn("Failed to authenticate user with JWT Session", err);
             // Remove the existing cookie, to force the user to re-sing in, and hence refresh it
             return undefined;
         }
     }
 
+    /**
+     * Verify the JWT session cookie.
+     * @throws only in case of programming errors, if the cookie is invalid, undefined is returned
+     * @param cookie - the cookie header value to verify
+     * @returns returns the claims of the JWT session cookie or undefined if the cookie is not present or invalid
+     */
+    async verifyJWTCookie(cookie: string): Promise<JwtPayload | undefined> {
+        const cookies = parseCookieHeader(cookie);
+        const jwtToken = cookies[this.getJWTCookieName(this.config)];
+        if (!jwtToken) {
+            log.debug("No JWT session present on request");
+            return undefined;
+        }
+        const claims = await this.authJWT.verify(jwtToken);
+        log.debug("JWT Session token verified", { claims });
+        return claims;
+    }
+
     public async createJWTSessionCookie(
         userID: string,
     ): Promise<{ name: string; value: string; opts: express.CookieOptions }> {