Add SafeDial to block connections (#270)

* Add SafeDial to block connections

Closes #255

* Add SafeTransport for convenience

Author: mrdg

http/http.go

@@ -2,6 +2,8 @@ package http
 
 import (
 	"context"
+	"errors"
+	"fmt"
 	"net"
 	"net/http"
 	"net/http/httptrace"
@@ -85,6 +87,8 @@ func (no noLocalTransport) RoundTrip(req *http.Request) (*http.Response, error)
 	return no.inner.RoundTrip(req)
 }
 
+// SafeRoundtripper blocks requests to private ip ranges
+// Deprecated: use SafeTransport instead
 func SafeRoundtripper(trans http.RoundTripper, log logrus.FieldLogger, allowedBlocks ...*net.IPNet) http.RoundTripper {
 	if trans == nil {
 		trans = http.DefaultTransport
@@ -99,8 +103,65 @@ func SafeRoundtripper(trans http.RoundTripper, log logrus.FieldLogger, allowedBl
 	return ret
 }
 
+// SafeHTTPClient blocks requests to private ip ranges
+// Deprecated: use SafeTransport instead
 func SafeHTTPClient(client *http.Client, log logrus.FieldLogger, allowedBlocks ...*net.IPNet) *http.Client {
 	client.Transport = SafeRoundtripper(client.Transport, log, allowedBlocks...)
 
 	return client
 }
+
+// SafeTransport blocks requests to private ip ranges
+func SafeTransport(allowedBlocks ...*net.IPNet) *http.Transport {
+	return &http.Transport{
+		DialContext: SafeDial(&net.Dialer{}, allowedBlocks...),
+	}
+}
+
+type DialFunc func(ctx context.Context, network, address string) (net.Conn, error)
+
+// SafeDial wraps a *net.Dialer and restricts connections to private ip ranges.
+func SafeDial(dialer *net.Dialer, allowedBlocks ...*net.IPNet) DialFunc {
+	d := &safeDialer{dialer: dialer, allowedBlocks: allowedBlocks}
+	return d.dialContext
+}
+
+type safeDialer struct {
+	allowedBlocks []*net.IPNet
+	dialer        *net.Dialer
+}
+
+func (d *safeDialer) dialContext(ctx context.Context, network, address string) (net.Conn, error) {
+	conn, err := d.dialer.DialContext(ctx, network, address)
+	if err != nil {
+		return nil, err
+	}
+	addr := conn.RemoteAddr().String()
+
+	host, _, err := net.SplitHostPort(addr)
+	if err != nil {
+		_ = conn.Close()
+		return nil, fmt.Errorf("safe dialer: invalid address: %w", err)
+	}
+
+	ip := net.ParseIP(host)
+	if ip == nil {
+		_ = conn.Close()
+		return nil, fmt.Errorf("safe dialer: invalid ip: %v", host)
+	}
+
+	for _, block := range d.allowedBlocks {
+		if block.Contains(ip) {
+			return conn, nil
+		}
+	}
+
+	for _, block := range privateIPBlocks {
+		if block.Contains(ip) {
+			_ = conn.Close()
+			return nil, errors.New("safe dialer: private ip not allowed")
+		}
+	}
+
+	return conn, nil
+}

http/http_test.go

@@ -34,7 +34,32 @@ func TestIsPrivateIP(t *testing.T) {
 	}
 }
 
-func TestSafeHTTPClient(t *testing.T) {
+func TestBlockList(t *testing.T) {
+	t.Run("safe http client", func(t *testing.T) {
+		client := SafeHTTPClient(&http.Client{}, logrus.New())
+		testBlockList(t, client)
+	})
+	t.Run("safe dial", func(t *testing.T) {
+		client := &http.Client{Transport: SafeTransport()}
+		testBlockList(t, client)
+	})
+}
+
+func TestAllowList(t *testing.T) {
+	_, local, err := net.ParseCIDR("127.0.0.1/8")
+	require.NoError(t, err)
+
+	t.Run("safe http client", func(t *testing.T) {
+		client := SafeHTTPClient(&http.Client{}, logrus.New(), local)
+		testAllowList(t, client)
+	})
+	t.Run("safe dial", func(t *testing.T) {
+		client := &http.Client{Transport: SafeTransport(local)}
+		testAllowList(t, client)
+	})
+}
+
+func testBlockList(t *testing.T, client *http.Client) {
 	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 		w.Write([]byte("Done"))
 	}))
@@ -44,8 +69,6 @@ func TestSafeHTTPClient(t *testing.T) {
 		t.Fatal(err)
 	}
 
-	client := SafeHTTPClient(&http.Client{}, logrus.New())
-
 	// It allows accessing non-local addresses
 	_, err = client.Get("https://google.com")
 	require.Nil(t, err)
@@ -64,26 +87,13 @@ func TestSafeHTTPClient(t *testing.T) {
 		assert.Nil(t, res)
 		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
+func testAllowList(t *testing.T, client *http.Client) {
+	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		w.Write([]byte("Done"))
+	}))
+	defer ts.Close()
+	_, err := client.Get(ts.URL)
+	assert.NoError(t, err)
 }