Allow whitelisting local IPs for development

Author: vbrown608

http/http.go

@@ -29,18 +29,23 @@ func init() {
 	}
 }
 
-func isPrivateIP(ip net.IP) bool {
-	for _, block := range privateIPBlocks {
+func blocksContain(blocks []*net.IPNet, ip net.IP) bool {
+	for _, block := range blocks {
 		if block.Contains(ip) {
 			return true
 		}
 	}
 	return false
 }
 
+func isPrivateIP(ip net.IP) bool {
+	return blocksContain(privateIPBlocks, ip)
+}
+
 type noLocalTransport struct {
-	inner  http.RoundTripper
-	errlog logrus.FieldLogger
+	inner         http.RoundTripper
+	errlog        logrus.FieldLogger
+	allowedBlocks []*net.IPNet
 }
 
 func (no noLocalTransport) RoundTrip(req *http.Request) (*http.Response, error) {
@@ -61,34 +66,38 @@ func (no noLocalTransport) RoundTrip(req *http.Request) (*http.Response, error)
 				return
 			}
 
+			if blocksContain(no.allowedBlocks, ip) {
+				return
+			}
+
 			if isPrivateIP(ip) {
 				cancel()
 				no.errlog.Error("Cancelled attempted request to ip in private range")
 				return
 			}
-
 		},
 	})
 
 	req = req.WithContext(ctx)
 	return no.inner.RoundTrip(req)
 }
 
-func SafeRountripper(trans http.RoundTripper, log logrus.FieldLogger) http.RoundTripper {
+func SafeRountripper(trans http.RoundTripper, log logrus.FieldLogger, allowedBlocks ...*net.IPNet) http.RoundTripper {
 	if trans == nil {
 		trans = http.DefaultTransport
 	}
 
 	ret := &noLocalTransport{
-		inner:  trans,
-		errlog: log.WithField("transport", "local_blocker"),
+		inner:         trans,
+		errlog:        log.WithField("transport", "local_blocker"),
+		allowedBlocks: allowedBlocks,
 	}
 
 	return ret
 }
 
-func SafeHTTPClient(client *http.Client, log logrus.FieldLogger) *http.Client {
-	client.Transport = SafeRountripper(client.Transport, log)
+func SafeHTTPClient(client *http.Client, log logrus.FieldLogger, allowedBlocks ...*net.IPNet) *http.Client {
+	client.Transport = SafeRountripper(client.Transport, log, allowedBlocks...)
 
 	return client
 }

http/http_test.go

@@ -42,21 +42,23 @@ func TestSafeHTTPClient(t *testing.T) {
 
 	client := SafeHTTPClient(&http.Client{}, logrus.New())
 
-	// It blocks the local IP
+	// It blocks the local IP.
 	_, err = client.Get(ts.URL)
 	assert.NotNil(t, err)
 
-	// It blocks localhost
+	// It blocks localhost.
 	_, err = client.Get("http://localhost:" + tsURL.Port())
 	assert.NotNil(t, err)
 
 	// It succeeds when the local IP range used by the testserver is removed from
 	// the blacklist.
 	ipNet := popMatchingBlock(net.ParseIP(tsURL.Hostname()))
-	defer func() {
-		privateIPBlocks = append(privateIPBlocks, ipNet)
-	}()
+	_, err = client.Get(ts.URL)
+	assert.Nil(t, err)
+	privateIPBlocks = append(privateIPBlocks, ipNet)
 
+	// It allows whitelisting for local development.
+	client = SafeHTTPClient(&http.Client{}, logrus.New(), ipNet)
 	_, err = client.Get(ts.URL)
 	assert.Nil(t, err)
 }