mse/mse_test.go

   1 package mse
   2
   3 import (
   4         "bytes"
   5         "crypto/rand"
   6         "io"
   7         "io/ioutil"
   8         "net"
   9         "sync"
  10
  11         "github.com/bradfitz/iter"
  12
  13         "testing"
  14 )
  15
  16 func TestReadUntil(t *testing.T) {
  17         test := func(data, until string, leftover int, expectedErr error) {
  18                 r := bytes.NewReader([]byte(data))
  19                 err := readUntil(r, []byte(until))
  20                 if err != expectedErr {
  21                         t.Fatal(err)
  22                 }
  23                 if r.Len() != leftover {
  24                         t.Fatal(r.Len())
  25                 }
  26         }
  27         test("feakjfeafeafegbaabc00", "abc", 2, nil)
  28         test("feakjfeafeafegbaadc00", "abc", 0, io.EOF)
  29 }
  30
  31 func TestSuffixMatchLen(t *testing.T) {
  32         test := func(a, b string, expected int) {
  33                 actual := suffixMatchLen([]byte(a), []byte(b))
  34                 if actual != expected {
  35                         t.Fatalf("expected %d, got %d for %q and %q", expected, actual, a, b)
  36                 }
  37         }
  38         test("hello", "world", 0)
  39         test("hello", "lo", 2)
  40         test("hello", "llo", 3)
  41         test("hello", "hell", 0)
  42         test("hello", "helloooo!", 5)
  43         test("hello", "lol!", 2)
  44         test("hello", "mondo", 0)
  45         test("mongo", "webscale", 0)
  46         test("sup", "person", 1)
  47 }
  48
  49 func handshakeTest(t testing.TB, ia []byte, aData, bData string) {
  50         a, b := net.Pipe()
  51         wg := sync.WaitGroup{}
  52         wg.Add(2)
  53         go func() {
  54                 defer wg.Done()
  55                 a, err := InitiateHandshake(a, []byte("yep"), ia)
  56                 if err != nil {
  57                         t.Fatal(err)
  58                         return
  59                 }
  60                 go a.Write([]byte(aData))
  61
  62                 var msg [20]byte
  63                 n, _ := a.Read(msg[:])
  64                 if n != len(bData) {
  65                         t.FailNow()
  66                 }
  67                 // t.Log(string(msg[:n]))
  68         }()
  69         go func() {
  70                 defer wg.Done()
  71                 b, err := ReceiveHandshake(b, [][]byte{[]byte("nope"), []byte("yep"), []byte("maybe")})
  72                 if err != nil {
  73                         t.Fatal(err)
  74                         return
  75                 }
  76                 go b.Write([]byte(bData))
  77                 // Need to be exact here, as there are several reads, and net.Pipe is
  78                 // most synchronous.
  79                 msg := make([]byte, len(ia)+len(aData))
  80                 n, _ := io.ReadFull(b, msg[:])
  81                 if n != len(msg) {
  82                         t.FailNow()
  83                 }
  84                 // t.Log(string(msg[:n]))
  85         }()
  86         wg.Wait()
  87         a.Close()
  88         b.Close()
  89 }
  90
  91 func allHandshakeTests(t testing.TB) {
  92         handshakeTest(t, []byte("jump the gun, "), "hello world", "yo dawg")
  93         handshakeTest(t, nil, "hello world", "yo dawg")
  94         handshakeTest(t, []byte{}, "hello world", "yo dawg")
  95 }
  96
  97 func TestHandshake(t *testing.T) {
  98         allHandshakeTests(t)
  99         t.Logf("crypto provides encountered: %s", cryptoProvidesCount)
 100 }
 101
 102 func BenchmarkHandshake(b *testing.B) {
 103         for range iter.N(b.N) {
 104                 allHandshakeTests(b)
 105         }
 106 }
 107
 108 type trackReader struct {
 109         r io.Reader
 110         n int64
 111 }
 112
 113 func (me *trackReader) Read(b []byte) (n int, err error) {
 114         n, err = me.r.Read(b)
 115         me.n += int64(n)
 116         return
 117 }
 118
 119 func TestReceiveRandomData(t *testing.T) {
 120         tr := trackReader{rand.Reader, 0}
 121         ReceiveHandshake(readWriter{&tr, ioutil.Discard}, nil)
 122 }