Added secure traffic access (#180)

* secure traffic access

* changed type syntax

* bugbot suggestions

* update test setup conf

* format

* cursor suggestion

* increase default timeout

* sandbox not callable

* nit

Author: mishushakov

.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,17 @@ export class Sandbox extends BaseSandbox {
    */
   async createCodeContext(opts?: CreateCodeContextOpts): Promise<Context> {
     try {
+      const headers: Record<string, string> = {
+        'Content-Type': 'application/json',
+      }
+
+      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 +340,17 @@ 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',
+      }
+
+      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
@@ -359,12 +373,17 @@ export class Sandbox extends BaseSandbox {
    */
   async listCodeContexts(): Promise<Context[]> {
     try {
+      const headers: Record<string, string> = {
+        'Content-Type': 'application/json',
+      }
+
+      if (this.trafficAccessToken) {
+        headers['E2B-Traffic-Access-Token'] = this.trafficAccessToken
+      }
+
       const res = await fetch(`${this.jupyterUrl}/contexts`, {
         method: 'GET',
-        headers: {
-          'Content-Type': 'application/json',
-          ...this.connectionConfig.headers,
-        },
+        headers,
         keepalive: true,
         signal: this.connectionConfig.getSignal(
           this.connectionConfig.requestTimeoutMs
@@ -392,12 +411,17 @@ 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',
+      }
+
+      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,14 @@
 import { expect } from 'vitest'
-
-import { sandboxTest } from './setup'
+import { isDebug, sandboxTest, secureSandboxTest } from './setup'
 
 sandboxTest('basic', async ({ sandbox }) => {
   const result = await sandbox.runCode('x =1; x')
 
   expect(result.text).toEqual('1')
 })
+
+secureSandboxTest.skipIf(isDebug)('secure access', async ({ sandbox }) => {
+  const result = await sandbox.runCode('x =1; x')
+
+  expect(result.text).toEqual('1')
+})

js/tests/contexts.test.ts

@@ -1,6 +1,6 @@
 import { expect } from 'vitest'
 
-import { sandboxTest } from './setup'
+import { isDebug, sandboxTest, secureSandboxTest } from './setup'
 
 sandboxTest('create context with no options', async ({ sandbox }) => {
   const context = await sandbox.createCodeContext()
@@ -61,3 +61,65 @@ sandboxTest('restart context', async ({ sandbox }) => {
   expect(execution.error?.name).toBe('NameError')
   expect(execution.error?.value).toBe("name 'x' is not defined")
 })
+
+secureSandboxTest.skipIf(isDebug)(
+  'create context (secure traffic)',
+  async ({ sandbox }) => {
+    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)
+  }
+)
+
+secureSandboxTest.skipIf(isDebug)(
+  'remove context (secure traffic)',
+  async ({ sandbox }) => {
+    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()
+  }
+)
+
+secureSandboxTest.skipIf(isDebug)(
+  'list contexts (secure traffic)',
+  async ({ sandbox }) => {
+    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()
+  }
+)
+
+secureSandboxTest.skipIf(isDebug)(
+  'restart context (secure traffic)',
+  async ({ sandbox }) => {
+    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")
+  }
+)

js/tests/setup.ts

@@ -1,21 +1,31 @@
-import { Sandbox } from '../src'
 import { test as base } from 'vitest'
-
-const timeoutMs = 60_000
-
-const template = process.env.E2B_TESTS_TEMPLATE || 'code-interpreter-v1'
+import { Sandbox, SandboxOpts } from '../src'
 
 interface SandboxFixture {
   sandbox: Sandbox
   template: string
+  sandboxTestId: string
+  sandboxOpts: Partial<SandboxOpts>
 }
 
+const template = process.env.E2B_TESTS_TEMPLATE || 'code-interpreter-v1'
+
 export const sandboxTest = base.extend<SandboxFixture>({
-  sandbox: [
+  template,
+  sandboxTestId: [
     // eslint-disable-next-line no-empty-pattern
     async ({}, use) => {
+      const id = `test-${generateRandomString()}`
+      await use(id)
+    },
+    { auto: true },
+  ],
+  sandboxOpts: {},
+  sandbox: [
+    async ({ sandboxTestId, sandboxOpts }, use) => {
       const sandbox = await Sandbox.create(template, {
-        timeoutMs,
+        metadata: { sandboxTestId },
+        ...sandboxOpts,
       })
       try {
         await use(sandbox)
@@ -31,13 +41,30 @@ export const sandboxTest = base.extend<SandboxFixture>({
         }
       }
     },
-    { auto: true },
+    { auto: false },
   ],
-  template,
 })
 
 export const isDebug = process.env.E2B_DEBUG !== undefined
+export const isIntegrationTest = process.env.E2B_INTEGRATION_TEST !== undefined
+
+export const secureSandboxTest = sandboxTest.extend({
+  sandboxOpts: {
+    secure: true,
+    network: {
+      allowPublicTraffic: false,
+    },
+  },
+})
+
+function generateRandomString(length: number = 8): string {
+  return Math.random()
+    .toString(36)
+    .substring(2, length + 2)
+}
 
 export async function wait(ms: number) {
   return new Promise((resolve) => setTimeout(resolve, ms))
 }
+
+export { template }

pnpm-lock.yaml

@@ -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,