From c053651f0b9fa3c56fc08e1452e212e164f8b1da Mon Sep 17 00:00:00 2001
From: Louis-Alexandre Duchesneau <la@duchesneau.dev>
Date: Wed, 27 Nov 2024 20:34:37 -0500
Subject: [PATCH] Fix Issue #151: Bind not restarting when user uses CloseBind
 (#152)

---
 receivable.go    | 38 +++++++++-----------------------------
 transceivable.go | 41 ++++++++++++++++++++++++-----------------
 2 files changed, 33 insertions(+), 46 deletions(-)

diff --git a/receivable.go b/receivable.go
index 1e6b5a3..1932e74 100644
--- a/receivable.go
+++ b/receivable.go
@@ -68,20 +68,6 @@ func (t *receivable) start() {
 	}()
 }
 
-// check error and do closing if needed
-func (t *receivable) check(err error) (closing bool) {
-	if err == nil {
-		return
-	}
-
-	if t.settings.OnReceivingError != nil {
-		t.settings.OnReceivingError(err)
-	}
-
-	closing = true
-	return
-}
-
 func (t *receivable) loop() {
 	var err error
 	for {
@@ -96,9 +82,13 @@ func (t *receivable) loop() {
 		if err = t.conn.SetReadTimeout(t.settings.ReadTimeout); err == nil {
 			p, err = pdu.Parse(t.conn)
 		}
-		closeOnError := t.check(err)
-		if closeOnError {
-			t.closing(InvalidStreaming)
+		if err != nil {
+			if atomic.LoadInt32(&t.aliveState) == Alive {
+				if t.settings.OnReceivingError != nil {
+					t.settings.OnReceivingError(err)
+				}
+				t.closing(InvalidStreaming)
+			}
 			return
 		}
 
@@ -150,13 +140,9 @@ func (t *receivable) handleWindowPdu(p pdu.PDU) (closing bool) {
 			if t.settings.EnableAutoRespond {
 				t.settings.response(pp.GetResponse())
 			} else if t.settings.OnReceivedPduRequest != nil {
-				r, closeBind := t.settings.OnReceivedPduRequest(p)
+				r, _ := t.settings.OnReceivedPduRequest(p)
 				t.settings.response(r)
-				if closeBind {
-					time.Sleep(50 * time.Millisecond)
-					closing = true
-					t.closing(UnbindClosing)
-				}
+
 			}
 		case *pdu.Unbind:
 			if t.settings.EnableAutoRespond {
@@ -165,15 +151,12 @@ func (t *receivable) handleWindowPdu(p pdu.PDU) (closing bool) {
 				// wait to send response before closing
 				time.Sleep(50 * time.Millisecond)
 				closing = true
-				t.closing(UnbindClosing)
 			} else if t.settings.OnReceivedPduRequest != nil {
 				r, closeBind := t.settings.OnReceivedPduRequest(p)
 				t.settings.response(r)
 				if closeBind {
 					time.Sleep(50 * time.Millisecond)
 					closing = true
-					t.closing(UnbindClosing)
-
 				}
 			}
 		default:
@@ -183,7 +166,6 @@ func (t *receivable) handleWindowPdu(p pdu.PDU) (closing bool) {
 				if closeBind {
 					time.Sleep(50 * time.Millisecond)
 					closing = true
-					t.closing(UnbindClosing)
 				}
 			}
 		}
@@ -198,7 +180,6 @@ func (t *receivable) handleAllPdu(p pdu.PDU) (closing bool) {
 		if closeBind {
 			time.Sleep(50 * time.Millisecond)
 			closing = true
-			t.closing(UnbindClosing)
 		}
 	}
 	return
@@ -216,7 +197,6 @@ func (t *receivable) handleOrClose(p pdu.PDU) (closing bool) {
 			time.Sleep(50 * time.Millisecond)
 
 			closing = true
-			t.closing(UnbindClosing)
 
 		default:
 			var responded bool
diff --git a/transceivable.go b/transceivable.go
index 586b03b..9f7e577 100644
--- a/transceivable.go
+++ b/transceivable.go
@@ -115,20 +115,7 @@ func (t *transceivable) SystemID() string {
 
 // Close transceiver and stop underlying daemons.
 func (t *transceivable) Close() (err error) {
-	if atomic.CompareAndSwapInt32(&t.aliveState, Alive, Closed) {
-		// closing input and output
-		_ = t.out.close(StoppingProcessOnly)
-		_ = t.in.close(StoppingProcessOnly)
-
-		// close underlying conn
-		err = t.conn.Close()
-
-		// notify transceiver closed
-		if t.settings.OnClosed != nil {
-			t.settings.OnClosed(ExplicitClosing)
-		}
-	}
-	return
+	return t.closing(ExplicitClosing)
 }
 
 // Submit a PDU.
@@ -159,9 +146,8 @@ func (t *transceivable) windowCleanup() {
 				if time.Since(request.TimeSent) > t.settings.PduExpireTimeOut {
 					_ = t.requestStore.Delete(ctx, request.GetSequenceNumber())
 					if t.settings.OnExpiredPduRequest != nil {
-						bindClose := t.settings.OnExpiredPduRequest(request.PDU)
-						if bindClose {
-							_ = t.Close()
+						if t.settings.OnExpiredPduRequest(request.PDU) {
+							_ = t.closing(ConnectionIssue)
 						}
 					}
 				}
@@ -170,3 +156,24 @@ func (t *transceivable) windowCleanup() {
 		}
 	}
 }
+
+func (t *transceivable) closing(state State) (err error) {
+	if atomic.CompareAndSwapInt32(&t.aliveState, Alive, Closed) {
+		t.cancel()
+
+		// closing input and output
+		_ = t.out.close(StoppingProcessOnly)
+		_ = t.in.close(StoppingProcessOnly)
+
+		// close underlying conn
+		err = t.conn.Close()
+
+		// notify transceiver closed
+		if t.settings.OnClosed != nil {
+			t.settings.OnClosed(state)
+		}
+
+		t.wg.Wait()
+	}
+	return
+}