chore: improve the length check for conn.Discard

panjf2000 · Nov 16, 2024 · 2d705af · 2d705af
1 parent adcbdd2
commit 2d705af
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 10 deletions.
diff --git a/connection_unix.go b/connection_unix.go
@@ -368,10 +368,9 @@ func (c *conn) Discard(n int) (int, error) {
 	}
 
 	inBufferLen := c.inboundBuffer.Buffered()
-	tempBufferLen := len(c.buffer)
-	if inBufferLen+tempBufferLen < n || n <= 0 {
+	if totalLen := inBufferLen + len(c.buffer); n >= totalLen || n <= 0 {
 		c.resetBuffer()
-		return inBufferLen + tempBufferLen, nil
+		return totalLen, nil
 	}
 
 	if c.inboundBuffer.IsEmpty() {

diff --git a/connection_windows.go b/connection_windows.go
@@ -186,18 +186,17 @@ func (c *conn) Peek(n int) (buf []byte, err error) {
 }
 
 func (c *conn) Discard(n int) (int, error) {
-	inBufferLen := c.inboundBuffer.Buffered()
-	tempBufferLen := c.buffer.Len()
-	if inBufferLen+tempBufferLen < n || n <= 0 {
-		c.resetBuffer()
-		return inBufferLen + tempBufferLen, nil
-	}
-
 	if len(c.cache) > 0 {
 		bsPool.Put(c.cache)
 		c.cache = nil
 	}
 
+	inBufferLen := c.inboundBuffer.Buffered()
+	if totalLen := inBufferLen + c.buffer.Len(); n >= totalLen || n <= 0 {
+		c.resetBuffer()
+		return totalLen, nil
+	}
+
 	if c.inboundBuffer.IsEmpty() {
 		c.buffer.B = c.buffer.B[n:]
 		return n, nil