From 41e16adf67ba1b61256a51a335224456b8e7c48a Mon Sep 17 00:00:00 2001
From: Marten Seemann <martenseemann@gmail.com>
Date: Thu, 20 Feb 2020 12:20:42 +0700
Subject: [PATCH] use the ipnet.PSK instead of the ipnet.Protector for private
 networks

---
 go.mod      |  1 +
 go.sum      |  2 ++
 upgrader.go | 19 ++++++++++---------
 3 files changed, 13 insertions(+), 9 deletions(-)

diff --git a/go.mod b/go.mod
index 07b697a..45d500a 100644
--- a/go.mod
+++ b/go.mod
@@ -5,6 +5,7 @@ require (
 	github.com/jbenet/go-temp-err-catcher v0.0.0-20150120210811-aac704a3f4f2
 	github.com/libp2p/go-libp2p-core v0.3.0
 	github.com/libp2p/go-libp2p-mplex v0.2.1
+	github.com/libp2p/go-libp2p-pnet v0.1.0
 	github.com/libp2p/go-maddr-filter v0.0.5
 	github.com/multiformats/go-multiaddr v0.2.0
 	github.com/multiformats/go-multiaddr-net v0.1.2
diff --git a/go.sum b/go.sum
index 0d99222..20e1b8b 100644
--- a/go.sum
+++ b/go.sum
@@ -79,6 +79,8 @@ github.com/libp2p/go-libp2p-mplex v0.2.0 h1:vnR+/fUnL5d2kTB1Lfce1KVSP+bCCy2aJIvq
 github.com/libp2p/go-libp2p-mplex v0.2.0/go.mod h1:Ejl9IyjvXJ0T9iqUTE1jpYATQ9NM3g+OtR+EMMODbKo=
 github.com/libp2p/go-libp2p-mplex v0.2.1 h1:E1xaJBQnbSiTHGI1gaBKmKhu1TUKkErKJnE8iGvirYI=
 github.com/libp2p/go-libp2p-mplex v0.2.1/go.mod h1:SC99Rxs8Vuzrf/6WhmH41kNn13TiYdAWNYHrwImKLnE=
+github.com/libp2p/go-libp2p-pnet v0.1.0 h1:kRUES28dktfnHNIRW4Ro78F7rKBHBiw5MJpl0ikrLIA=
+github.com/libp2p/go-libp2p-pnet v0.1.0/go.mod h1:ZkyZw3d0ZFOex71halXRihWf9WH/j3OevcJdTmD0lyE=
 github.com/libp2p/go-libp2p-testing v0.0.3 h1:bdij4bKaaND7tCsaXVjRfYkMpvoOeKj9AVQGJllA6jM=
 github.com/libp2p/go-libp2p-testing v0.0.3/go.mod h1:gvchhf3FQOtBdr+eFUABet5a4MBLK8jM3V4Zghvmi+E=
 github.com/libp2p/go-maddr-filter v0.0.4 h1:hx8HIuuwk34KePddrp2mM5ivgPkZ09JH4AvsALRbFUs=
diff --git a/upgrader.go b/upgrader.go
index 69fb90b..a464df5 100644
--- a/upgrader.go
+++ b/upgrader.go
@@ -8,9 +8,10 @@ import (
 
 	"github.com/libp2p/go-libp2p-core/mux"
 	"github.com/libp2p/go-libp2p-core/peer"
-	"github.com/libp2p/go-libp2p-core/pnet"
+	ipnet "github.com/libp2p/go-libp2p-core/pnet"
 	"github.com/libp2p/go-libp2p-core/sec"
 	"github.com/libp2p/go-libp2p-core/transport"
+	"github.com/libp2p/go-libp2p-pnet"
 
 	filter "github.com/libp2p/go-maddr-filter"
 	manet "github.com/multiformats/go-multiaddr-net"
@@ -26,10 +27,10 @@ var AcceptQueueLength = 16
 // Upgrader is a multistream upgrader that can upgrade an underlying connection
 // to a full transport connection (secure and multiplexed).
 type Upgrader struct {
-	Protector pnet.Protector
-	Secure    sec.SecureTransport
-	Muxer     mux.Multiplexer
-	Filters   *filter.Filters
+	PSK     ipnet.PSK
+	Secure  sec.SecureTransport
+	Muxer   mux.Multiplexer
+	Filters *filter.Filters
 }
 
 // UpgradeListener upgrades the passed multiaddr-net listener into a full libp2p-transport listener.
@@ -71,17 +72,17 @@ func (u *Upgrader) upgrade(ctx context.Context, t transport.Transport, maconn ma
 	}
 
 	var conn net.Conn = maconn
-	if u.Protector != nil {
-		pconn, err := u.Protector.Protect(conn)
+	if u.PSK != nil {
+		pconn, err := pnet.NewProtectedConn(u.PSK, conn)
 		if err != nil {
 			conn.Close()
 			return nil, fmt.Errorf("failed to setup private network protector: %s", err)
 		}
 		conn = pconn
-	} else if pnet.ForcePrivateNetwork {
+	} else if ipnet.ForcePrivateNetwork {
 		log.Error("tried to dial with no Private Network Protector but usage" +
 			" of Private Networks is forced by the enviroment")
-		return nil, pnet.ErrNotInPrivateNetwork
+		return nil, ipnet.ErrNotInPrivateNetwork
 	}
 	sconn, err := u.setupSecurity(ctx, conn, p)
 	if err != nil {