dht mode toggling (modulo dynamic switching)

dht.go

@@ -41,6 +41,13 @@ var logger = logging.Logger("dht")
 // collect members of the routing table.
 const NumBootstrapQueries = 5
 
+type DHTMode int
+
+const (
+	ModeServer = DHTMode(1)
+	ModeClient = DHTMode(2)
+)
+
 // IpfsDHT is an implementation of Kademlia with S/Kademlia modifications.
 // It is used to implement the base Routing module.
 type IpfsDHT struct {
@@ -66,6 +73,9 @@ type IpfsDHT struct {
 	plk sync.Mutex
 
 	protocols []protocol.ID // DHT protocols
+
+	mode   DHTMode
+	modeLk sync.Mutex
 }
 
 // Assert that IPFS assumptions about interfaces aren't broken. These aren't a
@@ -97,8 +107,10 @@ func New(ctx context.Context, h host.Host, options ...opts.Option) (*IpfsDHT, er
 
 	dht.proc.AddChild(dht.providers.Process())
 	dht.Validator = cfg.Validator
+	dht.mode = ModeClient
 
 	if !cfg.Client {
+		dht.mode = ModeServer
 		for _, p := range cfg.Protocols {
 			h.SetStreamHandler(p, dht.handleNewStream)
 		}
@@ -365,6 +377,61 @@ func (dht *IpfsDHT) betterPeersToQuery(pmes *pb.Message, p peer.ID, count int) [
 	return filtered
 }
 
+func (dht *IpfsDHT) SetMode(m DHTMode) error {
+	dht.modeLk.Lock()
+	defer dht.modeLk.Unlock()
+
+	if m == dht.mode {
+		return nil
+	}
+
+	switch m {
+	case ModeServer:
+		return dht.moveToServerMode()
+	case ModeClient:
+		return dht.moveToClientMode()
+	default:
+		return fmt.Errorf("unrecognized dht mode: %d", m)
+	}
+}
+
+func (dht *IpfsDHT) moveToServerMode() error {
+	dht.mode = ModeServer
+	for _, p := range dht.protocols {
+		dht.host.SetStreamHandler(p, dht.handleNewStream)
+	}
+	return nil
+}
+
+func (dht *IpfsDHT) moveToClientMode() error {
+	dht.mode = ModeClient
+	for _, p := range dht.protocols {
+		dht.host.RemoveStreamHandler(p)
+	}
+
+	pset := make(map[protocol.ID]bool)
+	for _, p := range dht.protocols {
+		pset[p] = true
+	}
+
+	for _, c := range dht.host.Network().Conns() {
+		for _, s := range c.GetStreams() {
+			if pset[s.Protocol()] {
+				if s.Stat().Direction == network.DirInbound {
+					s.Reset()
+				}
+			}
+		}
+	}
+	return nil
+}
+
+func (dht *IpfsDHT) getMode() DHTMode {
+	dht.modeLk.Lock()
+	defer dht.modeLk.Unlock()
+	return dht.mode
+}
+
 // Context return dht's context
 func (dht *IpfsDHT) Context() context.Context {
 	return dht.ctx

dht_net.go

@@ -14,6 +14,7 @@ import (
 
 	"github.com/libp2p/go-libp2p-kad-dht/metrics"
 	pb "github.com/libp2p/go-libp2p-kad-dht/pb"
+	msmux "github.com/multiformats/go-multistream"
 
 	ggio "github.com/gogo/protobuf/io"
 
@@ -79,6 +80,11 @@ func (dht *IpfsDHT) handleNewMessage(s network.Stream) bool {
 	defer timer.Stop()
 
 	for {
+		if dht.getMode() != ModeServer {
+			logger.Errorf("responding to dht message while not in server mode")
+			return false
+		}
+
 		var req pb.Message
 		switch err := r.ReadMsg(&req); err {
 		case io.EOF:
@@ -153,6 +159,9 @@ func (dht *IpfsDHT) sendRequest(ctx context.Context, p peer.ID, pmes *pb.Message
 
 	ms, err := dht.messageSenderForPeer(ctx, p)
 	if err != nil {
+		if err == msmux.ErrNotSupported {
+			dht.RoutingTable().Remove(p)
+		}
 		stats.Record(ctx, metrics.SentRequestErrors.M(1))
 		return nil, err
 	}
@@ -161,6 +170,9 @@ func (dht *IpfsDHT) sendRequest(ctx context.Context, p peer.ID, pmes *pb.Message
 
 	rpmes, err := ms.SendRequest(ctx, pmes)
 	if err != nil {
+		if err == msmux.ErrNotSupported {
+			dht.RoutingTable().Remove(p)
+		}
 		stats.Record(ctx, metrics.SentRequestErrors.M(1))
 		return nil, err
 	}
@@ -187,11 +199,17 @@ func (dht *IpfsDHT) sendMessage(ctx context.Context, p peer.ID, pmes *pb.Message
 
 	ms, err := dht.messageSenderForPeer(ctx, p)
 	if err != nil {
+		if err == msmux.ErrNotSupported {
+			dht.RoutingTable().Remove(p)
+		}
 		stats.Record(ctx, metrics.SentMessageErrors.M(1))
 		return err
 	}
 
 	if err := ms.SendMessage(ctx, pmes); err != nil {
+		if err == msmux.ErrNotSupported {
+			dht.RoutingTable().Remove(p)
+		}
 		stats.Record(ctx, metrics.SentMessageErrors.M(1))
 		return err
 	}