diff --git a/.changeset/network-transform-iam-placeholders.md b/.changeset/network-transform-iam-placeholders.md new file mode 100644 index 0000000000..b0d6e5004d --- /dev/null +++ b/.changeset/network-transform-iam-placeholders.md @@ -0,0 +1,56 @@ +--- +'e2b': minor +'@e2b/python-sdk': minor +--- + +Allow a network rule's `transform` to be a callback, so a workload identity token from the `iam` option can be injected into egress requests without the SDK ever seeing its value. The callback receives placeholder strings that the egress proxy resolves per request — `iam.tokens.aws` is `${e2b.identity.tokens.aws}` on the wire — and referencing a token that is not registered in `iam.tokens` fails with `InvalidArgumentError` / `InvalidArgumentException` instead of silently sending a placeholder no token will ever replace. + +`updateNetwork` / `update_network` accepts the same callbacks, but its payload carries no `iam` config, so token names cannot be checked there and every name resolves to its placeholder. + +```ts +import { Sandbox, Secret } from 'e2b' + +const sandbox = await Sandbox.create({ + iam: { + tokens: { + aws: Secret.iamToken({ audience: 'sts.amazonaws.com', tokenType: 'JWT-SVID' }), + }, + }, + network: { + allowOut: ({ rules }) => [...rules.keys()], + rules: { + 'api.internal.example.com': [ + { + transform: ({ iam }) => ({ + headers: { Authorization: `Bearer ${iam.tokens.aws}` }, + }), + }, + ], + }, + }, +}) +``` + +```python +from e2b import Sandbox, Secret + +sandbox = Sandbox.create( + iam={ + "tokens": { + "aws": Secret.iam_token(audience="sts.amazonaws.com", token_type="JWT-SVID"), + }, + }, + network={ + "allow_out": lambda ctx: list(ctx.rules.keys()), + "rules": { + "api.internal.example.com": [ + { + "transform": lambda ctx: { + "headers": {"Authorization": f"Bearer {ctx.iam.tokens['aws']}"}, + }, + }, + ], + }, + }, +) +``` diff --git a/packages/js-sdk/src/index.ts b/packages/js-sdk/src/index.ts index 54aa9a83ff..9227688f5b 100644 --- a/packages/js-sdk/src/index.ts +++ b/packages/js-sdk/src/index.ts @@ -73,6 +73,8 @@ export type { SandboxNetworkRuleInfo, SandboxNetworkRules, SandboxNetworkTransform, + SandboxNetworkTransformContext, + SandboxNetworkTransformResolver, SandboxNetworkUpdate, SandboxOnTimeout, SandboxLifecycle, diff --git a/packages/js-sdk/src/sandbox/sandboxApi.ts b/packages/js-sdk/src/sandbox/sandboxApi.ts index a11c955078..3c8b2023c3 100644 --- a/packages/js-sdk/src/sandbox/sandboxApi.ts +++ b/packages/js-sdk/src/sandbox/sandboxApi.ts @@ -52,14 +52,58 @@ export type SandboxNetworkTransform = { headers?: Record } +/** + * Context passed to a {@link SandboxNetworkRule} `transform` callback. Its + * values are literal placeholder strings that the egress proxy resolves per + * request, so the secret itself never leaves the platform. + */ +export type SandboxNetworkTransformContext = { + /** Workload identity placeholders. */ + iam: { + /** + * Placeholder for each workload token registered in + * {@link SandboxOpts.iam}, keyed by token name. `tokens.aws` is the string + * `'${e2b.identity.tokens.aws}'`, which the egress proxy replaces with a + * freshly minted token when it forwards the request. + * + * Reading a name that is not registered throws + * {@link InvalidArgumentError} — the proxy never turns an unregistered name + * into a token, so a typo would surface as a confusing auth failure at the + * destination. + */ + tokens: Record + } +} + +/** + * Callback form of {@link SandboxNetworkRule.transform}. Invoked once while the + * request is being built, with a context of placeholder strings. + */ +export type SandboxNetworkTransformResolver = ( + ctx: SandboxNetworkTransformContext +) => SandboxNetworkTransform + /** * Per-domain rule applied to egress requests. */ export type SandboxNetworkRule = { /** * Transform applied to requests matching this rule. + * + * Accepts either a static object or a callback that receives a + * {@link SandboxNetworkTransformContext} of placeholder strings — use the + * callback to inject a workload identity token the proxy mints per request. + * + * @example + * ```ts + * { + * transform: ({ iam }) => ({ + * headers: { Authorization: `Bearer ${iam.tokens.aws}` }, + * }), + * } + * ``` */ - transform?: SandboxNetworkTransform + transform?: SandboxNetworkTransform | SandboxNetworkTransformResolver } /** @@ -138,6 +182,11 @@ export type SandboxNetworkOpts = { * also appear in {@link allowOut}. Hosts registered here are exposed to the * `allowOut`/`denyOut` callbacks via `rules`. * + * A rule's `transform` can also be a callback receiving a + * {@link SandboxNetworkTransformContext}, which is how a workload identity + * token from {@link SandboxOpts.iam} gets injected without the SDK ever + * seeing its value. + * * @example * ```ts * await Sandbox.create({ @@ -147,6 +196,13 @@ export type SandboxNetworkOpts = { * 'api.openai.com': [ * { transform: { headers: { Authorization: `Bearer ${token}` } } }, * ], + * 'api.internal.example.com': [ + * { + * transform: ({ iam }) => ({ + * headers: { Authorization: `Bearer ${iam.tokens.aws}` }, + * }), + * }, + * ], * }, * }, * }) @@ -191,7 +247,12 @@ export type SandboxNetworkUpdate = { allowOut?: SandboxNetworkSelector /** See {@link SandboxNetworkOpts.denyOut}. */ denyOut?: SandboxNetworkSelector - /** See {@link SandboxNetworkOpts.rules}. */ + /** + * See {@link SandboxNetworkOpts.rules}. A `transform` callback works here + * too, but the update payload carries no `iam` config, so token names cannot + * be checked against the sandbox's registered tokens — every name resolves to + * its placeholder and a typo only surfaces at the destination. + */ rules?: SandboxNetworkRules /** * Allow sandbox to access the internet. When set to `false`, it behaves the @@ -443,6 +504,10 @@ export interface SandboxOpts extends ConnectionOpts { * Sandbox workload identity configuration. Providing a non-empty * `tokens` map enables workload identity for the sandbox. * + * Registered tokens are exposed to {@link SandboxNetworkOpts.rules} + * `transform` callbacks as `iam.tokens.` placeholders, which the egress + * proxy resolves per request. + * * @example * ```ts * const sandbox = await Sandbox.create({ @@ -735,14 +800,143 @@ function resolveNetworkSelector( return selector } +function iamTokenPlaceholder(name: string): string { + return `\${e2b.identity.tokens.${name}}` +} + +/** + * Properties the language and the runtime read off any object they serialize, + * await, or coerce to a string. A token is never named after them, so they + * resolve normally instead of counting as a token reference — otherwise + * `JSON.stringify(iam.tokens)` inside a callback would throw. + */ +const RUNTIME_PROBED_PROPS = new Set(['toJSON', 'then', 'toString', 'valueOf']) + +/** + * Build the context handed to `transform` callbacks. + * + * `tokenNames` are the workload tokens the request registers. Referencing any + * other name throws: the proxy never turns an unregistered name into a token, so + * a typo would surface as a confusing auth failure at the destination instead of + * an error here. + * + * `validate: false` is for the update-network endpoint, whose payload carries no + * `iam` config — the sandbox's registered token names are not known client-side + * there, so any name resolves to its placeholder. + */ +function buildTransformContext( + tokenNames: string[], + { validate }: { validate: boolean } +): SandboxNetworkTransformContext { + const tokens: Record = {} + for (const name of tokenNames) { + tokens[name] = iamTokenPlaceholder(name) + } + + return { + iam: { + tokens: new Proxy(tokens, { + get(target, prop, receiver) { + if ( + typeof prop === 'string' && + // Own keys only: a bare `in` also matches inherited + // `Object.prototype` members, so an unregistered token named + // `constructor` or `__proto__` would resolve to a built-in instead + // of being reported. Python's mapping treats them as missing too. + !Object.hasOwn(target, prop) && + !RUNTIME_PROBED_PROPS.has(prop) + ) { + if (!validate) { + return iamTokenPlaceholder(prop) + } + + const hint = + tokenNames.length === 0 + ? `Pass it to Sandbox.create as iam: { tokens: { '${prop}': Secret.iamToken({ audience, tokenType }) } }.` + : `Registered tokens: ${tokenNames.map((name) => `'${name}'`).join(', ')}.` + + throw new InvalidArgumentError( + `Network transform references iam token '${prop}', which is not registered. ${hint}` + ) + } + + return Reflect.get(target, prop, receiver) + }, + + // `name in iam.tokens` answers "is this token registered?", so it must + // agree with the `get` trap above and not report inherited + // `Object.prototype` members as tokens. Mirrors Python's + // `_IamTokenPlaceholders.__contains__`. + has(target, prop) { + return Object.hasOwn(target, prop) + }, + }), + }, + } +} + +function isPlainObject(value: unknown): value is Record { + if (typeof value !== 'object' || value === null) { + return false + } + + const proto = Object.getPrototypeOf(value) + return proto === Object.prototype || proto === null +} + +/** Name the shape a `transform` callback returned, for the error message. */ +function describeValue(value: unknown): string { + if (value === null) { + return 'null' + } + + if (typeof value !== 'object') { + return typeof value + } + + return Array.isArray(value) ? 'array' : (value.constructor?.name ?? 'object') +} + function resolveRulesForBody( - rules: Map + rules: Map, + ctx: SandboxNetworkTransformContext ): Record { const out: Record = {} for (const [host, hostRules] of rules) { - out[host] = hostRules.map((rule) => - rule.transform === undefined ? {} : { transform: rule.transform } - ) + out[host] = hostRules.map((rule) => { + // `== null` also covers an explicit `transform: null`, which Python's + // `rule.get('transform') is None` treats as no transform too. + if (rule.transform == null) { + return {} + } + + if (typeof rule.transform !== 'function') { + return { transform: rule.transform } + } + + const transform: unknown = rule.transform(ctx) + // A callback that returns something other than a transform resolves to no + // headers at all, which would silently create the rule without the headers + // it is for. + if (typeof (transform as PromiseLike)?.then === 'function') { + // Swallow a later rejection so the caller gets this error instead of an + // unhandled rejection. + void Promise.resolve(transform).catch(() => {}) + throw new InvalidArgumentError( + `Network transform callback for '${host}' must be synchronous, it returned a promise. Resolve the value before creating the sandbox.` + ) + } + + // Mirrors Python's `isinstance(transform, Mapping)`: an array, `Map`, + // `Date` or class instance serializes to something the API cannot read. + if (!isPlainObject(transform)) { + throw new InvalidArgumentError( + `Network transform callback for '${host}' must return a transform object, got ${describeValue(transform)}.` + ) + } + + return { transform: transform as SandboxNetworkTransform } + }) } return out } @@ -753,11 +947,14 @@ type NetworkEgressBody = { rules?: Record } -function buildNetworkEgress(network: { - allowOut?: SandboxNetworkSelector - denyOut?: SandboxNetworkSelector - rules?: SandboxNetworkRules -}): NetworkEgressBody { +function buildNetworkEgress( + network: { + allowOut?: SandboxNetworkSelector + denyOut?: SandboxNetworkSelector + rules?: SandboxNetworkRules + }, + transformContext: SandboxNetworkTransformContext +): NetworkEgressBody { const rules = network.rules instanceof Map ? network.rules @@ -769,20 +966,24 @@ function buildNetworkEgress(network: { ...(allowOut !== undefined ? { allowOut } : {}), ...(denyOut !== undefined ? { denyOut } : {}), ...(network.rules !== undefined - ? { rules: resolveRulesForBody(rules) } + ? { rules: resolveRulesForBody(rules, transformContext) } : {}), } } function buildNetworkBody( - network: SandboxNetworkOpts | undefined + network: SandboxNetworkOpts | undefined, + iam: components['schemas']['SandboxIam'] | undefined ): components['schemas']['SandboxNetworkConfig'] | undefined { if (!network) { return undefined } return { - ...buildNetworkEgress(network), + ...buildNetworkEgress( + network, + buildTransformContext(Object.keys(iam?.tokens ?? {}), { validate: true }) + ), ...(network.allowPublicTraffic !== undefined ? { allowPublicTraffic: network.allowPublicTraffic } : {}), @@ -831,7 +1032,10 @@ function buildNetworkUpdateBody( network: SandboxNetworkUpdate ): components['schemas']['SandboxNetworkUpdateConfig'] { return { - ...buildNetworkEgress(network), + ...buildNetworkEgress( + network, + buildTransformContext([], { validate: false }) + ), ...(network.allowInternetAccess !== undefined ? { allow_internet_access: network.allowInternetAccess } : {}), @@ -1291,6 +1495,10 @@ export class SandboxApi { ) } + // Built before the network config: `transform` callbacks are resolved + // against the workload tokens this request registers. + const iam = buildIamBody(opts?.iam) + const body: components['schemas']['NewSandbox'] = { templateID: template, metadata: opts?.metadata, @@ -1299,8 +1507,8 @@ export class SandboxApi { timeout: timeoutToSeconds(timeoutMs), secure: opts?.secure ?? true, allow_internet_access: opts?.allowInternetAccess ?? true, - network: buildNetworkBody(opts?.network), - iam: buildIamBody(opts?.iam), + network: buildNetworkBody(opts?.network, iam), + iam, autoPause: action === 'pause', autoPauseMemory: action === 'pause' ? keepMemory : undefined, autoResume: { enabled: autoResume }, diff --git a/packages/js-sdk/tests/sandbox/networkTransform.test.ts b/packages/js-sdk/tests/sandbox/networkTransform.test.ts new file mode 100644 index 0000000000..383df30a06 --- /dev/null +++ b/packages/js-sdk/tests/sandbox/networkTransform.test.ts @@ -0,0 +1,308 @@ +import { afterAll, afterEach, beforeAll, expect, test } from 'vitest' +import { http, HttpResponse } from 'msw' +import { setupServer } from 'msw/node' + +import { InvalidArgumentError, Sandbox, Secret } from '../../src' +import { TEST_API_KEY, apiUrl } from '../setup' + +const sandboxId = 'test-sandbox-id' + +let lastCreateBody: Record | undefined +let lastUpdateBody: Record | undefined + +const server = setupServer( + http.post(apiUrl('/sandboxes'), async ({ request }) => { + lastCreateBody = (await request.json()) as Record + return HttpResponse.json({ + sandboxID: sandboxId, + templateID: 'base', + envdVersion: '0.2.4', + }) + }), + http.put(apiUrl(`/sandboxes/${sandboxId}/network`), async ({ request }) => { + lastUpdateBody = (await request.json()) as Record + return new HttpResponse(null, { status: 204 }) + }) +) + +beforeAll(() => server.listen({ onUnhandledRequest: 'error' })) + +afterAll(() => server.close()) + +afterEach(() => { + lastCreateBody = undefined + lastUpdateBody = undefined + server.resetHandlers() +}) + +const awsToken = Secret.iamToken({ + audience: 'sts.amazonaws.com', + tokenType: 'JWT-SVID', +}) + +test('transform callback resolves an iam token to its placeholder', async () => { + await Sandbox.create('base', { + apiKey: TEST_API_KEY, + iam: { tokens: { aws: awsToken } }, + network: { + allowOut: ({ rules }) => [...rules.keys()], + rules: { + 'api.internal.example.com': [ + { + transform: ({ iam }) => ({ + headers: { Authorization: `Bearer ${iam.tokens.aws}` }, + }), + }, + ], + }, + }, + }) + + expect(lastCreateBody?.network).toEqual({ + allowOut: ['api.internal.example.com'], + rules: { + 'api.internal.example.com': [ + { + transform: { + headers: { + // The SDK never resolves the placeholder — the egress proxy + // substitutes a freshly minted token per request. + Authorization: 'Bearer ${e2b.identity.tokens.aws}', + }, + }, + }, + ], + }, + }) +}) + +test('transform callback sees every registered iam token', async () => { + await Sandbox.create('base', { + apiKey: TEST_API_KEY, + iam: { tokens: { aws: awsToken, gcp: awsToken } }, + network: { + rules: { + 'api.internal.example.com': [ + { + transform: ({ iam }) => ({ + headers: { + 'X-Tokens': Object.keys(iam.tokens).join(','), + // Membership answers "is it registered?" without throwing, so a + // callback can branch on it. + 'X-Has-Aws': String('aws' in iam.tokens), + 'X-Has-Gh': String('gh' in iam.tokens), + // Membership must agree with a lookup: an inherited object + // member is not a registered token either. + 'X-Has-Ctor': String('constructor' in iam.tokens), + // Serializing the context must not trip the unknown-token guard + // on the runtime's `toJSON` probe. + 'X-Json': JSON.stringify(iam.tokens), + }, + }), + }, + ], + }, + }, + }) + + expect( + lastCreateBody?.network.rules['api.internal.example.com'][0].transform + .headers + ).toEqual({ + 'X-Tokens': 'aws,gcp', + 'X-Has-Aws': 'true', + 'X-Has-Gh': 'false', + 'X-Has-Ctor': 'false', + 'X-Json': JSON.stringify({ + aws: '${e2b.identity.tokens.aws}', + gcp: '${e2b.identity.tokens.gcp}', + }), + }) +}) + +test('a static transform is still sent unchanged', async () => { + await Sandbox.create('base', { + apiKey: TEST_API_KEY, + network: { + rules: { + 'api.openai.com': [ + { transform: { headers: { Authorization: 'Bearer static' } } }, + ], + }, + }, + }) + + expect(lastCreateBody?.network.rules).toEqual({ + 'api.openai.com': [ + { transform: { headers: { Authorization: 'Bearer static' } } }, + ], + }) +}) + +test('transform callback rejects an unregistered iam token', async () => { + await expect( + Sandbox.create('base', { + apiKey: TEST_API_KEY, + iam: { tokens: { aws: awsToken } }, + network: { + rules: { + 'api.internal.example.com': [ + { + transform: ({ iam }) => ({ + headers: { Authorization: `Bearer ${iam.tokens.awz}` }, + }), + }, + ], + }, + }, + }) + ).rejects.toThrowError(/iam token 'awz'.*Registered tokens: 'aws'/s) + + expect(lastCreateBody).toBeUndefined() +}) + +test.for(['constructor', '__proto__', 'hasOwnProperty'])( + 'transform callback rejects an unregistered iam token named %s', + async (name: string) => { + // Inherited Object.prototype members are not registered tokens; resolving + // them would put a built-in function in the header. Python's mapping treats + // them as missing too. + await expect( + Sandbox.create('base', { + apiKey: TEST_API_KEY, + iam: { tokens: { aws: awsToken } }, + network: { + rules: { + 'api.internal.example.com': [ + { + transform: ({ iam }) => ({ + headers: { Authorization: `Bearer ${iam.tokens[name]}` }, + }), + }, + ], + }, + }, + }) + ).rejects.toThrowError(`iam token '${name}', which is not registered`) + + expect(lastCreateBody).toBeUndefined() + } +) + +test('transform callback rejects an iam token when no iam config is set', async () => { + await expect( + Sandbox.create('base', { + apiKey: TEST_API_KEY, + network: { + rules: { + 'api.internal.example.com': [ + { + transform: ({ iam }) => ({ + headers: { Authorization: `Bearer ${iam.tokens.aws}` }, + }), + }, + ], + }, + }, + }) + ).rejects.toThrowError(InvalidArgumentError) + + expect(lastCreateBody).toBeUndefined() +}) + +test.for([ + ['undefined', () => undefined], + ['string', () => 'headers'], + ['array', () => [{ headers: {} }]], + ['Map', () => new Map([['headers', {}]])], +])( + 'transform callback returning a %s is rejected', + async ([, transform]: [string, () => unknown]) => { + // Untyped callers can forget the return value or return the wrong shape; the + // rule would otherwise be created without the headers it exists for. + await expect( + Sandbox.create('base', { + apiKey: TEST_API_KEY, + network: { + rules: { + 'api.internal.example.com': [{ transform: transform as never }], + }, + }, + }) + ).rejects.toThrowError( + /must return a transform object, got (undefined|string|array|Map)/ + ) + + expect(lastCreateBody).toBeUndefined() + } +) + +test('async transform callback is rejected', async () => { + await expect( + Sandbox.create('base', { + apiKey: TEST_API_KEY, + network: { + rules: { + 'api.internal.example.com': [ + { + transform: (async () => ({ + headers: { Authorization: 'Bearer late' }, + })) as never, + }, + ], + }, + }, + }) + ).rejects.toThrowError(/must be synchronous/) + + expect(lastCreateBody).toBeUndefined() +}) + +test('an explicit null transform sends an empty rule', async () => { + // Rules built from parsed JSON spell "no transform" as null; Python treats it + // the same way. + await Sandbox.create('base', { + apiKey: TEST_API_KEY, + network: { + rules: { 'api.internal.example.com': [{ transform: null as never }] }, + }, + }) + + expect(lastCreateBody?.network.rules).toEqual({ + 'api.internal.example.com': [{}], + }) +}) + +test('updateNetwork resolves transform callbacks without an iam config', async () => { + // The update payload carries no iam config, so the sandbox's registered token + // names are unknown client-side and any name resolves to its placeholder. + await Sandbox.updateNetwork( + sandboxId, + { + allowOut: ({ rules }) => [...rules.keys()], + rules: { + 'api.internal.example.com': [ + { + transform: ({ iam }) => ({ + headers: { Authorization: `Bearer ${iam.tokens.aws}` }, + }), + }, + ], + }, + }, + { apiKey: TEST_API_KEY } + ) + + expect(lastUpdateBody).toEqual({ + allowOut: ['api.internal.example.com'], + rules: { + 'api.internal.example.com': [ + { + transform: { + headers: { Authorization: 'Bearer ${e2b.identity.tokens.aws}' }, + }, + }, + ], + }, + }) +}) diff --git a/packages/python-sdk/e2b/__init__.py b/packages/python-sdk/e2b/__init__.py index 86f912f0e0..11ae55d80b 100644 --- a/packages/python-sdk/e2b/__init__.py +++ b/packages/python-sdk/e2b/__init__.py @@ -91,6 +91,8 @@ SandboxNetworkSelector, SandboxNetworkSelectorContext, SandboxNetworkTransform, + SandboxNetworkTransformContext, + SandboxNetworkTransformResolver, SandboxNetworkUpdate, SandboxQuery, SandboxState, @@ -204,6 +206,8 @@ "SandboxNetworkRuleInfo", "SandboxNetworkRules", "SandboxNetworkTransform", + "SandboxNetworkTransformContext", + "SandboxNetworkTransformResolver", "SandboxNetworkUpdate", "SandboxLifecycle", "SandboxOnTimeout", diff --git a/packages/python-sdk/e2b/sandbox/sandbox_api.py b/packages/python-sdk/e2b/sandbox/sandbox_api.py index 699ec8040a..14a17c0fb2 100644 --- a/packages/python-sdk/e2b/sandbox/sandbox_api.py +++ b/packages/python-sdk/e2b/sandbox/sandbox_api.py @@ -1,9 +1,12 @@ +import inspect from dataclasses import dataclass, field from datetime import datetime from typing import ( Any, Callable, Dict, + Iterable, + Iterator, List, Literal, Mapping, @@ -101,14 +104,119 @@ class SandboxNetworkTransform(TypedDict): """ +def _iam_token_placeholder(name: str) -> str: + return f"${{e2b.identity.tokens.{name}}}" + + +class _IamTokenPlaceholders(Mapping[str, str]): + """ + Workload token placeholders keyed by token name, as exposed to a + ``transform`` callable. Every lookup — ``[name]``, ``get(name)`` — resolves + through :meth:`__getitem__`, so an unregistered name cannot slip through as + ``None``; iteration and ``in`` see only the registered names. + + ``validate=False`` is for the update-network endpoint, whose payload carries + no ``iam`` config — the sandbox's registered token names are not known + client-side there, so any name resolves to its placeholder. + """ + + def __init__(self, names: Iterable[str], *, validate: bool) -> None: + # dict.fromkeys keeps registration order while dropping duplicates. + self._names = tuple(dict.fromkeys(names)) + self._validate = validate + + def __getitem__(self, name: str) -> str: + # The proxy never turns an unregistered name into a token, so a typo + # would surface as a confusing auth failure at the destination instead + # of an error here. + if not self._validate or name in self._names: + return _iam_token_placeholder(name) + + hint = ( + f"Registered tokens: {', '.join(repr(known) for known in self._names)}." + if self._names + else ( + f"Pass it to Sandbox.create as iam={{'tokens': " + f"{{{name!r}: Secret.iam_token(audience=..., token_type=...)}}}}." + ) + ) + raise InvalidArgumentException( + f"Network transform references iam token {name!r}, which is not " + f"registered. {hint}" + ) + + def __contains__(self, name: object) -> bool: + # Membership answers "is this registered?" and never raises, so callables + # can branch on it; mirrors `name in iam.tokens` in the JS SDK. + return name in self._names + + def __iter__(self) -> Iterator[str]: + return iter(self._names) + + def __len__(self) -> int: + return len(self._names) + + +@dataclass(frozen=True) +class _SandboxNetworkTransformIam: + """Workload identity placeholders.""" + + tokens: Mapping[str, str] + """ + Placeholder for each workload token registered in + :attr:`SandboxOpts.iam`, keyed by token name. ``tokens["aws"]`` is the + string ``"${e2b.identity.tokens.aws}"``, which the egress proxy replaces + with a freshly minted token when it forwards the request. + + Reading a name that is not registered raises + :class:`InvalidArgumentException` — the proxy never turns an unregistered + name into a token, so a typo would surface as a confusing auth failure at + the destination. + """ + + +@dataclass(frozen=True) +class SandboxNetworkTransformContext: + """ + Context passed to a :class:`SandboxNetworkRule` ``transform`` callable. Its + values are literal placeholder strings that the egress proxy resolves per + request, so the secret itself never leaves the platform. + """ + + iam: _SandboxNetworkTransformIam + """Workload identity placeholders.""" + + +SandboxNetworkTransformResolver = Callable[ + [SandboxNetworkTransformContext], SandboxNetworkTransform +] +""" +Callable form of ``SandboxNetworkRule.transform``. Invoked once while the +request is being built, with a context of placeholder strings. +""" + + class SandboxNetworkRule(TypedDict): """ Per-domain rule applied to egress requests. """ - transform: NotRequired[SandboxNetworkTransform] + transform: NotRequired[ + Union[SandboxNetworkTransform, SandboxNetworkTransformResolver] + ] """ Transform applied to requests matching this rule. + + Accepts either a static :class:`SandboxNetworkTransform` or a callable that + receives a :class:`SandboxNetworkTransformContext` of placeholder strings — + use the callable to inject a workload identity token the proxy mints per + request:: + + { + "transform": lambda ctx: { + "headers": {"Authorization": f"Bearer {ctx.iam.tokens['aws']}"}, + }, + } """ @@ -195,6 +303,23 @@ class SandboxNetworkOpts(TypedDict): Registering a host here does not allow egress on its own — the host must also appear in ``allow_out``. Hosts registered here are exposed to the ``allow_out``/``deny_out`` callables via ``ctx.rules``. + + A rule's ``transform`` can also be a callable receiving a + :class:`SandboxNetworkTransformContext`, which is how a workload identity + token from :attr:`SandboxOpts.iam` gets injected without the SDK ever + seeing its value:: + + rules={ + "api.internal.example.com": [ + { + "transform": lambda ctx: { + "headers": { + "Authorization": f"Bearer {ctx.iam.tokens['aws']}", + }, + }, + }, + ], + } """ allow_public_traffic: NotRequired[bool] @@ -227,7 +352,12 @@ class SandboxNetworkUpdate(TypedDict, total=False): """See :attr:`SandboxNetworkOpts.deny_out`.""" rules: SandboxNetworkRules - """See :attr:`SandboxNetworkOpts.rules`.""" + """ + See :attr:`SandboxNetworkOpts.rules`. A ``transform`` callable works here + too, but the update payload carries no ``iam`` config, so token names cannot + be checked against the sandbox's registered tokens — every name resolves to + its placeholder and a typo only surfaces at the destination. + """ allow_internet_access: bool """ @@ -381,7 +511,28 @@ def _resolve_network_selector( return list(selector) -def _build_client_rules(rules: SandboxNetworkRules) -> SandboxNetworkConfigRules: +def _build_transform_context( + token_names: Iterable[str], + *, + validate: bool, +) -> SandboxNetworkTransformContext: + """ + Build the context handed to ``transform`` callables. + + ``token_names`` are the workload tokens the request registers; see + :class:`_IamTokenPlaceholders` for what ``validate`` controls. + """ + return SandboxNetworkTransformContext( + iam=_SandboxNetworkTransformIam( + tokens=_IamTokenPlaceholders(token_names, validate=validate), + ), + ) + + +def _build_client_rules( + rules: SandboxNetworkRules, + ctx: SandboxNetworkTransformContext, +) -> SandboxNetworkConfigRules: client_rules = SandboxNetworkConfigRules() for host, host_rules in rules.items(): converted: List[ClientSandboxNetworkRule] = [] @@ -391,6 +542,27 @@ def _build_client_rules(rules: SandboxNetworkRules) -> SandboxNetworkConfigRules converted.append(ClientSandboxNetworkRule()) continue + if callable(transform): + transform = transform(ctx) + # A callable that returns something other than a transform + # resolves to no headers at all, which would silently create the + # rule without the headers it is for. + if inspect.isawaitable(transform): + if inspect.iscoroutine(transform): + # Close it so the caller does not also get a + # "coroutine was never awaited" RuntimeWarning. + transform.close() + raise InvalidArgumentException( + f"Network transform callable for {host!r} must be " + "synchronous, it returned an awaitable. Resolve the " + "value before creating the sandbox." + ) + if not isinstance(transform, Mapping): + raise InvalidArgumentException( + f"Network transform callable for {host!r} must return a " + f"transform dict, got {type(transform).__name__}." + ) + client_transform = ClientSandboxNetworkTransform() headers = transform.get("headers") if headers: @@ -406,6 +578,7 @@ def _build_client_rules(rules: SandboxNetworkRules) -> SandboxNetworkConfigRules def _build_network_egress( network: Mapping[str, Any], + ctx: SandboxNetworkTransformContext, ) -> Dict[str, Any]: """ Resolve the shared egress fields (``allow_out`` / ``deny_out`` / per-host @@ -423,19 +596,31 @@ def _build_network_egress( if deny_out is not None: body["deny_out"] = deny_out if "rules" in network and network["rules"] is not None: - body["rules"] = _build_client_rules(network["rules"]).additional_properties + body["rules"] = _build_client_rules(network["rules"], ctx).additional_properties return body def build_network_config( network: Optional[SandboxNetworkOpts], + iam: Optional[ClientSandboxIam] = None, ) -> Optional[Dict[str, Any]]: - """Resolve a :class:`SandboxNetworkOpts` into the dict the API expects.""" + """Resolve a :class:`SandboxNetworkOpts` into the dict the API expects. + + ``iam`` is the built workload identity config of the same request — its + token names are what ``transform`` callables may reference. + """ if network is None: return None - body = _build_network_egress(network) + token_names: List[str] = [] + if iam is not None and not isinstance(iam.tokens, Unset): + token_names = iam.tokens.additional_keys + + body = _build_network_egress( + network, + _build_transform_context(token_names, validate=True), + ) if "rules" in body: client_rules = SandboxNetworkConfigRules() client_rules.additional_properties = body["rules"] @@ -493,7 +678,9 @@ def build_network_update_body( network: SandboxNetworkUpdate, ) -> SandboxNetworkUpdateConfig: """Resolve a :class:`SandboxNetworkUpdate` into the API client body.""" - egress = _build_network_egress(network) + egress = _build_network_egress( + network, _build_transform_context([], validate=False) + ) body = SandboxNetworkUpdateConfig() if "allow_out" in egress: diff --git a/packages/python-sdk/e2b/sandbox_async/main.py b/packages/python-sdk/e2b/sandbox_async/main.py index f22eaa729d..f2ed5fd61a 100644 --- a/packages/python-sdk/e2b/sandbox_async/main.py +++ b/packages/python-sdk/e2b/sandbox_async/main.py @@ -198,8 +198,8 @@ async def create( :param secure: Envd is secured with access token and cannot be used without it, defaults to `True`. :param allow_internet_access: Allow sandbox to access the internet, defaults to `True`. If set to `False`, it works the same as setting network `deny_out` to `[0.0.0.0/0]`. :param mcp: MCP server to enable in the sandbox - :param network: Sandbox network configuration. ``allow_out``/``deny_out`` may also be a callable receiving a :class:`SandboxNetworkSelectorContext` (``ctx.all_traffic``, ``ctx.rules``) and returning a list of strings. Per-host transform rules are nested under ``network.rules``. - :param iam: Sandbox workload identity configuration. A non-empty ``tokens`` map enables workload identity for the sandbox; token definitions can be created with :meth:`Secret.iam_token`. Example: ``{"tokens": {"aws": Secret.iam_token(audience="sts.amazonaws.com", token_type="JWT-SVID")}}`` + :param network: Sandbox network configuration. ``allow_out``/``deny_out`` may also be a callable receiving a :class:`SandboxNetworkSelectorContext` (``ctx.all_traffic``, ``ctx.rules``) and returning a list of strings. Per-host transform rules are nested under ``network.rules``; a rule's ``transform`` may be a callable receiving a :class:`SandboxNetworkTransformContext` of placeholder strings (``ctx.iam.tokens[name]``). + :param iam: Sandbox workload identity configuration. A non-empty ``tokens`` map enables workload identity for the sandbox; token definitions can be created with :meth:`Secret.iam_token`. Example: ``{"tokens": {"aws": Secret.iam_token(audience="sts.amazonaws.com", token_type="JWT-SVID")}}``. Registered tokens are exposed to ``network.rules`` ``transform`` callables as ``ctx.iam.tokens[name]`` placeholders, which the egress proxy resolves per request :param lifecycle: Sandbox lifecycle configuration — ``on_timeout``: ``"kill"`` (default) or ``"pause"``, or an object ``{"action": "pause"|"kill", "keep_memory": bool}`` where ``keep_memory`` (default ``True``) set to ``False`` makes a timeout auto-pause filesystem-only (cold-boots on resume; cannot be combined with ``auto_resume``); ``auto_resume``: ``False`` (default) or ``True`` (only when ``on_timeout`` action is ``"pause"``). Example: ``{"on_timeout": {"action": "pause", "keep_memory": False}}`` :param volume_mounts: Dictionary mapping mount paths to AsyncVolume instances or volume names :param logger: Logger used for request and response logging for this sandbox. Accepts any standard library `logging.Logger`. When omitted, no request/response logging is emitted. diff --git a/packages/python-sdk/e2b/sandbox_async/sandbox_api.py b/packages/python-sdk/e2b/sandbox_async/sandbox_api.py index a540e8c523..bebf686348 100644 --- a/packages/python-sdk/e2b/sandbox_async/sandbox_api.py +++ b/packages/python-sdk/e2b/sandbox_async/sandbox_api.py @@ -257,7 +257,10 @@ async def _create_sandbox( "must be resumed explicitly using Sandbox.connect()." ) - network_body = build_network_config(network) + # Built before the network config: ``transform`` callables are resolved + # against the workload tokens this request registers. + iam_body = build_iam_config(iam) + network_body = build_network_config(network, iam_body) body = NewSandbox( template_id=template, auto_pause=on_timeout == "pause", @@ -270,7 +273,7 @@ async def _create_sandbox( secure=secure, allow_internet_access=allow_internet_access, network=SandboxNetworkConfig(**network_body) if network_body else UNSET, - iam=build_iam_config(iam) or UNSET, + iam=iam_body or UNSET, volume_mounts=volume_mounts if volume_mounts else UNSET, ) diff --git a/packages/python-sdk/e2b/sandbox_sync/main.py b/packages/python-sdk/e2b/sandbox_sync/main.py index db4d46d3ba..6cd0666432 100644 --- a/packages/python-sdk/e2b/sandbox_sync/main.py +++ b/packages/python-sdk/e2b/sandbox_sync/main.py @@ -187,8 +187,8 @@ def create( :param secure: Envd is secured with access token and cannot be used without it, defaults to `True`. :param allow_internet_access: Allow sandbox to access the internet, defaults to `True`. If set to `False`, it works the same as setting network `deny_out` to `[0.0.0.0/0]`. :param mcp: MCP server to enable in the sandbox - :param network: Sandbox network configuration. ``allow_out``/``deny_out`` may also be a callable receiving a :class:`SandboxNetworkSelectorContext` (``ctx.all_traffic``, ``ctx.rules``) and returning a list of strings. Per-host transform rules are nested under ``network.rules``. - :param iam: Sandbox workload identity configuration. A non-empty ``tokens`` map enables workload identity for the sandbox; token definitions can be created with :meth:`Secret.iam_token`. Example: ``{"tokens": {"aws": Secret.iam_token(audience="sts.amazonaws.com", token_type="JWT-SVID")}}`` + :param network: Sandbox network configuration. ``allow_out``/``deny_out`` may also be a callable receiving a :class:`SandboxNetworkSelectorContext` (``ctx.all_traffic``, ``ctx.rules``) and returning a list of strings. Per-host transform rules are nested under ``network.rules``; a rule's ``transform`` may be a callable receiving a :class:`SandboxNetworkTransformContext` of placeholder strings (``ctx.iam.tokens[name]``). + :param iam: Sandbox workload identity configuration. A non-empty ``tokens`` map enables workload identity for the sandbox; token definitions can be created with :meth:`Secret.iam_token`. Example: ``{"tokens": {"aws": Secret.iam_token(audience="sts.amazonaws.com", token_type="JWT-SVID")}}``. Registered tokens are exposed to ``network.rules`` ``transform`` callables as ``ctx.iam.tokens[name]`` placeholders, which the egress proxy resolves per request :param lifecycle: Sandbox lifecycle configuration — ``on_timeout``: ``"kill"`` (default) or ``"pause"``, or an object ``{"action": "pause"|"kill", "keep_memory": bool}`` where ``keep_memory`` (default ``True``) set to ``False`` makes a timeout auto-pause filesystem-only (cold-boots on resume; cannot be combined with ``auto_resume``); ``auto_resume``: ``False`` (default) or ``True`` (only when ``on_timeout`` action is ``"pause"``). Example: ``{"on_timeout": {"action": "pause", "keep_memory": False}}`` :param volume_mounts: Dictionary mapping mount paths to Volume instances or volume names :param logger: Logger used for request and response logging for this sandbox. Accepts any standard library `logging.Logger`. When omitted, no request/response logging is emitted. diff --git a/packages/python-sdk/e2b/sandbox_sync/sandbox_api.py b/packages/python-sdk/e2b/sandbox_sync/sandbox_api.py index d11876f321..e4a8cc0afe 100644 --- a/packages/python-sdk/e2b/sandbox_sync/sandbox_api.py +++ b/packages/python-sdk/e2b/sandbox_sync/sandbox_api.py @@ -256,7 +256,10 @@ def _create_sandbox( "must be resumed explicitly using Sandbox.connect()." ) - network_body = build_network_config(network) + # Built before the network config: ``transform`` callables are resolved + # against the workload tokens this request registers. + iam_body = build_iam_config(iam) + network_body = build_network_config(network, iam_body) body = NewSandbox( template_id=template, auto_pause=on_timeout == "pause", @@ -269,7 +272,7 @@ def _create_sandbox( secure=secure, allow_internet_access=allow_internet_access, network=SandboxNetworkConfig(**network_body) if network_body else UNSET, - iam=build_iam_config(iam) or UNSET, + iam=iam_body or UNSET, volume_mounts=volume_mounts if volume_mounts else UNSET, ) diff --git a/packages/python-sdk/tests/shared/sandbox/test_network_transform.py b/packages/python-sdk/tests/shared/sandbox/test_network_transform.py new file mode 100644 index 0000000000..d15a7ec13f --- /dev/null +++ b/packages/python-sdk/tests/shared/sandbox/test_network_transform.py @@ -0,0 +1,234 @@ +from typing import Any, cast + +import pytest + +from e2b import SandboxNetworkOpts, Secret +from e2b.exceptions import InvalidArgumentException +from e2b.sandbox.sandbox_api import ( + build_iam_config, + build_network_config, + build_network_update_body, +) + + +def _aws_iam(): + return build_iam_config( + { + "tokens": { + "aws": Secret.iam_token( + audience="sts.amazonaws.com", token_type="JWT-SVID" + ), + }, + } + ) + + +def test_transform_callable_resolves_iam_token_placeholder(): + network: SandboxNetworkOpts = { + "allow_out": lambda ctx: list(ctx.rules.keys()), + "rules": { + "api.internal.example.com": [ + { + "transform": lambda ctx: { + "headers": { + "Authorization": f"Bearer {ctx.iam.tokens['aws']}", + }, + }, + }, + ], + }, + } + + body = build_network_config(network, _aws_iam()) + assert body is not None + assert body["allow_out"] == ["api.internal.example.com"] + # The SDK never resolves the placeholder — the egress proxy substitutes a + # freshly minted token per request. + assert body["rules"].to_dict() == { + "api.internal.example.com": [ + { + "transform": { + "headers": { + "Authorization": "Bearer ${e2b.identity.tokens.aws}", + }, + }, + }, + ], + } + + +def test_transform_callable_sees_every_registered_iam_token(): + token = Secret.iam_token(audience="sts.amazonaws.com", token_type="JWT-SVID") + iam = build_iam_config({"tokens": {"aws": token, "gcp": token}}) + + network: SandboxNetworkOpts = { + "rules": { + "api.internal.example.com": [ + { + "transform": lambda ctx: { + "headers": { + "X-Tokens": ",".join(ctx.iam.tokens), + # Membership answers "is it registered?" without + # raising, so a callable can branch on it. + "X-Has-Aws": str("aws" in ctx.iam.tokens), + "X-Has-Gh": str("gh" in ctx.iam.tokens), + # Membership must agree with a lookup: a name that + # shadows a mapping member is not a token either. + "X-Has-Keys": str("keys" in ctx.iam.tokens), + }, + }, + }, + ], + }, + } + + body = build_network_config(network, iam) + assert body is not None + assert body["rules"].to_dict() == { + "api.internal.example.com": [ + { + "transform": { + "headers": { + "X-Tokens": "aws,gcp", + "X-Has-Aws": "True", + "X-Has-Gh": "False", + "X-Has-Keys": "False", + }, + }, + }, + ], + } + + +def test_static_transform_is_sent_unchanged(): + network: SandboxNetworkOpts = { + "rules": { + "api.openai.com": [ + {"transform": {"headers": {"Authorization": "Bearer static"}}}, + ], + }, + } + + body = build_network_config(network) + assert body is not None + assert body["rules"].to_dict() == { + "api.openai.com": [ + {"transform": {"headers": {"Authorization": "Bearer static"}}}, + ], + } + + +def _typo_network(access) -> Any: + return cast( + Any, + { + "rules": { + "api.internal.example.com": [ + { + "transform": lambda ctx: { + "headers": {"Authorization": f"Bearer {access(ctx)}"}, + }, + }, + ], + }, + }, + ) + + +@pytest.mark.parametrize( + "access", + [ + # Every lookup form must reject an unregistered name — `.get()` returning + # None would ship the literal header "Bearer None". + pytest.param(lambda ctx: ctx.iam.tokens["awz"], id="getitem"), + pytest.param(lambda ctx: ctx.iam.tokens.get("awz"), id="get"), + # A name that shadows an object member is not a registered token either; + # the JS SDK checks own keys only for the same reason. + pytest.param(lambda ctx: ctx.iam.tokens["keys"], id="dunder-lookalike"), + ], +) +def test_transform_callable_rejects_unregistered_iam_token(access): + # The proxy never turns an unregistered name into a token, so a typo would + # surface as a confusing auth failure at the destination. + with pytest.raises(InvalidArgumentException, match="Registered tokens: 'aws'"): + build_network_config(_typo_network(access), _aws_iam()) + + # Same rule without any iam config at all. + with pytest.raises(InvalidArgumentException, match="not registered"): + build_network_config(_typo_network(access)) + + +@pytest.mark.parametrize( + "returned", + [ + pytest.param(None, id="none"), + pytest.param("headers", id="str"), + pytest.param([{"headers": {}}], id="list"), + ], +) +def test_transform_callable_returning_a_non_transform_is_rejected(returned): + # Untyped callers can forget the return value or return the wrong shape; the + # rule would otherwise be created without the headers it exists for. + network = cast( + Any, + { + "rules": { + "api.internal.example.com": [{"transform": lambda ctx: returned}] + }, + }, + ) + + with pytest.raises(InvalidArgumentException, match="must return a transform dict"): + build_network_config(network) + + +def test_async_transform_callable_is_rejected(): + async def transform(ctx): + return {"headers": {"Authorization": "Bearer late"}} + + network = cast( + Any, + {"rules": {"api.internal.example.com": [{"transform": transform}]}}, + ) + + with pytest.raises(InvalidArgumentException, match="must be synchronous"): + build_network_config(network) + + +def test_update_network_resolves_transform_callables_without_iam(): + # The update payload carries no iam config, so the sandbox's registered + # token names are unknown client-side and any name resolves to its + # placeholder. + body = build_network_update_body( + { + "allow_out": lambda ctx: list(ctx.rules.keys()), + "rules": { + "api.internal.example.com": [ + { + "transform": lambda ctx: { + "headers": { + "Authorization": f"Bearer {ctx.iam.tokens['aws']}", + "X-Also": f"{ctx.iam.tokens.get('gcp')}", + }, + }, + }, + ], + }, + } + ) + + assert body.to_dict() == { + "allowOut": ["api.internal.example.com"], + "rules": { + "api.internal.example.com": [ + { + "transform": { + "headers": { + "Authorization": "Bearer ${e2b.identity.tokens.aws}", + "X-Also": "${e2b.identity.tokens.gcp}", + }, + }, + }, + ], + }, + }