--- /dev/null
+package connection
+
+import (
+ "bytes"
+ "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()
+ fooPubKey := fooPrvKey.PubKey().Unwrap().(crypto.PubKeyEd25519)
+ barPrvKey := crypto.GenPrivKeyEd25519()
+ barPubKey := barPrvKey.PubKey().Unwrap().(crypto.PubKeyEd25519)
+
+ 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 !bytes.Equal(remotePubBytes[:], 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 !bytes.Equal(remotePubBytes[:], fooPubKey[:]) {
+ tb.Errorf("Unexpected barSecConn.RemotePubKey. Expected %v, got %v",
+ fooPubKey, barSecConn.RemotePubKey())
+ }
+ })
+
+ return
+}
+
+func TestSecretConnectionHandshake(t *testing.T) {
+ fooSecConn, barSecConn := makeSecretConnPair(t)
+ fooSecConn.Close()
+ barSecConn.Close()
+}
+
+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()
+ 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
+ }
+ }
+ nodeConn.PipeWriter.Close()
+ },
+ 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]))
+ }
+ nodeConn.PipeReader.Close()
+ })
+ }
+ }
+
+ // 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++
+ 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()
+
+ fooSecConn.Close()
+ //barSecConn.Close() race condition
+}
OSDN Git Service
