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
+  });
+}