Allowlist & RLS Application Layer

Allowlist & RLS Application Layer

Author: Brayden

install.sh

@@ -40,7 +40,7 @@ echo "Cloning the repository..."
 git clone https://github.com/outerbase/starbasedb.git > /dev/null 2>&1 || { echo "Error: Failed to clone the repository. Please check your internet connection and try again."; exit 1; }
 cd starbasedb || { echo "Error: Failed to change to the starbasedb directory. The clone might have failed."; exit 1; }
 
-# Step 3: Generate a secure AUTHORIZATION_TOKEN and update wrangler.toml
+# Step 3: Generate a secure ADMIN_AUTHORIZATION_TOKEN and update wrangler.toml
 os=$(uname -s)
 PLATFORM_SED="sed -i ''"
 
@@ -64,8 +64,8 @@ case "$os" in
         ;;
 esac
 
-AUTHORIZATION_TOKEN=$(openssl rand -hex 16)
-$PLATFORM_SED "s/AUTHORIZATION_TOKEN = \"[^\"]*\"/AUTHORIZATION_TOKEN = \"$AUTHORIZATION_TOKEN\"/" wrangler.toml
+ADMIN_AUTHORIZATION_TOKEN=$(openssl rand -hex 16)
+$PLATFORM_SED "s/ADMIN_AUTHORIZATION_TOKEN = \"[^\"]*\"/ADMIN_AUTHORIZATION_TOKEN = \"$ADMIN_AUTHORIZATION_TOKEN\"/" wrangler.toml
 
 # Step 4: Prompt the user for Cloudflare account_id (force interaction)
 echo " "

package.json

@@ -17,9 +17,10 @@
   "dependencies": {
     "@libsql/client": "^0.14.0",
     "@outerbase/sdk": "2.0.0-rc.3",
+    "jose": "^5.9.6",
     "mongodb": "^6.11.0",
     "mysql2": "^3.11.4",
-    "pg": "^8.13.1",
-    "node-sql-parser": "^4.18.0"
+    "node-sql-parser": "^4.18.0",
+    "pg": "^8.13.1"
   }
 }

pnpm-lock.yaml

@@ -0,0 +1,86 @@
+import { Env } from "..";
+import { DataSource } from "../types";
+
+const parser = new (require('node-sql-parser').Parser)();
+
+let allowlist: string[] | null = null;
+let normalizedAllowlist: any[] | null = null;
+
+function normalizeSQL(sql: string) {
+    // Remove trailing semicolon. This allows a user to send a SQL statement that has
+    // a semicolon where the allow list might not include it but both statements can
+    // equate to being the same. AST seems to have an issue with matching the difference
+    // when included in one query vs another.
+    return sql.trim().replace(/;\s*$/, '');
+}
+
+async function loadAllowlist(dataSource?: DataSource): Promise<string[]> {
+    try {
+        const statement = 'SELECT sql_statement FROM tmp_allowlist_queries'
+        const result = await dataSource?.internalConnection?.durableObject.executeQuery(statement, [], false) as any[];
+        return result.map((row: any) => row.sql_statement);
+    } catch (error) {
+        console.error('Error loading allowlist:', error);
+        return [];
+    }
+}
+
+export async function isQueryAllowed(sql: string, isEnabled: boolean, dataSource?: DataSource, env?: Env): Promise<boolean | Error> {
+    // If the feature is not turned on then by default the query is allowed
+    if (!isEnabled) return true;
+
+    // If we are using the administrative AUTHORIZATION token value, this request is allowed.
+    // We want database UI's to be able to have more free reign to run queries so we can load
+    // tables, run queries, and more. If you want to block queries with the allowlist then we
+    // advise you to do so by implementing user authentication with JWT.
+    if (dataSource?.request.headers.get('Authorization') === `Bearer ${env?.ADMIN_AUTHORIZATION_TOKEN}`) {
+        return true;
+    }
+
+    allowlist = await loadAllowlist(dataSource);
+    normalizedAllowlist = allowlist.map(query => parser.astify(normalizeSQL(query)));
+    
+    try {
+        if (!sql) {
+            return Error('No SQL provided for allowlist check')
+        }
+
+        const normalizedQuery = parser.astify(normalizeSQL(sql));
+        
+        // Compare ASTs while ignoring specific values
+        const isCurrentAllowed = normalizedAllowlist?.some(allowedQuery => {
+            // Create deep copies to avoid modifying original ASTs
+            const allowedAst = JSON.parse(JSON.stringify(allowedQuery));
+            const queryAst = JSON.parse(JSON.stringify(normalizedQuery));
+            
+            // Remove or normalize value fields from both ASTs
+            const normalizeAst = (ast: any) => {
+                if (Array.isArray(ast)) {
+                    ast.forEach(normalizeAst);
+                } else if (ast && typeof ast === 'object') {
+                    // Remove or normalize fields that contain specific values
+                    if ('value' in ast) {
+                        ast.value = '?';
+                    }
+                    
+                    Object.values(ast).forEach(normalizeAst);
+                }
+
+                return ast;
+            };
+
+            normalizeAst(allowedAst);
+            normalizeAst(queryAst);
+            
+            return JSON.stringify(allowedAst) === JSON.stringify(queryAst);
+        });
+
+        if (!isCurrentAllowed) {
+            throw new Error("Query not allowed");
+        }
+
+        return true;
+    } catch (error: any) {
+        throw new Error(error?.message ?? 'Error');
+    }
+}

