Merge pull request #83 from AikidoSec/AIK-3345

Add benchmarking/fix some performance issues

.github/workflows/benchmark.yml

@@ -0,0 +1,24 @@
+name: Benchmark
+on:
+  push: {}
+  workflow_call: {}
+jobs:
+  benchmark:
+    runs-on: ubuntu-latest
+    timeout-minutes: 10
+    steps:
+      - name: Checkout code
+        uses: actions/checkout@v2
+
+      - name: Set up Docker Buildx
+        uses: docker/setup-buildx-action@v1
+      - name: Run Docker Compose
+        working-directory: ./sample-apps/flask-mysql
+        run: |
+          docker compose -f docker-compose.yml -f docker-compose.benchmark.yml up --build -d
+      - name: Install K6
+        uses: grafana/setup-k6-action@v1
+      - name: Run flask-mysql k6 Benchmark
+        run: |
+          k6 run -q ./benchmarks/flask-mysql-benchmarks.js
+

aikido_firewall/background_process/__init__.py

@@ -3,15 +3,17 @@
 and listen for data sent by our sources and sinks
 """
 
+import os
 from aikido_firewall.helpers.token import get_token_from_env
+from aikido_firewall.helpers.get_temp_dir import get_temp_dir
 from aikido_firewall.helpers.logging import logger
 from aikido_firewall.background_process.comms import (
     AikidoIPCCommunications,
     get_comms,
     reset_comms,
 )
 
-IPC_ADDRESS = ("localhost", 9898)  # Specify the IP address and port
+IPC_ADDRESS = get_temp_dir() + "/aikido_python_socket.sock"
 
 
 def start_background_process():
@@ -22,5 +24,9 @@ def start_background_process():
     # Generate a secret key :
     secret_key_bytes = str.encode(str(get_token_from_env()))
 
+    # Remove the socket file if it already exists
+    if os.path.exists(IPC_ADDRESS):
+        os.remove(IPC_ADDRESS)
+
     comms = AikidoIPCCommunications(IPC_ADDRESS, secret_key_bytes)
     comms.start_aikido_listener()

aikido_firewall/helpers/get_temp_dir.py

@@ -0,0 +1,15 @@
+"""
+Helper function file, see function docstring
+"""
+
+import os
+
+
+def get_temp_dir():
+    """
+    Checks the environment variable "AIKIDO_TMP_DIR"
+    """
+    aikido_temp_dir_env = os.getenv("AIKIDO_TMP_DIR")
+    if aikido_temp_dir_env is not None:
+        return aikido_temp_dir_env
+    return "/tmp"

aikido_firewall/sources/flask.py

@@ -10,34 +10,46 @@
 from aikido_firewall.context import Context
 from aikido_firewall.background_process.packages import add_wrapped_package
 from .functions.request_handler import request_handler
+from aikido_firewall.context import get_current_context
 
 
-class AikidoMiddleware:
+def generate_aikido_view_func_wrapper(former_view_func):
     """
