Merge pull request #321 from AikidoSec/AIK-3229-2

Check if IP has access to route

library/agent/Config.ts

@@ -6,6 +6,7 @@ export type Endpoint = {
     type: "query" | "mutation";
     name: string;
   };
+  allowedIPAddresses?: string[];
   rateLimiting: {
     enabled: boolean;
     maxRequests: number;

library/agent/ServiceConfig.ts

@@ -1,4 +1,5 @@
 import { LimitedContext, matchEndpoint } from "../helpers/matchEndpoint";
+import { matchEndpoints } from "../helpers/matchEndpoints";
 import { Endpoint } from "./Config";
 
 export class ServiceConfig {
@@ -29,6 +30,10 @@ export class ServiceConfig {
     );
   }
 
+  getEndpoints(context: LimitedContext) {
+    return matchEndpoints(context, this.nonGraphQLEndpoints);
+  }
+
   getEndpoint(context: LimitedContext) {
     return matchEndpoint(context, this.nonGraphQLEndpoints);
   }

library/helpers/matchEndpoints.ts

@@ -0,0 +1,53 @@
+import { Endpoint } from "../agent/Config";
+import { Context } from "../agent/Context";
+import { tryParseURLPath } from "./tryParseURLPath";
+
+export type LimitedContext = Pick<Context, "url" | "method" | "route">;
+
+export function matchEndpoints(context: LimitedContext, endpoints: Endpoint[]) {
+  const matches: Endpoint[] = [];
+
+  if (!context.method) {
+    return matches;
+  }
+
+  const possible = endpoints.filter((endpoint) => {
+    if (endpoint.method === "*") {
+      return true;
+    }
+
+    return endpoint.method === context.method;
+  });
+
+  const exact = possible.find((endpoint) => endpoint.route === context.route);
+  if (exact) {
+    matches.push(exact);
+  }
+
+  if (context.url) {
+    // req.url is relative, so we need to prepend a host to make it absolute
+    // We just match the pathname, we don't use the host for matching
+    const path = tryParseURLPath(context.url);
+    const wildcards = possible
+      .filter((endpoint) => endpoint.route.includes("*"))
+      .sort((a, b) => {
+        // Sort endpoints based on the amount of * in the route
+        return b.route.split("*").length - a.route.split("*").length;
+      });
+
+    if (path) {
+      for (const wildcard of wildcards) {
+        const regex = new RegExp(
+          `^${wildcard.route.replace(/\*/g, "(.*)")}\/?$`,
+          "i"
+        );
+
+        if (regex.test(path)) {
+          matches.push(wildcard);
+        }
+      }
+    }
+  }
+
+  return matches;
+}

library/sources/HTTPServer.test.ts

@@ -29,12 +29,20 @@ const api = new ReportingAPIForTesting({
       route: "/rate-limited",
       method: "GET",
       forceProtectionOff: false,
+      allowedIPAddresses: [],
       rateLimiting: {
         enabled: true,
         maxRequests: 3,
         windowSizeInMS: 60 * 60 * 1000,
       },
     },
+    {
+      route: "/ip-allowed",
+      method: "GET",
+      forceProtectionOff: false,
+      allowedIPAddresses: ["8.8.8.8"],
+      rateLimiting: undefined,
+    },
   ],
   heartbeatIntervalInMS: 10 * 60 * 1000,
 });
@@ -47,8 +55,11 @@ const agent = new Agent(
 );
 agent.start([new HTTPServer()]);
 
