From 1891bdff1dd7cf0ac0874683a196db34169af25e Mon Sep 17 00:00:00 2001
From: xdrm-brackets <doowap31@gmail.com>
Date: Mon, 7 May 2018 19:25:54 +0200
Subject: [PATCH] fix send length + add test for simple/extended-16/extended-64
 length

---
 ws/message.go      |  4 ++--
 ws/message_test.go | 34 +++++++++++++++++++++++++++++++++-
 2 files changed, 35 insertions(+), 3 deletions(-)

diff --git a/ws/message.go b/ws/message.go
index ee80dda..6f300b2 100644
--- a/ws/message.go
+++ b/ws/message.go
@@ -178,7 +178,7 @@ func (m Message) Send(writer io.Writer) error {
 
 		header = append(header, byte(m.Size) )
 
-	} else if m.Size < 0xffff { // extended: 16 bits
+	} else if m.Size <= 0xffff { // extended: 16 bits
 
 		header = append(header, 126)
 
@@ -186,7 +186,7 @@ func (m Message) Send(writer io.Writer) error {
 		binary.BigEndian.PutUint16(buf, uint16(m.Size))
 		header = append(header, buf...)
 
-	} else if m.Size < 0xffffffffffffffff { // extended: 64 bits
+	} else if m.Size <= 0xffffffffffffffff { // extended: 64 bits
 
 		header = append(header, 127)
 
diff --git a/ws/message_test.go b/ws/message_test.go
index 04254c4..9e8db72 100644
--- a/ws/message_test.go
+++ b/ws/message_test.go
@@ -193,6 +193,14 @@ func TestReadEOF(t *testing.T) {
 
 func TestSimpleMessageSending(t *testing.T) {
 
+	m4b1 := make([]byte, 0x7e - 1)
+	m4b2 := make([]byte, 0x7e)
+	m4b3 := make([]byte, 0x7e + 1)
+
+	m16b1 := make([]byte, 0xffff - 1)
+	m16b2 := make([]byte, 0xffff)
+	m16b3 := make([]byte, 0xffff + 1)
+
 	cases := []struct{
 		Name     string
 		Base     Message
@@ -214,6 +222,30 @@ func TestSimpleMessageSending(t *testing.T) {
 			"non final frame",
 			Message{ false, TEXT, 2, []byte{0x01,0x02} },
 			[]byte{0x01,0x02,0x01,0x02},
+		}, {
+			"125 > normal length",
+			Message{ true, TEXT, uint(len(m4b1)), m4b1 },
+			append([]byte{0x81,0x7e-1}, m4b1...),
+		}, {
+			"126 > extended 16 bits length",
+			Message{ true, TEXT, uint(len(m4b2)), m4b2 },
+			append([]byte{0x81,126,0x00,0x7e}, m4b2...),
+		}, {
+			"127 > extended 16 bits length",
+			Message{ true, TEXT, uint(len(m4b3)), m4b3 },
+			append([]byte{0x81,126,0x00,0x7e+1}, m4b3...),
+		}, {
+			"fffe > extended 16 bits length",
+			Message{ true, TEXT, uint(len(m16b1)), m16b1 },
+			append([]byte{0x81,126, 0xff, 0xfe}, m16b1...),
+		}, {
+			"ffff > extended 16 bits length",
+			Message{ true, TEXT, uint(len(m16b2)), m16b2 },
+			append([]byte{0x81,126,0xff,0xff}, m16b2...),
+		}, {
+			"10000 > extended 64 bits length",
+			Message{ true, TEXT, uint(len(m16b3)), m16b3 },
+			append([]byte{0x81,127, 0x00,0x00,0x00,0x00,0x00,0x01,0x00,0x00,}, m16b3...),
 		},
 	}
 
@@ -232,7 +264,7 @@ func TestSimpleMessageSending(t *testing.T) {
 
 			// check buffer
 			if writer.String() != string(tc.Expected) {
-				t.Errorf("expected '%x', got '%x'", tc.Expected, writer.String())
+				t.Errorf("expected '%.20x', got '%.20x'", tc.Expected, writer.String())
 			}
 
 		})