You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

202 lines
5.2 KiB

9 years ago
  1. package p2p
  2. import (
  3. "bytes"
  4. "io"
  5. "testing"
  6. . "github.com/tendermint/go-common"
  7. "github.com/tendermint/go-crypto"
  8. )
  9. type dummyConn struct {
  10. *io.PipeReader
  11. *io.PipeWriter
  12. }
  13. func (drw dummyConn) Close() (err error) {
  14. err2 := drw.PipeWriter.CloseWithError(io.EOF)
  15. err1 := drw.PipeReader.Close()
  16. if err2 != nil {
  17. return err
  18. }
  19. return err1
  20. }
  21. // Each returned ReadWriteCloser is akin to a net.Connection
  22. func makeDummyConnPair() (fooConn, barConn dummyConn) {
  23. barReader, fooWriter := io.Pipe()
  24. fooReader, barWriter := io.Pipe()
  25. return dummyConn{fooReader, fooWriter}, dummyConn{barReader, barWriter}
  26. }
  27. func makeSecretConnPair(tb testing.TB) (fooSecConn, barSecConn *SecretConnection) {
  28. fooConn, barConn := makeDummyConnPair()
  29. fooPrvKey := crypto.GenPrivKeyEd25519()
  30. fooPubKey := fooPrvKey.PubKey().(crypto.PubKeyEd25519)
  31. barPrvKey := crypto.GenPrivKeyEd25519()
  32. barPubKey := barPrvKey.PubKey().(crypto.PubKeyEd25519)
  33. Parallel(
  34. func() {
  35. var err error
  36. fooSecConn, err = MakeSecretConnection(fooConn, fooPrvKey)
  37. if err != nil {
  38. tb.Errorf("Failed to establish SecretConnection for foo: %v", err)
  39. return
  40. }
  41. remotePubBytes := fooSecConn.RemotePubKey()
  42. if !bytes.Equal(remotePubBytes[:], barPubKey[:]) {
  43. tb.Errorf("Unexpected fooSecConn.RemotePubKey. Expected %v, got %v",
  44. barPubKey, fooSecConn.RemotePubKey())
  45. }
  46. },
  47. func() {
  48. var err error
  49. barSecConn, err = MakeSecretConnection(barConn, barPrvKey)
  50. if barSecConn == nil {
  51. tb.Errorf("Failed to establish SecretConnection for bar: %v", err)
  52. return
  53. }
  54. remotePubBytes := barSecConn.RemotePubKey()
  55. if !bytes.Equal(remotePubBytes[:], fooPubKey[:]) {
  56. tb.Errorf("Unexpected barSecConn.RemotePubKey. Expected %v, got %v",
  57. fooPubKey, barSecConn.RemotePubKey())
  58. }
  59. })
  60. return
  61. }
  62. func TestSecretConnectionHandshake(t *testing.T) {
  63. fooSecConn, barSecConn := makeSecretConnPair(t)
  64. fooSecConn.Close()
  65. barSecConn.Close()
  66. }
  67. func TestSecretConnectionReadWrite(t *testing.T) {
  68. fooConn, barConn := makeDummyConnPair()
  69. fooWrites, barWrites := []string{}, []string{}
  70. fooReads, barReads := []string{}, []string{}
  71. // Pre-generate the things to write (for foo & bar)
  72. for i := 0; i < 100; i++ {
  73. fooWrites = append(fooWrites, RandStr((RandInt()%(dataMaxSize*5))+1))
  74. barWrites = append(barWrites, RandStr((RandInt()%(dataMaxSize*5))+1))
  75. }
  76. // A helper that will run with (fooConn, fooWrites, fooReads) and vice versa
  77. genNodeRunner := func(nodeConn dummyConn, nodeWrites []string, nodeReads *[]string) func() {
  78. return func() {
  79. // Node handskae
  80. nodePrvKey := crypto.GenPrivKeyEd25519()
  81. nodeSecretConn, err := MakeSecretConnection(nodeConn, nodePrvKey)
  82. if err != nil {
  83. t.Errorf("Failed to establish SecretConnection for node: %v", err)
  84. return
  85. }
  86. // In parallel, handle reads and writes
  87. Parallel(
  88. func() {
  89. // Node writes
  90. for _, nodeWrite := range nodeWrites {
  91. n, err := nodeSecretConn.Write([]byte(nodeWrite))
  92. if err != nil {
  93. t.Errorf("Failed to write to nodeSecretConn: %v", err)
  94. return
  95. }
  96. if n != len(nodeWrite) {
  97. t.Errorf("Failed to write all bytes. Expected %v, wrote %v", len(nodeWrite), n)
  98. return
  99. }
  100. }
  101. nodeConn.PipeWriter.Close()
  102. },
  103. func() {
  104. // Node reads
  105. readBuffer := make([]byte, dataMaxSize)
  106. for {
  107. n, err := nodeSecretConn.Read(readBuffer)
  108. if err == io.EOF {
  109. return
  110. } else if err != nil {
  111. t.Errorf("Failed to read from nodeSecretConn: %v", err)
  112. return
  113. }
  114. *nodeReads = append(*nodeReads, string(readBuffer[:n]))
  115. }
  116. nodeConn.PipeReader.Close()
  117. })
  118. }
  119. }
  120. // Run foo & bar in parallel
  121. Parallel(
  122. genNodeRunner(fooConn, fooWrites, &fooReads),
  123. genNodeRunner(barConn, barWrites, &barReads),
  124. )
  125. // A helper to ensure that the writes and reads match.
  126. // Additionally, small writes (<= dataMaxSize) must be atomically read.
  127. compareWritesReads := func(writes []string, reads []string) {
  128. for {
  129. // Pop next write & corresponding reads
  130. var read, write string = "", writes[0]
  131. var readCount = 0
  132. for _, readChunk := range reads {
  133. read += readChunk
  134. readCount += 1
  135. if len(write) <= len(read) {
  136. break
  137. }
  138. if len(write) <= dataMaxSize {
  139. break // atomicity of small writes
  140. }
  141. }
  142. // Compare
  143. if write != read {
  144. t.Errorf("Expected to read %X, got %X", write, read)
  145. }
  146. // Iterate
  147. writes = writes[1:]
  148. reads = reads[readCount:]
  149. if len(writes) == 0 {
  150. break
  151. }
  152. }
  153. }
  154. compareWritesReads(fooWrites, barReads)
  155. compareWritesReads(barWrites, fooReads)
  156. }
  157. func BenchmarkSecretConnection(b *testing.B) {
  158. b.StopTimer()
  159. fooSecConn, barSecConn := makeSecretConnPair(b)
  160. fooWriteText := RandStr(dataMaxSize)
  161. // Consume reads from bar's reader
  162. go func() {
  163. readBuffer := make([]byte, dataMaxSize)
  164. for {
  165. _, err := barSecConn.Read(readBuffer)
  166. if err == io.EOF {
  167. return
  168. } else if err != nil {
  169. b.Fatalf("Failed to read from barSecConn: %v", err)
  170. }
  171. }
  172. }()
  173. b.StartTimer()
  174. for i := 0; i < b.N; i++ {
  175. _, err := fooSecConn.Write([]byte(fooWriteText))
  176. if err != nil {
  177. b.Fatalf("Failed to write to fooSecConn: %v", err)
  178. }
  179. }
  180. b.StopTimer()
  181. fooSecConn.Close()
  182. //barSecConn.Close() race condition
  183. }