add testcase to check Sec-WebSocket-Key header

server_test.go

@@ -9,6 +9,7 @@ import (
 	"bytes"
 	"net"
 	"net/http"
+	"net/http/httptest"
 	"reflect"
 	"strings"
 	"testing"
@@ -72,12 +73,12 @@ func TestCheckSameOrigin(t *testing.T) {
 	}
 }
 
-type reuseTestResponseWriter struct {
+type testResponseWriter struct {
 	brw *bufio.ReadWriter
 	http.ResponseWriter
 }
 
-func (resp *reuseTestResponseWriter) Hijack() (net.Conn, *bufio.ReadWriter, error) {
+func (resp *testResponseWriter) Hijack() (net.Conn, *bufio.ReadWriter, error) {
 	return fakeNetConn{strings.NewReader(""), &bytes.Buffer{}}, resp.brw, nil
 }
 
@@ -93,7 +94,7 @@ func TestBufioReuse(t *testing.T) {
 	for i, tt := range bufioReuseTests {
 		br := bufio.NewReaderSize(strings.NewReader(""), tt.n)
 		bw := bufio.NewWriterSize(&bytes.Buffer{}, tt.n)
-		resp := &reuseTestResponseWriter{
+		resp := &testResponseWriter{
 			brw: bufio.NewReadWriter(br, bw),
 		}
 		upgrader := Upgrader{}
@@ -117,3 +118,37 @@ func TestBufioReuse(t *testing.T) {
 		}
 	}
 }
+
+var secWebSocketKeyTest = []struct {
+	expect string
+	key    string
+}{
+	{"", "dGhlIHNhbXBsZSBub25jZQ=="},
+	{"websocket: not a websocket handshake: 'Sec-WebSocket-Key' header is missing or blank", ""},
+	{"websocket: not a websocket handshake: 'Sec-WebSocket-Key' header is illegal base64 data", "IllegalBase64Data"},
+	{"websocket: not a websocket handshake: 'Sec-WebSocket-Key' header must be 16-byte when decoded from base64", "WHQ4eXhscUtKYjBvOGN3WEdtOEQ="},
+}
+
+func TestSecWebSocketKeyTest(t *testing.T) {
+	for _, tt := range secWebSocketKeyTest {
+		br := bufio.NewReaderSize(strings.NewReader(""), 0)
+		bw := bufio.NewWriterSize(&bytes.Buffer{}, 0)
+		w := &testResponseWriter{
+			bufio.NewReadWriter(br, bw),
+			httptest.NewRecorder(),
+		}
+		upgrader := Upgrader{}
+		_, err := upgrader.Upgrade(w, &http.Request{
+			Method: http.MethodGet,
+			Header: http.Header{
+				"Upgrade":               []string{"websocket"},
+				"Connection":            []string{"upgrade"},
+				"Sec-Websocket-Key":     []string{tt.key},
+				"Sec-Websocket-Version": []string{"13"},
+			}}, nil)
+
+		if err != nil && err.Error() != tt.expect {
+			t.Errorf("Upgrade return error=%v, want=%v", err, tt.expect)
+		}
+	}
+}