add support for persistence to datastore

p2p/net/conngater/conngater.go

@@ -12,6 +12,9 @@ import (
 	ma "github.com/multiformats/go-multiaddr"
 	manet "github.com/multiformats/go-multiaddr/net"
 
+	"github.com/ipfs/go-datastore"
+	"github.com/ipfs/go-datastore/namespace"
+	"github.com/ipfs/go-datastore/query"
 	logging "github.com/ipfs/go-log"
 )
 
@@ -21,13 +24,91 @@ type BasicConnectionGater struct {
 	blockedPeers   map[peer.ID]struct{}
 	blockedAddrs   map[string]struct{}
 	blockedSubnets map[string]*net.IPNet
+
+	ds datastore.Datastore
 }
 
 var log = logging.Logger("net/conngater")
 
-func NewBasicConnectionGater() *BasicConnectionGater {
-	// XXX
-	return nil
+const ns = "/libp2p/net/conngater"
+const keyPeer = "/peer/"
+const keyAddr = "/addr/"
+const keySubnet = "/subnet/"
+
+// NewBasicConnectionGater creates a new connection gater.
+// The ds argument is an (optional, can be nil) datastore to persist the connection gater
+// filters
+func NewBasicConnectionGater(ds datastore.Datastore) *BasicConnectionGater {
+	cg := &BasicConnectionGater{
+		blockedPeers:   make(map[peer.ID]struct{}),
+		blockedAddrs:   make(map[string]struct{}),
+		blockedSubnets: make(map[string]*net.IPNet),
+	}
+
+	if ds != nil {
+		cg.ds = namespace.Wrap(ds, datastore.NewKey(ns))
+		cg.loadRules()
+	}
+
+	return cg
+}
+
+func (cg *BasicConnectionGater) loadRules() {
+	// load blocked peers
+	res, err := cg.ds.Query(query.Query{Prefix: keyPeer})
+	if err != nil {
+		log.Errorf("error querying datastore for blocked peers: %s", err)
+		return
+	}
+
+	for r := range res.Next() {
+		if r.Error != nil {
+			log.Errorf("query result error: %s", r.Error)
+			return
+		}
+
+		p := peer.ID(r.Entry.Value)
+		cg.blockedPeers[p] = struct{}{}
+	}
+
+	// load blocked addrs
+	res, err = cg.ds.Query(query.Query{Prefix: keyAddr})
+	if err != nil {
+		log.Errorf("error querying datastore for blocked addrs: %s", err)
+		return
+	}
+
+	for r := range res.Next() {
+		if r.Error != nil {
+			log.Errorf("query result error: %s", r.Error)
+			return
+		}
+
+		ip := net.IP(r.Entry.Value)
+		cg.blockedAddrs[ip.String()] = struct{}{}
+	}
+
+	// load blocked subnets
+	res, err = cg.ds.Query(query.Query{Prefix: keySubnet})
+	if err != nil {
+		log.Errorf("error querying datastore for blocked subnets: %s", err)
+		return
+	}
+
+	for r := range res.Next() {
+		if r.Error != nil {
+			log.Errorf("query result error: %s", r.Error)
+			return
+		}
+
+		ipnetStr := string(r.Entry.Value)
+		_, ipnet, err := net.ParseCIDR(ipnetStr)
+		if err != nil {
+			log.Errorf("error parsing CIDR subnet: %s", err)
+			return
+		}
+		cg.blockedSubnets[ipnetStr] = ipnet
+	}
 }
 
 // BlockPeer adds a peer to the set of blocked peers
@@ -36,6 +117,13 @@ func (cg *BasicConnectionGater) BlockPeer(p peer.ID) {
 	defer cg.Unlock()
 
 	cg.blockedPeers[p] = struct{}{}
+
+	if cg.ds != nil {
+		err := cg.ds.Put(datastore.NewKey(keyPeer+p.String()), []byte(p))
+		if err != nil {
+			log.Errorf("error writing blocked peer to datastore: %s", err)
+		}
+	}
 }
 
 // UnblockPeer removes a peer from the set of blocked peers
@@ -44,6 +132,13 @@ func (cg *BasicConnectionGater) UnblockPeer(p peer.ID) {
 	defer cg.Unlock()
 
 	delete(cg.blockedPeers, p)
+
+	if cg.ds != nil {
+		err := cg.ds.Delete(datastore.NewKey(keyPeer + p.String()))
+		if err != nil {
+			log.Errorf("error deleting blocked peer from datastore: %s", err)
+		}
+	}
 }
 
 // BlockAddr adds an IP address to the set of blocked addresses
@@ -52,6 +147,13 @@ func (cg *BasicConnectionGater) BlockAddr(ip net.IP) {
 	defer cg.Unlock()
 
 	cg.blockedAddrs[ip.String()] = struct{}{}
+
+	if cg.ds != nil {
+		err := cg.ds.Put(datastore.NewKey(keyAddr+ip.String()), []byte(ip))
+		if err != nil {
+			log.Errorf("error writing blocked addr to datastore: %s", err)
+		}
+	}
 }
 
 // UnblockAddr removes an IP address from the set of blocked addresses
@@ -60,6 +162,13 @@ func (cg *BasicConnectionGater) UnblockAddr(ip net.IP) {
 	defer cg.Unlock()
 
 	delete(cg.blockedAddrs, ip.String())
+
+	if cg.ds != nil {
+		err := cg.ds.Delete(datastore.NewKey(keyAddr + ip.String()))
+		if err != nil {
+			log.Errorf("error deleting blocked addr from datastore: %s", err)
+		}
+	}
 }
 
 // BlockSubnet adds an IP subnet to the set of blocked addresses
@@ -68,6 +177,13 @@ func (cg *BasicConnectionGater) BlockSubnet(ipnet *net.IPNet) {
 	defer cg.Unlock()
 
 	cg.blockedSubnets[ipnet.String()] = ipnet
+
+	if cg.ds != nil {
+		err := cg.ds.Put(datastore.NewKey(keySubnet+ipnet.String()), []byte(ipnet.String()))
+		if err != nil {
+			log.Errorf("error writing blocked addr to datastore: %s", err)
+		}
+	}
 }
 
 // UnblockSubnet removes an IP address from the set of blocked addresses
@@ -76,6 +192,13 @@ func (cg *BasicConnectionGater) UnblockSubnet(ipnet *net.IPNet) {
 	defer cg.Unlock()
 
 	delete(cg.blockedSubnets, ipnet.String())
+
+	if cg.ds != nil {
+		err := cg.ds.Delete(datastore.NewKey(keySubnet + ipnet.String()))
+		if err != nil {
+			log.Errorf("error deleting blocked subnet from datastore: %s", err)
+		}
+	}
 }
 
 // ConnectionGater interface