Skip to content

Commit

Permalink
feat: automatically set container proxy env variables [DET-3414] (#838)
Browse files Browse the repository at this point in the history
* go.sum addition

* pass proxy env variables to agent container

* clarifying some comments

* updating comments

* proxy vars from agent config

* formatting files

* linting and documentation

* proxy vars do not override existing vars

* added friendlier formatting and code cleanup
  • Loading branch information
eecsliu authored Jul 22, 2020
1 parent b75442a commit f11ef7f
Show file tree
Hide file tree
Showing 7 changed files with 145 additions and 0 deletions.
1 change: 1 addition & 0 deletions agent/go.mod
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,7 @@ require (
github.com/docker/go-connections v0.4.0
github.com/ghodss/yaml v1.0.1-0.20190212211648-25d852aebe32
github.com/go-ole/go-ole v1.2.4 // indirect
github.com/golang-collections/collections v0.0.0-20130729185459-604e922904d3
github.com/golangci/golangci-lint v1.27.0
github.com/google/uuid v1.1.1
github.com/goreleaser/goreleaser v0.133.0
Expand Down
13 changes: 13 additions & 0 deletions agent/go.sum
Original file line number Diff line number Diff line change
Expand Up @@ -28,22 +28,32 @@ github.com/Azure/go-ansiterm v0.0.0-20170929234023-d6e3b3328b78 h1:w+iIsaOQNcT7O
github.com/Azure/go-ansiterm v0.0.0-20170929234023-d6e3b3328b78/go.mod h1:LmzpDX56iTiv29bbRTIsUNlaFfuhWRQBWjQdVyAevI8=
github.com/Azure/go-autorest v12.0.0+incompatible h1:N+VqClcomLGD/sHb3smbSYYtNMgKpVV3Cd5r5i8z6bQ=
github.com/Azure/go-autorest v12.0.0+incompatible/go.mod h1:r+4oMnoxhatjLLJ6zxSWATqVooLgysK6ZNox3g/xq24=
github.com/Azure/go-autorest v12.2.0+incompatible h1:2Fxszbg492oAJrcvJlgyVaTqnQYRkxmEK6VPCLLVpBI=
github.com/Azure/go-autorest v12.2.0+incompatible/go.mod h1:r+4oMnoxhatjLLJ6zxSWATqVooLgysK6ZNox3g/xq24=
github.com/Azure/go-autorest/autorest v0.9.0/go.mod h1:xyHB1BMZT0cuDHU7I0+g046+BFDTQ8rEZB0s4Yfa6bI=
github.com/Azure/go-autorest/autorest v0.9.3 h1:OZEIaBbMdUE/Js+BQKlpO81XlISgipr6yDJ+PSwsgi4=
github.com/Azure/go-autorest/autorest v0.9.3/go.mod h1:GsRuLYvwzLjjjRoWEIyMUaYq8GNUx2nRB378IPt/1p0=
github.com/Azure/go-autorest/autorest/adal v0.5.0/go.mod h1:8Z9fGy2MpX0PvDjB1pEgQTmVqjGhiHBW7RJJEciWzS0=
github.com/Azure/go-autorest/autorest/adal v0.8.0/go.mod h1:Z6vX6WXXuyieHAXwMj0S6HY6e6wcHn37qQMBQlvY3lc=
github.com/Azure/go-autorest/autorest/adal v0.8.1 h1:pZdL8o72rK+avFWl+p9nE8RWi1JInZrWJYlnpfXJwHk=
github.com/Azure/go-autorest/autorest/adal v0.8.1/go.mod h1:ZjhuQClTqx435SRJ2iMlOxPYt3d2C/T/7TiQCVZSn3Q=
github.com/Azure/go-autorest/autorest/azure/auth v0.4.2 h1:iM6UAvjR97ZIeR93qTcwpKNMpV+/FTWjwEbuPD495Tk=
github.com/Azure/go-autorest/autorest/azure/auth v0.4.2/go.mod h1:90gmfKdlmKgfjUpnCEpOJzsUEjrWDSLwHIG73tSXddM=
github.com/Azure/go-autorest/autorest/azure/cli v0.3.1 h1:LXl088ZQlP0SBppGFsRZonW6hSvwgL5gRByMbvUbx8U=
github.com/Azure/go-autorest/autorest/azure/cli v0.3.1/go.mod h1:ZG5p860J94/0kI9mNJVoIoLgXcirM2gF5i2kWloofxw=
github.com/Azure/go-autorest/autorest/date v0.1.0/go.mod h1:plvfp3oPSKwf2DNjlBjWF/7vwR+cUD/ELuzDCXwHUVA=
github.com/Azure/go-autorest/autorest/date v0.2.0 h1:yW+Zlqf26583pE43KhfnhFcdmSWlm5Ew6bxipnr/tbM=
github.com/Azure/go-autorest/autorest/date v0.2.0/go.mod h1:vcORJHLJEh643/Ioh9+vPmf1Ij9AEBM5FuBIXLmIy0g=
github.com/Azure/go-autorest/autorest/mocks v0.1.0/go.mod h1:OTyCOPRA2IgIlWxVYxBee2F5Gr4kF2zd2J5cFRaIDN0=
github.com/Azure/go-autorest/autorest/mocks v0.2.0/go.mod h1:OTyCOPRA2IgIlWxVYxBee2F5Gr4kF2zd2J5cFRaIDN0=
github.com/Azure/go-autorest/autorest/mocks v0.3.0/go.mod h1:a8FDP3DYzQ4RYfVAxAN3SVSiiO77gL2j2ronKKP0syM=
github.com/Azure/go-autorest/autorest/to v0.3.0 h1:zebkZaadz7+wIQYgC7GXaz3Wb28yKYfVkkBKwc38VF8=
github.com/Azure/go-autorest/autorest/to v0.3.0/go.mod h1:MgwOyqaIuKdG4TL/2ywSsIWKAfJfgHDo8ObuUk3t5sA=
github.com/Azure/go-autorest/autorest/validation v0.2.0 h1:15vMO4y76dehZSq7pAaOLQxC6dZYsSrj2GQpflyM/L4=
github.com/Azure/go-autorest/autorest/validation v0.2.0/go.mod h1:3EEqHnBxQGHXRYq3HT1WyXAvT7LLY3tl70hw6tQIbjI=
github.com/Azure/go-autorest/logger v0.1.0 h1:ruG4BSDXONFRrZZJ2GUXDiUyVpayPmb1GnWeHDdaNKY=
github.com/Azure/go-autorest/logger v0.1.0/go.mod h1:oExouG+K6PryycPJfVSxi/koC6LSNgds39diKLz7Vrc=
github.com/Azure/go-autorest/tracing v0.5.0 h1:TRn4WjSnkcSy5AEG3pnbtFSwNtwzjr4VYyQflFE619k=
github.com/Azure/go-autorest/tracing v0.5.0/go.mod h1:r/s2XiOKccPW3HrqB+W0TQzfbtp2fGCgRFtBroKn4Dk=
github.com/BurntSushi/toml v0.3.1 h1:WXkYYl6Yr3qBf1K79EBnL4mak0OimBfB0XUf9Vl28OQ=
github.com/BurntSushi/toml v0.3.1/go.mod h1:xHWCNGjB5oqiDr8zfno3MHue2Ht5sIBksp03qcyfWMU=
Expand Down Expand Up @@ -222,6 +232,8 @@ github.com/gogo/protobuf v1.2.1 h1:/s5zKNz0uPFCZ5hddgPdo2TK2TVrUNMn0OOX8/aZMTE=
github.com/gogo/protobuf v1.2.1/go.mod h1:hp+jE20tsWTFYpLwKvXlhS1hjn+gTNwPg2I6zVXpSg4=
github.com/gogo/protobuf v1.2.2-0.20190723190241-65acae22fc9d h1:3PaI8p3seN09VjbTYC/QWlUZdZ1qS1zGjy7LH2Wt07I=
github.com/gogo/protobuf v1.2.2-0.20190723190241-65acae22fc9d/go.mod h1:SlYgWuQ5SjCEi6WLHjHCa1yvBfUnHcTbrrZtXPKa29o=
github.com/golang-collections/collections v0.0.0-20130729185459-604e922904d3 h1:zN2lZNZRflqFyxVaTIU61KNKQ9C0055u9CAfpmqUvo4=
github.com/golang-collections/collections v0.0.0-20130729185459-604e922904d3/go.mod h1:nPpo7qLxd6XL3hWJG/O60sR8ZKfMCiIoNap5GvD12KU=
github.com/golang-migrate/migrate v3.5.4+incompatible/go.mod h1:IsVUlFN5puWOmXrqjgGUfIRIbU7mr8oNBE2tyERd9Wk=
github.com/golang/glog v0.0.0-20160126235308-23def4e6c14b h1:VKtxabqXZkF25pY9ekfRL6a582T4P37/31XEstQ5p58=
github.com/golang/glog v0.0.0-20160126235308-23def4e6c14b/go.mod h1:SBH7ygxi8pfUlaOkMMuAQtPIUF8ecWP5IEl/CR7VP2Q=
Expand Down Expand Up @@ -542,6 +554,7 @@ github.com/sourcegraph/go-diff v0.5.1/go.mod h1:j2dHj3m8aZgQO8lMTcTnBcXkRRRqi34c
github.com/spaolacci/murmur3 v0.0.0-20180118202830-f09979ecbc72/go.mod h1:JwIasOWyU6f++ZhiEuf87xNszmSA2myDM2Kzu9HwQUA=
github.com/spf13/afero v1.1.2 h1:m8/z1t7/fwjysjQRYbP0RD+bUIF/8tJwPdEZsI83ACI=
github.com/spf13/afero v1.1.2/go.mod h1:j4pytiNVoe2o6bmDsKpLACNPDBIoEAkihy7loJ1B0CQ=
github.com/spf13/afero v1.2.2 h1:5jhuqJyZCZf2JRofRvN/nIFgIWNzPa3/Vz8mYylgbWc=
github.com/spf13/afero v1.2.2/go.mod h1:9ZxEEn6pIJ8Rxe320qSDBk6AsU0r9pR7Q4OcevTdifk=
github.com/spf13/cast v1.3.0 h1:oget//CVOEoFewqQxwr0Ej5yjygnqGkvggSE/gB35Q8=
github.com/spf13/cast v1.3.0/go.mod h1:Qx5cxh0v+4UWYiBimWS+eyWzqEqokIECu5etghLkUJE=
Expand Down
27 changes: 27 additions & 0 deletions agent/internal/agent.go
Original file line number Diff line number Diff line change
Expand Up @@ -9,8 +9,10 @@ import (
"net/http/pprof"
"os"
"runtime"
"strings"
"syscall"

"github.com/docker/docker/api/types/container"
"github.com/google/uuid"
"github.com/gorilla/websocket"
"github.com/labstack/echo"
Expand All @@ -36,6 +38,30 @@ type agent struct {
cm *actor.Ref
}

func (a *agent) addProxy(config *container.Config) {
addVars := map[string]string{
"HTTP_PROXY": a.Options.HTTPProxy,
"HTTPS_PROXY": a.Options.HTTPSProxy,
"FTP_PROXY": a.Options.FTPProxy,
"NO_PROXY": a.Options.NoProxy,
}

for _, v := range config.Env {
key := strings.SplitN(v, "=", 2)[0]
key = strings.ToUpper(key)
_, ok := addVars[key]
if ok {
delete(addVars, key)
}
}

for k, v := range addVars {
if v != "" {
config.Env = append(config.Env, k+"="+v)
}
}
}

func (a *agent) Receive(ctx *actor.Context) error {
switch msg := ctx.Message().(type) {
case actor.PreStart:
Expand All @@ -44,6 +70,7 @@ func (a *agent) Receive(ctx *actor.Context) error {
case proto.AgentMessage:
switch {
case msg.StartContainer != nil:
a.addProxy(&msg.StartContainer.Spec.RunSpec.ContainerConfig)
if !a.validateDevices(msg.StartContainer.Container.Devices) {
return errors.New("could not start container; devices specified in spec not found on agent")
}
Expand Down
92 changes: 92 additions & 0 deletions agent/internal/agent_test.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,92 @@
package internal

import (
"testing"

"github.com/docker/docker/api/types/container"
"github.com/golang-collections/collections/set"
)

func TestNoAddProxy(t *testing.T) {
inputEnv := container.Config{}
testAgent := agent{}

inputEnv.Env = []string{
"FIRST_VAR=1",
}

// InputEnv should not change because we didn't set any environment variables.
ans := append([]string{}, inputEnv.Env...)

testAgent.addProxy(&inputEnv)

if !compareSlices(inputEnv.Env, ans) {
t.Errorf("Expected: %v But got: %v", ans, inputEnv.Env)
}
}

func TestAddProxy(t *testing.T) {
inputEnv := container.Config{}
testAgent := agent{}

inputEnv.Env = []string{
"FIRST_VAR=1",
}

testAgent.Options.HTTPProxy = "192.168.1.1"
testAgent.Options.HTTPSProxy = "192.168.1.2"
testAgent.Options.FTPProxy = "192.168.1.3"
testAgent.Options.NoProxy = "*.test.com"

ans := append(inputEnv.Env, []string{
"HTTP_PROXY=192.168.1.1",
"HTTPS_PROXY=192.168.1.2",
"FTP_PROXY=192.168.1.3",
"NO_PROXY=*.test.com",
}...)

testAgent.addProxy(&inputEnv)

if !compareSlices(inputEnv.Env, ans) {
t.Errorf("Expected: %v But got: %v", ans, inputEnv.Env)
}
}

func TestAlreadyAddedProxy(t *testing.T) {
inputEnv := container.Config{}
testAgent := agent{}

inputEnv.Env = []string{
"FIRST_VAR=1",
"HTTP_PROXY=10.0.0.1",
}

testAgent.Options.HTTPProxy = "10.0.0.2"

// InputEnv should not change because existing config should not be overridden
ans := append([]string{}, inputEnv.Env...)

testAgent.addProxy(&inputEnv)

if !compareSlices(inputEnv.Env, ans) {
t.Errorf("Expected: %v But got: %v", ans, inputEnv.Env)
}
}

func compareSlices(env []string, ans []string) bool {
output := set.New()
correct := set.New()

for _, v := range env {
output.Insert(v)
}

for _, v := range ans {
correct.Insert(v)
}

if diff := output.Difference(correct); diff.Len() != 0 {
return false
}
return true
}
5 changes: 5 additions & 0 deletions agent/internal/options.go
Original file line number Diff line number Diff line change
Expand Up @@ -29,6 +29,11 @@ type Options struct {
TLS bool `json:"tls"`
CertFile string `json:"cert_file"`
KeyFile string `json:"key_file"`

HTTPProxy string `json:"http_proxy"`
HTTPSProxy string `json:"https_proxy"`
FTPProxy string `json:"ftp_proxy"`
NoProxy string `json:"no_proxy"`
}

// Validate validates the state of the Options struct.
Expand Down
3 changes: 3 additions & 0 deletions docs/how-to/custom-env.txt
Original file line number Diff line number Diff line change
Expand Up @@ -32,6 +32,9 @@ experiment config. The format is a list of strings in the format
Variables are set sequentially, which affect variables that depend
on expansion of other variables.

Proxy variables set in this way will take precedent over those set using the
:ref:`agent configuration <agent-configuration>`.

``startup-hook.sh``
~~~~~~~~~~~~~~~~~~~

Expand Down
4 changes: 4 additions & 0 deletions docs/reference/cluster-config.txt
Original file line number Diff line number Diff line change
Expand Up @@ -547,3 +547,7 @@ Agent Configuration
each specified by a 0-based index, UUID, PCI bus ID, or board serial number.
The 0-based index of NVIDIA GPUs can be obtained via the ``nvidia-smi``
command.
- ``http_proxy``: The HTTP proxy address for the agent's containers.
- ``https_proxy``: The HTTPS proxy address for the agent's containers.
- ``ftp_proxy``: The FTP proxy address for the agent's containers.
- ``no_proxy``: The addresses that the container should not proxy.

0 comments on commit f11ef7f

Please sign in to comment.