feat(core-p2p): rate limit plugin (#4193)

Author: air1one

__tests__/unit/core-p2p/socket-server/plugins/rate-limit.test.ts

@@ -0,0 +1,95 @@
+import { Server } from "@hapi/hapi";
+import Joi from "@hapi/joi";
+import { Container } from "@arkecosystem/core-kernel";
+
+import { RateLimitPlugin } from "@arkecosystem/core-p2p/src/socket-server/plugins/rate-limit";
+import * as utils from "@arkecosystem/core-p2p/src/utils/build-rate-limiter";
+
+afterEach(() => {
+    jest.clearAllMocks();
+});
+
+describe("RateLimitPlugin", () => {
+    let rateLimitPlugin: RateLimitPlugin;
+
+    const container = new Container.Container();
+
+    const responsePayload = { status: "ok" };
+    const mockRouteByPath = {
+        "/p2p/peer/mockroute": {
+            id: "p2p.peer.getPeers",
+            handler: () => responsePayload,
+            validation: Joi.object().max(0),
+        },
+    };
+    const mockRoute = {
+        method: "POST",
+        path: "/p2p/peer/mockroute",
+        config: {
+            id: mockRouteByPath["/p2p/peer/mockroute"].id,
+            handler: mockRouteByPath["/p2p/peer/mockroute"].handler,
+        },
+    };
+
+    const app = {
+        resolve: jest.fn().mockReturnValue({ getRoutesConfigByPath: () => mockRouteByPath }),
+    };
+    const pluginConfiguration = { getOptional: (id, defaultValue) => defaultValue };
+    const rateLimiter = {
+        hasExceededRateLimit: jest.fn().mockReturnValue(false),
+    };
+
+    beforeAll(() => {
+        container.unbindAll();
+        container.bind(Container.Identifiers.Application).toConstantValue(app);
+        container.bind(Container.Identifiers.PluginConfiguration).toConstantValue(pluginConfiguration);
+
+        jest.spyOn(utils, "buildRateLimiter").mockReturnValue(rateLimiter as any);
+    });
+
+    beforeEach(() => {
+        rateLimitPlugin = container.resolve<RateLimitPlugin>(RateLimitPlugin);
+    });
+
+    it("should register the plugin", async () => {
+        const server = new Server({ port: 4100 });
+        server.route(mockRoute);
+
+        const spyExt = jest.spyOn(server, "ext");
+
+        rateLimitPlugin.register(server);
+
+        expect(spyExt).toBeCalledWith(expect.objectContaining({ type: "onPreAuth" }));
+
+        // try the route with a valid payload
+        const remoteAddress = "187.166.55.44";
+        const responseValid = await server.inject({
+            method: "POST",
+            url: "/p2p/peer/mockroute",
+            payload: {},
+            remoteAddress,
+        });
+        expect(JSON.parse(responseValid.payload)).toEqual(responsePayload);
+        expect(responseValid.statusCode).toBe(200);
+        expect(rateLimiter.hasExceededRateLimit).toBeCalledTimes(1);
+    });
+
+    it("should return a tooManyRequests error when exceeded rate limit", async () => {
+        const server = new Server({ port: 4100 });
+        server.route(mockRoute);
+        rateLimitPlugin.register(server);
+
+        rateLimiter.hasExceededRateLimit.mockReturnValueOnce(true);
+
+        // try the route with a valid payload
+        const remoteAddress = "187.166.55.44";
+        const responseForbidden = await server.inject({
+            method: "POST",
+            url: "/p2p/peer/mockroute",
+            payload: {},
+            remoteAddress,
+        });
+        expect(responseForbidden.statusCode).toBe(429);
+        expect(rateLimiter.hasExceededRateLimit).toBeCalledTimes(1);
+    });
+});

packages/core-p2p/src/socket-server/plugins/rate-limit.ts

@@ -0,0 +1,47 @@
+import { Container, Contracts, Providers } from "@arkecosystem/core-kernel";
+import Boom from "@hapi/boom";
+import { RateLimiter } from "../../rate-limiter";
+import { buildRateLimiter } from "../../utils/build-rate-limiter";
+import { BlocksRoute } from "../routes/blocks";
+import { InternalRoute } from "../routes/internal";
+import { PeerRoute } from "../routes/peer";
+import { TransactionsRoute } from "../routes/transactions";
+
+@Container.injectable()
+export class RateLimitPlugin {
+    @Container.inject(Container.Identifiers.Application)
+    protected readonly app!: Contracts.Kernel.Application;
+
+    @Container.inject(Container.Identifiers.PluginConfiguration)
+    @Container.tagged("plugin", "@arkecosystem/core-p2p")
+    private readonly configuration!: Providers.PluginConfiguration;
+
+    private rateLimiter!: RateLimiter;
+
+    public register(server) {
+        this.rateLimiter = buildRateLimiter({
+            whitelist: [],
+            remoteAccess: this.configuration.getOptional<Array<string>>("remoteAccess", []),
+            rateLimit: this.configuration.getOptional<number>("rateLimit", 100),
+        });
+
+        const allRoutesConfigByPath = {
+            ...this.app.resolve(InternalRoute).getRoutesConfigByPath(),
+            ...this.app.resolve(PeerRoute).getRoutesConfigByPath(),
+            ...this.app.resolve(BlocksRoute).getRoutesConfigByPath(),
+            ...this.app.resolve(TransactionsRoute).getRoutesConfigByPath(),
+        };
+
+        server.ext({
+            type: "onPreAuth",
+            method: async (request, h) => {
+                const endpoint = allRoutesConfigByPath[request.path].id;
+
+                if (await this.rateLimiter.hasExceededRateLimit(request.info.remoteAddress, endpoint)) {
+                    return Boom.tooManyRequests("Rate limit exceeded");
+                }
+                return h.continue;
+            },
+        });
+    }
+}

packages/core-p2p/src/socket-server/server.ts

@@ -11,6 +11,7 @@ import { BlocksRoute } from "./routes/blocks";
 import { InternalRoute } from "./routes/internal";
 import { PeerRoute } from "./routes/peer";
 import { TransactionsRoute } from "./routes/transactions";
+import { RateLimitPlugin } from "./plugins/rate-limit";
 
 // todo: review the implementation
 @Container.injectable()
@@ -76,6 +77,7 @@ export class Server {
 
         // onPreAuth
         this.app.resolve(WhitelistForgerPlugin).register(this.server);
+        this.app.resolve(RateLimitPlugin).register(this.server);
 
         // onPostAuth
         this.app.resolve(CodecPlugin).register(this.server);