cmd/compile: optimize XOR masking code

[nhooyr]

In the WebSocket protocol, clients have to generate a random masking key for every frame sent to the server and XOR mask the frame payload with the key. The server will unmask it on retrieval.

The mask/unmask function looks like this https://github.com/gorilla/websocket/blob/0ec3d1bd7fe50c503d6df98ee649d81f4857c564/mask_safe.go#L9-L15

This is masking byte by byte but can be made significantly faster if the masking was done word by word. Unfortunately, the Go compiler does not do this automatically so you have to use unsafe and write this yourself: https://github.com/gorilla/websocket/blob/0ec3d1bd7fe50c503d6df98ee649d81f4857c564/mask.go#L13-L54

Benchmarks at gorilla/websocket#31

It would be nice if the Go compiler did this automatically.

[ALTree]

The two links you posted are the same (copy-paste mistake?). The 2nd one in particular does not use unsafe.

[nhooyr]

My bad, fixed.

The second link was supposed to be https://github.com/gorilla/websocket/blob/0ec3d1bd7fe50c503d6df98ee649d81f4857c564/mask.go#L13-L54

[odeke-em]

Thank you for filing this issue @nhooyr and @ALTree for the triaging!

Kindly paging @josharian @randall77 @martisch.

[crvv]

Please see #28113 #30553 and #28465

And

go/src/crypto/cipher/xor_generic.go

Line 45 in 68d4b12