-    Aikido WSGI Middleware for ratelimiting and route reporting
+    Generates our own wrapper for the function in self.view_functions[]
     """
 
-    def __init__(self, app, flask_app=None):
-        self.app = app
-        self.flask_app = flask_app
-
-    def __call__(self, environ, start_response):
-        response = request_handler(stage="pre_response")
-        if response:
-            from flask import jsonify  #  We don't want to install flask
-
-            with self.flask_app.app_context():
-                start_response(f"{response[1]} Aikido", [])
-                return [response[0].encode("utf-8")]
-
-        def custom_start_response(status, headers):
-            """Is current route useful snippet :"""
-            status_code = int(status.split(" ")[0])
+    def aikido_view_func(*args, **kwargs):
+        from werkzeug.exceptions import HTTPException
+        from flask.globals import request_ctx
+
+        req = request_ctx.request
+        # Set body :
+        context = get_current_context()
+        if context:
+            if req.is_json:
+                context.body = req.get_json()
+                context.set_as_current_context()
+            else:
+                context.body = req.form
+                context.set_as_current_context()
+
+        pre_response = request_handler(stage="pre_response")
+        if pre_response:
+            return pre_response[0], pre_response[1]
+        try:
+            res = former_view_func(*args, **kwargs)
+            status_code = 200
+            if isinstance(res, tuple):
+                status_code = res[1]
+            elif hasattr(res, "status_code"):
+                status_code = res.status_code
             request_handler(stage="post_response", status_code=status_code)
-            return start_response(status, headers)
+            return res
+        except HTTPException as e:
+            request_handler(stage="post_response", status_code=e.code)
+            raise e
 
-        response = self.app(environ, custom_start_response)
-        return response
+    return aikido_view_func
 
 
 def aikido___call__(flask_app, environ, start_response):
@@ -46,13 +58,7 @@ def aikido___call__(flask_app, environ, start_response):
     # pylint: disable=import-outside-toplevel
     try:
         request_handler(stage="init")
-        #  https://stackoverflow.com/a/11163649 :
-        length = int(environ.get("CONTENT_LENGTH") or 0)
-        body = environ["wsgi.input"].read(length)
-        # replace the stream since it was exhausted by read()
-        environ["wsgi.input"] = BytesIO(body)
-
-        context1 = Context(req=environ, raw_body=body.decode("utf-8"), source="flask")
+        context1 = Context(req=environ, raw_body={}, source="flask")
         logger.debug("Context : %s", json.dumps(context1.__dict__))
         context1.set_as_current_context()
     except Exception as e:
@@ -67,18 +73,24 @@ def on_flask_import(flask):
     Hook 'n wrap on `flask.app`
     Our goal is to wrap the __init__ function of the "Flask" class,
     so we can insert our middleware. Returns : Modified flask.app object
+
+    Flask class |-> App class |-> Scaffold class
+    @app.route is implemented in Scaffold and calls `add_url_rule` in App class
+    This function writes to self.view_functions[endpoint] = view_func
+    The only other reference where view_functions is called is on this line:
+    https://github.com/pallets/flask/blob/8a6cdf1e2a5efa81c30f6166602064ceefb0a35b/src/flask/app.py#L882
+    So we would have to wrap the `ensure_sync` function of the app object
     """
     modified_flask = importhook.copy_module(flask)
 
-    prev_flask_init = copy.deepcopy(flask.Flask.__init__)
-
-    def aikido_flask_init(_self, *args, **kwargs):
-        prev_flask_init(_self, *args, **kwargs)
-        setattr(_self, "__call__", aikido___call__)
-        _self.wsgi_app = AikidoMiddleware(_self.wsgi_app, _self)
+    def aikido_ensure_sync(_self, func):
+        """
+        We're wrapping this function, so we can wrap the passed along function `func`
+        https://github.com/pallets/flask/blob/8a6cdf1e2a5efa81c30f6166602064ceefb0a35b/src/flask/app.py#L946
+        """
+        return generate_aikido_view_func_wrapper(func)
 
-    # pylint: disable=no-member
-    setattr(modified_flask.Flask, "__init__", aikido_flask_init)
     setattr(modified_flask.Flask, "__call__", aikido___call__)
+    setattr(modified_flask.Flask, "ensure_sync", aikido_ensure_sync)
     add_wrapped_package("flask")
     return modified_flask

benchmarks/flask-mysql-benchmarks.js

