From 77f110791cfaf7893f72043b54f469352b836d57 Mon Sep 17 00:00:00 2001
From: Gary Burd <gary@beagledreams.com>
Date: Thu, 20 Oct 2016 16:25:46 -0700
Subject: [PATCH] Improve mask performance

---
 bench_test.go | 19 --------------
 conn.go       | 15 -----------
 mask.go       | 61 ++++++++++++++++++++++++++++++++++++++++++
 mask_test.go  | 73 +++++++++++++++++++++++++++++++++++++++++++++++++++
 4 files changed, 134 insertions(+), 34 deletions(-)
 delete mode 100644 bench_test.go
 create mode 100644 mask.go
 create mode 100644 mask_test.go

diff --git a/bench_test.go b/bench_test.go
deleted file mode 100644
index f66fc36..0000000
--- a/bench_test.go
+++ /dev/null
@@ -1,19 +0,0 @@
-// Copyright 2014 The Gorilla WebSocket Authors. All rights reserved.
-// Use of this source code is governed by a BSD-style
-// license that can be found in the LICENSE file.
-
-package websocket
-
-import (
-	"testing"
-)
-
-func BenchmarkMaskBytes(b *testing.B) {
-	var key [4]byte
-	data := make([]byte, 1024)
-	pos := 0
-	for i := 0; i < b.N; i++ {
-		pos = maskBytes(key, pos, data)
-	}
-	b.SetBytes(int64(len(data)))
-}
diff --git a/conn.go b/conn.go
index d916c1d..5b26b53 100644
--- a/conn.go
+++ b/conn.go
@@ -10,7 +10,6 @@ import (
 	"errors"
 	"io"
 	"io/ioutil"
-	"math/rand"
 	"net"
 	"strconv"
 	"time"
@@ -218,20 +217,6 @@ func isValidReceivedCloseCode(code int) bool {
 	return validReceivedCloseCodes[code] || (code >= 3000 && code <= 4999)
 }
 
-func maskBytes(key [4]byte, pos int, b []byte) int {
-	for i := range b {
-		b[i] ^= key[pos&3]
-		pos++
-	}
-	return pos & 3
-}
-
-func newMaskKey() [4]byte {
-	n := rand.Uint32()
-	return [4]byte{byte(n), byte(n >> 8), byte(n >> 16), byte(n >> 24)}
-}
-
-// Conn represents a WebSocket connection.
 type Conn struct {
 	conn        net.Conn
 	isServer    bool
diff --git a/mask.go b/mask.go
new file mode 100644
index 0000000..68861fc
--- /dev/null
+++ b/mask.go
@@ -0,0 +1,61 @@
+// Copyright 2016 The Gorilla WebSocket Authors. All rights reserved.  Use of
+// this source code is governed by a BSD-style license that can be found in the
+// LICENSE file.
+
+package websocket
+
+import (
+	"math/rand"
+	"unsafe"
+)
+
+const wordSize = int(unsafe.Sizeof(uintptr(0)))
+
+func newMaskKey() [4]byte {
+	n := rand.Uint32()
+	return [4]byte{byte(n), byte(n >> 8), byte(n >> 16), byte(n >> 24)}
+}
+
+func maskBytes(key [4]byte, pos int, b []byte) int {
+
+	// Mask one byte at a time for small buffers.
+	if len(b) < 2*wordSize {
+		for i := range b {
+			b[i] ^= key[pos&3]
+			pos++
+		}
+		return pos & 3
+	}
+
+	// Mask one byte at a time to word boundary.
+	if n := int(uintptr(unsafe.Pointer(&b))) % wordSize; n != 0 {
+		n = wordSize - n
+		for i := range b[:n] {
+			b[i] ^= key[pos&3]
+			pos++
+		}
+		b = b[n:]
+	}
+
+	// Create aligned word size key.
+	var k [wordSize]byte
+	for i := range k {
+		k[i] = key[(pos+i)&3]
+	}
+	kw := *(*uintptr)(unsafe.Pointer(&k))
+
+	// Mask one word at a time.
+	n := (len(b) / wordSize) * wordSize
+	for i := 0; i < n; i += wordSize {
+		*(*uintptr)(unsafe.Pointer(uintptr(unsafe.Pointer(&b[0])) + uintptr(i))) ^= kw
+	}
+
+	// Mask one byte at a time for remaining bytes.
+	b = b[n:]
+	for i := range b {
+		b[i] ^= key[pos&3]
+		pos++
+	}
+
+	return pos & 3
+}
diff --git a/mask_test.go b/mask_test.go
new file mode 100644
index 0000000..de06029
--- /dev/null
+++ b/mask_test.go
@@ -0,0 +1,73 @@
+// Copyright 2016 The Gorilla WebSocket Authors. All rights reserved.  Use of
+// this source code is governed by a BSD-style license that can be found in the
+// LICENSE file.
+
+// Require 1.7 for sub-bencmarks
+// +build go1.7
+
+package websocket
+
+import (
+	"fmt"
+	"testing"
+)
+
+func maskBytesByByte(key [4]byte, pos int, b []byte) int {
+	for i := range b {
+		b[i] ^= key[pos&3]
+		pos++
+	}
+	return pos & 3
+}
+
+func notzero(b []byte) int {
+	for i := range b {
+		if b[i] != 0 {
+			return i
+		}
+	}
+	return -1
+}
+
+func TestMaskBytes(t *testing.T) {
+	key := [4]byte{1, 2, 3, 4}
+	for size := 1; size <= 1024; size++ {
+		for align := 0; align < wordSize; align++ {
+			for pos := 0; pos < 4; pos++ {
+				b := make([]byte, size+align)[align:]
+				maskBytes(key, pos, b)
+				maskBytesByByte(key, pos, b)
+				if i := notzero(b); i >= 0 {
+					t.Errorf("size:%d, align:%d, pos:%d, offset:%d", size, align, pos, i)
+				}
+			}
+		}
+	}
+}
+
+func BenchmarkMaskBytes(b *testing.B) {
+	for _, size := range []int{2, 4, 8, 16, 32, 512, 1024} {
+		b.Run(fmt.Sprintf("size-%d", size), func(b *testing.B) {
+			for _, align := range []int{wordSize / 2} {
+				b.Run(fmt.Sprintf("align-%d", align), func(b *testing.B) {
+					for _, fn := range []struct {
+						name string
+						fn   func(key [4]byte, pos int, b []byte) int
+					}{
+						{"byte", maskBytesByByte},
+						{"word", maskBytes},
+					} {
+						b.Run(fn.name, func(b *testing.B) {
+							key := newMaskKey()
+							data := make([]byte, size+align)[align:]
+							for i := 0; i < b.N; i++ {
+								fn.fn(key, 0, data)
+							}
+							b.SetBytes(int64(len(data)))
+						})
+					}
+				})
+			}
+		})
+	}
+}