From f0d7cace9464fde1239eb979f93122323c08f02a Mon Sep 17 00:00:00 2001
From: Tully <tully@apiary.io>
Date: Tue, 2 Sep 2014 13:19:01 +0200
Subject: [PATCH] modules loading inside VM (just stdlib
 http://nodejs.org/api/)

---
 src/forkable.coffee      |  8 ++++++--
 src/pitboss.coffee       |  9 +++++----
 test/pitboss_test.coffee | 42 ++++++++++++++++++++++++++++++----------
 3 files changed, 43 insertions(+), 16 deletions(-)

diff --git a/src/forkable.coffee b/src/forkable.coffee
index 6d9ae7a..7ae8b3d 100644
--- a/src/forkable.coffee
+++ b/src/forkable.coffee
@@ -29,8 +29,12 @@ run = (msg) ->
   unless script
     error "No code to run"
     return false
-  if msg?.context?.debug is true
-    msg.context.console = console
+
+  msg.context ?= {}
+
+  if msg?.libraries
+    for lib in msg?.libraries
+      msg.context[lib] = eval lib
   try
     res =
       result: script.runInNewContext(msg.context || {}) || null # script can return undefined, ensure it's null
diff --git a/src/pitboss.coffee b/src/pitboss.coffee
index bae89c1..a3d57f1 100644
--- a/src/pitboss.coffee
+++ b/src/pitboss.coffee
@@ -8,15 +8,15 @@ exports.Pitboss = class Pitboss extends EventEmitter
       @next()
     @q = []
 
-  run: (context, callback) ->
-    @q.push({context:context,callback:callback})
+  run: ({context, libraries}, callback) ->
+    @q.push({context: context, libraries: libraries, callback: callback})
     @next()
 
   next: ->
     return false if @runner.running
     c = @q.shift()
     if c
-      @runner.run(c.context, c.callback)
+      @runner.run({context: c.context, libraries: c.libraries}, c.callback)
 
 # Can only run one at a time due to the blocking nature of VM
 # Need to queue this up outside of the process since it's over an async channel
@@ -34,11 +34,12 @@ exports.Runner = class Runner extends EventEmitter
     @proc.on 'exit', @failedForkHandler
     @proc.send {code:@code}
 
-  run: (context, callback) ->
+  run: ({context, libraries}, callback) ->
     return false if @running
     id = Date.now().toString() + Math.floor(Math.random() * 1000)
     msg =
       context: context
+      libraries: libraries
       id: id
     @callback = callback || false
     @startTimer()
diff --git a/test/pitboss_test.coffee b/test/pitboss_test.coffee
index 163828c..21813c9 100644
--- a/test/pitboss_test.coffee
+++ b/test/pitboss_test.coffee
@@ -14,12 +14,34 @@ describe "Pitboss running code", ->
 
   it "should take a JSON encodable message", (done) ->
     pitboss = new Pitboss(@code)
-    pitboss.run {data: "test"}, (err, result) ->
+    pitboss.run context: {data: "test"}, (err, result) ->
       assert.equal "test", result
-    pitboss.run {data: 456}, (err, result) ->
+    pitboss.run context: {data: 456}, (err, result) ->
       assert.equal 456, result
       done()
 
+
+describe "Pitboss modules loading code", ->
+  beforeEach ->
+    @code = """
+      console.error(data);
+      data;
+    """
+
+  it "should not return an error when loaded module is used", (done) ->
+    pitboss = new Pitboss(@code)
+    pitboss.run context: {data: "test"}, libraries: ['console'], (err, result) ->
+      assert.equal "test", result
+      assert.equal undefined, err
+      done()
+
+  it "should return an error when unknown module is used", (done) ->
+    pitboss = new Pitboss(@code)
+    pitboss.run context: {data: "test"}, libraries: [], (err, result) ->
+      assert.equal undefined, result
+      assert.equal 'VM Runtime Error: ReferenceError: console is not defined', err
+      done()
+
 describe "Running dubius code", ->
   beforeEach ->
     @code = """
@@ -32,7 +54,7 @@ describe "Running dubius code", ->
 
   it "should take a JSON encodable message", (done) ->
     pitboss = new Runner(@code)
-    pitboss.run {data: 123}, (err, result) ->
+    pitboss.run context: {data: 123}, (err, result) ->
       assert.equal 123, result
       done()
 
@@ -44,7 +66,7 @@ describe "Running shitty code", ->
 
   it "should return the error", (done) ->
     pitboss = new Runner(@code)
-    pitboss.run {data: 123}, (err, result) ->
+    pitboss.run context: {data: 123}, (err, result) ->
       assert.equal "VM Syntax Error: SyntaxError: Unexpected identifier", err
       assert.equal null, result
       done()
@@ -60,19 +82,19 @@ describe "Running infinite loop code", ->
 
   it "should timeout and restart fork", (done) ->
     pitboss = new Runner @code,
-      timeout: 100
-    pitboss.run {infinite: true}, (err, result) ->
+      timeout: 200
+    pitboss.run context: {infinite: true}, (err, result) ->
       assert.equal "Timedout", err
-      pitboss.run {infinite: false}, (err, result) ->
+      pitboss.run context: {infinite: false}, (err, result) ->
         assert.equal "OK", result
         done()
 
   it "should happily allow for process failure (e.g. ulimit kills)", (done) ->
     pitboss = new Runner @code,
-      timeout: 100
-    pitboss.run {infinite: true}, (err, result) ->
+      timeout: 200
+    pitboss.run context: {infinite: true}, (err, result) ->
       assert.equal "Process Failed", err
-      pitboss.run {infinite: false}, (err, result) ->
+      pitboss.run context: {infinite: false}, (err, result) ->
         assert.equal "OK", result
         done()
     pitboss.proc.kill()