+t.setTimeout(30 * 1000);
+
 t.beforeEach(() => {
   delete process.env.AIKIDO_MAX_BODY_SIZE_MB;
+  delete process.env.NODE_ENV;
 });
 
 t.test("it wraps the createServer function of http module", async () => {
@@ -289,9 +300,9 @@ t.test("it uses x-forwarded-for header", async () => {
   });
 
   await new Promise<void>((resolve) => {
-    server.listen(3316, () => {
+    server.listen(3350, () => {
       fetch({
-        url: new URL("http://localhost:3316"),
+        url: new URL("http://localhost:3350"),
         method: "GET",
         headers: {
           "x-forwarded-for": "203.0.113.195",
@@ -515,9 +526,9 @@ t.test("it wraps on request event of http", async () => {
   http.globalAgent = new http.Agent({ keepAlive: false });
 
   await new Promise<void>((resolve) => {
-    server.listen(3314, () => {
+    server.listen(3367, () => {
       fetch({
-        url: new URL("http://localhost:3314"),
+        url: new URL("http://localhost:3367"),
         method: "GET",
         headers: {},
         timeoutInMS: 500,
@@ -526,7 +537,7 @@ t.test("it wraps on request event of http", async () => {
         t.same(context, {
           url: "/",
           method: "GET",
-          headers: { host: "localhost:3314", connection: "close" },
+          headers: { host: "localhost:3367", connection: "close" },
           query: {},
           route: "/",
           source: "http.createServer",
@@ -564,9 +575,9 @@ t.test("it wraps on request event of https", async () => {
   https.globalAgent = new https.Agent({ keepAlive: false });
 
   await new Promise<void>((resolve) => {
-    server.listen(3315, () => {
+    server.listen(3361, () => {
       fetch({
-        url: new URL("https://localhost:3315"),
+        url: new URL("https://localhost:3361"),
         method: "GET",
         headers: {},
         timeoutInMS: 500,
@@ -575,7 +586,7 @@ t.test("it wraps on request event of https", async () => {
         t.same(context, {
           url: "/",
           method: "GET",
-          headers: { host: "localhost:3315", connection: "close" },
+          headers: { host: "localhost:3361", connection: "close" },
           query: {},
           route: "/",
           source: "https.createServer",
@@ -590,3 +601,52 @@ t.test("it wraps on request event of https", async () => {
     });
   });
 });
+
+t.test("it checks if IP can access route", async () => {
+  const http = require("http");
+  const server = http.createServer((req, res) => {
+    res.setHeader("Content-Type", "text/plain");
+    res.end("OK");
+  });
+
+  process.env.NODE_ENV = "production";
+
+  await new Promise<void>((resolve) => {
+    server.listen(3324, () => {
+      Promise.all([
+        fetch({
+          url: new URL("http://localhost:3324/ip-allowed"),
+          method: "GET",
+          headers: {
+            "x-forwarded-for": "8.8.8.8",
+          },
+          timeoutInMS: 500,
+        }),
+        fetch({
+          url: new URL("http://localhost:3324/ip-allowed"),
+          method: "GET",
+          headers: {},
+          timeoutInMS: 500,
+        }),
+        fetch({
+          url: new URL("http://localhost:3324/ip-allowed"),
+          method: "GET",
+          headers: {
+            "x-forwarded-for": "1.2.3.4",
+          },
+          timeoutInMS: 500,
+        }),
+      ]).then(([response1, response2, response3]) => {
+        t.equal(response1.statusCode, 200);
+        t.equal(response2.statusCode, 200);
+        t.equal(response3.statusCode, 403);
+        t.same(
+          response3.body,
+          "Your IP address is not allowed to access this resource. (Your IP: 1.2.3.4)"
+        );
+        server.close();
+        resolve();
+      });
+    });
+  });
+});

library/sources/HTTPServer.ts

@@ -39,9 +39,11 @@ export class HTTPServer implements Wrapper {
     ) {
       return args;
     }
+
     if (args[0] !== "request") {
       return args;
     }
+
     return this.wrapRequestListener(args, module, agent);
   }
 

library/sources/http-server/createRequestListener.ts

@@ -4,6 +4,7 @@ import { getContext, runWithContext } from "../../agent/Context";
 import { escapeHTML } from "../../helpers/escapeHTML";
 import { shouldRateLimitRequest } from "../../ratelimiting/shouldRateLimitRequest";
 import { contextFromRequest } from "./contextFromRequest";
+import { ipAllowedToAccessRoute } from "./ipAllowedToAccessRoute";
 import { readBodyStream } from "./readBodyStream";
 import { shouldDiscoverRoute } from "./shouldDiscoverRoute";
 
@@ -46,29 +47,19 @@ function callListenerWithContext(
   const context = contextFromRequest(req, body, module);
 
   return runWithContext(context, () => {
-    res.on("finish", () => {
-      const context = getContext();
-
-      if (
-        context &&
-        context.route &&
-        context.method &&
-        shouldDiscoverRoute({
-          statusCode: res.statusCode,
-          route: context.route,
-          method: context.method,
-        })
-      ) {
-        agent.onRouteExecute(context.method, context.route);
-      }
+    res.on("finish", createOnFinishRequestHandler(res, agent));
+
+    if (!ipAllowedToAccessRoute(context, agent)) {
+      res.statusCode = 403;
+      res.setHeader("Content-Type", "text/plain");
 
-      agent.getInspectionStatistics().onRequest();
-      if (context && context.attackDetected) {
-        agent.getInspectionStatistics().onDetectedAttack({
-          blocked: agent.shouldBlock(),
-        });
+      let message = "Your IP address is not allowed to access this resource.";
+      if (context.remoteAddress) {
+        message += ` (Your IP: ${escapeHTML(context.remoteAddress)})`;
       }
-    });
+
+      return res.end(message);
+    }
 
     const result = shouldRateLimitRequest(context, agent);
 
@@ -87,3 +78,32 @@ function callListenerWithContext(
     return listener(req, res);
   });
 }
+
+function createOnFinishRequestHandler(
+  res: ServerResponse<IncomingMessage>,
+  agent: Agent
+) {
+  return function onFinishRequest() {
+    const context = getContext();
+
+    if (
+      context &&
+      context.route &&
+      context.method &&
+      shouldDiscoverRoute({
+        statusCode: res.statusCode,
+        route: context.route,
+        method: context.method,
+      })
+    ) {
+      agent.onRouteExecute(context.method, context.route);
+    }
+
+    agent.getInspectionStatistics().onRequest();
+    if (context && context.attackDetected) {
+      agent.getInspectionStatistics().onDetectedAttack({
+        blocked: agent.shouldBlock(),
+      });
+    }
+  };
+}