Added secure traffic access

.changeset/fast-dogs-brush.md

@@ -0,0 +1,6 @@
+---
+'@e2b/code-interpreter-python': patch
+'@e2b/code-interpreter': patch
+---
+
+secure traffic access

js/src/sandbox.ts

@@ -223,6 +223,10 @@ export class Sandbox extends BaseSandbox {
     const headers: Record<string, string> = {
       'Content-Type': 'application/json',
     }
+
+    if (this.trafficAccessToken) {
+      headers['E2B-Traffic-Access-Token'] = this.trafficAccessToken
+    }
     if (this.envdAccessToken) {
       headers['X-Access-Token'] = this.envdAccessToken
     }
@@ -296,12 +300,18 @@ export class Sandbox extends BaseSandbox {
    */
   async createCodeContext(opts?: CreateCodeContextOpts): Promise<Context> {
     try {
+      const headers: Record<string, string> = {
+        'Content-Type': 'application/json',
+        ...this.connectionConfig.headers,
+      }
+
+      if (this.trafficAccessToken) {
+        headers['E2B-Traffic-Access-Token'] = this.trafficAccessToken
+      }
+
       const res = await fetch(`${this.jupyterUrl}/contexts`, {
         method: 'POST',
-        headers: {
-          'Content-Type': 'application/json',
-          ...this.connectionConfig.headers,
-        },
+        headers,
         body: JSON.stringify({
           language: opts?.language,
           cwd: opts?.cwd,
@@ -331,12 +341,18 @@ export class Sandbox extends BaseSandbox {
   async removeCodeContext(context: Context | string): Promise<void> {
     try {
       const id = typeof context === 'string' ? context : context.id
+      const headers: Record<string, string> = {
+        'Content-Type': 'application/json',
+        ...this.connectionConfig.headers,
+      }
+
+      if (this.trafficAccessToken) {
+        headers['E2B-Traffic-Access-Token'] = this.trafficAccessToken
+      }
+
       const res = await fetch(`${this.jupyterUrl}/contexts/${id}`, {
         method: 'DELETE',
-        headers: {
-          'Content-Type': 'application/json',
-          ...this.connectionConfig.headers,
-        },
+        headers,
         keepalive: true,
         signal: this.connectionConfig.getSignal(
           this.connectionConfig.requestTimeoutMs
@@ -364,6 +380,9 @@ export class Sandbox extends BaseSandbox {
         headers: {
           'Content-Type': 'application/json',
           ...this.connectionConfig.headers,
+          ...(this.trafficAccessToken
+            ? { 'E2B-Traffic-Access-Token': this.trafficAccessToken }
+            : {}),
         },
         keepalive: true,
         signal: this.connectionConfig.getSignal(
@@ -392,12 +411,18 @@ export class Sandbox extends BaseSandbox {
   async restartCodeContext(context: Context | string): Promise<void> {
     try {
       const id = typeof context === 'string' ? context : context.id
+      const headers: Record<string, string> = {
+        'Content-Type': 'application/json',
+        ...this.connectionConfig.headers,
+      }
+
+      if (this.trafficAccessToken) {
+        headers['E2B-Traffic-Access-Token'] = this.trafficAccessToken
+      }
+
       const res = await fetch(`${this.jupyterUrl}/contexts/${id}/restart`, {
         method: 'POST',
-        headers: {
-          'Content-Type': 'application/json',
-          ...this.connectionConfig.headers,
-        },
+        headers,
         keepalive: true,
         signal: this.connectionConfig.getSignal(
           this.connectionConfig.requestTimeoutMs

js/tests/basic.test.ts

@@ -1,9 +1,23 @@
 import { expect } from 'vitest'
-
-import { sandboxTest } from './setup'
+import { isDebug, sandboxTest } from './setup'
+import { Sandbox } from '../src'
 
 sandboxTest('basic', async ({ sandbox }) => {
   const result = await sandbox.runCode('x =1; x')
 
   expect(result.text).toEqual('1')
 })
+
+sandboxTest.skipIf(isDebug)('secure access', async ({ template }) => {
+  const sandbox = await Sandbox.create(template, {
+    network: {
+      allowPublicTraffic: false,
+    },
+  })
+
+  const result = await sandbox.runCode('x =1; x')
+
+  expect(result.text).toEqual('1')
+
+  await sandbox.kill()
+})

js/tests/contexts.test.ts

@@ -1,6 +1,7 @@
 import { expect } from 'vitest'
 
-import { sandboxTest } from './setup'
+import { isDebug, sandboxTest } from './setup'
+import Sandbox from '../src'
 
 sandboxTest('create context with no options', async ({ sandbox }) => {
   const context = await sandbox.createCodeContext()
@@ -61,3 +62,89 @@ sandboxTest('restart context', async ({ sandbox }) => {
   expect(execution.error?.name).toBe('NameError')
   expect(execution.error?.value).toBe("name 'x' is not defined")
 })
+
+sandboxTest.skipIf(isDebug)(
+  'create context (secure traffic)',
+  async ({ template }) => {
+    const sandbox = await Sandbox.create(template, {
+      network: {
+        allowPublicTraffic: false,
+      },
+    })
+    const context = await sandbox.createCodeContext()
+
+    const contexts = await sandbox.listCodeContexts()
+    const lastContext = contexts[contexts.length - 1]
+
+    expect(lastContext.id).toBe(context.id)
+    expect(lastContext.language).toBe(context.language)
+    expect(lastContext.cwd).toBe(context.cwd)
+
+    await sandbox.kill()
+  }
+)
+
+sandboxTest.skipIf(isDebug)(
+  'remove context (secure traffic)',
+  async ({ template }) => {
+    const sandbox = await Sandbox.create(template, {
+      network: {
+        allowPublicTraffic: false,
+      },
+    })
+    const context = await sandbox.createCodeContext()
+
+    await sandbox.removeCodeContext(context.id)
+    const contexts = await sandbox.listCodeContexts()
+
+    expect(contexts.map((context) => context.id)).not.toContain(context.id)
+
+    await sandbox.kill()
+  }
+)
+
+sandboxTest.skipIf(isDebug)(
+  'list contexts (secure traffic)',
+  async ({ template }) => {
+    const sandbox = await Sandbox.create(template, {
+      network: {
+        allowPublicTraffic: false,
+      },
+    })
+    const contexts = await sandbox.listCodeContexts()
+
+    // default contexts should include python and javascript
+    expect(contexts.map((context) => context.language)).toContain('python')
+    expect(contexts.map((context) => context.language)).toContain('javascript')
+
+    await sandbox.kill()
+  }
+)
+
+sandboxTest.skipIf(isDebug)(
+  'restart context (secure traffic)',
+  async ({ template }) => {
+    const sandbox = await Sandbox.create(template, {
+      network: {
+        allowPublicTraffic: false,
+      },
+    })
+    const context = await sandbox.createCodeContext()
+
+    // set a variable in the context
+    await sandbox.runCode('x = 1', { context: context })
+
+    // restart the context
+    await sandbox.restartCodeContext(context.id)
+
+    // check that the variable no longer exists
+    const execution = await sandbox.runCode('x', { context: context })
+
+    // check for an NameError with message "name 'x' is not defined"
+    expect(execution.error).toBeDefined()
+    expect(execution.error?.name).toBe('NameError')
+    expect(execution.error?.value).toBe("name 'x' is not defined")
+
+    await sandbox.kill()
+  }
+)

python/e2b_code_interpreter/code_interpreter_async.py

@@ -190,12 +190,15 @@ async def run_code(
         timeout = None if timeout == 0 else (timeout or DEFAULT_TIMEOUT)
         request_timeout = request_timeout or self.connection_config.request_timeout
         context_id = context.id if context else None
-
-        headers: Dict[str, str] = {}
-        if self._envd_access_token:
-            headers = {"X-Access-Token": self._envd_access_token}
-
         try:
+            headers = {
+                "Content-Type": "application/json",
+            }
+            if self._envd_access_token:
+                headers["X-Access-Token"] = self._envd_access_token
+            if self.traffic_access_token:
+                headers["E2B-Traffic-Access-Token"] = self.traffic_access_token
+
             async with self._client.stream(
                 "POST",
                 f"{self._jupyter_url}/execute",
@@ -253,11 +256,13 @@ async def create_code_context(
         if cwd:
             data["cwd"] = cwd
 
-        headers: Dict[str, str] = {}
-        if self._envd_access_token:
-            headers = {"X-Access-Token": self._envd_access_token}
-
         try:
+            headers = {
+                "Content-Type": "application/json",
+            }
+            if self.traffic_access_token:
+                headers["E2B-Traffic-Access-Token"] = self.traffic_access_token
+
             response = await self._client.post(
                 f"{self._jupyter_url}/contexts",
                 headers=headers,
@@ -287,11 +292,13 @@ async def remove_code_context(
         """
         context_id = context.id if isinstance(context, Context) else context
 
-        headers: Dict[str, str] = {}
-        if self._envd_access_token:
-            headers = {"X-Access-Token": self._envd_access_token}
-
         try:
+            headers = {
+                "Content-Type": "application/json",
+            }
+            if self.traffic_access_token:
+                headers["E2B-Traffic-Access-Token"] = self.traffic_access_token
+
             response = await self._client.delete(
                 f"{self._jupyter_url}/contexts/{context_id}",
                 headers=headers,
@@ -310,11 +317,13 @@ async def list_code_contexts(self) -> List[Context]:
 
         :return: List of contexts.
         """
-        headers: Dict[str, str] = {}
-        if self._envd_access_token:
-            headers = {"X-Access-Token": self._envd_access_token}
-
         try:
+            headers = {
+                "Content-Type": "application/json",
+            }
+            if self.traffic_access_token:
+                headers["E2B-Traffic-Access-Token"] = self.traffic_access_token
+
             response = await self._client.get(
                 f"{self._jupyter_url}/contexts",
                 headers=headers,
@@ -342,12 +351,13 @@ async def restart_code_context(
         :return: None
         """
         context_id = context.id if isinstance(context, Context) else context
-
-        headers: Dict[str, str] = {}
-        if self._envd_access_token:
-            headers = {"X-Access-Token": self._envd_access_token}
-
         try:
+            headers = {
+                "Content-Type": "application/json",
+            }
+            if self.traffic_access_token:
+                headers["E2B-Traffic-Access-Token"] = self.traffic_access_token
+
             response = await self._client.post(
                 f"{self._jupyter_url}/contexts/{context_id}/restart",
                 headers=headers,