更新

antlabs · Jul 18, 2023 · 1dec3bb · 1dec3bb
1 parent 32b9ee4
commit 1dec3bb
Show file tree

Hide file tree

Showing 9 changed files with 95 additions and 752 deletions.
diff --git a/benchmark_mask_test.go b/benchmark_mask_test.go
@@ -5,14 +5,6 @@ import (
 	"testing"
 )
 
-func Benchmark_Rand(t *testing.B) {
-	var maskValue [4]byte
-
-	for i := 0; i < t.N; i++ {
-		newMask(maskValue[:])
-	}
-}
-
 func Benchmark_Rand_Uint32(t *testing.B) {
 	for i := 0; i < t.N; i++ {
 		// newMask(maskValue[:])

diff --git a/benchmark_read_write_message_test.go b/benchmark_read_write_message_test.go
@@ -38,7 +38,6 @@ func (testconn *testConn) Read(b []byte) (n int, err error) {
 // time limit; see SetDeadline and SetWriteDeadline.
 func (testconn *testConn) Write(b []byte) (n int, err error) {
 	return testconn.buf.Write(b)
-	return
 }
 
 // Close closes the connection.
@@ -102,59 +101,52 @@ type testConn struct {
 	buf *bytes.Buffer
 }
 
-// func Benchmark_WriteMessage(b *testing.B) {
-// 	var c Conn
-// 	buf2 := bytes.NewBuffer(make([]byte, 0, 1024))
-// 	c.c = &testConn{buf: buf2}
-// 	buf := make([]byte, 1024)
-// 	for i := range buf {
-// 		buf[i] = 1
-// 	}
-
-// 	b.ResetTimer()
-// 	for i := 0; i < b.N; i++ {
-// 		c.WriteMessage(opcode.Binary, buf)
-// 		buf2.Reset()
-// 	}
-// }
-
 func Benchmark_WriteMessage(b *testing.B) {
-	var c Conn
-	buf2 := bytes.NewBuffer(make([]byte, 0, 1024))
-	c.c = &testConn{buf: buf2}
-	buf := make([]byte, 1024)
-	for i := range buf {
-		buf[i] = 1
-	}
+	b.Run("1.case", func(b *testing.B) {
+		var c Conn
+		buf2 := bytes.NewBuffer(make([]byte, 0, 1024))
+		c.c = &testConn{buf: buf2}
+		buf := make([]byte, 1024)
+		for i := range buf {
+			buf[i] = 1
+		}
 
-	b.ResetTimer()
-	for i := 0; i < b.N; i++ {
-		c.WriteMessage(opcode.Binary, buf)
-		buf2.Reset()
-	}
+		b.ResetTimer()
+		for i := 0; i < b.N; i++ {
+			c.WriteMessage(opcode.Binary, buf)
+			buf2.Reset()
+		}
+	})
 }
 
-// TODO
-// func Benchmark_ReadMessage(b *testing.B) {
-// 	var c Conn
-// 	buf2 := bytes.NewBuffer(make([]byte, 0, 1024))
-// 	c.c = &testConn{buf: buf2}
-// 	rv := make([]byte, 0, 1024)
-// 	c.fr.Init(c.c, &rv)
-// 	c.Callback = &DefCallback{}
-
-// 	wbuf := make([]byte, 1024)
-// 	for i := range wbuf {
-// 		wbuf[i] = 1
-// 	}
-
-// 	b.ResetTimer()
-// 	for i := 0; i < b.N; i++ {
-// 		c.WriteMessage(opcode.Binary, wbuf)
-// 		c.ReadLoop()
-// 		buf2.Reset()
-// 	}
-// }
+func Benchmark_ReadMessage(b *testing.B) {
+	b.Run("bufio-TODO", func(b *testing.B) {
+	})
+
+	b.Run("windows", func(b *testing.B) {
+		var c Conn
+		buf2 := bytes.NewBuffer(make([]byte, 0, 1024+enum.MaxFrameHeaderSize))
+
+		c.c = &testConn{buf: buf2}
+
+		windows := make([]byte, 0, 1024)
+
+		c.fr.Init(c.c, &windows)
+		c.Callback = &DefCallback{}
+
+		wbuf := make([]byte, 1024)
+		for i := range wbuf {
+			wbuf[i] = 1
+		}
+
+		b.ResetTimer()
+		for i := 0; i < b.N; i++ {
+			c.WriteMessage(opcode.Binary, wbuf)
+			c.ReadLoop()
+			buf2.Reset()
+		}
+	})
+}
 
 func Benchmark_ReadFrame(b *testing.B) {
 	r := bytes.NewReader(noMaskData)

diff --git a/conn.go b/conn.go
@@ -186,20 +186,22 @@ func (c *Conn) readLoop() error {
 			op = fragmentFrameHeader.Opcode
 		}
 
+		rsv1 := f.GetRsv1()
 		// 检查Rsv1 rsv2 Rsv3
-		if f.Rsv1 && c.failRsv1(op) || f.Rsv2 || f.Rsv3 {
-			err = fmt.Errorf("%w:Rsv1(%t) Rsv2(%t) rsv2(%t) compression:%t", ErrRsv123, f.Rsv1, f.Rsv2, f.Rsv3, c.compression)
+		if rsv1 && c.failRsv1(op) || f.GetRsv2() || f.GetRsv3() {
+			err = fmt.Errorf("%w:Rsv1(%t) Rsv2(%t) rsv2(%t) compression:%t", ErrRsv123, rsv1, f.GetRsv2(), f.GetRsv3(), c.compression)
 			return c.writeErrAndOnClose(ProtocolError, err)
 		}
 
+		fin := f.GetFin()
 		if fragmentFrameHeader != nil && !f.Opcode.IsControl() {
 			if f.Opcode == 0 {
 				fragmentFrameBuf = append(fragmentFrameBuf, f.Payload...)
 
 				// 分段的在这返回
-				if f.Fin {
+				if fin {
 					// 解压缩
-					if fragmentFrameHeader.Rsv1 && c.decompression {
+					if fragmentFrameHeader.GetRsv1() && c.decompression {
 						tmpeBuf, err := decode(fragmentFrameBuf)
 						if err != nil {
 							return err
@@ -224,10 +226,8 @@ func (c *Conn) readLoop() error {
 			return ErrFrameOpcode
 		}
 
-		// 检查Opcode
-		switch f.Opcode {
-		case opcode.Text, opcode.Binary:
-			if !f.Fin {
+		if f.Opcode == opcode.Text || f.Opcode == opcode.Binary {
+			if !fin {
 				prevFrame := f.FrameHeader
 				// 第一次分段
 				if len(fragmentFrameBuf) == 0 {
@@ -240,7 +240,7 @@ func (c *Conn) readLoop() error {
 				continue
 			}
 
-			if f.Rsv1 && c.decompression {
+			if rsv1 && c.decompression {
 				// 不分段的解压缩
 				f.Payload, err = decode(f.Payload)
 				if err != nil {
@@ -257,14 +257,17 @@ func (c *Conn) readLoop() error {
 			}
 
 			c.Callback.OnMessage(c, f.Opcode, f.Payload)
-		case Close, Ping, Pong:
+			continue
+		}
+
+		if f.Opcode == Close || f.Opcode == Ping || f.Opcode == Pong {
 			//  对方发的控制消息太大
 			if f.PayloadLen > maxControlFrameSize {
 				c.writeErrAndOnClose(ProtocolError, ErrMaxControlFrameSize)
 				return ErrMaxControlFrameSize
 			}
 			// Close, Ping, Pong 不能分片
-			if !f.Fin {
+			if !fin {
 				c.writeErrAndOnClose(ProtocolError, ErrNOTBeFragmented)
 				return ErrNOTBeFragmented
 			}
@@ -314,11 +317,11 @@ func (c *Conn) readLoop() error {
 			}
 
 			c.Callback.OnMessage(c, f.Opcode, nil)
-		default:
-			c.writeErrAndOnClose(ProtocolError, ErrOpcode)
-			return ErrOpcode
+			continue
 		}
-
+		// 检查Opcode
+		c.writeErrAndOnClose(ProtocolError, ErrOpcode)
+		return ErrOpcode
 	}
 }
 
@@ -361,39 +364,38 @@ func (c *Conn) WriteMessage(op Opcode, writeBuf []byte) (err error) {
 	return frame.WriteFrame(&c.fw, c.c, writeBuf, rsv1, c.client, op, maskValue)
 }
 
-// func (c *Conn) WriteMessageOld(op Opcode, writeBuf []byte) (err error) {
-// 	var f frame.FrameHeader
-
-// 	if op == opcode.Text {
-// 		if !c.utf8Check(writeBuf) {
-// 			return ErrTextNotUTF8
-// 		}
-// 	}
-
-// 	f.Fin = true
-// 	f.Rsv1 = c.compression && (op == opcode.Text || op == opcode.Binary)
-// 	if f.Rsv1 {
-// 		var out wrapBuffer
-// 		w := compressNoContextTakeover(&out, defaultCompressionLevel)
-// 		if _, err = io.Copy(w, bytes.NewReader(writeBuf)); err != nil {
-// 			return
-// 		}
-
-// 		if err = w.Close(); err != nil {
-// 			return
-// 		}
-// 		writeBuf = out.Bytes()
-// 	}
-
-// 	f.Opcode = op
-// 	f.PayloadLen = int64(len(writeBuf))
-// 	if c.client {
-// 		f.Mask = true
-// 		newMask(f.MaskValue[:])
-// 	}
-
-// 	return frame.WriteFrameOld(c.c, f, writeBuf, &c.fw)
-// }
+// 这是一个不安全的方法, writeBuf的格式必须是 14个字节的空白长度+需要写的payload组成
+func (c *Conn) WriteMessageUnsafe(op Opcode, writeBuf []byte) (err error) {
+	// if op == opcode.Text {
+	// 	if !c.utf8Check(writeBuf) {
+	// 		return ErrTextNotUTF8
+	// 	}
+	// }
+
+	// rsv1 := c.compression && (op == opcode.Text || op == opcode.Binary)
+	// if rsv1 {
+	// 	var out wrapBuffer
+	// 	w := compressNoContextTakeover(&out, defaultCompressionLevel)
+	// 	if _, err = io.Copy(w, bytes.NewReader(writeBuf)); err != nil {
+	// 		return
+	// 	}
+
+	// 	if err = w.Close(); err != nil {
+	// 		return
+	// 	}
+	// 	writeBuf = out.Bytes()
+	// }
+
+	// // f.Opcode = op
+	// // f.PayloadLen = int64(len(writeBuf))
+	// maskValue := uint32(0)
+	// if c.client {
+	// 	maskValue = rand.Uint32()
+	// }
+
+	// return frame.WriteFrame(&c.fw, c.c, writeBuf, rsv1, c.client, op, maskValue)
+	return
+}
 
 func (c *Conn) SetDeadline(t time.Time) error {
 	return c.c.SetDeadline(t)

diff --git a/conn_test.go b/conn_test.go
diff --git a/go.mod b/go.mod
@@ -2,4 +2,4 @@ module github.com/antlabs/quickws
 
 go 1.20
 
-require github.com/antlabs/wsutil v0.0.9
+require github.com/antlabs/wsutil v0.0.10
diff --git a/go.sum b/go.sum
@@ -1,2 +1,2 @@
-github.com/antlabs/wsutil v0.0.9 h1:CNURXjtBZ3O/5xp+uimBtAoxR309R4QS1uxGJ/QCp+c=
-github.com/antlabs/wsutil v0.0.9/go.mod h1:7ec5eUM7nmKW+Oi6F1I58iatOeL9k+yIsfOh1zh910g=
+github.com/antlabs/wsutil v0.0.10 h1:5O5KdRH5qPasEhyOkXA+x3OcX7ICYnjSVkUmSDAP+0s=
+github.com/antlabs/wsutil v0.0.10/go.mod h1:7ec5eUM7nmKW+Oi6F1I58iatOeL9k+yIsfOh1zh910g=