Fix TLS fragment tests

Author: anrossi

src/platform/tls_schannel.c

@@ -2445,8 +2445,8 @@ CxPlatTlsWriteDataToSchannel(
             //
             // Not all the input buffer was consumed. There is some 'extra' left over.
             //
-            CXPLAT_DBG_ASSERT(InSecBuffers[1].cbBuffer <= *InBufferLength);
-            *InBufferLength -= InSecBuffers[1].cbBuffer;
+            CXPLAT_DBG_ASSERT(ExtraBuffer->cbBuffer <= *InBufferLength);
+            *InBufferLength -= ExtraBuffer->cbBuffer;
         }
 
         QuicTraceLogConnInfo(
@@ -2699,6 +2699,19 @@ CxPlatTlsWriteDataToSchannel(
             TlsContext->PeerTransportParamsLength = 0;
         }
 
+        //
+        // Some or all of the input data was processed. There may or may not be
+        // corresponding output data to send in response.
+        //
+
+        if (ExtraBuffer != NULL && ExtraBuffer->cbBuffer > 0) {
+            //
+            // Not all the input buffer was consumed. There is some 'extra' left over.
+            //
+            CXPLAT_DBG_ASSERT(ExtraBuffer->cbBuffer <= *InBufferLength);
+            *InBufferLength -= ExtraBuffer->cbBuffer;
+        }
+
         break;
 
     case SEC_E_INCOMPLETE_MESSAGE:

src/platform/unittest/TlsTest.cpp

@@ -488,6 +488,9 @@ struct TlsTest : public ::testing::TestWithParam<bool>
                 if (ConsumedBuffer > 0) {
                     Buffer += ConsumedBuffer;
                     BufferLength -= ConsumedBuffer;
+                    if (ConsumedBuffer > BufferLength) {
+                        ConsumedBuffer = BufferLength;
+                    }
                 } else {
                     ConsumedBuffer = FragmentSize * ++Count;
                     ConsumedBuffer = CXPLAT_MIN(ConsumedBuffer, BufferLength);
@@ -1280,7 +1283,7 @@ TEST_F(TlsTest, HandshakeMultiAlpnBoth)
     DoHandshake(ServerContext, ClientContext);
 }
 
-TEST_F(TlsTest, DISABLED_HandshakeFragmented) // Remove DISABLED_ after fixing #6035
+TEST_F(TlsTest, HandshakeFragmented)
 {
     CxPlatClientSecConfig ClientConfig;
     CxPlatServerSecConfig ServerConfig;
@@ -1290,7 +1293,7 @@ TEST_F(TlsTest, DISABLED_HandshakeFragmented) // Remove DISABLED_ after fixing #
     DoHandshake(ServerContext, ClientContext, 200);
 }
 
-TEST_F(TlsTest, DISABLED_HandshakeVeryFragmented) // Remove DISABLED_ after fixing #6035
+TEST_F(TlsTest, HandshakeVeryFragmented)
 {
     CxPlatClientSecConfig ClientConfig;
     CxPlatServerSecConfig ServerConfig;