Skip to content

Commit 1589964

Browse files
ericcurtinericcurtin
committed
Add proxy support for model pulling behind firewall
- Pass proxy environment variables (HTTP_PROXY, HTTPS_PROXY, NO_PROXY) to docker-model-runner container - Configure HTTP transport to use ProxyFromEnvironment for model pulling - Add tests for proxy configuration in both container creation and transport Co-authored-by: ericcurtin <1694275+ericcurtin@users.noreply.github.com> Signed-off-by: Eric Curtin <eric.curtin@docker.com>
1 parent 1f7eb6d commit 1589964

4 files changed

Lines changed: 208 additions & 1 deletion

File tree

‎cmd/cli/pkg/standalone/containers.go‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -230,6 +230,14 @@ func CreateControllerContainer(ctx context.Context, dockerClient *client.Client,
230230
if doNotTrack {
231231
env = append(env, "DO_NOT_TRACK=1")
232232
}
233+
234+
// Pass proxy environment variables to the container if they are set
235+
proxyEnvVars := []string{"HTTP_PROXY", "HTTPS_PROXY", "NO_PROXY", "http_proxy", "https_proxy", "no_proxy"}
236+
for _, proxyVar := range proxyEnvVars {
237+
if value := os.Getenv(proxyVar); value != "" {
238+
env = append(env, proxyVar+"="+value)
239+
}
240+
}
233241
config := &container.Config{
234242
Image: imageName,
235243
Env: env,

‎cmd/cli/pkg/standalone/containers_test.go‎

Lines changed: 122 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,122 @@
1+
package standalone
2+
3+
import (
4+
"os"
5+
"testing"
6+
)
7+
8+
// TestProxyEnvironmentVariablesPassedToContainer verifies that proxy environment
9+
// variables are correctly identified and would be passed to the container.
10+
func TestProxyEnvironmentVariablesPassedToContainer(t *testing.T) {
11+
tests := []struct {
12+
name string
13+
envVars map[string]string
14+
expected []string
15+
}{
16+
{
17+
name: "All proxy variables set",
18+
envVars: map[string]string{
19+
"HTTP_PROXY": "http://proxy.example.com:8080",
20+
"HTTPS_PROXY": "http://proxy.example.com:8080",
21+
"NO_PROXY": "localhost,127.0.0.1",
22+
},
23+
expected: []string{
24+
"HTTP_PROXY=http://proxy.example.com:8080",
25+
"HTTPS_PROXY=http://proxy.example.com:8080",
26+
"NO_PROXY=localhost,127.0.0.1",
27+
},
28+
},
29+
{
30+
name: "Lowercase proxy variables set",
31+
envVars: map[string]string{
32+
"http_proxy": "http://proxy.example.com:8080",
33+
"https_proxy": "http://proxy.example.com:8080",
34+
"no_proxy": "localhost,127.0.0.1",
35+
},
36+
expected: []string{
37+
"http_proxy=http://proxy.example.com:8080",
38+
"https_proxy=http://proxy.example.com:8080",
39+
"no_proxy=localhost,127.0.0.1",
40+
},
41+
},
42+
{
43+
name: "Mixed case proxy variables",
44+
envVars: map[string]string{
45+
"HTTP_PROXY": "http://proxy.example.com:8080",
46+
"https_proxy": "http://proxy.example.com:8080",
47+
"NO_PROXY": "localhost",
48+
},
49+
expected: []string{
50+
"HTTP_PROXY=http://proxy.example.com:8080",
51+
"https_proxy=http://proxy.example.com:8080",
52+
"NO_PROXY=localhost",
53+
},
54+
},
55+
{
56+
name: "No proxy variables set",
57+
envVars: map[string]string{},
58+
expected: []string{},
59+
},
60+
{
61+
name: "Only HTTP_PROXY set",
62+
envVars: map[string]string{
63+
"HTTP_PROXY": "http://proxy.example.com:8080",
64+
},
65+
expected: []string{
66+
"HTTP_PROXY=http://proxy.example.com:8080",
67+
},
68+
},
69+
}
70+
71+
for _, tt := range tests {
72+
t.Run(tt.name, func(t *testing.T) {
73+
// Save original environment
74+
originalEnv := make(map[string]string)
75+
proxyEnvVars := []string{"HTTP_PROXY", "HTTPS_PROXY", "NO_PROXY", "http_proxy", "https_proxy", "no_proxy"}
76+
for _, key := range proxyEnvVars {
77+
if val, exists := os.LookupEnv(key); exists {
78+
originalEnv[key] = val
79+
}
80+
os.Unsetenv(key)
81+
}
82+
defer func() {
83+
// Restore original environment
84+
for _, key := range proxyEnvVars {
85+
os.Unsetenv(key)
86+
}
87+
for key, val := range originalEnv {
88+
os.Setenv(key, val)
89+
}
90+
}()
91+
92+
// Set test environment variables
93+
for key, val := range tt.envVars {
94+
os.Setenv(key, val)
95+
}
96+
97+
// Simulate the proxy environment variable collection logic
98+
var result []string
99+
for _, proxyVar := range proxyEnvVars {
100+
if value := os.Getenv(proxyVar); value != "" {
101+
result = append(result, proxyVar+"="+value)
102+
}
103+
}
104+
105+
// Verify results
106+
if len(result) != len(tt.expected) {
107+
t.Fatalf("Expected %d environment variables, got %d", len(tt.expected), len(result))
108+
}
109+
110+
expectedMap := make(map[string]bool)
111+
for _, e := range tt.expected {
112+
expectedMap[e] = true
113+
}
114+
115+
for _, r := range result {
116+
if !expectedMap[r] {
117+
t.Errorf("Unexpected environment variable: %s", r)
118+
}
119+
}
120+
})
121+
}
122+
}

‎main.go‎

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -70,12 +70,17 @@ func main() {
7070

7171
memEstimator := memory.NewEstimator(sysMemInfo)
7272

73+
// Create a proxy-aware HTTP transport
74+
// http.DefaultTransport already uses ProxyFromEnvironment by default
75+
baseTransport := http.DefaultTransport.(*http.Transport).Clone()
76+
baseTransport.Proxy = http.ProxyFromEnvironment
77+
7378
modelManager := models.NewManager(
7479
log,
7580
models.ClientConfig{
7681
StoreRootPath: modelPath,
7782
Logger: log.WithFields(logrus.Fields{"component": "model-manager"}),
78-
Transport: resumable.New(http.DefaultTransport),
83+
Transport: resumable.New(baseTransport),
7984
},
8085
nil,
8186
memEstimator,

‎main_test.go‎

Lines changed: 72 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
package main
22

33
import (
4+
"net/http"
45
"os"
56
"testing"
67

@@ -106,3 +107,74 @@ func TestCreateLlamaCppConfigFromEnv(t *testing.T) {
106107
})
107108
}
108109
}
110+
111+
// TestProxyTransportConfiguration verifies that the HTTP transport
112+
// is configured to use proxy settings from environment variables.
113+
func TestProxyTransportConfiguration(t *testing.T) {
114+
tests := []struct {
115+
name string
116+
envVars map[string]string
117+
}{
118+
{
119+
name: "HTTP_PROXY set",
120+
envVars: map[string]string{
121+
"HTTP_PROXY": "http://proxy.example.com:8080",
122+
},
123+
},
124+
{
125+
name: "HTTPS_PROXY set",
126+
envVars: map[string]string{
127+
"HTTPS_PROXY": "http://proxy.example.com:8080",
128+
},
129+
},
130+
{
131+
name: "Both HTTP_PROXY and HTTPS_PROXY set",
132+
envVars: map[string]string{
133+
"HTTP_PROXY": "http://proxy.example.com:8080",
134+
"HTTPS_PROXY": "https://proxy.example.com:8080",
135+
},
136+
},
137+
{
138+
name: "No proxy variables set",
139+
envVars: map[string]string{},
140+
},
141+
}
142+
143+
for _, tt := range tests {
144+
t.Run(tt.name, func(t *testing.T) {
145+
// Save original environment
146+
originalEnv := make(map[string]string)
147+
proxyEnvVars := []string{"HTTP_PROXY", "HTTPS_PROXY", "NO_PROXY", "http_proxy", "https_proxy", "no_proxy"}
148+
for _, key := range proxyEnvVars {
149+
if val, exists := os.LookupEnv(key); exists {
150+
originalEnv[key] = val
151+
}
152+
os.Unsetenv(key)
153+
}
154+
defer func() {
155+
// Restore original environment
156+
for _, key := range proxyEnvVars {
157+
os.Unsetenv(key)
158+
}
159+
for key, val := range originalEnv {
160+
os.Setenv(key, val)
161+
}
162+
}()
163+
164+
// Set test environment variables
165+
for key, val := range tt.envVars {
166+
os.Setenv(key, val)
167+
}
168+
169+
// Test that we can create a proxy-aware transport
170+
// This simulates what we do in main()
171+
baseTransport := http.DefaultTransport.(*http.Transport).Clone()
172+
baseTransport.Proxy = http.ProxyFromEnvironment
173+
174+
// Verify the transport has a Proxy function set
175+
if baseTransport.Proxy == nil {
176+
t.Error("Expected Proxy function to be set on transport")
177+
}
178+
})
179+
}
180+
}

0 commit comments

Comments
 (0)