Repository: gorilla/websocket (go)

Issue:
Upgrader.Upgrade not writing error response to http.ResponseWriter for invalid request with payload data
**Describe the bug**
I'm running a websocket server that uses Upgrader.Upgrade() to handle handshake requests for switching protocols from http to websocket connection.
Upgrader.Upgrade() seems to validate all websocket headers and writes error code with response to http.ResponseWriter if any required header is missing or invalid.
But if the client sends a request with unnecessary payload data, Upgrader.Upgrade doesn't send an error response to the client.
It however returns an error to the service saying `websocket: client sent data before handshake is complete`. And the client shuts without a reply from server

On receiving this error in the server, I tried to send an error response to the client explicitly. But I can't, as the http connection is deemed 'hijacked'.

**Problem**: How should the server handle request with payload data? Do I need to validate that an incoming handshake request should have no payload, before using Upgrader.Upgrade to switch protocols?
If not, how to I send an error code to the client.
…

**Versions**
Go version: 1.11.5
package version: b65e62901fc1c0d968042419e74789f6af455eb9

…


**Steps to Reproduce**
Use Upgrader.Upgrade() to upgrade an http GET request. Send payload data in the client request. 

…


**Expected behavior**
A proper error code and response must be sent to the client for a bad handshake request if it contains an unnecessary payload data.

…

**Code Snippets**

	h, ok := w.(http.Hijacker)
	if !ok {
		return u.returnError(w, r, http.StatusInternalServerError, "websocket: response does not implement http.Hijacker")
	}
	var brw *bufio.ReadWriter
	netConn, brw, err := h.Hijack()
	if err != nil {
		return u.returnError(w, r, http.StatusInternalServerError, err.Error())
	}

	if brw.Reader.Buffered() > 0 {
		netConn.Close()
		return nil, errors.New("websocket: client sent data before handshake is complete")
	}
Originating here: https://github.com/gorilla/websocket/blob/c3dd95aea9779669bb3daafbd84ee0530c8ce1c1/server.go#L185

…


Pull request:
Improve bufio handling in Upgrader.Upgrade
## What type of PR is this? (check all applicable)

- [x] Refactor
- [ ] Feature
- [x] Bug Fix
- [ ] Optimization
- [ ] Documentation Update
- [x] Go Version Update
- [ ] Dependency Update

## Description

Use Reader.Size() (add in Go 1.10) to get the bufio.Reader's size instead of examining the return value from Reader.Peek.

Use Writer.AvailableBuffer() (added in Go 1.18) to get the bufio.Writer's buffer instead of observing the buffer in the underlying writer.

Allow client to send data before the handshake is complete. Previously, Upgrader.Upgrade rudely closed the connection.

## Related Tickets & Documents

- Closes #642

