Merge pull request #251 from DataDog/talwai/proxy

config: respect proxy_* settings in agent

agent/endpoint.go

@@ -7,6 +7,7 @@ import (
 	"sync/atomic"
 	"time"
 
+	"github.com/DataDog/datadog-trace-agent/config"
 	"github.com/DataDog/datadog-trace-agent/model"
 	"github.com/DataDog/datadog-trace-agent/statsd"
 	log "github.com/cihub/seelog"
@@ -66,6 +67,7 @@ type APIEndpoint struct {
 	apiKeys []string
 	urls    []string
 	stats   endpointStats
+	client  *http.Client
 }
 
 // NewAPIEndpoint returns a new APIEndpoint from a given config
@@ -83,11 +85,26 @@ func NewAPIEndpoint(urls, apiKeys []string) *APIEndpoint {
 	a := APIEndpoint{
 		apiKeys: apiKeys,
 		urls:    urls,
+		client:  http.DefaultClient,
 	}
 	go a.logStats()
 	return &a
 }
 
+// SetProxy updates the http client used by APIEndpoint to report via the given proxy
+func (a *APIEndpoint) SetProxy(settings *config.ProxySettings) {
+	proxyPath, err := settings.URL()
+	if err != nil {
+		log.Errorf("failed to configure proxy: %v", err)
+		return
+	}
+	a.client = &http.Client{
+		Transport: &http.Transport{
+			Proxy: http.ProxyURL(proxyPath),
+		},
+	}
+}
+
 // Write writes the bucket to the API collector endpoint.
 func (a *APIEndpoint) Write(p model.AgentPayload) (int, error) {
 	data, err := model.EncodeAgentPayload(p)
@@ -124,7 +141,7 @@ func (a *APIEndpoint) Write(p model.AgentPayload) (int, error) {
 		req.URL.RawQuery = queryParams.Encode()
 		model.SetAgentPayloadHeaders(req.Header)
 
-		resp, err := http.DefaultClient.Do(req)
+		resp, err := a.client.Do(req)
 		if err != nil {
 			log.Errorf("error when requesting to endpoint %s: %v", url, err)
 			atomic.AddInt64(&a.stats.TracesPayloadError, 1)
@@ -190,7 +207,7 @@ func (a *APIEndpoint) WriteServices(s model.ServicesMetadata) {
 		req.URL.RawQuery = queryParams.Encode()
 		model.SetServicesPayloadHeaders(req.Header)
 
-		resp, err := http.DefaultClient.Do(req)
+		resp, err := a.client.Do(req)
 		if err != nil {
 			log.Errorf("error when requesting to endpoint %s: %v", url, err)
 			atomic.AddInt64(&a.stats.ServicesPayloadError, 1)

agent/writer.go

@@ -66,6 +66,11 @@ func NewWriter(conf *config.AgentConfig) *Writer {
 
 	if conf.APIEnabled {
 		endpoint = NewAPIEndpoint(conf.APIEndpoints, conf.APIKeys)
+		if conf.Proxy != nil {
+			// we have some kind of proxy configured.
+			// make sure our http client uses it
+			endpoint.(*APIEndpoint).SetProxy(conf.Proxy)
+		}
 	} else {
 		log.Info("API interface is disabled, flushing to /dev/null instead")
 		endpoint = NullEndpoint{}

config/agent.go

@@ -59,6 +59,9 @@ type AgentConfig struct {
 	MaxMemory        float64       // MaxMemory is the threshold (bytes allocated) above which program panics and exits, to be restarted
 	MaxConnections   int           // MaxConnections is the threshold (opened TCP connections) above which program panics and exits, to be restarted
 	WatchdogInterval time.Duration // WatchdogInterval is the delay between 2 watchdog checks
+
+	// http/s proxying
+	Proxy *ProxySettings
 }
 
 // mergeEnv applies overrides from environment variables to the trace agent configuration
@@ -222,6 +225,10 @@ func NewAgentConfig(conf *File, legacyConf *File) (*AgentConfig, error) {
 		if v := m.Key("log_level").MustString(""); v != "" {
 			c.LogLevel = v
 		}
+
+		if p := getProxySettings(m); p.Host != "" {
+			c.Proxy = p
+		}
 	}
 
 APM_CONF:

config/proxy.go

@@ -0,0 +1,68 @@
+package config
+
+import (
+	"fmt"
+	"github.com/go-ini/ini"
+	"net/url"
+	"strings"
+)
+
+// mirror default behavior of the infra agent
+const defaultProxyPort = 3128
+
+// ProxySettings contains configuration for http/https proxying
+type ProxySettings struct {
+	User     string
+	Password string
+	Host     string
+	Port     int
+	Scheme   string
+}
+
+func getProxySettings(m *ini.Section) *ProxySettings {
+	p := ProxySettings{Port: defaultProxyPort, Scheme: "http"}
+
+	if v := m.Key("proxy_host").MustString(""); v != "" {
+		// accept either http://myproxy.com or myproxy.com
+		if i := strings.Index(v, "://"); i != -1 {
+			// when available, parse the scheme from the url
+			p.Scheme = v[0:i]
+			p.Host = v[i+3:]
+		} else {
+			p.Host = v
+		}
+	}
+	if v := m.Key("proxy_port").MustInt(-1); v != -1 {
+		p.Port = v
+	}
+	if v := m.Key("proxy_user").MustString(""); v != "" {
+		p.User = v
+	}
+	if v := m.Key("proxy_password").MustString(""); v != "" {
+		p.Password = v
+	}
+
+	return &p
+}
+
+// URL turns ProxySettings into an idiomatic URL struct
+func (p *ProxySettings) URL() (*url.URL, error) {
+	// construct scheme://user:pass@host:port
+	var userpass *url.Userinfo
+	if p.User != "" {
+		if p.Password != "" {
+			userpass = url.UserPassword(p.User, p.Password)
+		} else {
+			userpass = url.User(p.User)
+		}
+	}
+
+	var path string
+	if userpass != nil {
+		path = fmt.Sprintf("%s://%s@%s:%v", p.Scheme, userpass.String(), p.Host, p.Port)
+	} else {
+		path = fmt.Sprintf("%s://%s:%v", p.Scheme, p.Host, p.Port)
+	}
+
+	return url.Parse(path)
+}

config/proxy_test.go

@@ -0,0 +1,116 @@
+package config
+
+import (
+	"net/url"
+	"strings"
+
+	"github.com/stretchr/testify/assert"
+
+	"testing"
+
+	"github.com/go-ini/ini"
+)
+
+func getURL(f *ini.File) (*url.URL, error) {
+	conf := File{
+		f,
+		"some/path",
+	}
+	m, _ := conf.GetSection("Main")
+	p := getProxySettings(m)
+	return p.URL()
+}
+
+func TestGetProxySettings(t *testing.T) {
+	assert := assert.New(t)
+
+	f, _ := ini.Load([]byte("[Main]\n\nproxy_host = myhost"))
+
+	s, err := getURL(f)
+	assert.Nil(err)
+	assert.Equal("http://myhost:3128", s.String())
+
+	f, _ = ini.Load([]byte("[Main]\n\nproxy_host = http://myhost"))
+
+	s, err = getURL(f)
+	assert.Nil(err)
+	assert.Equal("http://myhost:3128", s.String())
+
+	f, _ = ini.Load([]byte("[Main]\n\nproxy_host = https://myhost"))
+
+	s, err = getURL(f)
+	assert.Nil(err)
+	assert.Equal("https://myhost:3128", s.String())
+
+	// generic user name
+	f, _ = ini.Load([]byte(strings.Join([]string{
+		"[Main]",
+		"proxy_host = https://myhost",
+		"proxy_port = 3129",
+		"proxy_user = aaditya",
+	}, "\n")))
+
+	s, err = getURL(f)
+	assert.Nil(err)
+
+	assert.Equal("https://aaditya@myhost:3129", s.String())
+
+	// special char in user name <3
+	f, _ = ini.Load([]byte(strings.Join([]string{
+		"[Main]",
+		"proxy_host = myhost",
+		"proxy_port = 3129",
+		"proxy_user = léo",
+	}, "\n")))
+
+	s, err = getURL(f)
+	assert.Nil(err)
+
+	// user is url-encoded and decodes to original string
+	assert.Equal("http://l%C3%A9o@myhost:3129", s.String())
+	assert.Equal("léo", s.User.Username())
+
+	// generic  user-pass
+	f, _ = ini.Load([]byte(strings.Join([]string{
+		"[Main]",
+		"proxy_host = myhost",
+		"proxy_port = 3129",
+		"proxy_user = aaditya",
+		"proxy_password = password_12",
+	}, "\n")))
+
+	s, err = getURL(f)
+	assert.Nil(err)
+	assert.Equal("http://aaditya:password_12@myhost:3129", s.String())
+
+	// user-pass with schemed host
+	f, _ = ini.Load([]byte(strings.Join([]string{
+		"[Main]",
+		"proxy_host = https://myhost",
+		"proxy_port = 3129",
+		"proxy_user = aaditya",
+		"proxy_password = password_12",
+	}, "\n")))
+
+	s, err = getURL(f)
+	assert.Nil(err)
+	assert.Equal("https://aaditya:password_12@myhost:3129", s.String())
+
+	// special characters in password
+	f, _ = ini.Load([]byte(strings.Join([]string{
+		"[Main]",
+		"proxy_host = https://myhost",
+		"proxy_port = 3129",
+		"proxy_user = aaditya",
+		"proxy_password = /:!?&=@éÔγλῶσσα",
+	}, "\n")))
+
+	s, err = getURL(f)
+	assert.Nil(err)
+
+	// password is url-encoded and decodes to the original string
+	assert.Equal("https://aaditya:%2F%3A%21%3F&=%40%C3%A9%C3%94%CE%B3%CE%BB%E1%BF%B6%CF%83%CF%83%CE%B1@myhost:3129", s.String())
+
+	pass, _ := s.User.Password()
+	assert.Equal("/:!?&=@éÔγλῶσσα", pass)
+}