From d6fc4be820eb529112390d66d896384049d3d046 Mon Sep 17 00:00:00 2001
From: Marten Seemann <martenseemann@gmail.com>
Date: Fri, 4 Mar 2022 10:56:56 +0400
Subject: [PATCH] update mplex (#31)

---
 p2p/muxer/mplex/transport.go      |  6 ++++-
 p2p/muxer/mplex/transport_test.go | 41 +++++++++++++++++++++++++++++++
 2 files changed, 46 insertions(+), 1 deletion(-)

diff --git a/p2p/muxer/mplex/transport.go b/p2p/muxer/mplex/transport.go
index ac08b9b6..8f9e33a9 100644
--- a/p2p/muxer/mplex/transport.go
+++ b/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
 }
diff --git a/p2p/muxer/mplex/transport_test.go b/p2p/muxer/mplex/transport_test.go
index 134fd444..c224667c 100644
--- a/p2p/muxer/mplex/transport_test.go
+++ b/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,
+	})
+}