Fix Client Cert Tests (Windows Kernel)

src/platform/tls_schannel.c

@@ -396,6 +396,17 @@ typedef struct CXPLAT_SEC_CONFIG {
     //
     CXPLAT_TLS_CALLBACKS Callbacks;
 
+#ifdef _KERNEL_MODE
+    //
+    // Impersonation token from the original call to initialize.
+    //
+    PACCESS_TOKEN ImpersonationToken;
+    BOOLEAN IsPrimaryToken;
+    BOOLEAN CopyOnOpen;
+    BOOLEAN EffectiveOnly;
+    SECURITY_IMPERSONATION_LEVEL ImpersonationLevel;
+#endif // _KERNEL_MODE
+
 } CXPLAT_SEC_CONFIG;
 
 typedef struct QUIC_ACH_CONTEXT {
@@ -1249,62 +1260,27 @@ CxPlatTlsSecConfigCreate(
         goto Error;
     }
 
+    CXPLAT_DBG_ASSERT(AchContext->SecConfig != NULL);
+    AchContext->SecConfig->ImpersonationToken =
+        PsReferenceImpersonationToken(
+            PsGetCurrentThread(),
+            &AchContext->SecConfig->CopyOnOpen,
+            &AchContext->SecConfig->EffectiveOnly,
+            &AchContext->SecConfig->ImpersonationLevel);
+
+    if (AchContext->SecConfig->ImpersonationToken == NULL) {
+        AchContext->SecConfig->ImpersonationToken =
+            PsReferencePrimaryToken(
+                PsGetCurrentProcess());
+        AchContext->SecConfig->IsPrimaryToken = TRUE;
+    }
+
     QuicTraceLogVerbose(
         SchannelAchWorkerStart,
         "[ tls] Starting ACH worker");
 
     TLS_WORKER_CONTEXT ThreadContext = { STATUS_SUCCESS, AchContext };
-    if (IsClient) {
-        //
-        // For schannel resumption to work, we have to call the client side
-        // of this from a SYSTEM thread.
-        //
-        HANDLE ThreadHandle;
-        Status =
-            PsCreateSystemThread(
-                &ThreadHandle,
-                THREAD_ALL_ACCESS,
-                NULL,
-                NULL,
-                NULL,
-                CxPlatTlsAchWorker,
-                &ThreadContext);
-        if (QUIC_FAILED(Status)) {
-            QuicTraceEvent(
-                LibraryErrorStatus,
-                "[ lib] ERROR, %u, %s.",
-                Status,
-                "PsCreateSystemThread(CxPlatTlsAchWorker)");
-            goto Error;
-        }
-        void* Thread = NULL;
-        Status =
-            ObReferenceObjectByHandle(
-                ThreadHandle,
-                THREAD_ALL_ACCESS,
-                *PsThreadType,
-                KernelMode,
-                &Thread,
-                NULL);
-        ZwClose(ThreadHandle);
-        if (QUIC_FAILED(Status)) {
-            QuicTraceEvent(
-                LibraryErrorStatus,
-                "[ lib] ERROR, %u, %s.",
-                Status,
-                "ObReferenceObjectByHandle(CxPlatTlsAchWorker)");
-            goto Error;
-        }
-        KeWaitForSingleObject(Thread, Executive, KernelMode, FALSE, NULL);
-        ObDereferenceObject(Thread);
-
-    } else {
-        //
-        // For schannel to successfully load the certificate (even a machine
-        // one), this needs to be on the caller's thread.
-        //
-        CxPlatTlsAchHelper(&ThreadContext);
-    }
+    CxPlatTlsAchHelper(&ThreadContext);
 
     Status = ThreadContext.CompletionStatus;
     AchContext = ThreadContext.AchContext;
@@ -1384,6 +1360,16 @@ CxPlatTlsSecConfigDelete(
         FreeCredentialsHandle(&ServerConfig->CredentialHandle);
     }
 
+#ifdef _KERNEL_MODE
+    if (ServerConfig->ImpersonationToken) {
+        if (ServerConfig->IsPrimaryToken) {
+            PsDereferencePrimaryToken(ServerConfig->ImpersonationToken);
+        } else {
+            PsDereferenceImpersonationToken(ServerConfig->ImpersonationToken);
+        }
+    }
+#endif
+
     CXPLAT_FREE(ServerConfig, QUIC_POOL_TLS_SECCONF);
 }
 
@@ -1782,6 +1768,37 @@ CxPlatTlsWriteDataToSchannel(
     ULONG ContextAttr;
     SECURITY_STATUS SecStatus;
 
+#ifdef _KERNEL_MODE
+    if (TlsContext->SecConfig->ImpersonationToken) {
+        NTSTATUS Status;
+        if (TlsContext->SecConfig->IsPrimaryToken) {
+            Status =
+                PsImpersonateClient(
+                    PsGetCurrentThread(),
+                    TlsContext->SecConfig->ImpersonationToken,
+                    FALSE,
+                    FALSE,
+                    SecurityImpersonation);
+        } else {
+            Status =
+                PsImpersonateClient(
+                    PsGetCurrentThread(),
+                    TlsContext->SecConfig->ImpersonationToken,
+                    TlsContext->SecConfig->CopyOnOpen,
+                    TlsContext->SecConfig->EffectiveOnly,
+                    TlsContext->SecConfig->ImpersonationLevel);
+        }
+        if (!NT_SUCCESS(Status)) {
+            QuicTraceEvent(
+                TlsErrorStatus,
+                "[ tls][%p] ERROR, %u, %s.",
+                TlsContext->Connection,
+                Status,
+                "PsImpersonateClient failed");
+        }
+    }
+#endif
+
     if (TlsContext->IsServer) {
         CXPLAT_DBG_ASSERT(!(TlsContext->SecConfig->Flags & QUIC_CREDENTIAL_FLAG_CLIENT));
 
@@ -1816,6 +1833,16 @@ CxPlatTlsWriteDataToSchannel(
                 NULL);
     }
 
+#ifdef _KERNEL_MODE
+    if (TlsContext->SecConfig->ImpersonationToken) {
+        //
+        // This must only be called on a worker thread.
+        // Otherwise, this might ruin existing impersonation.
+        //
+        PsRevertToSelf();
+    }
+#endif
+
     CXPLAT_TLS_RESULT_FLAGS Result = 0;
 
     SecBuffer* ExtraBuffer = NULL;

src/test/bin/quic_gtest.cpp

@@ -647,12 +647,11 @@ TEST_P(WithHandshakeArgs5, CustomCertificateValidation) {
 TEST_P(WithHandshakeArgs6, ConnectClientCertificate) {
     TestLoggerT<ParamType> Logger("QuicTestConnectClientCertificate", GetParam());
     if (TestingKernelMode) {
-        /*QUIC_RUN_CONNECT_CLIENT_CERT Params = {
+        QUIC_RUN_CONNECT_CLIENT_CERT Params = {
             GetParam().Family,
             (uint8_t)GetParam().UseClientCertificate
         };
-        ASSERT_TRUE(DriverClient.Run(IOCTL_QUIC_RUN_CONNECT_CLIENT_CERT, Params));*/
-        printf("WARNING: ConnectClientCertificate not supported in Kernel Mode yet!\n");
+        ASSERT_TRUE(DriverClient.Run(IOCTL_QUIC_RUN_CONNECT_CLIENT_CERT, Params));
     } else {
         QuicTestConnectClientCertificate(GetParam().Family, GetParam().UseClientCertificate);
     }