update mplex (#31)

p2p/muxer/mplex/transport.go

@@ -18,5 +18,9 @@ var _ network.Multiplexer = &Transport{}
 type Transport struct{}
 
 func (t *Transport) NewConn(nc net.Conn, isServer bool, scope network.PeerScope) (network.MuxedConn, error) {
-	return (*conn)(mp.NewMultiplex(nc, isServer, scope)), nil
+	m, err := mp.NewMultiplex(nc, isServer, scope)
+	if err != nil {
+		return nil, err
+	}
+	return (*conn)(m), nil
 }

p2p/muxer/mplex/transport_test.go

@@ -1,11 +1,52 @@
 package peerstream_multiplex
 
 import (
+	"errors"
+	"net"
 	"testing"
 
+	"github.com/libp2p/go-libp2p-core/network"
 	test "github.com/libp2p/go-libp2p-testing/suites/mux"
 )
 
 func TestDefaultTransport(t *testing.T) {
 	test.SubtestAll(t, DefaultTransport)
 }
+
+type memoryScope struct {
+	network.PeerScope
+	limit    int
+	reserved int
+}
+
+func (m *memoryScope) ReserveMemory(size int, prio uint8) error {
+	if m.reserved+size > m.limit {
+		return errors.New("too much")
+	}
+	m.reserved += size
+	return nil
+}
+
+func (m *memoryScope) ReleaseMemory(size int) {
+	m.reserved -= size
+	if m.reserved < 0 {
+		panic("too much memory released")
+	}
+}
+
+type memoryLimitedTransport struct {
+	Transport
+}
+
+func (t *memoryLimitedTransport) NewConn(nc net.Conn, isServer bool, scope network.PeerScope) (network.MuxedConn, error) {
+	return t.Transport.NewConn(nc, isServer, &memoryScope{
+		limit:     3 * 1 << 20,
+		PeerScope: scope,
+	})
+}
+
+func TestDefaultTransportWithMemoryLimit(t *testing.T) {
+	test.SubtestAll(t, &memoryLimitedTransport{
+		Transport: *DefaultTransport,
+	})
+}