Merge pull request #821 from nats-io/ws_fix_close_partial

[FIXED] Websocket: possible panic when decoding CLOSE frame

ws.go

@@ -53,6 +53,7 @@ const (
 	wsContinuationFrame     = 0
 	wsMaxFrameHeaderSize    = 14
 	wsMaxControlPayloadSize = 125
+	wsCloseSatusSize        = 2
 
 	// From https://tools.ietf.org/html/rfc6455#section-11.7
 	wsCloseStatusNormalClosure      = 1000
@@ -372,7 +373,6 @@ func (r *websocketReader) handleControlFrame(frameType wsOpCode, buf []byte, pos
 	var payload []byte
 	var err error
 
-	statusPos := pos
 	if rem > 0 {
 		payload, pos, err = wsGet(r.r, buf, pos, rem)
 		if err != nil {
@@ -382,17 +382,24 @@ func (r *websocketReader) handleControlFrame(frameType wsOpCode, buf []byte, pos
 	switch frameType {
 	case wsCloseMessage:
 		status := wsCloseStatusNoStatusReceived
-		body := ""
-		// If there is a payload, it should contain 2 unsigned bytes
-		// that represent the status code and then optional payload.
-		if len(payload) >= 2 {
-			status = int(binary.BigEndian.Uint16(buf[statusPos : statusPos+2]))
-			body = string(buf[statusPos+2 : statusPos+len(payload)])
-			if body != "" && !utf8.ValidString(body) {
-				// https://tools.ietf.org/html/rfc6455#section-5.5.1
-				// If body is present, it must be a valid utf8
-				status = wsCloseStatusInvalidPayloadData
-				body = "invalid utf8 body in close frame"
+		var body string
+		lp := len(payload)
+		// If there is a payload, the status is represented as a 2-byte
+		// unsigned integer (in network byte order). Then, there may be an
+		// optional body.
+		hasStatus, hasBody := lp >= wsCloseSatusSize, lp > wsCloseSatusSize
+		if hasStatus {
+			// Decode the status
+			status = int(binary.BigEndian.Uint16(payload[:wsCloseSatusSize]))
+			// Now if there is a body, capture it and make sure this is a valid UTF-8.
+			if hasBody {
+				body = string(payload[wsCloseSatusSize:])
+				if !utf8.ValidString(body) {
+					// https://tools.ietf.org/html/rfc6455#section-5.5.1
+					// If body is present, it must be a valid utf8
+					status = wsCloseStatusInvalidPayloadData
+					body = "invalid utf8 body in close frame"
+				}
 			}
 		}
 		r.nc.wsEnqueueCloseMsg(status, body)

ws_test.go

@@ -226,6 +226,29 @@ func TestWSParseControlFrames(t *testing.T) {
 	if err != io.EOF || n != 0 {
 		t.Fatalf("Error on read: n=%v err=%v", n, err)
 	}
+
+	// Write a CLOSE without payload
+	mr.buf.Write([]byte{136, 2, 3, 232})
+	n, err = r.Read(p)
+	if err != io.EOF || n != 0 {
+		t.Fatalf("Error on read: n=%v err=%v", n, err)
+	}
+
+	// Write a CLOSE with invalid status
+	mr.buf.Write([]byte{136, 1, 100})
+	n, err = r.Read(p)
+	if err != io.EOF || n != 0 {
+		t.Fatalf("Error on read: n=%v err=%v", n, err)
+	}
+
+	// Write CLOSE with valid status and payload but call with a read buffer
+	// that has capacity of 1.
+	mr.buf.Write([]byte{136, 6, 3, 232, 't', 'e', 's', 't'})
+	pl := []byte{136}
+	n, err = r.Read(pl[:])
+	if err != io.EOF || n != 0 {
+		t.Fatalf("Error on read: n=%v err=%v", n, err)
+	}
 }
 
 func TestWSParseInvalidFrames(t *testing.T) {