From d26364a67c66f0ad4cd0d4e5cfb008e4804c0d9f Mon Sep 17 00:00:00 2001
From: xtaci <daniel820313@gmail.com>
Date: Mon, 11 Mar 2024 22:08:35 +0800
Subject: [PATCH] fix for issue: https://github.com/xtaci/kcp-go/issues/259

---
 sess.go | 16 +++++++++++++---
 1 file changed, 13 insertions(+), 3 deletions(-)

diff --git a/sess.go b/sess.go
index 227d99b..4e32399 100644
--- a/sess.go
+++ b/sess.go
@@ -195,6 +195,7 @@ func newUDPSession(conv uint32, dataShards, parityShards int, l *Listener, conn
 
 // Read implements net.Conn
 func (s *UDPSession) Read(b []byte) (n int, err error) {
+RESET_TIMER:
 	var timeout *time.Timer
 	// deadline for current reading operation
 	var c <-chan time.Time
@@ -243,6 +244,10 @@ func (s *UDPSession) Read(b []byte) (n int, err error) {
 		// wait for read event or timeout or error
 		select {
 		case <-s.chReadEvent:
+			if timeout != nil {
+				timeout.Stop()
+				goto RESET_TIMER
+			}
 		case <-c:
 			return 0, errors.WithStack(errTimeout)
 		case <-s.chSocketReadError:
@@ -258,6 +263,7 @@ func (s *UDPSession) Write(b []byte) (n int, err error) { return s.WriteBuffers(
 
 // WriteBuffers write a vector of byte slices to the underlying connection
 func (s *UDPSession) WriteBuffers(v [][]byte) (n int, err error) {
+RESET_TIMER:
 	var timeout *time.Timer
 	var c <-chan time.Time
 	if !s.wd.IsZero() {
@@ -308,6 +314,10 @@ func (s *UDPSession) WriteBuffers(v [][]byte) (n int, err error) {
 
 		select {
 		case <-s.chWriteEvent:
+			if timeout != nil {
+				timeout.Stop()
+				goto RESET_TIMER
+			}
 		case <-c:
 			return 0, errors.WithStack(errTimeout)
 		case <-s.chSocketWriteError:
@@ -375,9 +385,9 @@ func (s *UDPSession) RemoteAddr() net.Addr { return s.remote }
 // SetDeadline sets the deadline associated with the listener. A zero time value disables the deadline.
 func (s *UDPSession) SetDeadline(t time.Time) error {
 	s.mu.Lock()
-	defer s.mu.Unlock()
 	s.rd = t
 	s.wd = t
+	s.mu.Unlock()
 	s.notifyReadEvent()
 	s.notifyWriteEvent()
 	return nil
@@ -386,8 +396,8 @@ func (s *UDPSession) SetDeadline(t time.Time) error {
 // SetReadDeadline implements the Conn SetReadDeadline method.
 func (s *UDPSession) SetReadDeadline(t time.Time) error {
 	s.mu.Lock()
-	defer s.mu.Unlock()
 	s.rd = t
+	s.mu.Unlock()
 	s.notifyReadEvent()
 	return nil
 }
@@ -395,8 +405,8 @@ func (s *UDPSession) SetReadDeadline(t time.Time) error {
 // SetWriteDeadline implements the Conn SetWriteDeadline method.
 func (s *UDPSession) SetWriteDeadline(t time.Time) error {
 	s.mu.Lock()
-	defer s.mu.Unlock()
 	s.wd = t
+	s.mu.Unlock()
 	s.notifyWriteEvent()
 	return nil
 }