Merge branch 'maskrand' of https://github.com/HaloArrow/websocket int…

…o coreydaley-patch-1

conn.go

@@ -11,7 +11,6 @@ import (
 	"errors"
 	"io"
 	"log"
-	"math/big"
 	"net"
 	"strconv"
 	"strings"
@@ -182,13 +181,16 @@ var (
 	errInvalidControlFrame = errors.New("websocket: invalid control frame")
 )
 
+// maskRand is an io.Reader for generating mask bytes. The reader is initialized
+// to crypto/rand Reader. Tests swap the reader to a math/rand reader for
+// reproducible results.
+var maskRand = rand.Reader
+
+// newMaskKey returns a new 32 bit value for masking client frames.
 func newMaskKey() [4]byte {
-	nBig, err := rand.Int(rand.Reader, big.NewInt(int64(1<<32)))
-	if err != nil {
-		log.Println(err)
-	}
-	nb := nBig.Uint64()
-	return [4]byte{byte(nb), byte(nb >> 8), byte(nb >> 16), byte(nb >> 24)}
+	var k [4]byte
+	_, _ = io.ReadFull(maskRand, k[:])
+	return k
 }
 
 func hideTempErr(err error) error {

prepared_test.go

@@ -4,18 +4,10 @@
 
 package websocket
 
-/* TODO: Test fails after changing from math/rand to crypto/rand in conn.go since we are unable to provide the same seed to the random number generator.
-=== RUN   TestPreparedMessage
-    prepared_test.go:66: write message != prepared message, got "\x81\x8e~\x06h\x02\nn\x01q^o\x1b\"\x1f&\x1cg\rr", want "\x81\x8e\x9d\f\xa9\xfa\xe9d\xc0\x89\xbde\xda\xda\xfc,ݟ\xeex"
-    prepared_test.go:66: write message != prepared message, got "\xc1\x94\x18C\xe8\xfa\x18M\xe8\v\xe77\x80\x93kc\x81\x898\"Ȏ}0\x9c\xfa", want "\xc1\x94Mkw\x11Mew\xe0\xb2\x1f\x1fx>K\x1ebm\nWe(\x18\x03\x11"
-    prepared_test.go:66: write message != prepared message, got "\xc1\x90\x1e\xebV}4\"\x9eQH#z+V\xbf~43\xc5W}", want "\xc1\x90\xd8-\xb3\xfc\xf2\xe4{Ў埪\x90y\x9b\xb5\xf5\x03\xb2\xfc"
-    prepared_test.go:66: write message != prepared message, got "\x89\x8e\x9b\x02\x15\x81\xefj|\xf2\xbbkf\xa1\xfa\"a\xe4\xe8v", want "\x89\x8e\xea\xe4)Ӟ\x8c@\xa0ʍZ\xf3\x8b\xc4]\xb6\x99\x90"
-    prepared_test.go:66: write message != prepared message, got "\x89\x8e\xd0\x1a\xa5 \xa4r\xccS\xf0s\xd6\x00\xb1:\xd1E\xa3n", want "\x89\x8e\xd0\x1f'\x9d\xa4wN\xee\xf0vT\xbd\xb1?S\xf8\xa3k"
---- FAIL: TestPreparedMessage (0.01s)
-
 import (
 	"bytes"
 	"compress/flate"
+	"math/rand"
 	"testing"
 )
 
@@ -41,16 +33,22 @@ var preparedMessageTests = []struct {
 }
 
 func TestPreparedMessage(t *testing.T) {
+	testRand := rand.New(rand.NewSource(99))
+	prevMaskRand := maskRand
+	maskRand = testRand
+	defer func() { maskRand = prevMaskRand }()
+
 	for _, tt := range preparedMessageTests {
 		var data = []byte("this is a test")
 		var buf bytes.Buffer
 		c := newTestConn(nil, &buf, tt.isServer)
 		if tt.enableWriteCompression {
 			c.newCompressionWriter = compressNoContextTakeover
 		}
-		if err := c.SetCompressionLevel(tt.compressionLevel); err != nil {
-			t.Fatal(err)
-		}
+		c.SetCompressionLevel(tt.compressionLevel)
+
+		// Seed random number generator for consistent frame mask.
+		testRand.Seed(1234)
 
 		if err := c.WriteMessage(tt.messageType, data); err != nil {
 			t.Fatal(err)
@@ -65,15 +63,17 @@ func TestPreparedMessage(t *testing.T) {
 		// Scribble on data to ensure that NewPreparedMessage takes a snapshot.
 		copy(data, "hello world")
 
+		// Seed random number generator for consistent frame mask.
+		testRand.Seed(1234)
+
 		buf.Reset()
 		if err := c.WritePreparedMessage(pm); err != nil {
 			t.Fatal(err)
 		}
 		got := buf.String()
 
 		if got != want {
-			t.Errorf("write message != prepared message, got %#v, want %#v", got, want)
+			t.Errorf("write message != prepared message for %+v", tt)
 		}
 	}
 }
-*/