From 2d705afc8c984c0a7cf87a98e4805fb8a9a63c87 Mon Sep 17 00:00:00 2001 From: Andy Pan Date: Sat, 16 Nov 2024 14:38:05 +0800 Subject: [PATCH] chore: improve the length check for conn.Discard --- connection_unix.go | 5 ++--- connection_windows.go | 13 ++++++------- 2 files changed, 8 insertions(+), 10 deletions(-) diff --git a/connection_unix.go b/connection_unix.go index 449039cec..7babbeeac 100644 --- 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 index b8b350fa7..fe8ccb187 100644 --- 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