Merge pull request #28 from DefangLabs/linda-csrf-protection

Added CSRF protection and rate-limiting

Author: commit111

.github/workflows/deploy.yaml

@@ -34,9 +34,10 @@ jobs:
       - name: Deploy
         uses: DefangLabs/[email protected]
         with:
-          config-env-vars: OPENAI_API_KEY
+          config-env-vars: OPENAI_API_KEY SECRET_KEY
           mode: production
           provider: aws
 
         env:
           OPENAI_API_KEY: ${{ secrets.OPENAI_API_KEY }}
+          SECRET_KEY: ${{ secrets.SECRET_KEY }}

app/app.py

@@ -1,8 +1,17 @@
 from flask import Flask, request, jsonify, render_template, Response, stream_with_context
+from flask_wtf.csrf import CSRFProtect
 from rag_system import rag_system
 import subprocess
 app = Flask(__name__, static_folder='templates/images')
 
+import os
+
+app.config['SECRET_KEY'] = os.getenv('SECRET_KEY')
+app.config['SESSION_COOKIE_HTTPONLY'] = True
+app.config['SESSION_COOKIE_SECURE'] = bool(os.getenv('SESSION_COOKIE_SECURE'))
+
+csrf = CSRFProtect(app)
+
 @app.route('/', methods=['GET', 'POST'])
 def index():
     return render_template('index.html')

app/requirements.txt

@@ -1,4 +1,5 @@
 Flask==2.0.1
+Flask-WTF==1.2.2
 Werkzeug==2.0.3
 scikit-learn==0.24.2
 numpy==1.22.0

app/templates/index.html

@@ -207,6 +207,27 @@ <h2>Ask Defang</h2>
             }
         });
 
+        async function rateLimitingFetch(url, options = {}) {
+            if (window.crypto && window.crypto.subtle) {
+                let bodyHash, nonceArray = new Uint32Array(1), nonce = 0;
+                const body = new TextEncoder().encode(options.body);
+                const bodyWithNonce = new Uint8Array(nonceArray.byteLength + body.byteLength);
+                bodyWithNonce.set(body, nonceArray.byteLength);
+
+                do {
+                    nonceArray[0] = ++nonce;
+                    bodyWithNonce.set(new Uint8Array(nonceArray.buffer));
+                    bodyHash = await crypto.subtle.digest('SHA-256', bodyWithNonce);
+                } while(new DataView(bodyHash).getUint32(0) > 0x50000);
+
+                options.headers = {
+                    ...options.headers,
+                    'X-Nonce': nonce
+                }
+                return fetch(url, options);
+            }
+        }
+
         function sendQuery() {
             const query = queryInput.value.trim();
             if (query === '') return;
@@ -228,10 +249,11 @@ <h2>Ask Defang</h2>
             sendButton.disabled = true;
 
             // Send query to server
-            fetch('/ask', {
+            rateLimitingFetch('/ask', {
                 method: 'POST',
                 headers: {
                     'Content-Type': 'application/json',
+                    'X-CSRFToken': '{{ csrf_token() }}'  
                 },
                 body: JSON.stringify({ query: query }),
             })

compose.dev.yaml

@@ -10,6 +10,8 @@ services:
         mode: ingress
     environment:
       FLASK_APP: app.py
+      SECRET_KEY: supersecret
+      SESSION_COOKIE_SECURE: 0
       OPENAI_API_KEY: ${OPENAI_API_KEY} # Set your OpenAI API key here or in the .env file
     command: flask run --host=0.0.0.0
     deploy:

compose.yaml

@@ -13,6 +13,8 @@ services:
         mode: ingress
     environment:
       FLASK_APP: app.py
+      SECRET_KEY: 
+      SESSION_COOKIE_SECURE: 1
       OPENAI_API_KEY: ${OPENAI_API_KEY} # Set your OpenAI API key here or in the .env file
     command: uwsgi --http 0.0.0.0:5000 --wsgi-file app.py --callable app --processes 4 --threads 2
     deploy: