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