From 89c1c82865e14d228ef9bf63fdb910d226fca749 Mon Sep 17 00:00:00 2001
From: Jeromy <jeromyj@gmail.com>
Date: Tue, 16 Jun 2015 11:06:04 -0700
Subject: [PATCH] add in basic address dial filtering

License: MIT
Signed-off-by: Jeromy <jeromyj@gmail.com>
---
 net/swarm/swarm.go      |  8 ++++++++
 net/swarm/swarm_dial.go | 28 ++++++++++++++++++++++++++++
 2 files changed, 36 insertions(+)

diff --git a/net/swarm/swarm.go b/net/swarm/swarm.go
index 15a5bbddb..489d6fec3 100644
--- a/net/swarm/swarm.go
+++ b/net/swarm/swarm.go
@@ -4,6 +4,7 @@ package swarm
 
 import (
 	"fmt"
+	"net"
 	"sync"
 	"time"
 
@@ -50,6 +51,9 @@ type Swarm struct {
 	notifmu sync.RWMutex
 	notifs  map[inet.Notifiee]ps.Notifiee
 
+	// filters for addresses that shouldnt be dialed
+	filters []*net.IPNet
+
 	cg  ctxgroup.ContextGroup
 	bwc metrics.Reporter
 }
@@ -84,6 +88,10 @@ func (s *Swarm) teardown() error {
 	return s.swarm.Close()
 }
 
+func (s *Swarm) AddDialFilter(f *net.IPNet) {
+	s.filters = append(s.filters, f)
+}
+
 // CtxGroup returns the Context Group of the swarm
 func filterAddrs(listenAddrs []ma.Multiaddr) ([]ma.Multiaddr, error) {
 	if len(listenAddrs) > 0 {
diff --git a/net/swarm/swarm_dial.go b/net/swarm/swarm_dial.go
index aacee6ec4..534f916d1 100644
--- a/net/swarm/swarm_dial.go
+++ b/net/swarm/swarm_dial.go
@@ -303,6 +303,8 @@ func (s *Swarm) dial(ctx context.Context, p peer.ID) (*Conn, error) {
 	ila, _ := s.InterfaceListenAddresses()
 	remoteAddrs = addrutil.Subtract(remoteAddrs, ila)
 	remoteAddrs = addrutil.Subtract(remoteAddrs, s.peers.Addrs(s.local))
+	remoteAddrs = s.filterAddrs(remoteAddrs)
+
 	log.Debugf("%s swarm dialing %s -- local:%s remote:%s", s.local, p, s.ListenAddresses(), remoteAddrs)
 	if len(remoteAddrs) == 0 {
 		err := errors.New("peer has no addresses")
@@ -454,6 +456,32 @@ func (s *Swarm) dialAddr(ctx context.Context, d *conn.Dialer, p peer.ID, addr ma
 	return connC, nil
 }
 
+func (s *Swarm) filterAddrs(addrs []ma.Multiaddr) []ma.Multiaddr {
+	var out []ma.Multiaddr
+	for _, a := range addrs {
+		if !s.addrBlocked(a) {
+			out = append(out, a)
+		}
+	}
+	return out
+}
+
+func (s *Swarm) addrBlocked(a ma.Multiaddr) bool {
+	_, addr, err := manet.DialArgs(a)
+	if err != nil {
+		// if we cant parse it, its probably not blocked
+		return false
+	}
+
+	ip := net.ParseIP(addr)
+	for _, f := range s.filters {
+		if f.Contains(ip) {
+			return true
+		}
+	}
+	return false
+}
+
 // dialConnSetup is the setup logic for a connection from the dial side. it
 // needs to add the Conn to the StreamSwarm, then run newConnSetup
 func dialConnSetup(ctx context.Context, s *Swarm, connC conn.Conn) (*Conn, error) {