@@ -0,0 +1,130 @@
+import http from 'k6/http';
+import { check, sleep, fail } from 'k6';
+import exec from 'k6/execution';
+import { Trend } from 'k6/metrics';
+
+const BASE_URL_8086 = 'http://localhost:8086';
+const BASE_URL_8087 = 'http://localhost:8087';
+
+export const options = {
+    vus: 1, // Number of virtual users
+    thresholds: {
+        test_40mb_payload: [{
+            threshold: "avg<2000", // Temporary exagurated threshold until optimizations are finished 
+            abortOnFail: true,
+            delayAbortEval: '10s',
+        }],
+        test_multiple_queries: [{
+            threshold: "avg<2000", 
+            abortOnFail: true,
+            delayAbortEval: '10s',
+        }],
+        test_multiple_queries_with_big_body: [{
+            threshold: "avg<2000", 
+            abortOnFail: true,
+            delayAbortEval: '10s',
+        }],
+        test_create_with_big_body: [{
+            threshold: "avg<2000", 
+            abortOnFail: true,
+            delayAbortEval: '10s',
+        }],
+        test_normal_route: [{
+            threshold: "avg<2000", 
+            abortOnFail: true,
+            delayAbortEval: '10s',
+        }],
+        test_id_route: [{
+            threshold: "avg<2000", 
+            abortOnFail: true,
+            delayAbortEval: '10s',
+        }],
+        test_open_file: [{
+            threshold: "avg<2000", 
+            abortOnFail: true,
+            delayAbortEval: '10s',
+        }],
+        test_execute_shell: [{
+            threshold: "avg<2000", 
+            abortOnFail: true,
+            delayAbortEval: '10s',
+        }],
+
+    },
+};
+const default_headers = {
+    "User-Agent":
+        "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/127.0.0.0 Safari/537.36",
+};
+
+const default_payload = {
+    dog_name: "Pops",
+    other_dogs: Array(2000).fill("Lorem Ipsum"),
+    other_dogs2: Array.from({length: 5000}, () => Math.floor(Math.random() * 99999999)),
+    text_message: "Lorem ipsum dolor sit amet".repeat(3000)
+
+};
+function generateLargeJson(sizeInMB) {
+    const sizeInBytes = sizeInMB * 1024; // Convert MB to Kilobytes
+    let long_text = "b".repeat(sizeInBytes)
+    return {
+        dog_name: "test",
+        long_texts: new Array(1024).fill(long_text)
+    }
+}
+
+function measureRequest(url, method = 'GET', payload, status_code=200, headers=default_headers) {
+    let res;
+    if (method === 'POST') {
+        res = http.post(url, payload, {
+            headers: headers
+        }
+        );
+    } else {
+        res = http.get(url, {
+            headers: headers
+        });
+    }
+    check(res, {
+        'status is correct': (r) => r.status === status_code,
+    });
+    return res.timings.duration; // Return the duration of the request
+}
+
+function route_test(trend, amount, route, method="GET", data=default_payload, status=200) {
+    for (let i = 0; i < amount; i++) {
+        let time_with_fw = measureRequest(BASE_URL_8086 + route, method, data, status)
+        let time_without_fw = measureRequest(BASE_URL_8087 + route, method, data, status)
+        trend.add(time_with_fw - time_without_fw)
+    }
+}
+
+export function handleSummary(data) {
+    for (const [metricName, metricValue] of Object.entries(data.metrics)) {
+        if(!metricName.startsWith('test_') || metricValue.values.avg == 0) {
+            continue
+        }
+        let values = metricValue.values
+        console.log(`\x1b[35m 🚅 ${metricName}\x1b[0m: ΔAverage is \x1b[4m${values.avg.toFixed(2)}ms\x1b[0m | ΔMedian is \x1b[4m${values.med.toFixed(2)}ms\x1b[0m`);
+    }
+    return {stdout: ""};
+}
+
+let test_40mb_payload = new Trend('test_40mb_payload')
+let test_multiple_queries = new Trend("test_multiple_queries")
+let test_multiple_queries_with_big_body = new Trend("test_multiple_queries_with_big_body")
+let test_create_with_big_body = new Trend("test_create_with_big_body")
+let test_normal_route = new Trend("test_normal_route")
+let test_id_route = new Trend("test_id_route")
+let test_open_file = new Trend("test_open_file")
+let test_execute_shell = new Trend("test_execute_shell")
+export default function () {
+    route_test(test_40mb_payload, 30, "/create", "POST", generateLargeJson(40)) // 40 Megabytes
+    route_test(test_multiple_queries, 50, "/multiple_queries", "POST", {dog_name: "W"})
+    route_test(test_multiple_queries_with_big_body, 50, "/multiple_queries", "POST")
+    route_test(test_create_with_big_body, 500, "/create", "POST")
+    route_test(test_normal_route, 500, "/")
+    route_test(test_id_route, 500, "/dogpage/1")
+    route_test(test_open_file, 500, "/open_file", 'POST', { filepath: '.env.example' })
+    route_test(test_execute_shell, 500, "/shell", "POST", { command: 'xyzwh'})
+}

