From 2c9dcca969e2b7c06ace72f658832a5555a8d42d Mon Sep 17 00:00:00 2001
From: Dmitry Meyer <me@undef.im>
Date: Fri, 29 May 2026 12:49:19 +0000
Subject: [PATCH] [shim] Pass proxy variables to the container

Fixes: https://github.com/dstackai/dstack/issues/3906
---
 runner/cmd/shim/main.go             | 14 ++++++++++++++
 runner/internal/shim/docker.go      | 26 ++++++++++++++++++++++++--
 runner/internal/shim/docker_test.go |  4 ++++
 runner/internal/shim/models.go      |  2 ++
 4 files changed, 44 insertions(+), 2 deletions(-)

diff --git a/runner/cmd/shim/main.go b/runner/cmd/shim/main.go
index 2bf9549ba2..6926488455 100644
--- a/runner/cmd/shim/main.go
+++ b/runner/cmd/shim/main.go
@@ -27,6 +27,13 @@ import (
 // Version is a build-time variable. The value is overridden by ldflags.
 var Version string
 
+// https://everything.curl.dev/usingcurl/proxies/env.html
+// https://cs.opensource.google/go/x/net/+/657eb1317b5dd33038d683297c6be9cae05fa97d:http/httpproxy/proxy.go
+// We accept HTTP_PROXY in upper case without additional checks as it's unlikely that
+// the shim is running in the CGI context
+// The lower case form should be used as some applications ignore the upper case form, e.g., curl, apt
+const defaultPassEnv = "http_proxy,https_proxy,no_proxy,HTTP_PROXY,HTTPS_PROXY,NO_PROXY"
+
 func main() {
 	os.Exit(mainInner())
 }
@@ -147,6 +154,13 @@ func mainInner() int {
 				Sources:     cli.EnvVars("DSTACK_DCGM_ADDRESS"),
 			},
 			/* Docker Parameters */
+			&cli.StringFlag{
+				Name:        "pass-env",
+				Usage:       "Environment variables to pass on to the container, a comma-separated list of names",
+				Value:       defaultPassEnv,
+				Destination: &args.Docker.PassEnv,
+				Sources:     cli.EnvVars("DSTACK_DOCKER_PASS_ENV"),
+			},
 			&cli.BoolFlag{
 				Name:        "privileged",
 				Usage:       "Give extended privileges to the container",
diff --git a/runner/internal/shim/docker.go b/runner/internal/shim/docker.go
index 4a43427e66..49a11c7e7e 100644
--- a/runner/internal/shim/docker.go
+++ b/runner/internal/shim/docker.go
@@ -135,6 +135,7 @@ type DockerRunner struct {
 	client       *docker.Client
 	dockerParams DockerParameters
 	dockerInfo   dockersystem.Info
+	baseEnv      []string
 	gpus         []host.GpuInfo
 	gpuVendor    gpu.GpuVendor
 	gpuLock      *GpuLock
@@ -151,6 +152,15 @@ func NewDockerRunner(ctx context.Context, dockerParams DockerParameters) (*Docke
 		return nil, fmt.Errorf("get docker info: %w", err)
 	}
 
+	// Copy variables once rather than on a per-task basis
+	// We don't expect variables to change during the shim's lifetime
+	baseEnv := []string{}
+	for _, name := range dockerParams.DockerPassEnv() {
+		if value, ok := os.LookupEnv(name); ok {
+			baseEnv = append(baseEnv, fmt.Sprintf("%s=%s", name, value))
+		}
+	}
+
 	var gpuVendor gpu.GpuVendor
 	gpus := host.GetGpuInfo(ctx)
 	if len(gpus) > 0 {
@@ -167,6 +177,7 @@ func NewDockerRunner(ctx context.Context, dockerParams DockerParameters) (*Docke
 		client:       client,
 		dockerParams: dockerParams,
 		dockerInfo:   dockerInfo,
+		baseEnv:      baseEnv,
 		gpus:         gpus,
 		gpuVendor:    gpuVendor,
 		gpuLock:      gpuLock,
@@ -861,8 +872,9 @@ func (d *DockerRunner) createContainer(ctx context.Context, task *Task) error {
 
 	// Set the environment variables
 	envVars := []string{}
-	if d.dockerParams.DockerPJRTDevice() != "" {
-		envVars = append(envVars, fmt.Sprintf("PJRT_DEVICE=%s", d.dockerParams.DockerPJRTDevice()))
+	envVars = append(envVars, d.baseEnv...)
+	if pjrtDevice := d.dockerParams.DockerPJRTDevice(); pjrtDevice != "" {
+		envVars = append(envVars, fmt.Sprintf("PJRT_DEVICE=%s", pjrtDevice))
 	}
 
 	// Override /dev/shm with tmpfs mount with `exec` option (the default is `noexec`)
@@ -1235,6 +1247,16 @@ func getContainerLastLogs(ctx context.Context, client docker.APIClient, containe
 
 /* DockerParameters interface implementation for CLIArgs */
 
+func (c *CLIArgs) DockerPassEnv() []string {
+	names := []string{}
+	for _, name := range strings.Split(c.Docker.PassEnv, ",") {
+		if name = strings.TrimSpace(name); name != "" {
+			names = append(names, name)
+		}
+	}
+	return names
+}
+
 func (c *CLIArgs) DockerPrivileged() bool {
 	return c.Docker.Privileged
 }
diff --git a/runner/internal/shim/docker_test.go b/runner/internal/shim/docker_test.go
index 68e45ed1cf..259c76dcd7 100644
--- a/runner/internal/shim/docker_test.go
+++ b/runner/internal/shim/docker_test.go
@@ -119,6 +119,10 @@ func (c *dockerParametersMock) DockerShellCommands(authorizedKeys []string, runn
 	return commands
 }
 
+func (c *dockerParametersMock) DockerPassEnv() []string {
+	return []string{}
+}
+
 func (c *dockerParametersMock) DockerPorts() []int {
 	return []int{}
 }
diff --git a/runner/internal/shim/models.go b/runner/internal/shim/models.go
index 549b7df7fa..6cc0935053 100644
--- a/runner/internal/shim/models.go
+++ b/runner/internal/shim/models.go
@@ -5,6 +5,7 @@ import (
 )
 
 type DockerParameters interface {
+	DockerPassEnv() []string
 	DockerPrivileged() bool
 	DockerShellCommands(authorizedKeys []string, runnerHttpAddress string) []string
 	DockerMounts(string) ([]mount.Mount, error)
@@ -40,6 +41,7 @@ type CLIArgs struct {
 	}
 
 	Docker struct {
+		PassEnv    string
 		Privileged bool
 		PJRTDevice string
 	}