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.

211 lines
5.4 KiB

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