From 6e0b53537c9a1e78ea148899a7f0a24fcd6b4029 Mon Sep 17 00:00:00 2001
From: ami-GS <1991.daiki@gmail.com>
Date: Wed, 29 May 2024 18:35:09 -0700
Subject: [PATCH] retry connection if timed out

---
 src/test/lib/HandshakeTest.cpp  | 73 ++++++++++++++++++++-------------
 src/test/lib/TestConnection.cpp |  1 +
 src/test/lib/TestConnection.h   |  2 +
 3 files changed, 47 insertions(+), 29 deletions(-)

diff --git a/src/test/lib/HandshakeTest.cpp b/src/test/lib/HandshakeTest.cpp
index aec150c9e3..5e74162e9e 100644
--- a/src/test/lib/HandshakeTest.cpp
+++ b/src/test/lib/HandshakeTest.cpp
@@ -284,42 +284,57 @@ QuicTestConnect(
                     TEST_QUIC_SUCCEEDED(Client.SetRemoteAddr(RemoteAddr));
                 }
 
-                TEST_QUIC_SUCCEEDED(
-                    Client.Start(
-                        ClientConfiguration,
-                        QuicAddrFamily,
-                        QUIC_LOCALHOST_FOR_AF(QuicAddrFamily),
-                        ServerLocalAddr.GetPort()));
+                do {
+                    if (Client.GetIsConnected()) {
+                        Client.Shutdown(QUIC_CONNECTION_SHUTDOWN_FLAG_NONE, QUIC_TEST_NO_ERROR);
+                        if (!Client.WaitForShutdownComplete()) {
+                            return;
+                        }
+                    }
+                    if (Server->GetIsConnected()) {
+                        Server->Shutdown(QUIC_CONNECTION_SHUTDOWN_FLAG_NONE, QUIC_TEST_NO_ERROR);
+                        if (!Server->WaitForShutdownComplete()) {
+                            return;
+                        }
+                    }
 
-                if (AsyncConfiguration || AsyncTicketValidation) {
-                    if (!CxPlatEventWaitWithTimeout(ServerAcceptCtx.NewConnectionReady, TestWaitTimeout)) {
-                        TEST_FAILURE("Timed out waiting for server accept.");
-                    } else if (Server == nullptr) {
-                        TEST_FAILURE("Failed to accept server connection.");
-                    } else {
-                        if (AsyncConfiguration) {
-                            if (AsyncConfiguration == QUIC_TEST_ASYNC_CONFIG_DELAYED) {
+                    TEST_QUIC_SUCCEEDED(
+                        Client.Start(
+                            ClientConfiguration,
+                            QuicAddrFamily,
+                            QUIC_LOCALHOST_FOR_AF(QuicAddrFamily),
+                            ServerLocalAddr.GetPort()));
+
+                    if (AsyncConfiguration || AsyncTicketValidation) {
+                        if (!CxPlatEventWaitWithTimeout(ServerAcceptCtx.NewConnectionReady, TestWaitTimeout)) {
+                            TEST_FAILURE("Timed out waiting for server accept.");
+                        } else if (Server == nullptr) {
+                            TEST_FAILURE("Failed to accept server connection.");
+                        } else {
+                            if (AsyncConfiguration) {
+                                if (AsyncConfiguration == QUIC_TEST_ASYNC_CONFIG_DELAYED) {
+                                    CxPlatSleep(1000);
+                                }
+                                TEST_QUIC_SUCCEEDED(
+                                    Server->SetConfiguration(ServerConfiguration));
+                            }
+                            if (AsyncTicketValidation) {
                                 CxPlatSleep(1000);
+                                TEST_QUIC_SUCCEEDED(Server->SetCustomTicketValidationResult(SessionResumption == QUIC_TEST_RESUMPTION_ENABLED_ASYNC));
                             }
-                            TEST_QUIC_SUCCEEDED(
-                                Server->SetConfiguration(ServerConfiguration));
-                        }
-                        if (AsyncTicketValidation) {
-                            CxPlatSleep(1000);
-                            TEST_QUIC_SUCCEEDED(Server->SetCustomTicketValidationResult(SessionResumption == QUIC_TEST_RESUMPTION_ENABLED_ASYNC));
                         }
                     }
-                }
 
-                if (!Client.WaitForConnectionComplete()) {
-                    return;
-                }
-                TEST_TRUE(Client.GetIsConnected());
+                    if (!Client.WaitForConnectionComplete()) {
+                        return;
+                    }
 
-                TEST_NOT_EQUAL(nullptr, Server);
-                if (!Server->WaitForConnectionComplete()) {
-                    return;
-                }
+                    TEST_NOT_EQUAL(nullptr, Server);
+                    if (!Server->WaitForConnectionComplete()) {
+                        return;
+                    }
+                } while (Server->GetConnectionTimeout() || Client.GetConnectionTimeout());
+                TEST_TRUE(Client.GetIsConnected());
                 TEST_TRUE(Server->GetIsConnected());
 
                 TEST_EQUAL(
diff --git a/src/test/lib/TestConnection.cpp b/src/test/lib/TestConnection.cpp
index 058c01dbb7..d13ae8cd07 100644
--- a/src/test/lib/TestConnection.cpp
+++ b/src/test/lib/TestConnection.cpp
@@ -820,6 +820,7 @@ TestConnection::HandleConnectionEvent(
                 QuicTraceLogInfo(
                     TestIgnoreConnectionTimeout,
                     "[test] Ignoring timeout unexpected status because of random loss");
+                ConnectionTimeout = true;
             } else {
                 TEST_FAILURE(
                     "Unexpected transport Close Error, expected=0x%x, actual=0x%x",
diff --git a/src/test/lib/TestConnection.h b/src/test/lib/TestConnection.h
index 1beeb9c3d3..3a85d11913 100644
--- a/src/test/lib/TestConnection.h
+++ b/src/test/lib/TestConnection.h
@@ -63,6 +63,7 @@ class TestConnection
     bool HasRandomLoss      : 1;
     bool AsyncCustomValidation : 1;
     bool CustomValidationResultSet : 1;
+    bool ConnectionTimeout : 1;
 
     bool ExpectedResumed    : 1;
     QUIC_STATUS ExpectedCustomTicketValidationResult;
@@ -193,6 +194,7 @@ class TestConnection
     bool GetTransportClosed() const { return TransportClosed; }
     bool GetIsShutdown() const { return IsShutdown; }
     bool GetShutdownTimedOut() const { return ShutdownTimedOut; }
+    bool GetConnectionTimeout() const { return ConnectionTimeout; }
 
     bool GetExpectedResumed() const { return ExpectedResumed; };
     void SetExpectedResumed(bool Value) { ExpectedResumed = Value; }