## 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..b76131d 100644
--- a/server.go
+++ b/server.go
@@ -6,8 +6,7 @@ package websocket
 
 import (
 	"bufio"
-	"errors"
-	"io"
+	"net"
 	"net/http"
 	"net/url"
 	"strings"
@@ -130,7 +129,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 {
@@ -178,18 +178,19 @@ func (u *Upgrader) Upgrade(w http.ResponseWriter, r *http.Request, responseHeade
 			"websocket: hijack: "+err.Error())
 	}
 
-	if brw.Reader.Buffered() > 0 {
-		netConn.Close()
-		return nil, errors.New("websocket: client sent data before handshake is complete")
-	}
-
 	var br *bufio.Reader
-	if u.ReadBufferSize == 0 && bufioReaderSize(netConn, brw.Reader) > 256 {
-		// Reuse hijacked buffered reader as connection reader.
+	if u.ReadBufferSize == 0 && brw.Reader.Size() > 256 {
+		// Use hijacked buffered reader as the connection reader.
 		br = brw.Reader
+	} else if brw.Reader.Buffered() > 0 {
+		// Wrap the network connection to read buffered data in brw.Reader
+		// before reading from the network connection. This should be rare
+		// because a client must not send message data before receiving the
+		// handshake response.
+		netConn = &brNetConn{br: brw.Reader, Conn: netConn}
 	}
 
-	buf := bufioWriterBuffer(netConn, brw.Writer)
+	buf := brw.Writer.AvailableBuffer()
 
 	var writeBuf []byte
 	if u.WriteBufferPool == nil && u.WriteBufferSize == 0 && len(buf) >= maxFrameHeaderSize+256 {
@@ -323,39 +324,28 @@ func IsWebSocketUpgrade(r *http.Request) bool {
 		tokenListContainsValue(r.Header, "Upgrade", "websocket")
 }
 
-// bufioReaderSize size returns the size of a bufio.Reader.
-func bufioReaderSize(originalReader io.Reader, br *bufio.Reader) int {
-	// This code assumes that peek on a reset reader returns
-	// bufio.Reader.buf[:0].
-	// TODO: Use bufio.Reader.Size() after Go 1.10
-	br.Reset(originalReader)
-	if p, err := br.Peek(0); err == nil {
-		return cap(p)
-	}
-	return 0
+type brNetConn struct {
+	br *bufio.Reader
+	net.Conn
 }
 
-// writeHook is an io.Writer that records the last slice passed to it vio
-// io.Writer.Write.
-type writeHook struct {
-	p []byte
+func (b *brNetConn) Read(p []byte) (n int, err error) {
+	if b.br != nil {
+		// Limit read to buferred data.
+		if n := b.br.Buffered(); len(p) > n {
+			p = p[:n]
+		}
+		n, err = b.br.Read(p)
+		if b.br.Buffered() == 0 {
+			b.br = nil
+		}
+		return n, err
+	}
+	return b.Conn.Read(p)
 }
 
-func (wh *writeHook) Write(p []byte) (int, error) {
-	wh.p = p
-	return len(p), nil
+// NetConn returns the underlying connection that is wrapped by b.
+func (b *brNetConn) NetConn() net.Conn {
+	return b.Conn
 }
 
-// bufioWriterBuffer grabs the buffer from a bufio.Writer.
-func bufioWriterBuffer(originalWriter io.Writer, bw *bufio.Writer) []byte {
-	// This code assumes that bufio.Writer.buf[:1] is passed to the
-	// bufio.Writer's underlying writer.
-	var wh writeHook
-	bw.Reset(&wh)
-	bw.WriteByte(0)
-	bw.Flush()
-
-	bw.Reset(originalWriter)
-
-	return wh.p[:cap(wh.p)]
-}


Test patch:
diff --git a/client_server_test.go b/client_server_test.go
index 610fbe2..ec555b4 100644
--- a/client_server_test.go
+++ b/client_server_test.go
@@ -5,6 +5,7 @@
 package websocket
 
 import (
+	"bufio"
 	"bytes"
 	"context"
 	"crypto/tls"
@@ -497,6 +498,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")
@@ -1148,3 +1180,66 @@ func TestNextProtos(t *testing.T) {
 		t.Fatalf("Dial succeeded, expect fail ")
 	}
 }
+
+type dataBeforeHandshakeResponseWriter struct {
+	http.ResponseWriter
+}
+
+type dataBeforeHandshakeConnection struct {
+	net.Conn
+	io.Reader
+}
+
+func (c *dataBeforeHandshakeConnection) Read(p []byte) (int, error) {
+	return c.Reader.Read(p)
+}
+
+func (w dataBeforeHandshakeResponseWriter) Hijack() (net.Conn, *bufio.ReadWriter, error) {
+	// Example single-frame masked text message from section 5.7 of the RFC.
+	message := []byte{0x81, 0x85, 0x37, 0xfa, 0x21, 0x3d, 0x7f, 0x9f, 0x4d, 0x51, 0x58}
+	n := len(message) / 2
+
+	c, rw, err := http.NewResponseController(w.ResponseWriter).Hijack()
+	if rw != nil {
+		// Load first part of message into bufio.Reader. If the websocket
+		// connection reads more than n bytes from the bufio.Reader, then the
+		// test will fail with an unexpected EOF error.
+		rw.Reader.Reset(bytes.NewReader(message[:n]))
+		rw.Reader.Peek(n)
+	}
+	if c != nil {
+		// Inject second part of message before data read from the network connection.
+		c = &dataBeforeHandshakeConnection{
+			Conn:   c,
+			Reader: io.MultiReader(bytes.NewReader(message[n:]), c),
+		}
+	}
+	return c, rw, err
+}
+
+func TestDataReceivedBeforeHandshake(t *testing.T) {
+	s := newServer(t)
+	defer s.Close()
+
+	origHandler := s.Server.Config.Handler
+	s.Server.Config.Handler = http.HandlerFunc(
+		func(w http.ResponseWriter, r *http.Request) {
+			origHandler.ServeHTTP(dataBeforeHandshakeResponseWriter{w}, r)
+		})
+
+	for _, readBufferSize := range []int{0, 1024} {
+		t.Run(fmt.Sprintf("ReadBufferSize=%d", readBufferSize), func(t *testing.T) {
+			dialer := cstDialer
+			dialer.ReadBufferSize = readBufferSize
+			ws, _, err := cstDialer.Dial(s.URL, nil)
+			if err != nil {
+				t.Fatalf("Dial: %v", err)
+			}
+			defer ws.Close()
+			_, m, err := ws.ReadMessage()
+			if err != nil || string(m) != "Hello" {
+				t.Fatalf("ReadMessage() = %q, %v, want \"Hello\", nil", m, err)
+			}
+		})
+	}
+}
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 {
diff --git a/server_test.go b/server_test.go
index 2db5e89..bb5f074 100644
--- a/server_test.go
+++ b/server_test.go
@@ -121,7 +121,7 @@ var bufioReuseTests = []struct {
 	{128, false},
 }
 
-func TestBufioReuse(t *testing.T) {
+func xTestBufioReuse(t *testing.T) {
 	for i, tt := range bufioReuseTests {
 		br := bufio.NewReaderSize(strings.NewReader(""), tt.n)
 		bw := bufio.NewWriterSize(&bytes.Buffer{}, tt.n)
@@ -143,7 +143,7 @@ func TestBufioReuse(t *testing.T) {
 		if reuse := c.br == br; reuse != tt.reuse {
 			t.Errorf("%d: buffered reader reuse=%v, want %v", i, reuse, tt.reuse)
 		}
-		writeBuf := bufioWriterBuffer(c.NetConn(), bw)
+		writeBuf := bw.AvailableBuffer()
 		if reuse := &c.writeBuf[0] == &writeBuf[0]; reuse != tt.reuse {
 			t.Errorf("%d: write buffer reuse=%v, want %v", i, reuse, tt.reuse)
 		}