sample-apps/README.md

@@ -1,14 +1,20 @@
 # Aikido's sample apps
 Overview : 
-- `django-mysql/` is a Django app using MySQL. 
+- `django-mysql/` is a Django app using MySQL.
   - It runs **multi-threaded**
+  - Runs on 8080. Without Aikido runs on 8081
 - `django-mysql-gunicorn/` is a Django app using MySQL and runnin with a Gunicorn backend.
   - it runs 4 processes, called workers, (**multi-process**) which handle requests using 2 threads (**multi-threaded**)
+  - Runs on 8082. Without Aikido runs on 8083
 - `flask-mongo/` is a Flask app using MongoDB.
   - It runs **multi-threaded**
+  - Runs on 8084. Without Aikido runs on 8085
 - `flask-mysql/` is a Flask app using MySQL.
   - It runs **single-threaded**
+  - Runs on 8086. Without Aikido runs on 8087
 - `flask-mysql-uwsgi/` is a Flask app using Mysql and running with a uWSGI backend.
   - It runs 4 processes (**multi-process**) which handle requests **multi-threaded**
+  - Runs on 8088. Without aikido runs on 8089
 - `flask-postres/` is a Flask app using Postgres
   - It runs **multi-threaded**
+  - Runs on 8090. Without aikido runs on 8091

sample-apps/django-mysql-gunicorn/.env.example

@@ -7,3 +7,4 @@ MYSQL_ROOT_PASSWORD="password"
 # Aikido keys
 AIKIDO_DEBUG=true
 AIKIDO_TOKEN="AIK_secret_token"
+AIKIDO_BLOCKING=true

sample-apps/django-mysql-gunicorn/README.md

@@ -6,11 +6,11 @@ With docker-compose installed run
 ```bash
 docker-compose up --build
 ```
-This will expose a Django web server at [localhost:8080](http://localhost:8080)
+This will expose a Django web server at [localhost:8082](http://localhost:8082)
 
 ## URLS : 
-- Homepage : `http://localhost:8080/app`
-- Create a dog : `http://localhost:8080/app/create/`
+- Homepage : `http://localhost:8082/app`
+- Create a dog : `http://localhost:8082/app/create/`
 - MySQL attack : Enter `Malicious dog", "Injected wrong boss name"); -- `
 
 To verify your attack was successfully note that the boss_name usualy is 'N/A', if you open the dog page (you can do this from the home page). You should see a "malicious dog" with a boss name that is not permitted.

sample-apps/django-mysql-gunicorn/docker-compose.benchmark.yml

@@ -0,0 +1,24 @@
+version: "3"
+services:
+  backend_firewall_disabled:
+    build: 
+      context: ./../../
+      dockerfile: ./sample-apps/django-mysql-gunicorn/Dockerfile
+    command: sh -c "python manage.py migrate && gunicorn -c gunicorn_config.py --workers 4 --threads 2 --log-level debug --access-logfile '-' --error-logfile '-' --bind 0.0.0.0:8000 sample-django-mysql-gunicorn-app.wsgi"
+    restart: always
+    volumes:
+      - .:/app
+    ports:
+      - "8083:8000"
+    env_file:
+      - .env
+    depends_on:
+      db:
+        condition: service_healthy
+    extra_hosts:
+      - "app.local.aikido.io:host-gateway"
+    environment:
+      - FIREWALL_DISABLED=1
+  backend:
+    environment:
+      - AIKIDO_TOKEN="test_aikido_token"