Merge pull request #547 from jbenet/fix/mux-privatize

fix: privatize Mux fields

p2p/host/basic/basic_host.go

@@ -16,7 +16,7 @@ var log = eventlog.Logger("p2p/host/basic")
 
 type BasicHost struct {
 	network inet.Network
-	mux     protocol.Mux
+	mux     *protocol.Mux
 	ids     *identify.IDService
 	relay   *relay.RelayService
 }
@@ -25,7 +25,7 @@ type BasicHost struct {
 func New(net inet.Network) *BasicHost {
 	h := &BasicHost{
 		network: net,
-		mux:     protocol.Mux{Handlers: protocol.StreamHandlerMap{}},
+		mux:     protocol.NewMux(),
 	}
 
 	// setup host services
@@ -65,7 +65,7 @@ func (h *BasicHost) Network() inet.Network {
 
 // Mux returns the Mux multiplexing incoming streams to protocol handlers
 func (h *BasicHost) Mux() *protocol.Mux {
-	return &h.mux
+	return h.mux
 }
 
 func (h *BasicHost) IDService() *identify.IDService {

p2p/protocol/mux.go

@@ -14,7 +14,7 @@ import (
 
 var log = eventlog.Logger("net/mux")
 
-type StreamHandlerMap map[ID]inet.StreamHandler
+type streamHandlerMap map[ID]inet.StreamHandler
 
 // Mux provides simple stream multixplexing.
 // It helps you precisely when:
@@ -23,24 +23,28 @@ type StreamHandlerMap map[ID]inet.StreamHandler
 //
 // It contains the handlers for each protocol accepted.
 // It dispatches handlers for streams opened by remote peers.
-//
-// WARNING: this datastructure IS NOT threadsafe.
-// do not modify it once the network is using it.
 type Mux struct {
-	Default  inet.StreamHandler // handles unknown protocols.
-	Handlers StreamHandlerMap
+	// defaultHandler handles unknown protocols. Callers modify at your own risk.
+	defaultHandler inet.StreamHandler
+
+	lock     sync.RWMutex
+	handlers streamHandlerMap
+}
 
-	sync.RWMutex
+func NewMux() *Mux {
+	return &Mux{
+		handlers: streamHandlerMap{},
+	}
 }
 
 // Protocols returns the list of protocols this muxer has handlers for
 func (m *Mux) Protocols() []ID {
-	m.RLock()
-	l := make([]ID, 0, len(m.Handlers))
-	for p := range m.Handlers {
+	m.lock.RLock()
+	l := make([]ID, 0, len(m.handlers))
+	for p := range m.handlers {
 		l = append(l, p)
 	}
-	m.RUnlock()
+	m.lock.RUnlock()
 	return l
 }
 
@@ -54,14 +58,14 @@ func (m *Mux) readHeader(s io.Reader) (ID, inet.StreamHandler, error) {
 	}
 
 	// log.Debug("readHeader got:", p)
-	m.RLock()
-	h, found := m.Handlers[p]
-	m.RUnlock()
+	m.lock.RLock()
+	h, found := m.handlers[p]
+	m.lock.RUnlock()
 
 	switch {
-	case !found && m.Default != nil:
-		return p, m.Default, nil
-	case !found && m.Default == nil:
+	case !found && m.defaultHandler != nil:
+		return p, m.defaultHandler, nil
+	case !found && m.defaultHandler == nil:
 		return p, nil, fmt.Errorf("%s no handler with name: %s (%d)", m, p, len(p))
 	default:
 		return p, h, nil
@@ -70,18 +74,24 @@ func (m *Mux) readHeader(s io.Reader) (ID, inet.StreamHandler, error) {
 
 // String returns the muxer's printing representation
 func (m *Mux) String() string {
-	m.RLock()
-	defer m.RUnlock()
-	return fmt.Sprintf("<Muxer %p %d>", m, len(m.Handlers))
+	m.lock.RLock()
+	defer m.lock.RUnlock()
+	return fmt.Sprintf("<Muxer %p %d>", m, len(m.handlers))
+}
+
+func (m *Mux) SetDefaultHandler(h inet.StreamHandler) {
+	m.lock.Lock()
+	m.defaultHandler = h
+	m.lock.Unlock()
 }
 
 // SetHandler sets the protocol handler on the Network's Muxer.
 // This operation is threadsafe.
 func (m *Mux) SetHandler(p ID, h inet.StreamHandler) {
 	log.Debugf("%s setting handler for protocol: %s (%d)", m, p, len(p))
-	m.Lock()
-	m.Handlers[p] = h
-	m.Unlock()
+	m.lock.Lock()
+	m.handlers[p] = h
+	m.lock.Unlock()
 }
 
 // Handle reads the next name off the Stream, and calls a handler function

p2p/protocol/mux_test.go

@@ -38,12 +38,12 @@ func TestHandler(t *testing.T) {
 		}
 	}
 
-	m := Mux{Handlers: StreamHandlerMap{}}
-	m.Default = h("default")
-	m.Handlers["/dht"] = h("bitswap")
+	m := NewMux()
+	m.SetDefaultHandler(h("default"))
+	m.SetHandler("/dht", h("bitswap"))
 	// m.Handlers["/ipfs"] = h("bitswap") // default!
-	m.Handlers["/bitswap"] = h("bitswap")
-	m.Handlers["/ipfs/dksnafkasnfkdajfkdajfdsjadosiaaodj"] = h("bitswap")
+	m.SetHandler("/bitswap", h("bitswap"))
+	m.SetHandler("/ipfs/dksnafkasnfkdajfkdajfdsjadosiaaodj", h("bitswap"))
 
 	for k, v := range testCases {
 		var buf bytes.Buffer