Skip to content

Commit 74dfba5

Browse files
cooper-grccooper-grc
committed
Networking Controls
1 parent 8b90eca commit 74dfba5

5 files changed

Lines changed: 213 additions & 3 deletions

File tree

‎lib/resources/abstraction/sandbox.ts‎

Lines changed: 42 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@ import type {
77
PodSandboxCreateImageFromFilesystemResponse,
88
PodSandboxUpdateTtlResponse,
99
PodSandboxExposePortResponse,
10+
PodSandboxUpdateNetworkPermissionsResponse,
1011
PodSandboxExecResponse,
1112
PodSandboxListFilesResponse,
1213
PodSandboxCreateDirectoryResponse,
@@ -51,6 +52,10 @@ function shellQuote(arg: string): string {
5152
* for sandboxes that never timeout.
5253
* - volumes (Volume[]): The volumes and/or cloud buckets to mount into the sandbox container.
5354
* - secrets (SecretVar[]): Secrets to pass to the sandbox, e.g. [{ name: "API_KEY" }].
55+
* - blockNetwork (boolean): Whether to block all outbound network access. Cannot be combined with
56+
* `allowList`.
57+
* - allowList (string[]): CIDR ranges that are allowed for outbound network access. When specified,
58+
* all other outbound traffic is blocked.
5459
* - authorized (boolean): Ignored for sandboxes (forced to false).
5560
*/
5661
export class Sandbox extends Pod {
@@ -373,17 +378,51 @@ export class SandboxInstance extends PodInstance {
373378
}
374379

375380
/**
376-
* List all exposed URLs in the sandbox.
381+
* Dynamically update outbound network permissions for the sandbox.
377382
*/
378-
public async listUrls(): Promise<string[]> {
383+
public async updateNetworkPermissions(
384+
blockNetwork: boolean = false,
385+
allowList?: string[]
386+
): Promise<void> {
387+
if (blockNetwork && allowList && allowList.length > 0) {
388+
throw new Error("Cannot specify both blockNetwork=true and allowList");
389+
}
390+
391+
const resp = await beamClient.request({
392+
method: "POST",
393+
url: `/api/v1/gateway/pods/${this.containerId}/network/update`,
394+
data: {
395+
stubId: this.stubId,
396+
blockNetwork,
397+
allowList: allowList ?? [],
398+
},
399+
});
400+
const data = resp.data as PodSandboxUpdateNetworkPermissionsResponse;
401+
if (!data.ok) {
402+
throw new SandboxProcessError(
403+
data.errorMsg || "Failed to update network permissions"
404+
);
405+
}
406+
}
407+
408+
/**
409+
* List all exposed URLs in the sandbox, keyed by port.
410+
*/
411+
public async listUrls(): Promise<Record<number, string>> {
379412
const resp = await beamClient.request({
380413
method: "GET",
381414
url: `/api/v1/gateway/pods/${this.containerId}/urls`,
382415
});
383416
const data = resp.data as PodSandboxListUrlsResponse;
384417
if (!data.ok)
385418
throw new SandboxProcessError(data.errorMsg || "Failed to list URLs");
386-
return Object.values(data.urls || {});
419+
420+
const urlsByPort: Record<number, string> = {};
421+
for (const [port, url] of Object.entries(data.urls || {})) {
422+
urlsByPort[Number(port)] = url;
423+
}
424+
425+
return urlsByPort;
387426
}
388427

389428
/**

‎lib/resources/abstraction/stub.ts‎

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -52,6 +52,8 @@ export interface StubConfig {
5252
inputs?: Schema;
5353
outputs?: Schema;
5454
tcp: boolean;
55+
blockNetwork: boolean;
56+
allowList?: string[];
5557
}
5658

5759
export interface CreateStubConfig extends Partial<StubConfig> {
@@ -110,6 +112,8 @@ export class StubBuilder {
110112
inputs = undefined,
111113
outputs = undefined,
112114
tcp = false,
115+
blockNetwork = false,
116+
allowList = undefined,
113117
}: CreateStubConfig) {
114118
this.config = {} as StubConfig;
115119
this.config.name = name;
@@ -137,6 +141,14 @@ export class StubBuilder {
137141
this.config.pricing = pricing;
138142
this.config.inputs = inputs;
139143
this.config.outputs = outputs;
144+
this.config.blockNetwork = blockNetwork;
145+
this.config.allowList = allowList;
146+
147+
if (this.config.blockNetwork && this.config.allowList !== undefined) {
148+
throw new Error(
149+
"Cannot specify both 'blockNetwork=true' and 'allowList'. Use 'allowList' with CIDR notation to allow specific ranges, or use 'blockNetwork=true' to block all outbound traffic."
150+
);
151+
}
140152

141153
// Set GPU count if GPU specified but count is 0
142154
if (
@@ -340,6 +352,8 @@ export class StubBuilder {
340352
inputs,
341353
outputs,
342354
tcp: this.config.tcp,
355+
blockNetwork: this.config.blockNetwork,
356+
allowList: this.config.allowList,
343357
};
344358

345359
try {

‎lib/types/pod.ts‎

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -217,6 +217,18 @@ export interface PodSandboxExposePortResponse {
217217
errorMsg: string;
218218
}
219219

220+
export interface PodSandboxUpdateNetworkPermissionsRequest {
221+
containerId: string;
222+
stubId: string;
223+
blockNetwork: boolean;
224+
allowList: string[];
225+
}
226+
227+
export interface PodSandboxUpdateNetworkPermissionsResponse {
228+
ok: boolean;
229+
errorMsg: string;
230+
}
231+
220232
export interface PodSandboxListUrlsResponse {
221233
ok: boolean;
222234
urls: Record<string, string>;

‎lib/types/stub.ts‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -72,6 +72,8 @@ export interface GetOrCreateStubRequest {
7272
inputs?: Schema;
7373
outputs?: Schema;
7474
tcp: boolean;
75+
blockNetwork: boolean;
76+
allowList?: string[];
7577
}
7678

7779
export interface GetOrCreateStubResponse {

‎tests/sandbox-network.test.ts‎

Lines changed: 143 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,143 @@
1+
import beamClient from "../lib";
2+
import { Sandbox, SandboxInstance } from "../lib/resources/abstraction/sandbox";
3+
import { EStubType } from "../lib/types/stub";
4+
5+
describe("Sandbox network parity", () => {
6+
beforeEach(() => {
7+
jest.spyOn(console, "log").mockImplementation(() => undefined);
8+
jest.spyOn(console, "warn").mockImplementation(() => undefined);
9+
jest.spyOn(console, "error").mockImplementation(() => undefined);
10+
});
11+
12+
afterEach(() => {
13+
jest.restoreAllMocks();
14+
});
15+
16+
test("rejects sandbox configs that set both blockNetwork and allowList", () => {
17+
expect(() => {
18+
new Sandbox({
19+
name: "networked-sandbox",
20+
blockNetwork: true,
21+
allowList: ["8.8.8.8/32"],
22+
});
23+
}).toThrow(
24+
"Cannot specify both 'blockNetwork=true' and 'allowList'. Use 'allowList' with CIDR notation to allow specific ranges, or use 'blockNetwork=true' to block all outbound traffic."
25+
);
26+
});
27+
28+
test("includes allowList in stub creation requests", async () => {
29+
const requestMock = jest.spyOn(beamClient, "request").mockResolvedValue({
30+
data: {
31+
ok: true,
32+
stubId: "stub-123",
33+
},
34+
});
35+
36+
const sandbox = new Sandbox({
37+
name: "networked-sandbox",
38+
allowList: ["8.8.8.8/32"],
39+
});
40+
41+
sandbox.stub.imageAvailable = true;
42+
sandbox.stub.filesSynced = true;
43+
sandbox.stub.objectId = "object-123";
44+
sandbox.stub.config.image.id = "image-123";
45+
46+
await expect(
47+
sandbox.stub.prepareRuntime(undefined, EStubType.Sandbox, true, ["*"])
48+
).resolves.toBe(true);
49+
50+
expect(requestMock).toHaveBeenCalledWith(
51+
expect.objectContaining({
52+
method: "POST",
53+
url: "/api/v1/gateway/stubs",
54+
data: expect.objectContaining({
55+
block_network: false,
56+
allow_list: ["8.8.8.8/32"],
57+
}),
58+
})
59+
);
60+
});
61+
62+
test("updates network permissions with the sandbox update endpoint", async () => {
63+
const requestMock = jest.spyOn(beamClient, "request").mockResolvedValue({
64+
data: {
65+
ok: true,
66+
errorMsg: "",
67+
},
68+
});
69+
70+
const instance = new SandboxInstance(
71+
{
72+
containerId: "sandbox-123",
73+
stubId: "stub-123",
74+
url: "",
75+
ok: true,
76+
errorMsg: "",
77+
},
78+
new Sandbox({ name: "networked-sandbox" })
79+
);
80+
81+
await expect(instance.updateNetworkPermissions(true)).resolves.toBeUndefined();
82+
83+
expect(requestMock).toHaveBeenCalledWith({
84+
method: "POST",
85+
url: "/api/v1/gateway/pods/sandbox-123/network/update",
86+
data: {
87+
stubId: "stub-123",
88+
blockNetwork: true,
89+
allowList: [],
90+
},
91+
});
92+
});
93+
94+
test("rejects conflicting network permission updates before making a request", async () => {
95+
const requestMock = jest.spyOn(beamClient, "request");
96+
97+
const instance = new SandboxInstance(
98+
{
99+
containerId: "sandbox-123",
100+
stubId: "stub-123",
101+
url: "",
102+
ok: true,
103+
errorMsg: "",
104+
},
105+
new Sandbox({ name: "networked-sandbox" })
106+
);
107+
108+
await expect(
109+
instance.updateNetworkPermissions(true, ["8.8.8.8/32"])
110+
).rejects.toThrow("Cannot specify both blockNetwork=true and allowList");
111+
112+
expect(requestMock).not.toHaveBeenCalled();
113+
});
114+
115+
test("returns exposed URLs keyed by port", async () => {
116+
jest.spyOn(beamClient, "request").mockResolvedValue({
117+
data: {
118+
ok: true,
119+
urls: {
120+
"3000": "https://3000.example.com",
121+
"8080": "https://8080.example.com",
122+
},
123+
errorMsg: "",
124+
},
125+
});
126+
127+
const instance = new SandboxInstance(
128+
{
129+
containerId: "sandbox-123",
130+
stubId: "stub-123",
131+
url: "",
132+
ok: true,
133+
errorMsg: "",
134+
},
135+
new Sandbox({ name: "networked-sandbox" })
136+
);
137+
138+
await expect(instance.listUrls()).resolves.toEqual({
139+
3000: "https://3000.example.com",
140+
8080: "https://8080.example.com",
141+
});
142+
});
143+
});

0 commit comments

Comments
 (0)