From 5e1612100b42b7ce4a713c4d2715b689a117a551 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Viktor=20Tigerstr=C3=B6m?= <vtigerstrom@gmail.com>
Date: Thu, 4 Jan 2024 20:39:59 +0100
Subject: [PATCH] gbn+mailbox: force connection status change on FIN

When the client receives+processes a FIN packet, we force the connection
to status ClientStatusSessionNotFound, as in rare occasions the
corresponding server may have time to close the connection before we've
already processed the sent FIN packet by the server. In that case, if we
didn't force a new status, the client would never mark the connection as
status ClientStatusSessionNotFound.

This solves a bug for the testServerReconnect itest, as that test would
sometimes fail due to the FIN packet being processed but without
changing the connection status, leading to a test predicate never being
satisfied.
---
 gbn/config.go          |  4 ++++
 gbn/gbn_conn.go        |  4 ++++
 gbn/options.go         |  8 ++++++++
 mailbox/client_conn.go | 28 ++++++++++++++++++++--------
 4 files changed, 36 insertions(+), 8 deletions(-)

diff --git a/gbn/config.go b/gbn/config.go
index 4d14f7f2..6a027b4d 100644
--- a/gbn/config.go
+++ b/gbn/config.go
@@ -38,6 +38,10 @@ type config struct {
 	// packet.
 	sendToStream sendBytesFunc
 
+	// onFIN is a callback that if set, will be called once a FIN packet has
+	// been received and processed.
+	onFIN func()
+
 	// handshakeTimeout is the time after which the server or client
 	// will abort and restart the handshake if the expected response is
 	// not received from the peer.
diff --git a/gbn/gbn_conn.go b/gbn/gbn_conn.go
index 4645fa82..50a0cc93 100644
--- a/gbn/gbn_conn.go
+++ b/gbn/gbn_conn.go
@@ -649,6 +649,10 @@ func (g *GoBackNConn) receivePacketsForever() error { // nolint:gocyclo
 
 			close(g.remoteClosed)
 
+			if g.cfg.onFIN != nil {
+				g.cfg.onFIN()
+			}
+
 			return errTransportClosing
 
 		default:
diff --git a/gbn/options.go b/gbn/options.go
index 2c1e8df1..4adcb3a0 100644
--- a/gbn/options.go
+++ b/gbn/options.go
@@ -43,3 +43,11 @@ func WithKeepalivePing(ping, pong time.Duration) Option {
 		conn.pongTime = pong
 	}
 }
+
+// WithOnFIN is used to set the onFIN callback that will be called once a FIN
+// packet has been received and processed.
+func WithOnFIN(fn func()) Option {
+	return func(conn *config) {
+		conn.onFIN = fn
+	}
+}
diff --git a/mailbox/client_conn.go b/mailbox/client_conn.go
index 6355a942..445daa3c 100644
--- a/mailbox/client_conn.go
+++ b/mailbox/client_conn.go
@@ -157,20 +157,32 @@ func NewClientConn(ctx context.Context, sid [64]byte, serverHost string,
 	}
 
 	c := &ClientConn{
-		transport: transport,
-		gbnOptions: []gbn.Option{
-			gbn.WithTimeout(gbnTimeout),
-			gbn.WithHandshakeTimeout(gbnHandshakeTimeout),
-			gbn.WithKeepalivePing(
-				gbnClientPingTimeout, gbnPongTimeout,
-			),
-		},
+		transport:   transport,
 		status:      ClientStatusNotConnected,
 		onNewStatus: onNewStatus,
 		quit:        make(chan struct{}),
 		cancel:      cancel,
 		log:         logger,
 	}
+
+	c.gbnOptions = []gbn.Option{
+		gbn.WithTimeout(gbnTimeout),
+		gbn.WithHandshakeTimeout(gbnHandshakeTimeout),
+		gbn.WithKeepalivePing(
+			gbnClientPingTimeout, gbnPongTimeout,
+		),
+		gbn.WithOnFIN(func() {
+			// We force the connection to set a new status after
+			// processing a FIN packet, as in rare occasions the
+			// corresponding server may have time to close the
+			// connection before we've already processed the sent
+			// FIN packet by the server. In that case, if we didn't
+			// force a new status, the client would never mark the
+			// connection as status ClientStatusSessionNotFound.
+			c.setStatus(ClientStatusSessionNotFound)
+		}),
+	}
+
 	c.connKit = &connKit{
 		ctx:        ctxc,
 		impl:       c,