Merge pull request #82 from netlify/safe-http-client

Add a SafeHttpClient method

Author: vbrown608

http/http.go

@@ -2,8 +2,11 @@ package http
 
 import (
 	"context"
-	"errors"
 	"net"
+	"net/http"
+	"net/http/httptrace"
+
+	"github.com/sirupsen/logrus"
 )
 
 var privateIPBlocks []*net.IPNet
@@ -26,27 +29,75 @@ 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 isLocalAddress(addr string) bool {
-	ip := net.ParseIP(addr)
-	return isPrivateIP(ip)
+func isPrivateIP(ip net.IP) bool {
+	return blocksContain(privateIPBlocks, ip)
 }
 
-// SafeDialContext exchanges a DialContext for a SafeDialContext that will never dial a reserved IP range
-func SafeDialContext(dialContext func(ctx context.Context, network, addr string) (net.Conn, error)) func(ctx context.Context, network, addr string) (net.Conn, error) {
-	return func(ctx context.Context, network, addr string) (net.Conn, error) {
-		if isLocalAddress(addr) {
-			return nil, errors.New("Connection to local network address denied")
-		}
+type noLocalTransport struct {
+	inner         http.RoundTripper
+	errlog        logrus.FieldLogger
+	allowedBlocks []*net.IPNet
+}
+
+func (no noLocalTransport) RoundTrip(req *http.Request) (*http.Response, error) {
+	ctx, cancel := context.WithCancel(req.Context())
+
+	ctx = httptrace.WithClientTrace(ctx, &httptrace.ClientTrace{
+		ConnectStart: func(network, addr string) {
+			host, _, err := net.SplitHostPort(addr)
+			if err != nil {
+				cancel()
+				no.errlog.WithError(err).Error("Cancelled request due to error in address parsing")
+				return
+			}
+			ip := net.ParseIP(host)
+			if ip == nil {
+				cancel()
+				no.errlog.WithError(err).Error("Cancelled request due to error in ip parsing")
+				return
+			}
+
+			if blocksContain(no.allowedBlocks, ip) {
+				return
+			}
+
+			if isPrivateIP(ip) {
+				cancel()
+				no.errlog.Error("Cancelled attempted request to ip in private range")
+				return
+			}
+		},
+	})
 
-		return dialContext(ctx, network, addr)
+	req = req.WithContext(ctx)
+	return no.inner.RoundTrip(req)
+}
+
+func SafeRoundtripper(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"),
+		allowedBlocks: allowedBlocks,
 	}
+
+	return ret
+}
+
+func SafeHTTPClient(client *http.Client, log logrus.FieldLogger, allowedBlocks ...*net.IPNet) *http.Client {
+	client.Transport = SafeRoundtripper(client.Transport, log, allowedBlocks...)
+
+	return client
 }

http/http_test.go

@@ -1,16 +1,74 @@
 package http
 
 import (
+	"net"
+	"net/http"
+	"net/http/httptest"
+	"net/url"
 	"testing"
 
+	"github.com/sirupsen/logrus"
 	"github.com/stretchr/testify/assert"
 )
 
-func TestIsLocalAddress(t *testing.T) {
-	assert.False(t, isLocalAddress("216.58.194.206"))
-	assert.True(t, isLocalAddress("127.0.0.1"))
-	assert.True(t, isLocalAddress("10.0.0.1"))
-	assert.True(t, isLocalAddress("192.168.0.1"))
-	assert.True(t, isLocalAddress("172.16.0.0"))
-	assert.True(t, isLocalAddress("169.254.169.254"))
+func TestIsPrivateIP(t *testing.T) {
+	tests := []struct {
+		ip       string
+		expected bool
+	}{
+		{"216.58.194.206", false},
+		{"127.0.0.1", true},
+		{"10.0.0.1", true},
+		{"192.168.0.1", true},
+		{"172.16.0.0", true},
+		{"169.254.169.254", true},
+	}
+
+	for _, tt := range tests {
+		ip := net.ParseIP(tt.ip)
+		assert.Equal(t, tt.expected, isPrivateIP(ip))
+	}
+}
+
+func TestSafeHTTPClient(t *testing.T) {
+	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		w.Write([]byte("Done"))
+	}))
+	defer ts.Close()
+	tsURL, err := url.Parse(ts.URL)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	client := SafeHTTPClient(&http.Client{}, logrus.New())
+
+	// It blocks the local IP.
+	_, err = client.Get(ts.URL)
+	assert.NotNil(t, err)
+
+	// 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()))
+	_, 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)
+}
+
+func popMatchingBlock(ip net.IP) *net.IPNet {
+	for i, ipNet := range privateIPBlocks {
+		if ipNet.Contains(ip) {
+			privateIPBlocks = append(privateIPBlocks[:i], privateIPBlocks[i+1:]...)
+			return ipNet
+		}
+	}
+	return nil
 }