1  package brontide
  2  
  3  import (
  4  	"bytes"
  5  	"io"
  6  	"math"
  7  	"math/rand"
  8  	"testing"
  9  	"time"
 10  
 11  	"github.com/stretchr/testify/require"
 12  )
 13  
 14  func BenchmarkReadHeaderAndBody(t *testing.B) {
 15  	// Create a test connection, grabbing either side of the connection
 16  	// into local variables. If the initial crypto handshake fails, then
 17  	// we'll get a non-nil error here.
 18  	localConn, remoteConn, err := establishTestConnection(t)
 19  	require.NoError(t, err, "unable to establish test connection: %v", err)
 20  
 21  	rand.Seed(time.Now().Unix())
 22  
 23  	noiseRemoteConn := remoteConn.(*Conn)
 24  	noiseLocalConn := localConn.(*Conn)
 25  
 26  	// Now that we have a local and remote side (to set up the initial
 27  	// handshake state, we'll have the remote side write out something
 28  	// similar to a large message in the protocol.
 29  	const pktSize = 60_000
 30  	msg := bytes.Repeat([]byte("a"), pktSize)
 31  	err = noiseRemoteConn.WriteMessage(msg)
 32  	require.NoError(t, err, "unable to write encrypted message: %v", err)
 33  
 34  	cipherHeader := noiseRemoteConn.noise.nextHeaderSend
 35  	cipherMsg := noiseRemoteConn.noise.nextBodySend
 36  
 37  	var (
 38  		benchErr error
 39  		msgBuf   [math.MaxUint16]byte
 40  	)
 41  
 42  	t.ReportAllocs()
 43  	t.ResetTimer()
 44  
 45  	nonceValue := noiseLocalConn.noise.recvCipher.nonce
 46  	for i := 0; i < t.N; i++ {
 47  		pktLen, benchErr := noiseLocalConn.noise.ReadHeader(
 48  			bytes.NewReader(cipherHeader),
 49  		)
 50  		require.NoError(
 51  			t, benchErr, "#%v: failed decryption: %v", i, benchErr,
 52  		)
 53  		_, benchErr = noiseLocalConn.noise.ReadBody(
 54  			bytes.NewReader(cipherMsg), msgBuf[:pktLen],
 55  		)
 56  		require.NoError(
 57  			t, benchErr, "#%v: failed decryption: %v", i, benchErr,
 58  		)
 59  
 60  		// We reset the internal nonce each time as otherwise, we'd
 61  		// continue to increment it which would cause a decryption
 62  		// failure.
 63  		noiseLocalConn.noise.recvCipher.nonce = nonceValue
 64  	}
 65  	require.NoError(t, benchErr)
 66  }
 67  
 68  // BenchmarkWriteMessage benchmarks the performance of writing a maximum-sized
 69  // message and flushing it to an io.Discard to measure the allocation and CPU
 70  // overhead of the encryption and writing logic.
 71  func BenchmarkWriteMessage(b *testing.B) {
 72  	localConn, remoteConn, err := establishTestConnection(b)
 73  	require.NoError(b, err, "unable to establish test connection: %v", err)
 74  
 75  	noiseLocalConn, ok := localConn.(*Conn)
 76  	require.True(b, ok, "expected *Conn type for localConn")
 77  
 78  	// Create the largest possible message we can write (MaxUint16 bytes).
 79  	// This is the maximum message size allowed by the protocol.
 80  	const maxMsgSize = math.MaxUint16
 81  	largeMsg := bytes.Repeat([]byte("a"), maxMsgSize)
 82  
 83  	// Use io.Discard to simulate writing to a network connection that
 84  	// continuously accepts data without needing resets.
 85  	discard := io.Discard
 86  
 87  	b.ReportAllocs()
 88  	b.ResetTimer()
 89  
 90  	for i := 0; i < b.N; i++ {
 91  		// Write our massive message, then call flush to actually write
 92  		// the encrypted message This simulates a full write operation
 93  		// to a network.
 94  		err := noiseLocalConn.noise.WriteMessage(largeMsg)
 95  		if err != nil {
 96  			b.Fatalf("WriteMessage failed: %v", err)
 97  		}
 98  		_, err = noiseLocalConn.noise.Flush(discard)
 99  		if err != nil {
100  			b.Fatalf("Flush failed: %v", err)
101  		}
102  	}
103  
104  	// We'll make sure to clean up the connections at the end of the
105  	// benchmark.
106  	b.Cleanup(func() {
107  		localConn.Close()
108  		remoteConn.Close()
109  	})
110  }