From 34c0789b9de1128edfbd225c5c5c1864dd7e4bd6 Mon Sep 17 00:00:00 2001
From: Steven Allen <steven@stebalien.com>
Date: Thu, 18 Jan 2018 13:31:54 -0800
Subject: [PATCH] wrap net conns, not transport conns

We now "protect" below the transport layer.

Also, make the tests work without the dummy conn package (one fewer deps).
---
 p2p/net/pnet/protector.go     |  4 ++--
 p2p/net/pnet/psk_conn.go      |  8 ++++----
 p2p/net/pnet/psk_conn_test.go | 36 +++++++++++++++++++++--------------
 3 files changed, 28 insertions(+), 20 deletions(-)

diff --git a/p2p/net/pnet/protector.go b/p2p/net/pnet/protector.go
index 72446c2c6e..38706817ab 100644
--- a/p2p/net/pnet/protector.go
+++ b/p2p/net/pnet/protector.go
@@ -3,9 +3,9 @@ package pnet
 import (
 	"fmt"
 	"io"
+	"net"
 
 	ipnet "github.com/libp2p/go-libp2p-interface-pnet"
-	tconn "github.com/libp2p/go-libp2p-transport"
 )
 
 var _ ipnet.Protector = (*protector)(nil)
@@ -33,7 +33,7 @@ type protector struct {
 	fingerprint []byte
 }
 
-func (p protector) Protect(in tconn.Conn) (tconn.Conn, error) {
+func (p protector) Protect(in net.Conn) (net.Conn, error) {
 	return newPSKConn(p.psk, in)
 }
 func (p protector) Fingerprint() []byte {
diff --git a/p2p/net/pnet/psk_conn.go b/p2p/net/pnet/psk_conn.go
index b9fc750b3c..18dbada1fe 100644
--- a/p2p/net/pnet/psk_conn.go
+++ b/p2p/net/pnet/psk_conn.go
@@ -4,10 +4,10 @@ import (
 	"crypto/cipher"
 	"crypto/rand"
 	"io"
+	"net"
 
 	salsa20 "github.com/davidlazar/go-crypto/salsa20"
 	ipnet "github.com/libp2p/go-libp2p-interface-pnet"
-	tconn "github.com/libp2p/go-libp2p-transport"
 	mpool "github.com/libp2p/go-msgio/mpool"
 )
 
@@ -20,7 +20,7 @@ var (
 )
 
 type pskConn struct {
-	tconn.Conn
+	net.Conn
 	psk *[32]byte
 
 	writeS20 cipher.Stream
@@ -73,9 +73,9 @@ func (c *pskConn) Write(in []byte) (int, error) {
 	return c.Conn.Write(out) // send
 }
 
-var _ tconn.Conn = (*pskConn)(nil)
+var _ net.Conn = (*pskConn)(nil)
 
-func newPSKConn(psk *[32]byte, insecure tconn.Conn) (tconn.Conn, error) {
+func newPSKConn(psk *[32]byte, insecure net.Conn) (net.Conn, error) {
 	if insecure == nil {
 		return nil, errInsecureNil
 	}
diff --git a/p2p/net/pnet/psk_conn_test.go b/p2p/net/pnet/psk_conn_test.go
index c76f474af7..23108f4072 100644
--- a/p2p/net/pnet/psk_conn_test.go
+++ b/p2p/net/pnet/psk_conn_test.go
@@ -4,19 +4,14 @@ import (
 	"bytes"
 	"context"
 	"math/rand"
+	"net"
 	"testing"
-
-	dconn "github.com/libp2p/go-libp2p-dummy-conn"
-	tconn "github.com/libp2p/go-libp2p-transport"
 )
 
 var testPSK = [32]byte{} // null bytes are as good test key as any other key
 
-func setupPSKConns(ctx context.Context, t *testing.T) (tconn.Conn, tconn.Conn) {
-	conn1, conn2, err := dconn.NewDummyConnPair()
-	if err != nil {
-		t.Fatal(err)
-	}
+func setupPSKConns(ctx context.Context, t *testing.T) (net.Conn, net.Conn) {
+	conn1, conn2 := net.Pipe()
 
 	psk1, err := newPSKConn(&testPSK, conn1)
 	if err != nil {
@@ -37,14 +32,21 @@ func TestPSKSimpelMessges(t *testing.T) {
 	msg1 := []byte("hello world")
 	out1 := make([]byte, len(msg1))
 
-	_, err := psk1.Write(msg1)
+	wch := make(chan error)
+	go func() {
+		_, err := psk1.Write(msg1)
+		wch <- err
+	}()
+	n, err := psk2.Read(out1)
 	if err != nil {
 		t.Fatal(err)
 	}
-	n, err := psk2.Read(out1)
+
+	err = <-wch
 	if err != nil {
 		t.Fatal(err)
 	}
+
 	if n != len(out1) {
 		t.Fatalf("expected to read %d bytes, read: %d", len(out1), n)
 	}
@@ -68,10 +70,11 @@ func TestPSKFragmentation(t *testing.T) {
 
 	out := make([]byte, 100)
 
-	_, err = psk1.Write(in)
-	if err != nil {
-		t.Fatal(err)
-	}
+	wch := make(chan error)
+	go func() {
+		_, err := psk1.Write(in)
+		wch <- err
+	}()
 
 	for i := 0; i < 10; i++ {
 		_, err = psk2.Read(out)
@@ -80,4 +83,9 @@ func TestPSKFragmentation(t *testing.T) {
 		}
 		in = in[100:]
 	}
+
+	err = <-wch
+	if err != nil {
+		t.Fatal(err)
+	}
 }