From 228ab5a934c3ec92796fbee767ffd8467f147151 Mon Sep 17 00:00:00 2001
From: Nick Banks <nibanks@microsoft.com>
Date: Fri, 10 Sep 2021 15:04:28 -0400
Subject: [PATCH] Fix Test Failures in SendResumption Tests (#1970)

---
 src/test/lib/ApiTest.cpp | 56 ++++++++++++++++++----------------------
 1 file changed, 25 insertions(+), 31 deletions(-)

diff --git a/src/test/lib/ApiTest.cpp b/src/test/lib/ApiTest.cpp
index 1b1647b77f..2c437ae03b 100644
--- a/src/test/lib/ApiTest.cpp
+++ b/src/test/lib/ApiTest.cpp
@@ -478,24 +478,10 @@ DummyConnectionCallback(
 }
 
 #ifndef QUIC_DISABLE_0RTT_TESTS
-static
-QUIC_STATUS
-AutoShutdownConnectionCallback(
-    MsQuicConnection* Connection,
-    void* Context,
-    QUIC_CONNECTION_EVENT* Event
-    )
-{
-    if (Event->Type == QUIC_CONNECTION_EVENT_CONNECTED) {
-        if (Context != nullptr) {
-            if (!((CxPlatEvent*)Context)->WaitTimeout(1000)) {
-                TEST_FAILURE("Peer never signaled connected event");
-            }
-        }
-        Connection->Shutdown(0);
-    }
-    return QUIC_STATUS_SUCCESS;
-}
+struct QuicServerSendResumeState {
+    CxPlatEvent ListenerAcceptEvent;
+    CxPlatEvent HandshakeCompleteEvent;
+};
 
 static
 _Function_class_(QUIC_CONNECTION_CALLBACK)
@@ -520,7 +506,7 @@ ResumptionFailConnectionCallback(
                 QUIC_STATUS_INVALID_STATE,
                 Status);
         }
-        CxPlatEventSet(*(CXPLAT_EVENT*)Context);
+        ((QuicServerSendResumeState*)Context)->HandshakeCompleteEvent.Set();
         return QUIC_STATUS_SUCCESS;
     } else if (Event->Type == QUIC_CONNECTION_EVENT_SHUTDOWN_COMPLETE) {
         MsQuic->ConnectionClose(Connection);
@@ -533,7 +519,7 @@ _Function_class_(NEW_CONNECTION_CALLBACK)
 static
 bool
 ListenerFailSendResumeCallback(
-    _In_ TestListener*  Listener,
+    _In_ TestListener* Listener,
     _In_ HQUIC ConnectionHandle
     )
 {
@@ -554,18 +540,21 @@ ListenerFailSendResumeCallback(
         return false;
     }
     MsQuic->SetCallbackHandler(ConnectionHandle, (void*)ResumptionFailConnectionCallback, Listener->Context);
-    CxPlatEventSet(*(CXPLAT_EVENT*)Listener->Context);
+    ((QuicServerSendResumeState*)Listener->Context)->ListenerAcceptEvent.Set();
     return true;
 }
 #endif
 
 void QuicTestValidateConnection()
 {
-    MsQuicRegistration Registration;
+    MsQuicRegistration Registration(true);
     TEST_TRUE(Registration.IsValid());
 
     MsQuicAlpn Alpn("MsQuicTest");
 
+    MsQuicConfiguration ServerConfigurationNoResumption(Registration, Alpn, ServerSelfSignedCredConfig);
+    TEST_TRUE(ServerConfigurationNoResumption.IsValid());
+
     MsQuicSettings Settings;
     Settings.SetServerResumptionLevel(QUIC_SERVER_RESUME_ONLY);
     MsQuicConfiguration ServerConfiguration(Registration, Alpn, Settings, ServerSelfSignedCredConfig);
@@ -939,25 +928,26 @@ void QuicTestValidateConnection()
     //
 #ifndef QUIC_DISABLE_0RTT_TESTS
     {
-        TestListener MyListener(Registration, ListenerFailSendResumeCallback, ServerConfiguration);
+        TestListener MyListener(Registration, ListenerFailSendResumeCallback, ServerConfigurationNoResumption);
         TEST_TRUE(MyListener.IsValid());
 
         TEST_QUIC_SUCCEEDED(MyListener.Start(Alpn, Alpn.Length()));
         QuicAddr ServerLocalAddr;
         TEST_QUIC_SUCCEEDED(MyListener.GetLocalAddr(ServerLocalAddr));
 
-        CxPlatEvent Event;
-        MyListener.Context = &Event;
+        QuicServerSendResumeState ListenerContext;
+        MyListener.Context = &ListenerContext;
 
         {
             //
             // Validate that the resumption ticket call fails in the listener.
             //
             {
-            MsQuicConnection Connection(Registration, CleanUpManual, AutoShutdownConnectionCallback);
+            TestScopeLogger logScope("SendResumption in Listener callback");
+            MsQuicConnection Connection(Registration);
             TEST_QUIC_SUCCEEDED(Connection.GetInitStatus());
             TEST_QUIC_SUCCEEDED(Connection.StartLocalhost(ClientConfiguration, ServerLocalAddr));
-            TEST_TRUE(Event.WaitTimeout(1000));
+            TEST_TRUE(ListenerContext.ListenerAcceptEvent.WaitTimeout(1000));
             }
 
             //
@@ -965,10 +955,12 @@ void QuicTestValidateConnection()
             // because resumption is not enabled.
             //
             {
-            MsQuicConnection Connection(Registration, CleanUpManual, AutoShutdownConnectionCallback, &Event);
+            TestScopeLogger logScope("SendResumption with resumption disabled");
+            MsQuicConnection Connection(Registration);
             TEST_QUIC_SUCCEEDED(Connection.GetInitStatus());
             TEST_QUIC_SUCCEEDED(Connection.StartLocalhost(ClientConfiguration, ServerLocalAddr));
-            TEST_TRUE(Event.WaitTimeout(1000));
+            TEST_TRUE(ListenerContext.ListenerAcceptEvent.WaitTimeout(1000));
+            TEST_TRUE(ListenerContext.HandshakeCompleteEvent.WaitTimeout(1000)); // Wait for server to get connected
             }
 
             //
@@ -976,10 +968,12 @@ void QuicTestValidateConnection()
             // isn't in connected state yet.
             //
             {
-            MsQuicConnection Connection(Registration, CleanUpManual, AutoShutdownConnectionCallback);
+            TestScopeLogger logScope("SendResumption handshake not complete");
+            MsQuicConnection Connection(Registration);
             TEST_QUIC_SUCCEEDED(Connection.GetInitStatus());
             TEST_QUIC_SUCCEEDED(Connection.StartLocalhost(ClientConfiguration, ServerLocalAddr));
-            TEST_TRUE(Event.WaitTimeout(1000));
+            TEST_TRUE(ListenerContext.ListenerAcceptEvent.WaitTimeout(1000));
+            TEST_TRUE(Connection.HandshakeCompleteEvent.WaitTimeout(1000)); // Wait for client to get connected
 
             //
             // TODO: add test case to validate ConnectionSendResumptionTicket: