Repository: gorilla/websocket
Pull request #943 (2024-06-23 22:44:16)

Use crypto/rand for mask key
## What type of PR is this? (check all applicable)

- [ ] Refactor
- [ ] Feature
- [x] Bug Fix
- [ ] Optimization
- [ ] Documentation Update
- [ ] Go Version Update
- [ ] Dependency Update

## Description

Silence a lint warning by using crypto/rand instead of math/rand to generate frame masks.

## Related Tickets & Documents

- Related Issue #
- Closes #

## Added/updated tests?

- [x] Yes
- [ ] No, and this is why: _please replace this line with details on why tests
      have not been included_
- [ ] I need help with writing tests

## Run verifications and test

- [ ] `make verify` is passing
- [ ] `make test` is passing



Patch:
diff --git a/conn.go b/conn.go
index 9353252..44b1aed 100644
--- a/conn.go
+++ b/conn.go
@@ -6,10 +6,10 @@ package websocket
 
 import (
 	"bufio"
+	"crypto/rand"
 	"encoding/binary"
 	"errors"
 	"io"
-	"math/rand"
 	"net"
 	"strconv"
 	"strings"
@@ -180,9 +180,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 {
-	n := rand.Uint32()
-	return [4]byte{byte(n), byte(n >> 8), byte(n >> 16), byte(n >> 24)}
+	var k [4]byte
+	_, _ = io.ReadFull(maskRand, k[:])
+	return k
 }
 
 func hideTempErr(err error) error {
diff --git a/server.go b/server.go
index ff7d03a..28f2013 100644
--- a/server.go
+++ b/server.go
@@ -130,7 +130,8 @@ func (u *Upgrader) Upgrade(w http.ResponseWriter, r *http.Request, responseHeade
 	}
 
 	if !tokenListContainsValue(r.Header, "Upgrade", "websocket") {
-		return u.returnError(w, r, http.StatusBadRequest, badHandshake+"'websocket' token not found in 'Upgrade' header")
+		w.Header().Set("Upgrade", "websocket")
+		return u.returnError(w, r, http.StatusUpgradeRequired, badHandshake+"'websocket' token not found in 'Upgrade' header")
 	}
 
 	if r.Method != http.MethodGet {


Test patch:
diff --git a/client_server_test.go b/client_server_test.go
index 610fbe2..c5c08d7 100644
--- a/client_server_test.go
+++ b/client_server_test.go
@@ -497,6 +497,37 @@ func TestBadMethod(t *testing.T) {
 	}
 }
 
+func TestNoUpgrade(t *testing.T) {
+	t.Parallel()
+	s := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		ws, err := cstUpgrader.Upgrade(w, r, nil)
+		if err == nil {
+			t.Errorf("handshake succeeded, expect fail")
+			ws.Close()
+		}
+	}))
+	defer s.Close()
+
+	req, err := http.NewRequest(http.MethodGet, s.URL, strings.NewReader(""))
+	if err != nil {
+		t.Fatalf("NewRequest returned error %v", err)
+	}
+	req.Header.Set("Connection", "upgrade")
+	req.Header.Set("Sec-Websocket-Version", "13")
+
+	resp, err := http.DefaultClient.Do(req)
+	if err != nil {
+		t.Fatalf("Do returned error %v", err)
+	}
+	resp.Body.Close()
+	if u := resp.Header.Get("Upgrade"); u != "websocket" {
+		t.Errorf("Uprade response header is %q, want %q", u, "websocket")
+	}
+	if resp.StatusCode != http.StatusUpgradeRequired {
+		t.Errorf("Status = %d, want %d", resp.StatusCode, http.StatusUpgradeRequired)
+	}
+}
+
 func TestDialExtraTokensInRespHeaders(t *testing.T) {
 	s := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 		challengeKey := r.Header.Get("Sec-Websocket-Key")
diff --git a/prepared_test.go b/prepared_test.go
index 2297802..536d58d 100644
--- a/prepared_test.go
+++ b/prepared_test.go
@@ -33,6 +33,11 @@ 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
@@ -43,7 +48,7 @@ func TestPreparedMessage(t *testing.T) {
 		c.SetCompressionLevel(tt.compressionLevel)
 
 		// Seed random number generator for consistent frame mask.
-		rand.Seed(1234)
+		testRand.Seed(1234)
 
 		if err := c.WriteMessage(tt.messageType, data); err != nil {
 			t.Fatal(err)
@@ -59,7 +64,7 @@ func TestPreparedMessage(t *testing.T) {
 		copy(data, "hello world")
 
 		// Seed random number generator for consistent frame mask.
-		rand.Seed(1234)
+		testRand.Seed(1234)
 
 		buf.Reset()
 		if err := c.WritePreparedMessage(pm); err != nil {