Store allowlist rejections

Author: Brayden

src/allowlist/index.ts

@@ -16,17 +16,38 @@ function normalizeSQL(sql: string) {
 
 async function loadAllowlist(dataSource: DataSource): Promise<string[]> {
     try {
-        const statement = 'SELECT sql_statement FROM tmp_allowlist_queries'
+        const statement =
+            'SELECT sql_statement, source FROM tmp_allowlist_queries'
         const result = (await dataSource.rpc.executeQuery({
             sql: statement,
         })) as QueryResult[]
-        return result.map((row) => String(row.sql_statement))
+        return result
+            .filter((row) => row.source === dataSource.source)
+            .map((row) => String(row.sql_statement))
     } catch (error) {
         console.error('Error loading allowlist:', error)
         return []
     }
 }
 
+async function addRejectedQuery(
+    query: string,
+    dataSource: DataSource
+): Promise<string[]> {
+    try {
+        const statement =
+            'INSERT INTO tmp_allowlist_rejections (sql_statement, source) VALUES (?, ?)'
+        const result = (await dataSource.rpc.executeQuery({
+            sql: statement,
+            params: [query, dataSource.source],
+        })) as QueryResult[]
+        return result.map((row) => String(row.sql_statement))
+    } catch (error) {
+        console.error('Error inserting rejected allowlist query:', error)
+        return []
+    }
+}
+
 export async function isQueryAllowed(opts: {
     sql: string
     isEnabled: boolean
@@ -59,34 +80,78 @@ export async function isQueryAllowed(opts: {
         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
+        // 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) {
+        //                 // Preserve the value for specific clauses like LIMIT
+        //                 if (ast.as === 'limit' || ast.type === 'limit') {
+        //                     // Do not normalize LIMIT values
+        //                     return;
+        //                 }
+        //                 ast.value = '?'; // Normalize other values
+        //             }
+
+        //             // Recursively normalize all other fields
+        //             Object.values(ast).forEach(normalizeAst)
+        //         }
+
+        //         return ast;
+        //     };
+
+        //     normalizeAst(allowedAst)
+        //     normalizeAst(queryAst)
+
+        //     return JSON.stringify(allowedAst) === JSON.stringify(queryAst)
+        // })
+
+        const deepCompareAst = (allowedAst: any, queryAst: any): boolean => {
+            if (typeof allowedAst !== typeof queryAst) return false
+
+            if (Array.isArray(allowedAst) && Array.isArray(queryAst)) {
+                if (allowedAst.length !== queryAst.length) return false
+                return allowedAst.every((item, index) =>
+                    deepCompareAst(item, queryAst[index])
+                )
+            } else if (
+                typeof allowedAst === 'object' &&
+                allowedAst !== null &&
+                queryAst !== null
+            ) {
+                const allowedKeys = Object.keys(allowedAst)
+                const queryKeys = Object.keys(queryAst)
+
+                if (allowedKeys.length !== queryKeys.length) return false
+
+                return allowedKeys.every((key) =>
+                    deepCompareAst(allowedAst[key], queryAst[key])
+                )
             }
 
-            normalizeAst(allowedAst)
-            normalizeAst(queryAst)
+            // Base case: Primitive value comparison
+            return allowedAst === queryAst
+        }
 
-            return JSON.stringify(allowedAst) === JSON.stringify(queryAst)
-        })
+        const isCurrentAllowed = normalizedAllowlist?.some((allowedQuery) =>
+            deepCompareAst(allowedQuery, normalizedQuery)
+        )
 
         if (!isCurrentAllowed) {
+            // For any rejected query, we can add it to a table of rejected queries
+            // to act both as an audit log as well as an easy way to see recent queries
+            // that may need to be added to the allowlist in an easy way via a user
+            // interface.
+            addRejectedQuery(sql, dataSource)
+
+            // Then throw the appropriate error to the user.
             throw new Error('Query not allowed')
         }
 

src/do.ts

@@ -38,7 +38,15 @@ export class StarbaseDBDurableObject extends DurableObject {
         const allowlistStatement = `
         CREATE TABLE IF NOT EXISTS tmp_allowlist_queries (
             id INTEGER PRIMARY KEY AUTOINCREMENT,
-            sql_statement TEXT NOT NULL
+            sql_statement TEXT NOT NULL,
+            source TEXT DEFAULT 'external'
+        )`
+        const allowlistRejectedStatement = `
+        CREATE TABLE IF NOT EXISTS tmp_allowlist_rejections (
+            id INTEGER PRIMARY KEY AUTOINCREMENT,
+            sql_statement TEXT NOT NULL,
+            source TEXT DEFAULT 'external',
+            created_at TEXT DEFAULT (datetime('now'))
         )`
 
         const rlsStatement = `
@@ -55,6 +63,7 @@ export class StarbaseDBDurableObject extends DurableObject {
 
         this.executeQuery({ sql: cacheStatement })
         this.executeQuery({ sql: allowlistStatement })
+        this.executeQuery({ sql: allowlistRejectedStatement })
         this.executeQuery({ sql: rlsStatement })
     }
 

worker-configuration.d.ts

@@ -6,7 +6,7 @@ interface Env {
     REGION: 'auto'
     STUDIO_USER: 'admin'
     STUDIO_PASS: '123456'
-    ENABLE_ALLOWLIST: 0
+    ENABLE_ALLOWLIST: 1
     ENABLE_RLS: 0
     EXTERNAL_DB_TYPE: 'postgresql'
     AUTH_ALGORITHM: 'RS256'

wrangler.toml

@@ -43,7 +43,7 @@ STUDIO_USER = "admin"
 STUDIO_PASS = "123456"
 
 # Toggle to enable default features 
-ENABLE_ALLOWLIST = 0
+ENABLE_ALLOWLIST = 1
 ENABLE_RLS = 0
 
 # External database source details