func fastXORBytes(dst, a, b []byte, n int) {

If you want fast xor, there are 4 different implementations in crypto/cipher.
I don't know whether the compiler can do so much optimization.

[josharian]

The compiler already recognizes sequences of byte loads and coalesces them into a single multi-byte loads. See e.g. the amd64 rules starting at

go/src/cmd/compile/internal/ssa/gen/AMD64.rules

Line 1531 in e6ae4e8

// Combining byte loads into larger (unaligned) loads.

. We could in theory write similar rules that convert a sequence of byte-wise load/xor/store operations to an XORQmodify.

In fact, I believe there is a way of writing this that would allow the compiler to do this optimization now. This code:

import "encoding/binary"

func f(key [4]byte, b []byte) {
	k := binary.LittleEndian.Uint32(key[:])
	v := binary.LittleEndian.Uint32(b)
	binary.LittleEndian.PutUint32(b, v^k)
}

Generates (at its core, surrounded by other stuff):

MOVQ	"".b+40(SP), AX
XORL	(AX), DX
MOVL	DX, (AX)

That seems pretty good.

What do you think, @nhooyr ?

[josharian]

(And if this is seriously performance-critical, you might want to manually inline the LittleEndian rountines and also duplicate key into a uint64 so that you can do 8 bytes at a time.)

[nhooyr]

That seems very reasonable @josharian

Will give this is a shot and let you know how the performance compares.

[nhooyr]

So benchmarked and here are the results on my machine (Intel(R) Core(TM) i7-7820HQ CPU @ 2.90GHz):

$ go test -bench=BenchmarkMaskBytes
goos: darwin
goarch: amd64
pkg: github.com/gorilla/websocket
BenchmarkMaskBytes/size-2/align-4/byte-8         	500000000	         3.74 ns/op	 535.10 MB/s
BenchmarkMaskBytes/size-2/align-4/gorilla-8      	300000000	         3.93 ns/op	 508.44 MB/s
BenchmarkMaskBytes/size-2/align-4/gobwas-8       	300000000	         4.67 ns/op	 428.21 MB/s
BenchmarkMaskBytes/size-2/align-4/nhooyr-8       	300000000	         4.73 ns/op	 422.43 MB/s
BenchmarkMaskBytes/size-2/align-4/crypto/cipher-8         	100000000	        16.9 ns/op	 118.23 MB/s
BenchmarkMaskBytes/size-4/align-4/byte-8                  	300000000	         5.58 ns/op	 716.61 MB/s
BenchmarkMaskBytes/size-4/align-4/gorilla-8               	300000000	         5.56 ns/op	 719.80 MB/s
BenchmarkMaskBytes/size-4/align-4/gobwas-8                	200000000	         7.27 ns/op	 550.31 MB/s
BenchmarkMaskBytes/size-4/align-4/nhooyr-8                	200000000	         7.04 ns/op	 567.92 MB/s
BenchmarkMaskBytes/size-4/align-4/crypto/cipher-8         	100000000	        19.6 ns/op	 204.31 MB/s
BenchmarkMaskBytes/size-16/align-4/byte-8                 	100000000	        13.7 ns/op	1170.31 MB/s
BenchmarkMaskBytes/size-16/align-4/gorilla-8              	100000000	        20.2 ns/op	 793.15 MB/s
BenchmarkMaskBytes/size-16/align-4/gobwas-8               	200000000	         6.59 ns/op	2426.43 MB/s
BenchmarkMaskBytes/size-16/align-4/nhooyr-8               	100000000	        13.6 ns/op	1174.09 MB/s
BenchmarkMaskBytes/size-16/align-4/crypto/cipher-8        	100000000	        16.2 ns/op	 985.94 MB/s
BenchmarkMaskBytes/size-32/align-4/byte-8                 	100000000	        23.9 ns/op	1336.39 MB/s
BenchmarkMaskBytes/size-32/align-4/gorilla-8              	100000000	        20.6 ns/op	1551.76 MB/s
BenchmarkMaskBytes/size-32/align-4/gobwas-8               	200000000	         8.13 ns/op	3934.47 MB/s
BenchmarkMaskBytes/size-32/align-4/nhooyr-8               	100000000	        15.6 ns/op	2054.30 MB/s
BenchmarkMaskBytes/size-32/align-4/crypto/cipher-8        	100000000	        24.7 ns/op	1298.15 MB/s
BenchmarkMaskBytes/size-512/align-4/byte-8                	 5000000	       345 ns/op	1480.68 MB/s
BenchmarkMaskBytes/size-512/align-4/gorilla-8             	30000000	        53.8 ns/op	9525.00 MB/s
BenchmarkMaskBytes/size-512/align-4/gobwas-8              	50000000	        39.9 ns/op	12840.94 MB/s
BenchmarkMaskBytes/size-512/align-4/nhooyr-8              	20000000	        83.6 ns/op	6126.90 MB/s
BenchmarkMaskBytes/size-512/align-4/crypto/cipher-8       	10000000	       241 ns/op	2122.82 MB/s
BenchmarkMaskBytes/size-4096/align-4/byte-8               	  500000	      2674 ns/op	1531.63 MB/s
BenchmarkMaskBytes/size-4096/align-4/gorilla-8            	 5000000	       285 ns/op	14322.91 MB/s
BenchmarkMaskBytes/size-4096/align-4/gobwas-8             	 5000000	       282 ns/op	14478.90 MB/s
BenchmarkMaskBytes/size-4096/align-4/nhooyr-8             	 3000000	       539 ns/op	7596.43 MB/s
BenchmarkMaskBytes/size-4096/align-4/crypto/cipher-8      	 1000000	      1818 ns/op	2252.51 MB/s
PASS
ok  	github.com/gorilla/websocket	58.645s

Source code at https://github.com/nhooyr/websocket/blob/40b4f235f2e2730ad1d8b81852e7610a8251d080/mask_test.go#L136-L174

Interestingly, the crypto/cipher assembly XOR was the slowest but I might be using it wrong.

Your solution is 2x slower than gorilla and gobwas but fast enough for me as I do not want to use unsafe.

Thanks @josharian.

I'll leave this issue open in case you think there is still improvement that can be made.

[nhooyr]

How I implemented what you described:

// xor applies the WebSocket masking algorithm to p
// with the given key where the first 3 bits of pos
// are the starting position in the key.
// See https://tools.ietf.org/html/rfc6455#section-5.3
//
// The returned value is the position of the next byte
// to be used for masking in the key. This is so that
// unmasking can be performed without the entire frame.
func nhooyr(key [4]byte, keyPos int, b []byte) int {
	// If the payload is greater than 16 bytes, then it's worth
	// masking 8 bytes at a time.
	// Optimization from https://github.com/golang/go/issues/31586#issuecomment-485530859
	if len(b) > 16 {
		// We first create a key that is 8 bytes long
		// and is aligned on the keyPos correctly.
		var alignedKey [8]byte
		for i := range alignedKey {
			alignedKey[i] = key[(i+keyPos)&3]
		}
		k := binary.LittleEndian.Uint64(alignedKey[:])

		// Then we xor until b is less than 8 bytes.
		for len(b) >= 8 {
			v := binary.LittleEndian.Uint64(b)
			binary.LittleEndian.PutUint64(b, v^k)
			b = b[8:]
		}
	}

	// xor remaining bytes.
	for i := range b {
		b[i] ^= key[keyPos&3]
		keyPos++
	}
	return keyPos & 3
}