Merge pull request #248 from kumavis/secure-eval

feat: more safely evaluate llm generated code in a SES Compartment

bots/template.js

@@ -1,10 +1,6 @@
-import * as skills from '../../../src/agent/library/skills.js';
-import * as world from '../../../src/agent/library/world.js';
-import Vec3 from 'vec3';
+(async (bot) => {
 
-const log = skills.log;
+/* CODE HERE */
+log(bot, 'Code finished.');
 
-export async function main(bot) {
-    /* CODE HERE */
-    log(bot, 'Code finished.');
-}
+})

package.json

@@ -18,6 +18,7 @@
         "prismarine-item": "^1.15.0",
         "prismarine-viewer": "^1.28.0",
         "replicate": "^0.29.4",
+        "ses": "^1.9.1",
         "vec3": "^0.1.10",
         "yargs": "^17.7.2"
     },

src/agent/coder.js

@@ -1,6 +1,9 @@
 import { writeFile, readFile, mkdirSync } from 'fs';
-import { checkSafe } from '../utils/safety.js';
 import settings from '../../settings.js';
+import { makeCompartment } from './library/lockdown.js';
+import * as skills from './library/skills.js';
+import * as world from './library/world.js';
+import { Vec3 } from 'vec3';
 
 export class Coder {
     constructor(agent) {
@@ -21,7 +24,7 @@ export class Coder {
         mkdirSync('.' + this.fp, { recursive: true });
     }
 
-    // write custom code to file and import it
+    // write custom code to file and prepare for evaluation
     async stageCode(code) {
         code = this.sanitizeCode(code);
         let src = '';
@@ -47,13 +50,24 @@ export class Coder {
         // } commented for now, useful to keep files for debugging
         this.file_counter++;
 
-        let write_result = await this.writeFilePromise('.' + this.fp + filename, src)
+        let write_result = await this.writeFilePromise('.' + this.fp + filename, src);
+        // This is where we determine the environment the agent's code should be exposed to.
+        // It will only have access to these things, (in addition to basic javascript objects like Array, Object, etc.)
+        // Note that the code may be able to modify the exposed objects.
+        const compartment = makeCompartment({
+            skills,
+            log: skills.log,
+            world,
+            Vec3,
+        });
+        const mainFn = compartment.evaluate(src);
         
         if (write_result) {
             console.error('Error writing code execution file: ' + result);
             return null;
         }
-        return await import('../..' + this.fp + filename);
+
+        return { main: mainFn };
     }
 
     sanitizeCode(code) {
@@ -130,21 +144,17 @@ export class Coder {
             }
             code = res.substring(res.indexOf('```')+3, res.lastIndexOf('```'));
 
-            if (!checkSafe(code)) {
-                console.warn(`Detected insecure generated code, not executing. Insecure code: \n\`${code}\``);
-                const message = 'Error: Code insecurity detected. Do not import, read/write files, execute dynamic code, or access the internet. Please try again:';
-                messages.push({ role: 'system', content: message });
-                continue;
-            }
-
-            const execution_file = await this.stageCode(code);
-            if (!execution_file) {
+            let codeStagingResult;
+            try {
+                codeStagingResult = await this.stageCode(code);
+            } catch (err) {
+                console.error('Error staging code:', err);
                 agent_history.add('system', 'Failed to stage code, something is wrong.');
                 return {success: false, message: null, interrupted: false, timedout: false};
             }
             
             code_return = await this.execute(async ()=>{
-                return await execution_file.main(this.agent.bot);
+                return await codeStagingResult.main(this.agent.bot);
             }, settings.code_timeout_mins);
             if (code_return.interrupted && !code_return.timedout)
                 return {success: false, message: null, interrupted: true, timedout: false};

src/agent/library/lockdown.js

@@ -0,0 +1,26 @@
+import 'ses';
+
+// This sets up the secure environment
+// We disable some of the taming to allow for more flexibility
+
+// For configuration, see https://github.com/endojs/endo/blob/master/packages/ses/docs/lockdown.md
+lockdown({
+  // basic devex and quality of life improvements
+  localeTaming: 'unsafe',
+  consoleTaming: 'unsafe',
+  errorTaming: 'unsafe',
+  stackFiltering: 'verbose',
+  // allow eval outside of created compartments
+  // (mineflayer dep "protodef" uses eval)
+  evalTaming: 'unsafeEval',
+});
+
+export const makeCompartment = (endowments = {}) => {
+  return new Compartment({
+    // provide untamed Math, Date, etc
+    Math,
+    Date,
+    // standard endowments
+    ...endowments
+  });
+}

src/utils/safety.js

@@ -1,38 +0,0 @@
-export function checkSafe(code) {
-    const dangerousPatterns = [
-        // Dynamic imports
-        /\bimport\s*\(/,
-        // Access to process and global
-        /\bprocess\b/,
-        /\bglobal\b/,
-        // Module manipulation
-        /\bmodule\b/,
-        /\bexports\b/,
-        // Require usage
-        /\brequire\s*\(/,
-        // Function constructors
-        /\bFunction\s*\(/,
-        /\beval\s*\(/,
-        // Access to __dirname and __filename
-        /\b__dirname\b/,
-        /\b__filename\b/,
-
-        // fetch
-        /\bfetch\s*\(/,
-        // XMLHttpRequest
-        /\bXMLHttpRequest\b/,
-        // Websockets
-        /\bWebSocket\b/,
-    ];
-
-    for (const pattern of dangerousPatterns) {
-        if (pattern.test(code)) {
-            return false;
-        }
-    }
-    return true;
-}
-
-// generated by o1
-// Basic check for malicious code like dynamic imports, code exec, disk access, internet access, etc.
-// Will not catch all, and can be bypassed by obfuscation.