p2p: tmconn->conn and types->p2p

p2p/conn/secret_connection_test.go

@@ -0,0 +1,211 @@
+package conn
+
+import (
+	"io"
+	"testing"
+
+	"github.com/tendermint/go-crypto"
+	cmn "github.com/tendermint/tmlibs/common"
+)
+
+type dummyConn struct {
+	*io.PipeReader
+	*io.PipeWriter
+}
+
+func (drw dummyConn) Close() (err error) {
+	err2 := drw.PipeWriter.CloseWithError(io.EOF)
+	err1 := drw.PipeReader.Close()
+	if err2 != nil {
+		return err
+	}
+	return err1
+}
+
+// Each returned ReadWriteCloser is akin to a net.Connection
+func makeDummyConnPair() (fooConn, barConn dummyConn) {
+	barReader, fooWriter := io.Pipe()
+	fooReader, barWriter := io.Pipe()
+	return dummyConn{fooReader, fooWriter}, dummyConn{barReader, barWriter}
+}
+
+func makeSecretConnPair(tb testing.TB) (fooSecConn, barSecConn *SecretConnection) {
+	fooConn, barConn := makeDummyConnPair()
+	fooPrvKey := crypto.GenPrivKeyEd25519().Wrap()
+	fooPubKey := fooPrvKey.PubKey()
+	barPrvKey := crypto.GenPrivKeyEd25519().Wrap()
+	barPubKey := barPrvKey.PubKey()
+
+	cmn.Parallel(
+		func() {
+			var err error
+			fooSecConn, err = MakeSecretConnection(fooConn, fooPrvKey)
+			if err != nil {
+				tb.Errorf("Failed to establish SecretConnection for foo: %v", err)
+				return
+			}
+			remotePubBytes := fooSecConn.RemotePubKey()
+			if !remotePubBytes.Equals(barPubKey) {
+				tb.Errorf("Unexpected fooSecConn.RemotePubKey.  Expected %v, got %v",
+					barPubKey, fooSecConn.RemotePubKey())
+			}
+		},
+		func() {
+			var err error
+			barSecConn, err = MakeSecretConnection(barConn, barPrvKey)
+			if barSecConn == nil {
+				tb.Errorf("Failed to establish SecretConnection for bar: %v", err)
+				return
+			}
+			remotePubBytes := barSecConn.RemotePubKey()
+			if !remotePubBytes.Equals(fooPubKey) {
+				tb.Errorf("Unexpected barSecConn.RemotePubKey.  Expected %v, got %v",
+					fooPubKey, barSecConn.RemotePubKey())
+			}
+		})
+
+	return
+}
+
+func TestSecretConnectionHandshake(t *testing.T) {
+	fooSecConn, barSecConn := makeSecretConnPair(t)
+	if err := fooSecConn.Close(); err != nil {
+		t.Error(err)
+	}
+	if err := barSecConn.Close(); err != nil {
+		t.Error(err)
+	}
+}
+
+func TestSecretConnectionReadWrite(t *testing.T) {
+	fooConn, barConn := makeDummyConnPair()
+	fooWrites, barWrites := []string{}, []string{}
+	fooReads, barReads := []string{}, []string{}
+
+	// Pre-generate the things to write (for foo & bar)
+	for i := 0; i < 100; i++ {
+		fooWrites = append(fooWrites, cmn.RandStr((cmn.RandInt()%(dataMaxSize*5))+1))
+		barWrites = append(barWrites, cmn.RandStr((cmn.RandInt()%(dataMaxSize*5))+1))
+	}
+
+	// A helper that will run with (fooConn, fooWrites, fooReads) and vice versa
+	genNodeRunner := func(nodeConn dummyConn, nodeWrites []string, nodeReads *[]string) func() {
+		return func() {
+			// Node handskae
+			nodePrvKey := crypto.GenPrivKeyEd25519().Wrap()
+			nodeSecretConn, err := MakeSecretConnection(nodeConn, nodePrvKey)
+			if err != nil {
+				t.Errorf("Failed to establish SecretConnection for node: %v", err)
+				return
+			}
+			// In parallel, handle reads and writes
+			cmn.Parallel(
+				func() {
+					// Node writes
+					for _, nodeWrite := range nodeWrites {
+						n, err := nodeSecretConn.Write([]byte(nodeWrite))
+						if err != nil {
+							t.Errorf("Failed to write to nodeSecretConn: %v", err)
+							return
+						}
+						if n != len(nodeWrite) {
+							t.Errorf("Failed to write all bytes. Expected %v, wrote %v", len(nodeWrite), n)
+							return
+						}
+					}
+					if err := nodeConn.PipeWriter.Close(); err != nil {
+						t.Error(err)
+					}
+				},
+				func() {
+					// Node reads
+					readBuffer := make([]byte, dataMaxSize)
+					for {
+						n, err := nodeSecretConn.Read(readBuffer)
+						if err == io.EOF {
+							return
+						} else if err != nil {
+							t.Errorf("Failed to read from nodeSecretConn: %v", err)
+							return
+						}
+						*nodeReads = append(*nodeReads, string(readBuffer[:n]))
+					}
+					if err := nodeConn.PipeReader.Close(); err != nil {
+						t.Error(err)
+					}
+				})
+		}
+	}
+
+	// Run foo & bar in parallel
+	cmn.Parallel(
+		genNodeRunner(fooConn, fooWrites, &fooReads),
+		genNodeRunner(barConn, barWrites, &barReads),
+	)
+
+	// A helper to ensure that the writes and reads match.
+	// Additionally, small writes (<= dataMaxSize) must be atomically read.
+	compareWritesReads := func(writes []string, reads []string) {
+		for {
+			// Pop next write & corresponding reads
+			var read, write string = "", writes[0]
+			var readCount = 0
+			for _, readChunk := range reads {
+				read += readChunk
+				readCount += 1
+				if len(write) <= len(read) {
+					break
+				}
+				if len(write) <= dataMaxSize {
+					break // atomicity of small writes
+				}
+			}
+			// Compare
+			if write != read {
+				t.Errorf("Expected to read %X, got %X", write, read)
+			}
+			// Iterate
+			writes = writes[1:]
+			reads = reads[readCount:]
+			if len(writes) == 0 {
+				break
+			}
+		}
+	}
+
+	compareWritesReads(fooWrites, barReads)
+	compareWritesReads(barWrites, fooReads)
+
+}
+
+func BenchmarkSecretConnection(b *testing.B) {
+	b.StopTimer()
+	fooSecConn, barSecConn := makeSecretConnPair(b)
+	fooWriteText := cmn.RandStr(dataMaxSize)
+	// Consume reads from bar's reader
+	go func() {
+		readBuffer := make([]byte, dataMaxSize)
+		for {
+			_, err := barSecConn.Read(readBuffer)
+			if err == io.EOF {
+				return
+			} else if err != nil {
+				b.Fatalf("Failed to read from barSecConn: %v", err)
+			}
+		}
+	}()
+
+	b.StartTimer()
+	for i := 0; i < b.N; i++ {
+		_, err := fooSecConn.Write([]byte(fooWriteText))
+		if err != nil {
+			b.Fatalf("Failed to write to fooSecConn: %v", err)
+		}
+	}
+	b.StopTimer()
+
+	if err := fooSecConn.Close(); err != nil {
+		b.Error(err)
+	}
+	//barSecConn.Close() race condition
+}