Merge pull request #22 from stripe/xieyuxi/custom-transport-proxy-addr

Configurable http and https proxy addrs

https.go

@@ -118,7 +118,7 @@ func (proxy *ProxyHttpServer) handleHttps(w http.ResponseWriter, r *http.Request
 			host += ":80"
 		}
 
-		httpsProxy, err := httpsProxyFromEnv(r.URL)
+		httpsProxy, err := httpsProxyAddr(r.URL, proxy.HttpsProxyAddr)
 		if err != nil {
 			ctx.Warnf("Error configuring HTTPS proxy err=%q url=%q", err, r.URL.String())
 		}
@@ -398,14 +398,20 @@ func copyAndClose(ctx *ProxyCtx, dst, src *net.TCPConn) {
 	src.CloseRead()
 }
 
-func dialerFromEnv(proxy *ProxyHttpServer) func(network, addr string) (net.Conn, error) {
-	https_proxy := os.Getenv("HTTPS_PROXY")
+// dialerFromProxy gets the HttpsProxyAddr from proxy to create a dialer.
+// When the HttpsProxyAddr from proxy is empty, use the HTTPS_PROXY, https_proxy from environment variables.
+func dialerFromProxy(proxy *ProxyHttpServer) func(network, addr string) (net.Conn, error) {
+	https_proxy := proxy.HttpsProxyAddr
 	if https_proxy == "" {
-		https_proxy = os.Getenv("https_proxy")
-	}
-	if https_proxy == "" {
-		return nil
+		https_proxy = os.Getenv("HTTPS_PROXY")
+		if https_proxy == "" {
+			https_proxy = os.Getenv("https_proxy")
+		}
+		if https_proxy == "" {
+			return nil
+		}
 	}
+
 	return proxy.NewConnectDialToProxy(https_proxy)
 }
 
@@ -559,10 +565,17 @@ func (proxy *ProxyHttpServer) connectDialProxyWithContext(ctx *ProxyCtx, proxyHo
 	return c, nil
 }
 
-// httpsProxyFromEnv allows goproxy to respect no_proxy env vars
+// httpsProxyAddr function uses the address in httpsProxy parameter. 
+// When the httpProxyAddr parameter is empty, uses the HTTPS_PROXY, https_proxy from environment variables.
+// httpsProxyAddr function allows goproxy to respect no_proxy env vars
 // https://github.com/stripe/goproxy/pull/5
-func httpsProxyFromEnv(reqURL *url.URL) (string, error) {
+func httpsProxyAddr(reqURL *url.URL, httpsProxy string) (string, error) {
 	cfg := httpproxy.FromEnvironment()
+
+	if httpsProxy != "" {
+		cfg.HTTPSProxy = httpsProxy
+	}
+
 	// We only use this codepath for HTTPS CONNECT proxies so we shouldn't
 	// return anything from HTTPProxy
 	cfg.HTTPProxy = ""

https_test.go

@@ -7,24 +7,26 @@ import (
 )
 
 var proxytests = map[string]struct {
-	noProxy     string
-	httpsProxy  string
-	url         string
-	expectProxy string
+	noProxy          string
+	envHttpsProxy    string
+	customHttpsProxy string
+	url              string
+	expectProxy      string
 }{
-	"do not proxy without a proxy configured":   {"", "", "https://foo.bar/baz", ""},
-	"proxy with a proxy configured":             {"", "daproxy", "https://foo.bar/baz", "http://daproxy:http"},
-	"proxy without a scheme":                    {"", "daproxy", "//foo.bar/baz", "http://daproxy:http"},
-	"proxy with a proxy configured with a port": {"", "http://daproxy:123", "https://foo.bar/baz", "http://daproxy:123"},
-	"proxy with an https proxy configured":      {"", "https://daproxy", "https://foo.bar/baz", "https://daproxy:https"},
-	"proxy with a non-matching no_proxy":        {"other.bar", "daproxy", "https://foo.bar/baz", "http://daproxy:http"},
-	"do not proxy with a full no_proxy match":   {"foo.bar", "daproxy", "https://foo.bar/baz", ""},
-	"do not proxy with a suffix no_proxy match": {".bar", "daproxy", "https://foo.bar/baz", ""},
+	"do not proxy without a proxy configured":   {"", "", "", "https://foo.bar/baz", ""},
+	"proxy with a proxy configured":             {"", "daproxy", "", "https://foo.bar/baz", "http://daproxy:http"},
+	"proxy without a scheme":                    {"", "daproxy", "", "//foo.bar/baz", "http://daproxy:http"},
+	"proxy with a proxy configured with a port": {"", "http://daproxy:123", "", "https://foo.bar/baz", "http://daproxy:123"},
+	"proxy with an https proxy configured":      {"", "https://daproxy", "", "https://foo.bar/baz", "https://daproxy:https"},
+	"proxy with a non-matching no_proxy":        {"other.bar", "daproxy", "", "https://foo.bar/baz", "http://daproxy:http"},
+	"do not proxy with a full no_proxy match":   {"foo.bar", "daproxy", "", "https://foo.bar/baz", ""},
+	"do not proxy with a suffix no_proxy match": {".bar", "daproxy", "", "https://foo.bar/baz", ""},
+	"proxy with an custom https proxy":          {"", "https://daproxy", "https://customproxy", "https://foo.bar/baz", "https://customproxy:https"},
 }
 
 var envKeys = []string{"no_proxy", "http_proxy", "https_proxy", "NO_PROXY", "HTTP_PROXY", "HTTPS_PROXY"}
 
-func TestHttpsProxyFromEnv(t *testing.T) {
+func TestHttpsProxyAddr(t *testing.T) {
 	for _, k := range envKeys {
 		v, ok := os.LookupEnv(k)
 		if ok {
@@ -43,14 +45,14 @@ func TestHttpsProxyFromEnv(t *testing.T) {
 	for name, spec := range proxytests {
 		t.Run(name, func(t *testing.T) {
 			os.Setenv("no_proxy", spec.noProxy)
-			os.Setenv("https_proxy", spec.httpsProxy)
+			os.Setenv("https_proxy", spec.envHttpsProxy)
 
 			url, err := url.Parse(spec.url)
 			if err != nil {
 				t.Fatalf("bad test input URL %s: %v", spec.url, err)
 			}
 
-			actual, err := httpsProxyFromEnv(url)
+			actual, err := httpsProxyAddr(url, spec.customHttpsProxy)
 			if err != nil {
 				t.Fatalf("unexpected error parsing proxy from env: %#v", err)
 			}

proxy.go

@@ -6,9 +6,12 @@ import (
 	"log"
 	"net"
 	"net/http"
+	"net/url"
 	"os"
 	"regexp"
 	"sync/atomic"
+
+	"golang.org/x/net/http/httpproxy"
 )
 
 // The basic proxy type. Implements http.Handler.
@@ -54,6 +57,10 @@ type ProxyHttpServer struct {
 	// ConnectRespHandler allows users to mutate the response to the CONNECT request before it
 	// is returned to the client.
 	ConnectRespHandler func(ctx *ProxyCtx, resp *http.Response) error
+
+	// HTTP and HTTPS proxy addresses
+	HttpProxyAddr  string
+	HttpsProxyAddr string
 }
 
 var hasPort = regexp.MustCompile(`:\d+$`)
@@ -179,8 +186,40 @@ func (proxy *ProxyHttpServer) ServeHTTP(w http.ResponseWriter, r *http.Request)
 	}
 }
 
+type options struct {
+	httpProxyAddr  string
+	httpsProxyAddr string
+}
+type fnOption func(*options)
+
+func (fn fnOption) apply(opts *options) { fn(opts) }
+
+type ProxyHttpServerOptions interface {
+	apply(*options)
+}
+
+func WithHttpProxyAddr(httpProxyAddr string) ProxyHttpServerOptions {
+	return fnOption(func(opts *options) {
+		opts.httpProxyAddr = httpProxyAddr
+	})
+}
+
+func WithHttpsProxyAddr(httpsProxyAddr string) ProxyHttpServerOptions {
+	return fnOption(func(opts *options) {
+		opts.httpsProxyAddr = httpsProxyAddr
+	})
+}
+
 // NewProxyHttpServer creates and returns a proxy server, logging to stderr by default
-func NewProxyHttpServer() *ProxyHttpServer {
+func NewProxyHttpServer(opts ...ProxyHttpServerOptions) *ProxyHttpServer {
+	appliedOpts := &options{
+		httpProxyAddr: "",
+		httpsProxyAddr:  "",
+	}
+	for _, opt := range opts {
+		opt.apply(appliedOpts)
+	}
+
 	proxy := ProxyHttpServer{
 		Logger:        log.New(os.Stderr, "", log.LstdFlags),
 		reqHandlers:   []ReqHandler{},
@@ -191,7 +230,27 @@ func NewProxyHttpServer() *ProxyHttpServer {
 		}),
 		Tr: &http.Transport{TLSClientConfig: tlsClientSkipVerify, Proxy: http.ProxyFromEnvironment},
 	}
-	proxy.ConnectDial = dialerFromEnv(&proxy)
+
+	// httpProxyCfg holds configuration for HTTP proxy settings. See FromEnvironment for details.
+	httpProxyCfg := httpproxy.FromEnvironment()
+
+	if appliedOpts.httpProxyAddr != "" {
+		proxy.HttpProxyAddr = appliedOpts.httpProxyAddr
+		httpProxyCfg.HTTPProxy = appliedOpts.httpProxyAddr
+	}
+
+	if appliedOpts.httpsProxyAddr != "" {
+		proxy.HttpsProxyAddr = appliedOpts.httpsProxyAddr
+		httpProxyCfg.HTTPSProxy = appliedOpts.httpsProxyAddr
+	}
+
+	proxy.ConnectDial = dialerFromProxy(&proxy)
+
+	if appliedOpts.httpProxyAddr != "" || appliedOpts.httpsProxyAddr != "" {
+		proxy.Tr.Proxy = func(req *http.Request) (*url.URL, error) {
+			return httpProxyCfg.ProxyFunc()(req.URL)
+		}
+	}
 
 	return &proxy
 }

proxy_test.go

@@ -703,7 +703,57 @@ func TestGoproxyThroughProxy(t *testing.T) {
 	if r := string(getOrFail(https.URL+"/bobo", client, t)); r != "bobo bobo" {
 		t.Error("Expected bobo doubled twice, got", r)
 	}
+}
+
+func TestHttpProxyAddrsFromEnv(t *testing.T) {
+	proxy := goproxy.NewProxyHttpServer()
+	doubleString := func(resp *http.Response, ctx *goproxy.ProxyCtx) *http.Response {
+		b, err := ioutil.ReadAll(resp.Body)
+		panicOnErr(err, "readAll resp")
+		resp.Body = ioutil.NopCloser(bytes.NewBufferString(string(b) + " " + string(b)))
+		return resp
+	}
+	proxy.OnRequest().HandleConnect(goproxy.AlwaysMitm)
+	proxy.OnResponse().DoFunc(doubleString)
+
+	_, l := oneShotProxy(proxy, t)
+	defer l.Close()
+
+	os.Setenv("http_proxy", l.URL)
+	os.Setenv("https_proxy", l.URL)
+	proxy2 := goproxy.NewProxyHttpServer()
+
+	client, l2 := oneShotProxy(proxy2, t)
+	defer l2.Close()
+	if r := string(getOrFail(https.URL+"/bobo", client, t)); r != "bobo bobo" {
+		t.Error("Expected bobo doubled twice, got", r)
+	}
 
+	os.Unsetenv("http_proxy")
+	os.Unsetenv("https_proxy")
+}
+
+func TestCustomHttpProxyAddrs(t *testing.T) {
+	proxy := goproxy.NewProxyHttpServer()
+	doubleString := func(resp *http.Response, ctx *goproxy.ProxyCtx) *http.Response {
+		b, err := ioutil.ReadAll(resp.Body)
+		panicOnErr(err, "readAll resp")
+		resp.Body = ioutil.NopCloser(bytes.NewBufferString(string(b) + " " + string(b)))
+		return resp
+	}
+	proxy.OnRequest().HandleConnect(goproxy.AlwaysMitm)
+	proxy.OnResponse().DoFunc(doubleString)
+
+	_, l := oneShotProxy(proxy, t)
+	defer l.Close()
+
+	proxy2 := goproxy.NewProxyHttpServer(goproxy.WithHttpProxyAddr(l.URL), goproxy.WithHttpsProxyAddr(l.URL))
+
+	client, l2 := oneShotProxy(proxy2, t)
+	defer l2.Close()
+	if r := string(getOrFail(https.URL+"/bobo", client, t)); r != "bobo bobo" {
+		t.Error("Expected bobo doubled twice, got", r)
+	}
 }
 
 func TestGoproxyHijackConnect(t *testing.T) {