13 _ "github.com/anacrolix/envpprof"
14 "github.com/stretchr/testify/assert"
15 "github.com/stretchr/testify/require"
18 func sliceIter(skeys [][]byte) SecretKeyIter {
19 return func(callback func([]byte) bool) {
20 for _, sk := range skeys {
28 func TestReadUntil(t *testing.T) {
29 test := func(data, until string, leftover int, expectedErr error) {
30 r := bytes.NewReader([]byte(data))
31 err := readUntil(r, []byte(until))
32 if err != expectedErr {
35 if r.Len() != leftover {
39 test("feakjfeafeafegbaabc00", "abc", 2, nil)
40 test("feakjfeafeafegbaadc00", "abc", 0, io.EOF)
43 func TestSuffixMatchLen(t *testing.T) {
44 test := func(a, b string, expected int) {
45 actual := suffixMatchLen([]byte(a), []byte(b))
46 if actual != expected {
47 t.Fatalf("expected %d, got %d for %q and %q", expected, actual, a, b)
50 test("hello", "world", 0)
51 test("hello", "lo", 2)
52 test("hello", "llo", 3)
53 test("hello", "hell", 0)
54 test("hello", "helloooo!", 5)
55 test("hello", "lol!", 2)
56 test("hello", "mondo", 0)
57 test("mongo", "webscale", 0)
58 test("sup", "person", 1)
61 func handshakeTest(t testing.TB, ia []byte, aData, bData string, cryptoProvides CryptoMethod, cryptoSelect CryptoSelector) {
63 wg := sync.WaitGroup{}
67 a, cm, err := InitiateHandshake(a, []byte("yep"), ia, cryptoProvides)
68 require.NoError(t, err)
69 assert.Equal(t, cryptoSelect(cryptoProvides), cm)
70 go a.Write([]byte(aData))
73 n, _ := a.Read(msg[:])
77 // t.Log(string(msg[:n]))
81 res := ReceiveHandshakeEx(b, sliceIter([][]byte{[]byte("nope"), []byte("yep"), []byte("maybe")}), cryptoSelect)
82 require.NoError(t, res.error)
83 assert.EqualValues(t, "yep", res.SecretKey)
85 assert.Equal(t, cryptoSelect(cryptoProvides), res.CryptoMethod)
86 go b.Write([]byte(bData))
87 // Need to be exact here, as there are several reads, and net.Pipe is most synchronous.
88 msg := make([]byte, len(ia)+len(aData))
89 n, _ := io.ReadFull(b, msg)
93 // t.Log(string(msg[:n]))
100 func allHandshakeTests(t testing.TB, provides CryptoMethod, selector CryptoSelector) {
101 handshakeTest(t, []byte("jump the gun, "), "hello world", "yo dawg", provides, selector)
102 handshakeTest(t, nil, "hello world", "yo dawg", provides, selector)
103 handshakeTest(t, []byte{}, "hello world", "yo dawg", provides, selector)
106 func TestHandshakeDefault(t *testing.T) {
107 allHandshakeTests(t, AllSupportedCrypto, DefaultCryptoSelector)
108 t.Logf("crypto provides encountered: %s", cryptoProvidesCount)
111 func TestHandshakeSelectPlaintext(t *testing.T) {
112 allHandshakeTests(t, AllSupportedCrypto, func(CryptoMethod) CryptoMethod { return CryptoMethodPlaintext })
115 func BenchmarkHandshakeDefault(b *testing.B) {
116 for i := 0; i < b.N; i += 1 {
117 allHandshakeTests(b, AllSupportedCrypto, DefaultCryptoSelector)
121 type trackReader struct {
126 func (tr *trackReader) Read(b []byte) (n int, err error) {
127 n, err = tr.r.Read(b)
132 func TestReceiveRandomData(t *testing.T) {
133 tr := trackReader{rand.Reader, 0}
134 _, _, err := ReceiveHandshake(readWriter{&tr, ioutil.Discard}, nil, DefaultCryptoSelector)
136 require.Error(t, err)
137 // Establishing S, and then reading the maximum padding for giving up on
139 require.EqualValues(t, 96+532, tr.n)
142 func fillRand(t testing.TB, bs ...[]byte) {
143 for _, b := range bs {
144 _, err := rand.Read(b)
145 require.NoError(t, err)
149 func readAndWrite(rw io.ReadWriter, r, w []byte) error {
150 var wg sync.WaitGroup
155 _, wErr = rw.Write(w)
157 _, err := io.ReadFull(rw, r)
165 func benchmarkStream(t *testing.B, crypto CryptoMethod) {
166 ia := make([]byte, 0x1000)
167 a := make([]byte, 1<<20)
168 b := make([]byte, 1<<20)
169 fillRand(t, ia, a, b)
171 t.SetBytes(int64(len(ia) + len(a) + len(b)))
173 for i := 0; i < t.N; i += 1 {
175 ar := make([]byte, len(b))
176 br := make([]byte, len(ia)+len(a))
178 var wg sync.WaitGroup
183 rw, _, err := InitiateHandshake(ac, []byte("cats"), ia, crypto)
184 require.NoError(t, err)
185 require.NoError(t, readAndWrite(rw, ar, a))
189 rw, _, err := ReceiveHandshake(bc, sliceIter([][]byte{[]byte("cats")}), func(CryptoMethod) CryptoMethod { return crypto })
190 require.NoError(t, err)
191 require.NoError(t, readAndWrite(rw, br, b))
195 if !bytes.Equal(ar, b) {
196 t.Fatalf("A read the wrong bytes")
198 if !bytes.Equal(br[:len(ia)], ia) {
199 t.Fatalf("B read the wrong IA")
201 if !bytes.Equal(br[len(ia):], a) {
202 t.Fatalf("B read the wrong A")
204 // require.Equal(t, b, ar)
205 // require.Equal(t, ia, br[:len(ia)])
206 // require.Equal(t, a, br[len(ia):])
210 func BenchmarkStreamRC4(t *testing.B) {
211 benchmarkStream(t, CryptoMethodRC4)
214 func BenchmarkStreamPlaintext(t *testing.B) {
215 benchmarkStream(t, CryptoMethodPlaintext)
218 func BenchmarkPipeRC4(t *testing.B) {
219 key := make([]byte, 20)
220 n, _ := rand.Read(key)
221 require.Equal(t, len(key), n)
223 c, err := rc4.NewCipher(key)
224 require.NoError(t, err)
229 c, err = rc4.NewCipher(key)
230 require.NoError(t, err)
235 a := make([]byte, 0x1000)
236 n, _ = io.ReadFull(rand.Reader, a)
237 require.Equal(t, len(a), n)
238 b := make([]byte, len(a))
239 t.SetBytes(int64(len(a)))
241 for i := 0; i < t.N; i += 1 {
250 if !bytes.Equal(a, b) {
256 func BenchmarkSkeysReceive(b *testing.B) {
258 for i := 0; i < 100000; i += 1 {
259 skeys = append(skeys, make([]byte, 20))
261 fillRand(b, skeys...)
262 initSkey := skeys[len(skeys)/2]
266 for i := 0; i < b.N; i += 1 {
267 initiator, receiver := net.Pipe()
269 _, _, err := InitiateHandshake(initiator, initSkey, nil, AllSupportedCrypto)
274 res := ReceiveHandshakeEx(receiver, sliceIter(skeys), DefaultCryptoSelector)
275 if res.error != nil {