Skip to content

Commit a98b799

Browse files
mishushakovmishushakov
committed
secure traffic access
1 parent 6acf33a commit a98b799

File tree

12 files changed

+350
-63
lines changed

12 files changed

+350
-63
lines changed

‎.changeset/fast-dogs-brush.md‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,6 @@
1+
---
2+
'@e2b/code-interpreter-python': patch
3+
'@e2b/code-interpreter': patch
4+
---
5+
6+
secure traffic access

‎js/src/sandbox.ts‎

Lines changed: 38 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -220,8 +220,12 @@ export class Sandbox extends BaseSandbox {
220220
}, requestTimeout)
221221
: undefined
222222

223-
const headers: Record<string, string> = {
223+
const headers = {
224224
'Content-Type': 'application/json',
225+
} as Record<string, string>
226+
227+
if (this.trafficAccessToken) {
228+
headers['E2B-Traffic-Access-Token'] = this.trafficAccessToken
225229
}
226230
if (this.envdAccessToken) {
227231
headers['X-Access-Token'] = this.envdAccessToken
@@ -296,12 +300,18 @@ export class Sandbox extends BaseSandbox {
296300
*/
297301
async createCodeContext(opts?: CreateCodeContextOpts): Promise<Context> {
298302
try {
303+
const headers = {
304+
'Content-Type': 'application/json',
305+
...this.connectionConfig.headers,
306+
} as Record<string, string>
307+
308+
if (this.trafficAccessToken) {
309+
headers['E2B-Traffic-Access-Token'] = this.trafficAccessToken
310+
}
311+
299312
const res = await fetch(`${this.jupyterUrl}/contexts`, {
300313
method: 'POST',
301-
headers: {
302-
'Content-Type': 'application/json',
303-
...this.connectionConfig.headers,
304-
},
314+
headers,
305315
body: JSON.stringify({
306316
language: opts?.language,
307317
cwd: opts?.cwd,
@@ -331,12 +341,18 @@ export class Sandbox extends BaseSandbox {
331341
async removeCodeContext(context: Context | string): Promise<void> {
332342
try {
333343
const id = typeof context === 'string' ? context : context.id
344+
const headers = {
345+
'Content-Type': 'application/json',
346+
...this.connectionConfig.headers,
347+
} as Record<string, string>
348+
349+
if (this.trafficAccessToken) {
350+
headers['E2B-Traffic-Access-Token'] = this.trafficAccessToken
351+
}
352+
334353
const res = await fetch(`${this.jupyterUrl}/contexts/${id}`, {
335354
method: 'DELETE',
336-
headers: {
337-
'Content-Type': 'application/json',
338-
...this.connectionConfig.headers,
339-
},
355+
headers,
340356
keepalive: true,
341357
signal: this.connectionConfig.getSignal(
342358
this.connectionConfig.requestTimeoutMs
@@ -364,6 +380,9 @@ export class Sandbox extends BaseSandbox {
364380
headers: {
365381
'Content-Type': 'application/json',
366382
...this.connectionConfig.headers,
383+
...(this.trafficAccessToken
384+
? { 'E2B-Traffic-Access-Token': this.trafficAccessToken }
385+
: {}),
367386
},
368387
keepalive: true,
369388
signal: this.connectionConfig.getSignal(
@@ -392,12 +411,18 @@ export class Sandbox extends BaseSandbox {
392411
async restartCodeContext(context: Context | string): Promise<void> {
393412
try {
394413
const id = typeof context === 'string' ? context : context.id
414+
const headers = {
415+
'Content-Type': 'application/json',
416+
...this.connectionConfig.headers,
417+
} as Record<string, string>
418+
419+
if (this.trafficAccessToken) {
420+
headers['E2B-Traffic-Access-Token'] = this.trafficAccessToken
421+
}
422+
395423
const res = await fetch(`${this.jupyterUrl}/contexts/${id}/restart`, {
396424
method: 'POST',
397-
headers: {
398-
'Content-Type': 'application/json',
399-
...this.connectionConfig.headers,
400-
},
425+
headers,
401426
keepalive: true,
402427
signal: this.connectionConfig.getSignal(
403428
this.connectionConfig.requestTimeoutMs

‎js/tests/basic.test.ts‎

Lines changed: 16 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,9 +1,23 @@
11
import { expect } from 'vitest'
2-
3-
import { sandboxTest } from './setup'
2+
import { isDebug, sandboxTest } from './setup'
3+
import { Sandbox } from '../src'
44

55
sandboxTest('basic', async ({ sandbox }) => {
66
const result = await sandbox.runCode('x =1; x')
77

88
expect(result.text).toEqual('1')
99
})
10+
11+
sandboxTest.skipIf(isDebug)('secure access', async ({ template }) => {
12+
const sandbox = await Sandbox.create(template, {
13+
network: {
14+
allowPublicTraffic: false,
15+
},
16+
})
17+
18+
const result = await sandbox.runCode('x =1; x')
19+
20+
expect(result.text).toEqual('1')
21+
22+
await sandbox.kill()
23+
})

‎js/tests/contexts.test.ts‎

Lines changed: 88 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
import { expect } from 'vitest'
22

3-
import { sandboxTest } from './setup'
3+
import { isDebug, sandboxTest } from './setup'
4+
import Sandbox from '../src'
45

56
sandboxTest('create context with no options', async ({ sandbox }) => {
67
const context = await sandbox.createCodeContext()
@@ -61,3 +62,89 @@ sandboxTest('restart context', async ({ sandbox }) => {
6162
expect(execution.error?.name).toBe('NameError')
6263
expect(execution.error?.value).toBe("name 'x' is not defined")
6364
})
65+
66+
sandboxTest.skipIf(isDebug)(
67+
'create context (secure traffic)',
68+
async ({ template }) => {
69+
const sandbox = await Sandbox.create(template, {
70+
network: {
71+
allowPublicTraffic: false,
72+
},
73+
})
74+
const context = await sandbox.createCodeContext()
75+
76+
const contexts = await sandbox.listCodeContexts()
77+
const lastContext = contexts[contexts.length - 1]
78+
79+
expect(lastContext.id).toBe(context.id)
80+
expect(lastContext.language).toBe(context.language)
81+
expect(lastContext.cwd).toBe(context.cwd)
82+
83+
await sandbox.kill()
84+
}
85+
)
86+
87+
sandboxTest.skipIf(isDebug)(
88+
'remove context (secure traffic)',
89+
async ({ template }) => {
90+
const sandbox = await Sandbox.create(template, {
91+
network: {
92+
allowPublicTraffic: false,
93+
},
94+
})
95+
const context = await sandbox.createCodeContext()
96+
97+
await sandbox.removeCodeContext(context.id)
98+
const contexts = await sandbox.listCodeContexts()
99+
100+
expect(contexts.map((context) => context.id)).not.toContain(context.id)
101+
102+
await sandbox.kill()
103+
}
104+
)
105+
106+
sandboxTest.skipIf(isDebug)(
107+
'list contexts (secure traffic)',
108+
async ({ template }) => {
109+
const sandbox = await Sandbox.create(template, {
110+
network: {
111+
allowPublicTraffic: false,
112+
},
113+
})
114+
const contexts = await sandbox.listCodeContexts()
115+
116+
// default contexts should include python and javascript
117+
expect(contexts.map((context) => context.language)).toContain('python')
118+
expect(contexts.map((context) => context.language)).toContain('javascript')
119+
120+
await sandbox.kill()
121+
}
122+
)
123+
124+
sandboxTest.skipIf(isDebug)(
125+
'restart context (secure traffic)',
126+
async ({ template }) => {
127+
const sandbox = await Sandbox.create(template, {
128+
network: {
129+
allowPublicTraffic: false,
130+
},
131+
})
132+
const context = await sandbox.createCodeContext()
133+
134+
// set a variable in the context
135+
await sandbox.runCode('x = 1', { context: context })
136+
137+
// restart the context
138+
await sandbox.restartCodeContext(context.id)
139+
140+
// check that the variable no longer exists
141+
const execution = await sandbox.runCode('x', { context: context })
142+
143+
// check for an NameError with message "name 'x' is not defined"
144+
expect(execution.error).toBeDefined()
145+
expect(execution.error?.name).toBe('NameError')
146+
expect(execution.error?.value).toBe("name 'x' is not defined")
147+
148+
await sandbox.kill()
149+
}
150+
)

‎pnpm-lock.yaml‎

Lines changed: 1 addition & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎python/e2b_code_interpreter/code_interpreter_async.py‎

Lines changed: 32 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -190,12 +190,15 @@ async def run_code(
190190
timeout = None if timeout == 0 else (timeout or DEFAULT_TIMEOUT)
191191
request_timeout = request_timeout or self.connection_config.request_timeout
192192
context_id = context.id if context else None
193-
194-
headers: Dict[str, str] = {}
195-
if self._envd_access_token:
196-
headers = {"X-Access-Token": self._envd_access_token}
197-
198193
try:
194+
headers = {
195+
"Content-Type": "application/json",
196+
}
197+
if self._envd_access_token:
198+
headers["X-Access-Token"] = self._envd_access_token
199+
if self.traffic_access_token:
200+
headers["E2B-Traffic-Access-Token"] = self.traffic_access_token
201+
199202
async with self._client.stream(
200203
"POST",
201204
f"{self._jupyter_url}/execute",
@@ -253,11 +256,13 @@ async def create_code_context(
253256
if cwd:
254257
data["cwd"] = cwd
255258

256-
headers: Dict[str, str] = {}
257-
if self._envd_access_token:
258-
headers = {"X-Access-Token": self._envd_access_token}
259-
260259
try:
260+
headers = {
261+
"Content-Type": "application/json",
262+
}
263+
if self.traffic_access_token:
264+
headers["E2B-Traffic-Access-Token"] = self.traffic_access_token
265+
261266
response = await self._client.post(
262267
f"{self._jupyter_url}/contexts",
263268
headers=headers,
@@ -287,11 +292,13 @@ async def remove_code_context(
287292
"""
288293
context_id = context.id if isinstance(context, Context) else context
289294

290-
headers: Dict[str, str] = {}
291-
if self._envd_access_token:
292-
headers = {"X-Access-Token": self._envd_access_token}
293-
294295
try:
296+
headers = {
297+
"Content-Type": "application/json",
298+
}
299+
if self.traffic_access_token:
300+
headers["E2B-Traffic-Access-Token"] = self.traffic_access_token
301+
295302
response = await self._client.delete(
296303
f"{self._jupyter_url}/contexts/{context_id}",
297304
headers=headers,
@@ -310,11 +317,13 @@ async def list_code_contexts(self) -> List[Context]:
310317
311318
:return: List of contexts.
312319
"""
313-
headers: Dict[str, str] = {}
314-
if self._envd_access_token:
315-
headers = {"X-Access-Token": self._envd_access_token}
316-
317320
try:
321+
headers = {
322+
"Content-Type": "application/json",
323+
}
324+
if self.traffic_access_token:
325+
headers["E2B-Traffic-Access-Token"] = self.traffic_access_token
326+
318327
response = await self._client.get(
319328
f"{self._jupyter_url}/contexts",
320329
headers=headers,
@@ -342,12 +351,13 @@ async def restart_code_context(
342351
:return: None
343352
"""
344353
context_id = context.id if isinstance(context, Context) else context
345-
346-
headers: Dict[str, str] = {}
347-
if self._envd_access_token:
348-
headers = {"X-Access-Token": self._envd_access_token}
349-
350354
try:
355+
headers = {
356+
"Content-Type": "application/json",
357+
}
358+
if self.traffic_access_token:
359+
headers["E2B-Traffic-Access-Token"] = self.traffic_access_token
360+
351361
response = await self._client.post(
352362
f"{self._jupyter_url}/contexts/{context_id}/restart",
353363
headers=headers,

0 commit comments

Comments
 (0)