mse/mse_test.go

   1 package mse
   2
   3 import (
   4         "bytes"
   5         "crypto/rand"
   6         "crypto/rc4"
   7         "io"
   8         "net"
   9         "sync"
  10         "testing"
  11
  12         _ "github.com/anacrolix/envpprof"
  13         "github.com/stretchr/testify/assert"
  14         "github.com/stretchr/testify/require"
  15 )
  16
  17 func sliceIter(skeys [][]byte) SecretKeyIter {
  18         return func(callback func([]byte) bool) {
  19                 for _, sk := range skeys {
  20                         if !callback(sk) {
  21                                 break
  22                         }
  23                 }
  24         }
  25 }
  26
  27 func TestReadUntil(t *testing.T) {
  28         test := func(data, until string, leftover int, expectedErr error) {
  29                 r := bytes.NewReader([]byte(data))
  30                 err := readUntil(r, []byte(until))
  31                 if err != expectedErr {
  32                         t.Fatal(err)
  33                 }
  34                 if r.Len() != leftover {
  35                         t.Fatal(r.Len())
  36                 }
  37         }
  38         test("feakjfeafeafegbaabc00", "abc", 2, nil)
  39         test("feakjfeafeafegbaadc00", "abc", 0, io.EOF)
  40 }
  41
  42 func TestSuffixMatchLen(t *testing.T) {
  43         test := func(a, b string, expected int) {
  44                 actual := suffixMatchLen([]byte(a), []byte(b))
  45                 if actual != expected {
  46                         t.Fatalf("expected %d, got %d for %q and %q", expected, actual, a, b)
  47                 }
  48         }
  49         test("hello", "world", 0)
  50         test("hello", "lo", 2)
  51         test("hello", "llo", 3)
  52         test("hello", "hell", 0)
  53         test("hello", "helloooo!", 5)
  54         test("hello", "lol!", 2)
  55         test("hello", "mondo", 0)
  56         test("mongo", "webscale", 0)
  57         test("sup", "person", 1)
  58 }
  59
  60 func handshakeTest(t testing.TB, ia []byte, aData, bData string, cryptoProvides CryptoMethod, cryptoSelect CryptoSelector) {
  61         a, b := net.Pipe()
  62         wg := sync.WaitGroup{}
  63         wg.Add(2)
  64         go func() {
  65                 defer wg.Done()
  66                 a, cm, err := InitiateHandshake(a, []byte("yep"), ia, cryptoProvides)
  67                 require.NoError(t, err)
  68                 assert.Equal(t, cryptoSelect(cryptoProvides), cm)
  69                 go a.Write([]byte(aData))
  70
  71                 var msg [20]byte
  72                 n, _ := a.Read(msg[:])
  73                 if n != len(bData) {
  74                         t.FailNow()
  75                 }
  76                 // t.Log(string(msg[:n]))
  77         }()
  78         go func() {
  79                 defer wg.Done()
  80                 res := ReceiveHandshakeEx(b, sliceIter([][]byte{[]byte("nope"), []byte("yep"), []byte("maybe")}), cryptoSelect)
  81                 require.NoError(t, res.error)
  82                 assert.EqualValues(t, "yep", res.SecretKey)
  83                 b := res.ReadWriter
  84                 assert.Equal(t, cryptoSelect(cryptoProvides), res.CryptoMethod)
  85                 go b.Write([]byte(bData))
  86                 // Need to be exact here, as there are several reads, and net.Pipe is most synchronous.
  87                 msg := make([]byte, len(ia)+len(aData))
  88                 n, _ := io.ReadFull(b, msg)
  89                 if n != len(msg) {
  90                         t.FailNow()
  91                 }
  92                 // t.Log(string(msg[:n]))
  93         }()
  94         wg.Wait()
  95         a.Close()
  96         b.Close()
  97 }
  98
  99 func allHandshakeTests(t testing.TB, provides CryptoMethod, selector CryptoSelector) {
 100         handshakeTest(t, []byte("jump the gun, "), "hello world", "yo dawg", provides, selector)
 101         handshakeTest(t, nil, "hello world", "yo dawg", provides, selector)
 102         handshakeTest(t, []byte{}, "hello world", "yo dawg", provides, selector)
 103 }
 104
 105 func TestHandshakeDefault(t *testing.T) {
 106         allHandshakeTests(t, AllSupportedCrypto, DefaultCryptoSelector)
 107         t.Logf("crypto provides encountered: %s", cryptoProvidesCount)
 108 }
 109
 110 func TestHandshakeSelectPlaintext(t *testing.T) {
 111         allHandshakeTests(t, AllSupportedCrypto, func(CryptoMethod) CryptoMethod { return CryptoMethodPlaintext })
 112 }
 113
 114 func BenchmarkHandshakeDefault(b *testing.B) {
 115         for i := 0; i < b.N; i += 1 {
 116                 allHandshakeTests(b, AllSupportedCrypto, DefaultCryptoSelector)
 117         }
 118 }
 119
 120 type trackReader struct {
 121         r io.Reader
 122         n int64
 123 }
 124
 125 func (tr *trackReader) Read(b []byte) (n int, err error) {
 126         n, err = tr.r.Read(b)
 127         tr.n += int64(n)
 128         return
 129 }
 130
 131 func TestReceiveRandomData(t *testing.T) {
 132         tr := trackReader{rand.Reader, 0}
 133         _, _, err := ReceiveHandshake(readWriter{&tr, io.Discard}, nil, DefaultCryptoSelector)
 134         // No skey matches
 135         require.Error(t, err)
 136         // Establishing S, and then reading the maximum padding for giving up on
 137         // synchronizing.
 138         require.EqualValues(t, 96+532, tr.n)
 139 }
 140
 141 func fillRand(t testing.TB, bs ...[]byte) {
 142         for _, b := range bs {
 143                 _, err := rand.Read(b)
 144                 require.NoError(t, err)
 145         }
 146 }
 147
 148 func readAndWrite(rw io.ReadWriter, r, w []byte) error {
 149         var wg sync.WaitGroup
 150         wg.Add(1)
 151         var wErr error
 152         go func() {
 153                 defer wg.Done()
 154                 _, wErr = rw.Write(w)
 155         }()
 156         _, err := io.ReadFull(rw, r)
 157         if err != nil {
 158                 return err
 159         }
 160         wg.Wait()
 161         return wErr
 162 }
 163
 164 func benchmarkStream(t *testing.B, crypto CryptoMethod) {
 165         ia := make([]byte, 0x1000)
 166         a := make([]byte, 1<<20)
 167         b := make([]byte, 1<<20)
 168         fillRand(t, ia, a, b)
 169         t.StopTimer()
 170         t.SetBytes(int64(len(ia) + len(a) + len(b)))
 171         t.ResetTimer()
 172         for i := 0; i < t.N; i += 1 {
 173                 ac, bc := net.Pipe()
 174                 ar := make([]byte, len(b))
 175                 br := make([]byte, len(ia)+len(a))
 176                 t.StartTimer()
 177                 var wg sync.WaitGroup
 178                 wg.Add(1)
 179                 go func() {
 180                         defer ac.Close()
 181                         defer wg.Done()
 182                         rw, _, err := InitiateHandshake(ac, []byte("cats"), ia, crypto)
 183                         require.NoError(t, err)
 184                         require.NoError(t, readAndWrite(rw, ar, a))
 185                 }()
 186                 func() {
 187                         defer bc.Close()
 188                         rw, _, err := ReceiveHandshake(bc, sliceIter([][]byte{[]byte("cats")}), func(CryptoMethod) CryptoMethod { return crypto })
 189                         require.NoError(t, err)
 190                         require.NoError(t, readAndWrite(rw, br, b))
 191                 }()
 192                 wg.Wait()
 193                 t.StopTimer()
 194                 if !bytes.Equal(ar, b) {
 195                         t.Fatalf("A read the wrong bytes")
 196                 }
 197                 if !bytes.Equal(br[:len(ia)], ia) {
 198                         t.Fatalf("B read the wrong IA")
 199                 }
 200                 if !bytes.Equal(br[len(ia):], a) {
 201                         t.Fatalf("B read the wrong A")
 202                 }
 203                 // require.Equal(t, b, ar)
 204                 // require.Equal(t, ia, br[:len(ia)])
 205                 // require.Equal(t, a, br[len(ia):])
 206         }
 207 }
 208
 209 func BenchmarkStreamRC4(t *testing.B) {
 210         benchmarkStream(t, CryptoMethodRC4)
 211 }
 212
 213 func BenchmarkStreamPlaintext(t *testing.B) {
 214         benchmarkStream(t, CryptoMethodPlaintext)
 215 }
 216
 217 func BenchmarkPipeRC4(t *testing.B) {
 218         key := make([]byte, 20)
 219         n, _ := rand.Read(key)
 220         require.Equal(t, len(key), n)
 221         var buf bytes.Buffer
 222         c, err := rc4.NewCipher(key)
 223         require.NoError(t, err)
 224         r := cipherReader{
 225                 c: c,
 226                 r: &buf,
 227         }
 228         c, err = rc4.NewCipher(key)
 229         require.NoError(t, err)
 230         w := cipherWriter{
 231                 c: c,
 232                 w: &buf,
 233         }
 234         a := make([]byte, 0x1000)
 235         n, _ = io.ReadFull(rand.Reader, a)
 236         require.Equal(t, len(a), n)
 237         b := make([]byte, len(a))
 238         t.SetBytes(int64(len(a)))
 239         t.ResetTimer()
 240         for i := 0; i < t.N; i += 1 {
 241                 n, _ = w.Write(a)
 242                 if n != len(a) {
 243                         t.FailNow()
 244                 }
 245                 n, _ = r.Read(b)
 246                 if n != len(b) {
 247                         t.FailNow()
 248                 }
 249                 if !bytes.Equal(a, b) {
 250                         t.FailNow()
 251                 }
 252         }
 253 }
 254
 255 func BenchmarkSkeysReceive(b *testing.B) {
 256         var skeys [][]byte
 257         for i := 0; i < 100000; i += 1 {
 258                 skeys = append(skeys, make([]byte, 20))
 259         }
 260         fillRand(b, skeys...)
 261         initSkey := skeys[len(skeys)/2]
 262         // c := qt.New(b)
 263         b.ReportAllocs()
 264         b.ResetTimer()
 265         for i := 0; i < b.N; i += 1 {
 266                 initiator, receiver := net.Pipe()
 267                 go func() {
 268                         _, _, err := InitiateHandshake(initiator, initSkey, nil, AllSupportedCrypto)
 269                         if err != nil {
 270                                 panic(err)
 271                         }
 272                 }()
 273                 res := ReceiveHandshakeEx(receiver, sliceIter(skeys), DefaultCryptoSelector)
 274                 if res.error != nil {
 275                         panic(res.error)
 276                 }
 277         }
 278 }