Add block_network to Sandboxes

docs/openapi/gateway.swagger.json

@@ -1480,6 +1480,9 @@
         },
         "tcp": {
           "type": "boolean"
+        },
+        "blockNetwork": {
+          "type": "boolean"
         }
       }
     },

pkg/abstractions/pod/pod.go

@@ -334,6 +334,7 @@ func (s *GenericPodService) run(ctx context.Context, authInfo *auth.AuthInfo, st
 		Ports:             ports,
 		CheckpointEnabled: checkpointEnabled,
 		Checkpoint:        checkpoint,
+		BlockNetwork:      stubConfig.BlockNetwork,
 	})
 	if err != nil {
 		return "", err

pkg/gateway/gateway.proto

@@ -445,6 +445,7 @@ message GetOrCreateStubRequest {
   Schema inputs = 35;
   Schema outputs = 36;
   bool tcp = 37;
+  bool block_network = 38;
 }
 
 message GetOrCreateStubResponse {

pkg/gateway/services/stub.go

@@ -158,6 +158,7 @@ func (gws *GatewayService) GetOrCreateStub(ctx context.Context, in *pb.GetOrCrea
 		Inputs:             inputs,
 		Outputs:            outputs,
 		TCP:                in.Tcp,
+		BlockNetwork:       in.BlockNetwork,
 	}
 
 	// Ensure GPU count is at least 1 if a GPU is required

pkg/types/backend.go

@@ -460,6 +460,7 @@ type StubConfigV1 struct {
 	Inputs             *Schema         `json:"inputs"`
 	Outputs            *Schema         `json:"outputs"`
 	TCP                bool            `json:"tcp"`
+	BlockNetwork       bool            `json:"block_network"`
 }
 
 type StubConfigLimitedValues struct {

pkg/types/scheduler.go

@@ -225,6 +225,7 @@ type ContainerRequest struct {
 	AppId             string          `json:"app_id"`
 	Checkpoint        *Checkpoint     `json:"checkpoint"`
 	ConfigPath        string          `json:"config_path"`
+	BlockNetwork      bool            `json:"block_network"`
 }
 
 func (c *ContainerRequest) RequiresGPU() bool {

pkg/types/types.proto

@@ -88,6 +88,7 @@ message ContainerRequest {
   string app_id = 23;
   Checkpoint checkpoint = 24;
   string config_path = 25;
+  bool block_network = 26;
 }
 
 message ContainerState {

pkg/worker/network.go

@@ -509,6 +509,31 @@ func (m *ContainerNetworkManager) configureContainerNetwork(opts *containerNetwo
 		m.containerInstances.Set(opts.containerId, containerInstance)
 	}
 
+	// Block outbound network access if requested
+	if opts.request.BlockNetwork {
+		// Block IPv4 outbound traffic
+		err = m.ipt.InsertUnique("filter", "FORWARD", 1, "-s", ipAddr.IP.String(), "-o", m.defaultLink.Attrs().Name, "-j", "DROP")
+		if err != nil {
+			return err
+		}
+
+		// Block IPv6 outbound traffic if enabled
+		if m.ipt6 != nil {
+			// Calculate the corresponding IPv6 address using the last octet of the IPv4 address
+			ipv4LastOctet := int(ipAddr.IP.To4()[3])
+			_, ipv6Net, _ := net.ParseCIDR(containerSubnetIPv6)
+			ipv6Prefix := ipv6Net.IP.String()
+			ipv6Address := fmt.Sprintf("%s%x", ipv6Prefix, ipv4LastOctet)
+
+			err = m.ipt6.InsertUnique("filter", "FORWARD", 1, "-s", ipv6Address, "-o", m.defaultLink.Attrs().Name, "-j", "DROP")
+			if err != nil {
+				return err
+			}
+		}
+
+		log.Info().Str("container_id", opts.containerId).Str("ip_address", ipAddr.IP.String()).Msg("outbound network access blocked for container")
+	}
+
 	return nil
 }
 

sdk/src/beta9/abstractions/base/runner.py

@@ -117,6 +117,7 @@ def __init__(
         inputs: Optional[Schema] = None,
         outputs: Optional[Schema] = None,
         tcp: bool = False,
+        block_network: bool = False,
     ) -> None:
         super().__init__()
 
@@ -163,6 +164,7 @@ def __init__(
         self.extra: dict = {}
         self.entrypoint: Optional[List[str]] = entrypoint
         self.tcp = tcp
+        self.block_network = block_network
 
         if (self.gpu != "" or len(self.gpu) > 0) and self.gpu_count == 0:
             self.gpu_count = 1
@@ -518,6 +520,7 @@ def prepare_runtime(
                 inputs=inputs,
                 outputs=outputs,
                 tcp=self.tcp,
+                block_network=self.block_network,
             )
 
             if _is_stub_created_for_workspace():

sdk/src/beta9/abstractions/pod.py

@@ -136,6 +136,7 @@ def __init__(
         keep_warm_seconds: int = 600,
         authorized: bool = False,
         tcp: bool = False,
+        block_network: bool = False,
     ) -> None:
         super().__init__(
             cpu=cpu,
@@ -148,6 +149,7 @@ def __init__(
             env=env,
             entrypoint=entrypoint,
             ports=ports,
+            block_network=block_network,
             name=name,
             authorized=authorized,
             keep_warm_seconds=keep_warm_seconds,

sdk/src/beta9/abstractions/sandbox.py

@@ -87,6 +87,10 @@ class Sandbox(Pod):
             A dictionary of environment variables to be injected into each Sandbox container. Default is {}.
         sync_local_dir (bool):
             Whether to sync the local directory to the sandbox filesystem on creation. Default is False.
+        block_network (bool):
+            Whether to block all outbound network access from the sandbox. When enabled, the sandbox cannot
+            make outbound connections to external services, but inbound connections to exposed ports are still
+            allowed. Default is False.
 
     Example:
         ```python
@@ -126,6 +130,7 @@ def __init__(
         secrets: Optional[List[str]] = None,
         env: Optional[Dict[str, str]] = {},
         sync_local_dir: bool = False,
+        block_network: bool = False,
     ):
         self.debug_buffer = io.StringIO()
         self.sync_local_dir = sync_local_dir
@@ -142,6 +147,7 @@ def __init__(
             volumes=volumes,
             secrets=secrets,
             env=env,
+            block_network=block_network,
         )
 
     def debug(self):