Handle multiple sec-websocket-keys

Author: Emyrk

accept.go

@@ -185,13 +185,17 @@ func verifyClientRequest(w http.ResponseWriter, r *http.Request) (errCode int, _
 		return http.StatusBadRequest, fmt.Errorf("unsupported WebSocket protocol version (only 13 is supported): %q", r.Header.Get("Sec-WebSocket-Version"))
 	}
 
-	websocketSecKey := r.Header.Get("Sec-WebSocket-Key")
-	// The RFC states to remove any leading or trailing whitespace.
-	websocketSecKey = strings.TrimSpace(websocketSecKey)
-	if websocketSecKey == "" {
+	websocketSecKeys := r.Header.Values("Sec-WebSocket-Key")
+	if len(websocketSecKeys) == 0 {
 		return http.StatusBadRequest, errors.New("WebSocket protocol violation: missing Sec-WebSocket-Key")
 	}
 
+	if len(websocketSecKeys) > 1 {
+		return http.StatusBadRequest, errors.New("WebSocket protocol violation: multiple Sec-WebSocket-Key headers")
+	}
+
+	// The RFC states to remove any leading or trailing whitespace.
+	websocketSecKey := strings.TrimSpace(websocketSecKeys[0])
 	if v, err := base64.StdEncoding.DecodeString(websocketSecKey); err != nil || len(v) != 16 {
 		return http.StatusBadRequest, fmt.Errorf("WebSocket protocol violation: invalid Sec-WebSocket-Key %q, must be a 16 byte base64 encoded string", websocketSecKey)
 	}

accept_test.go

@@ -185,6 +185,14 @@ func Test_verifyClientHandshake(t *testing.T) {
 		},
 		{
 			name: "missingWebSocketKey",
+			h: map[string]string{
+				"Connection":            "Upgrade",
+				"Upgrade":               "websocket",
+				"Sec-WebSocket-Version": "13",
+			},
+		},
+		{
+			name: "emptyWebSocketKey",
 			h: map[string]string{
 				"Connection":            "Upgrade",
 				"Upgrade":               "websocket",
@@ -210,6 +218,18 @@ func Test_verifyClientHandshake(t *testing.T) {
 				"Sec-WebSocket-Key":     "notbase64",
 			},
 		},
+		{
+			name: "extraWebSocketKey",
+			h: map[string]string{
+				"Connection":            "Upgrade",
+				"Upgrade":               "websocket",
+				"Sec-WebSocket-Version": "13",
+				// Kinda cheeky, but http headers are case-insensitive.
+				// If 2 sec keys are present, this is a failure condition.
+				"Sec-WebSocket-Key": xrand.Base64(16),
+				"sec-webSocket-key": xrand.Base64(16),
+			},
+		},
 		{
 			name: "badHTTPVersion",
 			h: map[string]string{
@@ -256,7 +276,7 @@ func Test_verifyClientHandshake(t *testing.T) {
 			}
 
 			for k, v := range tc.h {
-				r.Header.Set(k, v)
+				r.Header.Add(k, v)
 			}
 
 			_, err := verifyClientRequest(httptest.NewRecorder(), r)