280 lines
6.6 KiB
Go
280 lines
6.6 KiB
Go
package mse
|
|
|
|
import (
|
|
"bytes"
|
|
"crypto/rand"
|
|
"crypto/rc4"
|
|
"io"
|
|
"io/ioutil"
|
|
"net"
|
|
"sync"
|
|
"testing"
|
|
|
|
_ "github.com/anacrolix/envpprof"
|
|
"github.com/stretchr/testify/assert"
|
|
"github.com/stretchr/testify/require"
|
|
)
|
|
|
|
func sliceIter(skeys [][]byte) SecretKeyIter {
|
|
return func(callback func([]byte) bool) {
|
|
for _, sk := range skeys {
|
|
if !callback(sk) {
|
|
break
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
func TestReadUntil(t *testing.T) {
|
|
test := func(data, until string, leftover int, expectedErr error) {
|
|
r := bytes.NewReader([]byte(data))
|
|
err := readUntil(r, []byte(until))
|
|
if err != expectedErr {
|
|
t.Fatal(err)
|
|
}
|
|
if r.Len() != leftover {
|
|
t.Fatal(r.Len())
|
|
}
|
|
}
|
|
test("feakjfeafeafegbaabc00", "abc", 2, nil)
|
|
test("feakjfeafeafegbaadc00", "abc", 0, io.EOF)
|
|
}
|
|
|
|
func TestSuffixMatchLen(t *testing.T) {
|
|
test := func(a, b string, expected int) {
|
|
actual := suffixMatchLen([]byte(a), []byte(b))
|
|
if actual != expected {
|
|
t.Fatalf("expected %d, got %d for %q and %q", expected, actual, a, b)
|
|
}
|
|
}
|
|
test("hello", "world", 0)
|
|
test("hello", "lo", 2)
|
|
test("hello", "llo", 3)
|
|
test("hello", "hell", 0)
|
|
test("hello", "helloooo!", 5)
|
|
test("hello", "lol!", 2)
|
|
test("hello", "mondo", 0)
|
|
test("mongo", "webscale", 0)
|
|
test("sup", "person", 1)
|
|
}
|
|
|
|
func handshakeTest(t testing.TB, ia []byte, aData, bData string, cryptoProvides CryptoMethod, cryptoSelect CryptoSelector) {
|
|
a, b := net.Pipe()
|
|
wg := sync.WaitGroup{}
|
|
wg.Add(2)
|
|
go func() {
|
|
defer wg.Done()
|
|
a, cm, err := InitiateHandshake(a, []byte("yep"), ia, cryptoProvides)
|
|
require.NoError(t, err)
|
|
assert.Equal(t, cryptoSelect(cryptoProvides), cm)
|
|
go a.Write([]byte(aData))
|
|
|
|
var msg [20]byte
|
|
n, _ := a.Read(msg[:])
|
|
if n != len(bData) {
|
|
t.FailNow()
|
|
}
|
|
// t.Log(string(msg[:n]))
|
|
}()
|
|
go func() {
|
|
defer wg.Done()
|
|
res := ReceiveHandshakeEx(b, sliceIter([][]byte{[]byte("nope"), []byte("yep"), []byte("maybe")}), cryptoSelect)
|
|
require.NoError(t, res.error)
|
|
assert.EqualValues(t, "yep", res.SecretKey)
|
|
b := res.ReadWriter
|
|
assert.Equal(t, cryptoSelect(cryptoProvides), res.CryptoMethod)
|
|
go b.Write([]byte(bData))
|
|
// Need to be exact here, as there are several reads, and net.Pipe is most synchronous.
|
|
msg := make([]byte, len(ia)+len(aData))
|
|
n, _ := io.ReadFull(b, msg)
|
|
if n != len(msg) {
|
|
t.FailNow()
|
|
}
|
|
// t.Log(string(msg[:n]))
|
|
}()
|
|
wg.Wait()
|
|
a.Close()
|
|
b.Close()
|
|
}
|
|
|
|
func allHandshakeTests(t testing.TB, provides CryptoMethod, selector CryptoSelector) {
|
|
handshakeTest(t, []byte("jump the gun, "), "hello world", "yo dawg", provides, selector)
|
|
handshakeTest(t, nil, "hello world", "yo dawg", provides, selector)
|
|
handshakeTest(t, []byte{}, "hello world", "yo dawg", provides, selector)
|
|
}
|
|
|
|
func TestHandshakeDefault(t *testing.T) {
|
|
allHandshakeTests(t, AllSupportedCrypto, DefaultCryptoSelector)
|
|
t.Logf("crypto provides encountered: %s", cryptoProvidesCount)
|
|
}
|
|
|
|
func TestHandshakeSelectPlaintext(t *testing.T) {
|
|
allHandshakeTests(t, AllSupportedCrypto, func(CryptoMethod) CryptoMethod { return CryptoMethodPlaintext })
|
|
}
|
|
|
|
func BenchmarkHandshakeDefault(b *testing.B) {
|
|
for i := 0; i < b.N; i += 1 {
|
|
allHandshakeTests(b, AllSupportedCrypto, DefaultCryptoSelector)
|
|
}
|
|
}
|
|
|
|
type trackReader struct {
|
|
r io.Reader
|
|
n int64
|
|
}
|
|
|
|
func (tr *trackReader) Read(b []byte) (n int, err error) {
|
|
n, err = tr.r.Read(b)
|
|
tr.n += int64(n)
|
|
return
|
|
}
|
|
|
|
func TestReceiveRandomData(t *testing.T) {
|
|
tr := trackReader{rand.Reader, 0}
|
|
_, _, err := ReceiveHandshake(readWriter{&tr, ioutil.Discard}, nil, DefaultCryptoSelector)
|
|
// No skey matches
|
|
require.Error(t, err)
|
|
// Establishing S, and then reading the maximum padding for giving up on
|
|
// synchronizing.
|
|
require.EqualValues(t, 96+532, tr.n)
|
|
}
|
|
|
|
func fillRand(t testing.TB, bs ...[]byte) {
|
|
for _, b := range bs {
|
|
_, err := rand.Read(b)
|
|
require.NoError(t, err)
|
|
}
|
|
}
|
|
|
|
func readAndWrite(rw io.ReadWriter, r []byte, w []byte) error {
|
|
var wg sync.WaitGroup
|
|
wg.Add(1)
|
|
var wErr error
|
|
go func() {
|
|
defer wg.Done()
|
|
_, wErr = rw.Write(w)
|
|
}()
|
|
_, err := io.ReadFull(rw, r)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
wg.Wait()
|
|
return wErr
|
|
}
|
|
|
|
func benchmarkStream(t *testing.B, crypto CryptoMethod) {
|
|
ia := make([]byte, 0x1000)
|
|
a := make([]byte, 1<<20)
|
|
b := make([]byte, 1<<20)
|
|
fillRand(t, ia, a, b)
|
|
t.StopTimer()
|
|
t.SetBytes(int64(len(ia) + len(a) + len(b)))
|
|
t.ResetTimer()
|
|
for i := 0; i < t.N; i += 1 {
|
|
ac, bc := net.Pipe()
|
|
ar := make([]byte, len(b))
|
|
br := make([]byte, len(ia)+len(a))
|
|
t.StartTimer()
|
|
var wg sync.WaitGroup
|
|
wg.Add(1)
|
|
go func() {
|
|
defer ac.Close()
|
|
defer wg.Done()
|
|
rw, _, err := InitiateHandshake(ac, []byte("cats"), ia, crypto)
|
|
require.NoError(t, err)
|
|
require.NoError(t, readAndWrite(rw, ar, a))
|
|
}()
|
|
func() {
|
|
defer bc.Close()
|
|
rw, _, err := ReceiveHandshake(bc, sliceIter([][]byte{[]byte("cats")}), func(CryptoMethod) CryptoMethod { return crypto })
|
|
require.NoError(t, err)
|
|
require.NoError(t, readAndWrite(rw, br, b))
|
|
}()
|
|
wg.Wait()
|
|
t.StopTimer()
|
|
if !bytes.Equal(ar, b) {
|
|
t.Fatalf("A read the wrong bytes")
|
|
}
|
|
if !bytes.Equal(br[:len(ia)], ia) {
|
|
t.Fatalf("B read the wrong IA")
|
|
}
|
|
if !bytes.Equal(br[len(ia):], a) {
|
|
t.Fatalf("B read the wrong A")
|
|
}
|
|
// require.Equal(t, b, ar)
|
|
// require.Equal(t, ia, br[:len(ia)])
|
|
// require.Equal(t, a, br[len(ia):])
|
|
}
|
|
}
|
|
|
|
func BenchmarkStreamRC4(t *testing.B) {
|
|
benchmarkStream(t, CryptoMethodRC4)
|
|
}
|
|
|
|
func BenchmarkStreamPlaintext(t *testing.B) {
|
|
benchmarkStream(t, CryptoMethodPlaintext)
|
|
}
|
|
|
|
func BenchmarkPipeRC4(t *testing.B) {
|
|
key := make([]byte, 20)
|
|
n, _ := rand.Read(key)
|
|
require.Equal(t, len(key), n)
|
|
var buf bytes.Buffer
|
|
c, err := rc4.NewCipher(key)
|
|
require.NoError(t, err)
|
|
r := cipherReader{
|
|
c: c,
|
|
r: &buf,
|
|
}
|
|
c, err = rc4.NewCipher(key)
|
|
require.NoError(t, err)
|
|
w := cipherWriter{
|
|
c: c,
|
|
w: &buf,
|
|
}
|
|
a := make([]byte, 0x1000)
|
|
n, _ = io.ReadFull(rand.Reader, a)
|
|
require.Equal(t, len(a), n)
|
|
b := make([]byte, len(a))
|
|
t.SetBytes(int64(len(a)))
|
|
t.ResetTimer()
|
|
for i := 0; i < t.N; i += 1 {
|
|
n, _ = w.Write(a)
|
|
if n != len(a) {
|
|
t.FailNow()
|
|
}
|
|
n, _ = r.Read(b)
|
|
if n != len(b) {
|
|
t.FailNow()
|
|
}
|
|
if !bytes.Equal(a, b) {
|
|
t.FailNow()
|
|
}
|
|
}
|
|
}
|
|
|
|
func BenchmarkSkeysReceive(b *testing.B) {
|
|
var skeys [][]byte
|
|
for i := 0; i < 100000; i += 1 {
|
|
skeys = append(skeys, make([]byte, 20))
|
|
}
|
|
fillRand(b, skeys...)
|
|
initSkey := skeys[len(skeys)/2]
|
|
//c := qt.New(b)
|
|
b.ReportAllocs()
|
|
b.ResetTimer()
|
|
for i := 0; i < b.N; i += 1 {
|
|
initiator, receiver := net.Pipe()
|
|
go func() {
|
|
_, _, err := InitiateHandshake(initiator, initSkey, nil, AllSupportedCrypto)
|
|
if err != nil {
|
|
panic(err)
|
|
}
|
|
}()
|
|
res := ReceiveHandshakeEx(receiver, sliceIter(skeys), DefaultCryptoSelector)
|
|
if res.error != nil {
|
|
panic(res.error)
|
|
}
|
|
}
|
|
}
|