新v3函数+修改limit逻辑

Author: guonaihong

frame/read_frame_bufio_v2.go

@@ -14,6 +14,7 @@
 package frame
 
 import (
+	"fmt"
 	"io"
 
 	"github.com/antlabs/wsutil/enum"
@@ -22,15 +23,46 @@ import (
 
 func ReadFrameFromReaderV2(r io.Reader, headArray *[enum.MaxFrameHeaderSize]byte, buf *[]byte) (f Frame2, err error) {
 	h, _, err := ReadHeader(r, headArray)
+	if err != nil {
+		return f, fmt.Errorf("ReadFrameFromReaderV2:%w", err)
+	}
+
+	if cap(*buf) < int(h.PayloadLen) {
+		// TODO sync.Pool 处理
+		*buf = make([]byte, h.PayloadLen)
+	}
+	*buf = (*buf)[:h.PayloadLen]
+	n1, err := io.ReadFull(r, *buf)
 	if err != nil {
 		return f, err
 	}
+	if n1 != int(h.PayloadLen) {
+		return f, io.ErrUnexpectedEOF
+	}
+	f.Payload = buf
+	f.FrameHeader = h
+	if h.Mask {
+		mask.Mask(*f.Payload, h.MaskKey)
+	}
+
+	return f, nil
+}
+
+func ReadFrameFromReaderV3(r io.Reader, lr io.Reader, headArray *[enum.MaxFrameHeaderSize]byte, buf *[]byte) (f Frame2, err error) {
+	h, _, err := ReadHeader(r, headArray)
+	if err != nil {
+		return f, fmt.Errorf("ReadFrameFromReaderV2:%w", err)
+	}
 
 	if cap(*buf) < int(h.PayloadLen) {
 		// TODO sync.Pool 处理
 		*buf = make([]byte, h.PayloadLen)
 	}
 	*buf = (*buf)[:h.PayloadLen]
+	if lr != nil {
+		r = lr
+	}
+
 	n1, err := io.ReadFull(r, *buf)
 	if err != nil {
 		return f, err

limitreader/limitreader.go

@@ -17,12 +17,27 @@ func NewLimitReader(r io.Reader, m int64) *limitReader {
 	return &limitReader{r: r, m: m}
 }
 
+// 目前go.mod使用的是go1.20版本， go1.21才有min函数
+func minSize(n, m int) int {
+	if n < m {
+		return n
+	}
+	return m
+}
+
 // 实现io.Reader接口
 func (l *limitReader) Read(p []byte) (n int, err error) {
 	if l.m < 0 {
 		return 0, ErrTooBigMessage
 	}
-	n, err = l.r.Read(p)
+	rn := minSize(int(l.m), len(p))
+	if rn == 0 && len(p) > 0 {
+		rn = 1
+	}
+	n, err = l.r.Read(p[:rn])
 	l.m -= int64(n)
+	if l.m < 0 {
+		return 0, ErrTooBigMessage
+	}
 	return
 }