Fix request timeout handling (go-ldap#388)

* Fix request timeout handling * Ensure 64-bit alignment
inv2004 · Jan 17, 2023 · 45d75f5 · 45d75f5
1 parent 34362d5
commit 45d75f5
Show file tree

Hide file tree

Showing 2 changed files with 30 additions and 20 deletions.
diff --git a/conn.go b/conn.go
@@ -87,6 +87,7 @@ const (
 type Conn struct {
 	// requestTimeout is loaded atomically
 	// so we need to ensure 64-bit alignment on 32-bit platforms.
+	// https://github.com/go-ldap/ldap/pull/199
 	requestTimeout      int64
 	conn                net.Conn
 	isTLS               bool
@@ -281,9 +282,7 @@ func (l *Conn) Close() {
 
 // SetTimeout sets the time after a request is sent that a MessageTimeout triggers
 func (l *Conn) SetTimeout(timeout time.Duration) {
-	if timeout > 0 {
-		atomic.StoreInt64(&l.requestTimeout, int64(timeout))
-	}
+	atomic.StoreInt64(&l.requestTimeout, int64(timeout))
 }
 
 // Returns the next available messageID
@@ -486,20 +485,26 @@ func (l *Conn) processMessages() {
 				l.messageContexts[message.MessageID] = message.Context
 
 				// Add timeout if defined
-				requestTimeout := time.Duration(atomic.LoadInt64(&l.requestTimeout))
-				if requestTimeout > 0 {
+				if l.requestTimeout > 0 {
 					go func() {
+						timer := time.NewTimer(time.Duration(l.requestTimeout))
 						defer func() {
 							if err := recover(); err != nil {
 								logger.Printf("ldap: recovered panic in RequestTimeout: %v", err)
 							}
+
+							timer.Stop()
 						}()
-						time.Sleep(requestTimeout)
-						timeoutMessage := &messagePacket{
-							Op:        MessageTimeout,
-							MessageID: message.MessageID,
+
+						select {
+						case <-timer.C:
+							timeoutMessage := &messagePacket{
+								Op:        MessageTimeout,
+								MessageID: message.MessageID,
+							}
+							l.sendProcessMessage(timeoutMessage)
+						case <-message.Context.done:
 						}
-						l.sendProcessMessage(timeoutMessage)
 					}()
 				}
 			case MessageResponse:

diff --git a/v3/conn.go b/v3/conn.go
@@ -87,6 +87,7 @@ const (
 type Conn struct {
 	// requestTimeout is loaded atomically
 	// so we need to ensure 64-bit alignment on 32-bit platforms.
+	// https://github.com/go-ldap/ldap/pull/199
 	requestTimeout      int64
 	conn                net.Conn
 	isTLS               bool
@@ -281,9 +282,7 @@ func (l *Conn) Close() {
 
 // SetTimeout sets the time after a request is sent that a MessageTimeout triggers
 func (l *Conn) SetTimeout(timeout time.Duration) {
-	if timeout > 0 {
-		atomic.StoreInt64(&l.requestTimeout, int64(timeout))
-	}
+	atomic.StoreInt64(&l.requestTimeout, int64(timeout))
 }
 
 // Returns the next available messageID
@@ -486,20 +485,26 @@ func (l *Conn) processMessages() {
 				l.messageContexts[message.MessageID] = message.Context
 
 				// Add timeout if defined
-				requestTimeout := time.Duration(atomic.LoadInt64(&l.requestTimeout))
-				if requestTimeout > 0 {
+				if l.requestTimeout > 0 {
 					go func() {
+						timer := time.NewTimer(time.Duration(l.requestTimeout))
 						defer func() {
 							if err := recover(); err != nil {
 								logger.Printf("ldap: recovered panic in RequestTimeout: %v", err)
 							}
+
+							timer.Stop()
 						}()
-						time.Sleep(requestTimeout)
-						timeoutMessage := &messagePacket{
-							Op:        MessageTimeout,
-							MessageID: message.MessageID,
+
+						select {
+						case <-timer.C:
+							timeoutMessage := &messagePacket{
+								Op:        MessageTimeout,
+								MessageID: message.MessageID,
+							}
+							l.sendProcessMessage(timeoutMessage)
+						case <-message.Context.done:
 						}
-						l.sendProcessMessage(timeoutMessage)
 					}()
 				}
 			case MessageResponse: