From 6c29558900497ca12a2b5aa6d031cf685fbb853a Mon Sep 17 00:00:00 2001
From: Marten Seemann <martenseemann@gmail.com>
Date: Mon, 17 Jan 2022 07:19:10 -0800
Subject: [PATCH] update mplex

---
 go.mod            |  4 ++--
 go.sum            | 11 +++++++----
 transport.go      |  6 +++++-
 transport_test.go | 41 +++++++++++++++++++++++++++++++++++++++++
 4 files changed, 55 insertions(+), 7 deletions(-)

diff --git a/go.mod b/go.mod
index 1b239be..d23af31 100644
--- a/go.mod
+++ b/go.mod
@@ -4,6 +4,6 @@ go 1.16
 
 require (
 	github.com/libp2p/go-libp2p-core v0.14.0
-	github.com/libp2p/go-libp2p-testing v0.7.0
-	github.com/libp2p/go-mplex v0.4.0
+	github.com/libp2p/go-libp2p-testing v0.8.0
+	github.com/libp2p/go-mplex v0.6.0
 )
diff --git a/go.sum b/go.sum
index 0444337..c98fd60 100644
--- a/go.sum
+++ b/go.sum
@@ -55,8 +55,10 @@ github.com/kisielk/errcheck v1.2.0/go.mod h1:/BMXB+zMLi60iA8Vv6Ksmxu/1UDYcXs4uQL
 github.com/kisielk/gotool v1.0.0/go.mod h1:XhKaO+MFFWcvkIS/tQcRk01m1F5IRFswLeQ+oQHNcck=
 github.com/kkdai/bstream v0.0.0-20161212061736-f391b8402d23/go.mod h1:J+Gs4SYgM6CZQHDETBtE9HaSEkGmuNXF86RwHhHUvq4=
 github.com/konsorten/go-windows-terminal-sequences v1.0.1/go.mod h1:T0+1ngSBFLxvqU3pZ+m/2kptfBszLMUkC4ZK/EgS/cQ=
+github.com/kr/pretty v0.1.0 h1:L/CwN0zerZDmRFUapSPitk6f+Q3+0za1rQkzVuMiMFI=
 github.com/kr/pretty v0.1.0/go.mod h1:dAy3ld7l9f0ibDNOQOHHMYYIIbhfbHSm3C4ZsoJORNo=
 github.com/kr/pty v1.1.1/go.mod h1:pFQYn66WHrOpPYNljwOMqo10TkYh1fy3cYio2l3bCsQ=
+github.com/kr/text v0.1.0 h1:45sCR5RtlFHMR4UwH9sdQ5TC8v0qDQCHnXt+kaKSTVE=
 github.com/kr/text v0.1.0/go.mod h1:4Jbv+DJW3UT/LiOwJeYQe1efqtUx/iVham/4vfdArNI=
 github.com/libp2p/go-buffer-pool v0.0.2 h1:QNK2iAFa8gjAe1SPz6mHSMuCcjs+X1wlHzeOSqcmlfs=
 github.com/libp2p/go-buffer-pool v0.0.2/go.mod h1:MvaB6xw5vOrDl8rYZGLFdKAuk/hRoRZd1Vi32+RXyFM=
@@ -65,10 +67,10 @@ github.com/libp2p/go-libp2p-core v0.3.0/go.mod h1:ACp3DmS3/N64c2jDzcV429ukDpicbL
 github.com/libp2p/go-libp2p-core v0.14.0 h1:0kYSgiK/D7Eo28GTuRXo5YHsWwAisVpFCqCVPUd/vJs=
 github.com/libp2p/go-libp2p-core v0.14.0/go.mod h1:tLasfcVdTXnixsLB0QYaT1syJOhsbrhG7q6pGrHtBg8=
 github.com/libp2p/go-libp2p-testing v0.1.2-0.20200422005655-8775583591d8/go.mod h1:Qy8sAncLKpwXtS2dSnDOP8ktexIAHKu+J+pnZOFZLTc=
-github.com/libp2p/go-libp2p-testing v0.7.0 h1:9bfyhNINizxuLrKsenzGaZalXRXIaAEmx1BP/PzF1gM=
-github.com/libp2p/go-libp2p-testing v0.7.0/go.mod h1:OLbdn9DbgdMwv00v+tlp1l3oe2Cl+FAjoWIA2pa0X6E=
-github.com/libp2p/go-mplex v0.4.0 h1:Ukkez9/4EOX5rTw4sHefNJp10dksftAA05ZgyjplUbM=
-github.com/libp2p/go-mplex v0.4.0/go.mod h1:y26Lx+wNVtMYMaPu300Cbot5LkEZ4tJaNYeHeT9dh6E=
+github.com/libp2p/go-libp2p-testing v0.8.0 h1:/te8SOIyj5sGH5Jr1Uoo+qYB00aK8O4+yHGzLgfE3kc=
+github.com/libp2p/go-libp2p-testing v0.8.0/go.mod h1:gRdsNxQSxAZowTgcLY7CC33xPmleZzoBpqSYbWenqPc=
+github.com/libp2p/go-mplex v0.6.0 h1:5kKp029zrsLVJT5q6ASt4LwuZFxj3B13wXXaGmFrWg0=
+github.com/libp2p/go-mplex v0.6.0/go.mod h1:y26Lx+wNVtMYMaPu300Cbot5LkEZ4tJaNYeHeT9dh6E=
 github.com/libp2p/go-msgio v0.0.6/go.mod h1:4ecVB6d9f4BDSL5fqvPiC4A3KivjWn+Venn/1ALLMWA=
 github.com/libp2p/go-openssl v0.0.4/go.mod h1:unDrJpgy3oFr+rqXsarWifmJuNnJR4chtO1HmaZjggc=
 github.com/libp2p/go-openssl v0.0.7 h1:eCAzdLejcNVBzP/iZM9vqHnQm+XyCEbSSIheIPRGNsw=
@@ -214,6 +216,7 @@ google.golang.org/grpc v1.28.1/go.mod h1:rpkK4SK4GF4Ach/+MFLZUBavHOvF2JJB5uozKKa
 google.golang.org/grpc v1.31.1 h1:SfXqXS5hkufcdZ/mHtYCh53P2b+92WQq/DZcKLgsFRs=
 google.golang.org/grpc v1.31.1/go.mod h1:N36X2cJ7JwdamYAgDz+s+rVMFjt3numwzf/HckM8pak=
 gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
+gopkg.in/check.v1 v1.0.0-20180628173108-788fd7840127 h1:qIbj1fsPNlZgppZ+VLlY7N33q108Sa+fhmuc+sWQYwY=
 gopkg.in/check.v1 v1.0.0-20180628173108-788fd7840127/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
 gopkg.in/fsnotify.v1 v1.4.7/go.mod h1:Tz8NjZHkW78fSQdbUxIjBTcgA1z1m8ZHf0WmKUhAMys=
 gopkg.in/src-d/go-cli.v0 v0.0.0-20181105080154-d492247bbc0d/go.mod h1:z+K8VcOYVYcSwSjGebuDL6176A1XskgbtNl64NSg+n8=
diff --git a/transport.go b/transport.go
index ac08b9b..8f9e33a 100644
--- a/transport.go
+++ b/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/transport_test.go b/transport_test.go
index 134fd44..c224667 100644
--- a/transport_test.go
+++ b/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,
+	})
+}