From 75f568a1eca74699a06330e8068ebe5de19160e3 Mon Sep 17 00:00:00 2001
From: xdrm-brackets <doowap31@gmail.com>
Date: Sun, 6 May 2018 01:37:02 +0200
Subject: [PATCH] manage invalid Close codes + close() now can work without
 sending the CLOSE frame (if status=NONE)

[TODO] manage fragmentation
---
 ws/client.go | 67 ++++++++++++++++++++++++++++++++++------------------
 1 file changed, 44 insertions(+), 23 deletions(-)

diff --git a/ws/client.go b/ws/client.go
index 7a42604..7a544ff 100644
--- a/ws/client.go
+++ b/ws/client.go
@@ -164,6 +164,13 @@ func clientReader(c *client){
 			if msg.Size > 2 && !utf8.Valid(msg.Data[2:]) {
 				errorCode = INVALID_PAYLOAD
 			}
+			// invalid code
+			if msg.Size >= 2 {
+				cCode := binary.BigEndian.Uint16(msg.Data[0:2])
+				if invalidCloseCode(cCode) {
+					errorCode = PROTOCOL_ERR
+				}
+			}
 			clientAck = false
 			break
 
@@ -171,7 +178,7 @@ func clientReader(c *client){
 
 		/* (5) PING size error */
 		if msg.Type == PING && msg.Size > 125 {
-			fmt.Printf(" [reader] PING payload too big\n")
+			// fmt.Printf(" [reader] PING payload too big\n")
 			// fmt.Printf("[reader] PING err\n")
 			errorCode = PROTOCOL_ERR
 			break
@@ -190,14 +197,14 @@ func clientReader(c *client){
 
 		/* (7) Invalid UTF8 */
 		if msg.Type == TEXT && !utf8.Valid(msg.Data) {
-			fmt.Printf(" [reader] invalid utf-8\n")
+			// fmt.Printf(" [reader] invalid utf-8\n")
 			errorCode = INVALID_PAYLOAD
 			break
 		}
 
 		/* (8) Unknown opcode */
 		if msg.Type != TEXT && msg.Type != BINARY {
-			fmt.Printf(" [reader] unknown OpCode %d\n", msg.Type)
+			// fmt.Printf(" [reader] unknown OpCode %d\n", msg.Type)
 			errorCode = PROTOCOL_ERR
 			break
 		}
@@ -249,10 +256,10 @@ func clientWriter(c *client){
 
 
 
-// close writes the error message (if needed)
-// and it closes the socket
-// if 'clientACK' is true, reads the next message (CLOSE acknowledge)
-// before closing the socket
+// closes the connection
+// send CLOSE frame is 'status' is not NONE
+// wait for the next message (CLOSE acknowledge) if 'clientACK'
+// then delete client
 func (c *client) close(status MessageError, clientACK bool){
 
 	/* (1) Fail if already closing */
@@ -277,25 +284,24 @@ func (c *client) close(status MessageError, clientACK bool){
 
 
 
-	if status == NONE {
-		status = NORMAL
-	}
+	if status != NONE {
 
-	/* (3) Build message */
-	msg := &Message{
-		Final: true,
-		Type: CLOSE,
-		Size: 2,
-		Data: make([]byte, 2),
-	}
-	binary.BigEndian.PutUint16(msg.Data, uint16(status))
+		/* (3) Build message */
+		msg := &Message{
+			Final: true,
+			Type: CLOSE,
+			Size: 2,
+			Data: make([]byte, 2),
+		}
+		binary.BigEndian.PutUint16(msg.Data, uint16(status))
+
+		/* (4) Send message */
+		err := msg.Send(c.io.sock)
+		if err != nil {
+			fmt.Printf("[close] send error (%s0\n", err)
+		}
 
-	/* (4) Send message */
-	err := msg.Send(c.io.sock)
-	if err != nil {
-		fmt.Printf("[close] send error (%s0\n", err)
 	}
-	// fmt.Printf("[close] frame sent\n")
 
 
 	/* (2) Wait for client CLOSE if needed */
@@ -326,4 +332,19 @@ func (c *client) close(status MessageError, clientACK bool){
 
 	return
 
+}
+
+
+
+
+
+func invalidCloseCode(code uint16) bool{
+
+	tolow := code < 1000
+	badrange1 := code >= 1004 && code <= 1006
+	badrange2 := code >= 1012 && code <= 1016
+	badspecific := code == 1100 || code == 2000 || code == 2999
+
+	return tolow || badrange1 || badrange2 || badspecific
+
 }
\ No newline at end of file