src/allowlist/index.ts

@@ -1,20 +1,6 @@
 import { DataSource, Source } from "../types";
 const parser = new (require('node-sql-parser').Parser)();
 
-async function createCacheTable(dataSource?: DataSource) {
-    const statement = `
-    CREATE TABLE IF NOT EXISTS "main"."tmp_cache"(
-        "id" INTEGER PRIMARY KEY AUTOINCREMENT,
-        "timestamp" REAL NOT NULL,
-        "ttl" INTEGER NOT NULL,
-        "query" TEXT UNIQUE NOT NULL,
-        "results" TEXT
-    );
-    `
-
-    await dataSource?.internalConnection?.durableObject.executeQuery(statement, undefined, false)
-}
-
 function hasModifyingStatement(ast: any): boolean {
     // Check if current node is a modifying statement
     if (ast.type && ['insert', 'update', 'delete'].includes(ast.type.toLowerCase())) {
@@ -33,28 +19,30 @@ function hasModifyingStatement(ast: any): boolean {
             }
         }
     }
-    
+
     return false;
 }
 
-export async function beforeQueryCache(sql: string, params?: any[], dataSource?: DataSource): Promise<any | null> {
+export async function beforeQueryCache(sql: string, params?: any[], dataSource?: DataSource, dialect?: string): Promise<any | null> {
     // Currently we do not support caching queries that have dynamic parameters
     if (params?.length) return null
+    if (dataSource?.source === Source.internal || !dataSource?.request.headers.has('X-Starbase-Cache')) return null
 
-    let ast = parser.astify(sql);
+    if (!dialect) dialect = 'sqlite'
+    if (dialect.toLowerCase() === 'postgres') dialect = 'postgresql'
 
-    if (!hasModifyingStatement(ast) && dataSource?.source === Source.external && dataSource?.request.headers.has('X-Starbase-Cache')) {
-        await createCacheTable(dataSource)
-        const fetchCacheStatement = `SELECT timestamp, ttl, query, results FROM tmp_cache WHERE query = ?`
-        const result = await dataSource.internalConnection?.durableObject.executeQuery(fetchCacheStatement, [sql], false) as any[];
+    let ast = parser.astify(sql, { database: dialect });
+    if (hasModifyingStatement(ast)) return null
+    
+    const fetchCacheStatement = `SELECT timestamp, ttl, query, results FROM tmp_cache WHERE query = ?`
+    const result = await dataSource.internalConnection?.durableObject.executeQuery(fetchCacheStatement, [sql], false) as any[];
 
-        if (result?.length) {
-            const { timestamp, ttl, results } = result[0];
-            const expirationTime = new Date(timestamp).getTime() + (ttl * 1000);
-            
-            if (Date.now() < expirationTime) {
-                return JSON.parse(results)
-            }
+    if (result?.length) {
+        const { timestamp, ttl, results } = result[0];
+        const expirationTime = new Date(timestamp).getTime() + (ttl * 1000);
+        
+        if (Date.now() < expirationTime) {
+            return JSON.parse(results)
         }
     }
 
@@ -67,16 +55,19 @@ export async function beforeQueryCache(sql: string, params?: any[], dataSource?:
 // to look into include using Cloudflare Cache but need to find a good way to cache the
 // response in a safe way for our use case. Another option is another service for queues
 // or another way to ingest it directly to the Durable Object.
-export async function afterQueryCache(sql: string, params: any[] | undefined, result: any, dataSource?: DataSource) {
+export async function afterQueryCache(sql: string, params: any[] | undefined, result: any, dataSource?: DataSource, dialect?: string) {
     // Currently we do not support caching queries that have dynamic parameters
     if (params?.length) return;
+    if (dataSource?.source === Source.internal || !dataSource?.request.headers.has('X-Starbase-Cache')) return null
 
     try {
-        let ast = parser.astify(sql);
+        if (!dialect) dialect = 'sqlite'
+        if (dialect.toLowerCase() === 'postgres') dialect = 'postgresql'
+
+        let ast = parser.astify(sql, { database: dialect });
 
         // If any modifying query exists within our SQL statement then we shouldn't proceed
-        if (hasModifyingStatement(ast) ||
-            !(dataSource?.source === Source.external && dataSource?.request.headers.has('X-Starbase-Cache'))) return;
+        if (hasModifyingStatement(ast)) return;
 
         const timestamp = Date.now();
         const results = JSON.stringify(result);

src/cache/index.ts

@@ -24,6 +24,38 @@ export class DatabaseDurableObject extends DurableObject {
         super(ctx, env);
         this.sql = ctx.storage.sql;
         this.storage = ctx.storage;
+
+        // Install default necessary `tmp_` tables for various features here.
+        const cacheStatement = `
+        CREATE TABLE IF NOT EXISTS tmp_cache (
+            "id" INTEGER PRIMARY KEY AUTOINCREMENT,
+            "timestamp" REAL NOT NULL,
+            "ttl" INTEGER NOT NULL,
+            "query" TEXT UNIQUE NOT NULL,
+            "results" TEXT
+        );`
+
+        const allowlistStatement = `
+        CREATE TABLE IF NOT EXISTS tmp_allowlist_queries (
+            id INTEGER PRIMARY KEY AUTOINCREMENT,
+            sql_statement TEXT NOT NULL
+        )`
+
+        const rlsStatement = `
+        CREATE TABLE IF NOT EXISTS tmp_rls_policies (
+            "id" INTEGER PRIMARY KEY AUTOINCREMENT,
+            "actions" TEXT NOT NULL CHECK(actions IN ('SELECT', 'UPDATE', 'INSERT', 'DELETE')),
+            "schema" TEXT,
+            "table" TEXT NOT NULL,
+            "column" TEXT NOT NULL,
+            "value" TEXT NOT NULL,
+            "value_type" TEXT NOT NULL DEFAULT 'string',
+            "operator" TEXT DEFAULT '='
+        )`
+
+        this.executeQuery(cacheStatement, undefined, false)
+        this.executeQuery(allowlistStatement, undefined, false)
+        this.executeQuery(rlsStatement, undefined, false)
     }
 
     /**

src/do.ts

@@ -99,7 +99,7 @@ export class Handler {
             }
 
             const { sql, params, transaction } = await request.json() as QueryRequest & QueryTransactionRequest;
-
+            
             if (Array.isArray(transaction) && transaction.length) {
                 const queries = transaction.map((queryObj: any) => {
                     const { sql, params } = queryObj;
@@ -125,7 +125,7 @@ export class Handler {
             return createResponse(response, undefined, 200);
         } catch (error: any) {
             console.error('Query Route Error:', error);
-            return createResponse(undefined, error || 'An unexpected error occurred.', 500);
+            return createResponse(undefined, error?.message || 'An unexpected error occurred.', 500);
         }
     }