feat(sdk): support structured network rules with per-host transforms (#1286)

Author: mishushakov

.changeset/network-rules-transform.md

@@ -0,0 +1,6 @@
+---
+'e2b': minor
+'@e2b/python-sdk': minor
+---
+
+Support structured network rules with per-host transforms

packages/js-sdk/src/api/schema.gen.ts

@@ -57,6 +57,13 @@ export type {
   SandboxListOpts,
   SandboxPaginator,
   SandboxNetworkOpts,
+  SandboxNetworkInfo,
+  SandboxNetworkSelector,
+  SandboxNetworkSelectorContext,
+  SandboxNetworkRule,
+  SandboxNetworkRuleInfo,
+  SandboxNetworkRules,
+  SandboxNetworkTransform,
   SandboxLifecycle,
   SandboxInfoLifecycle,
   SnapshotInfo,

packages/js-sdk/src/index.ts

@@ -5,6 +5,7 @@ import {
   DEFAULT_SANDBOX_TIMEOUT_MS,
 } from '../connectionConfig'
 import { compareVersions } from 'compare-versions'
+import { ALL_TRAFFIC } from './network'
 import {
   InvalidArgumentError,
   SandboxNotFoundError,
@@ -37,23 +38,118 @@ export type GitHubMcpServer = {
   }
 }
 
+/**
+ * Transform applied to egress requests matching a {@link SandboxNetworkRule}.
+ */
+export type SandboxNetworkTransform = {
+  /**
+   * Headers to inject into the outbound request. Values override any headers
+   * already present on the request.
+   */
+  headers?: Record<string, string>
+}
+
+/**
+ * Per-domain rule applied to egress requests.
+ */
+export type SandboxNetworkRule = {
+  /**
+   * Transform applied to requests matching this rule.
+   */
+  transform?: SandboxNetworkTransform
+}
+
+/**
+ * Map of host (or CIDR / IP) to ordered list of rules applied to outbound
+ * requests for that host. Accepts either a plain object or a `Map`.
+ * Registering a host here does not allow egress on its own — the host must
+ * also appear in {@link SandboxNetworkOpts.allowOut}.
+ */
+export type SandboxNetworkRules =
+  | Record<string, SandboxNetworkRule[]>
+  | Map<string, SandboxNetworkRule[]>
+
+/**
+ * Per-domain rule as returned by the sandbox info endpoint. Mirrors
+ * {@link SandboxNetworkRule} but with `transform` always materialized to the
+ * static {@link SandboxNetworkTransform} shape — no callback variant.
+ */
+export type SandboxNetworkRuleInfo = {
+  transform?: SandboxNetworkTransform
+}
+
+/**
+ * Context passed to {@link SandboxNetworkOpts.allowOut} and
+ * {@link SandboxNetworkOpts.denyOut} when they are defined as functions.
+ */
+export type SandboxNetworkSelectorContext = {
+  /** All traffic sentinel — equivalent to `'0.0.0.0/0'`. */
+  allTraffic: string
+  /** Rules registered in {@link SandboxNetworkOpts.rules}. */
+  rules: Map<string, SandboxNetworkRule[]>
+}
+
+/**
+ * Egress rule list, either a static array of CIDR blocks / IP addresses /
+ * hostnames, or a callback that receives `{ allTraffic, rules }` and returns
+ * the same.
+ */
+export type SandboxNetworkSelector =
+  | string[]
+  | ((ctx: SandboxNetworkSelectorContext) => string[])
+
 export type SandboxNetworkOpts = {
   /**
    * Allow outbound traffic from the sandbox to the specified addresses.
    * If `allowOut` is not specified, all outbound traffic is allowed.
    *
+   * Accepts either a static array of CIDR blocks, IP addresses, or hostnames,
+   * or a callback that receives `{ allTraffic, rules }` and returns the same.
+   * `allTraffic` is `'0.0.0.0/0'`; `rules` is a `Map` view of
+   * {@link SandboxNetworkOpts.rules}.
+   *
    * Examples:
-   * - To allow traffic to a specific addresses: `["1.1.1.1", "8.8.8.0/24"]`
+   * - Static list: `["1.1.1.1", "8.8.8.0/24"]`
+   * - Allow only rule-registered hosts:
+   *   `({ rules }) => [...rules.keys()]`
    */
-  allowOut?: string[]
+  allowOut?: SandboxNetworkSelector
 
   /**
    * Deny outbound traffic from the sandbox to the specified addresses.
    *
+   * Accepts the same shapes as {@link allowOut}.
+   *
    * Examples:
-   * - To deny traffic to a specific addresses: `["1.1.1.1", "8.8.8.0/24"]`
+   * - Static list: `["1.1.1.1", "8.8.8.0/24"]`
+   * - Block all egress: `({ allTraffic }) => [allTraffic]`
    */
-  denyOut?: string[]
+  denyOut?: SandboxNetworkSelector
+
+  /**
+   * Per-domain transform rules applied to matching egress HTTP/HTTPS
+   * requests. Keys are domains (e.g. `"api.example.com"`); values are
+   * ordered lists of rules.
+   *
+   * Registering a host here does not allow egress on its own — the host must
+   * also appear in {@link allowOut}. Hosts registered here are exposed to the
+   * `allowOut`/`denyOut` callbacks via `rules`.
+   *
+   * @example
+   * ```ts
+   * await Sandbox.create({
+   *   network: {
+   *     allowOut: ({ rules }) => [...rules.keys()],
+   *     rules: {
+   *       'api.openai.com': [
+   *         { transform: { headers: { Authorization: `Bearer ${token}` } } },
+   *       ],
+   *     },
+   *   },
+   * })
+   * ```
+   */
+  rules?: SandboxNetworkRules
 
   /**
    * Specify if the sandbox URLs should be accessible only with authentication.
@@ -69,6 +165,19 @@ export type SandboxNetworkOpts = {
   maskRequestHost?: string
 }
 
+/**
+ * Network configuration as returned by the sandbox info endpoint. Mirrors
+ * {@link SandboxNetworkOpts} but with `allowOut`/`denyOut` always materialized
+ * to plain string arrays.
+ */
+export type SandboxNetworkInfo = {
+  allowOut?: string[]
+  denyOut?: string[]
+  rules?: Record<string, SandboxNetworkRuleInfo[]>
+  allowPublicTraffic?: boolean
+  maskRequestHost?: string
+}
+
 export type SandboxLifecycle = {
   /**
    * Action to take when sandbox timeout is reached.
@@ -360,7 +469,7 @@ export interface SandboxInfo {
   /**
    * Sandbox network configuration.
    */
-  network?: SandboxNetworkOpts
+  network?: SandboxNetworkInfo
 
   /**
    * Sandbox lifecycle configuration.
@@ -413,6 +522,62 @@ export interface SandboxMetrics {
   diskTotal: number
 }
 
+function resolveNetworkSelector(
+  selector: SandboxNetworkSelector | undefined,
+  rules: Map<string, SandboxNetworkRule[]>
+): string[] | undefined {
+  if (selector === undefined) {
+    return undefined
+  }
+
+  if (typeof selector === 'function') {
+    return selector({ allTraffic: ALL_TRAFFIC, rules })
+  }
+
+  return selector
+}
+
+function resolveRulesForBody(
+  rules: Map<string, SandboxNetworkRule[]>
+): Record<string, { transform?: SandboxNetworkTransform }[]> {
+  const out: Record<string, { transform?: SandboxNetworkTransform }[]> = {}
+  for (const [host, hostRules] of rules) {
+    out[host] = hostRules.map((rule) =>
+      rule.transform === undefined ? {} : { transform: rule.transform }
+    )
+  }
+  return out
+}
+
+function buildNetworkBody(
+  network: SandboxNetworkOpts | undefined
+): components['schemas']['SandboxNetworkConfig'] | undefined {
+  if (!network) {
+    return undefined
+  }
+
+  const rules =
+    network.rules instanceof Map
+      ? network.rules
+      : new Map(Object.entries(network.rules ?? {}))
+  const allowOut = resolveNetworkSelector(network.allowOut, rules)
+  const denyOut = resolveNetworkSelector(network.denyOut, rules)
+
+  return {
+    ...(allowOut !== undefined ? { allowOut } : {}),
+    ...(denyOut !== undefined ? { denyOut } : {}),
+    ...(network.rules !== undefined
+      ? { rules: resolveRulesForBody(rules) }
+      : {}),
+    ...(network.allowPublicTraffic !== undefined
+      ? { allowPublicTraffic: network.allowPublicTraffic }
+      : {}),
+    ...(network.maskRequestHost !== undefined
+      ? { maskRequestHost: network.maskRequestHost }
+      : {}),
+  }
+}
+
 export class SandboxApi {
   protected constructor() {}
 
@@ -605,6 +770,7 @@ export class SandboxApi {
         ? {
             allowOut: res.data.network.allowOut,
             denyOut: res.data.network.denyOut,
+            rules: res.data.network.rules ?? undefined,
             allowPublicTraffic: res.data.network.allowPublicTraffic,
             maskRequestHost: res.data.network.maskRequestHost,
           }
@@ -786,7 +952,7 @@ export class SandboxApi {
       timeout: timeoutToSeconds(timeoutMs),
       secure: opts?.secure ?? true,
       allow_internet_access: opts?.allowInternetAccess ?? true,
-      network: opts?.network,
+      network: buildNetworkBody(opts?.network),
       autoPause: onTimeout === 'pause',
       autoResume: { enabled: autoResume },
     }