[PATCH] secur32: Handle incomplete messages in schan_InitializeSecurityContextW().

[Henri Verbeet]

This supersedes patch 87616. Compared to the previous version this removes a
failing ok() on NT4, since that is already handled by the code below anyway.
---
 dlls/secur32/schannel.c       |   24 +++++++++++++++++++
 dlls/secur32/tests/schannel.c |   51 ++++++++++++++++++++++++++++++++++++++--
 2 files changed, 72 insertions(+), 3 deletions(-)

diff --git a/dlls/secur32/schannel.c b/dlls/secur32/schannel.c
index f703ab7..af34e66 100644
--- a/dlls/secur32/schannel.c
+++ b/dlls/secur32/schannel.c
@@ -714,6 +714,30 @@ static SECURITY_STATUS SEC_ENTRY schan_InitializeSecurityContextW(
     }
     else
     {
+        unsigned int expected_size;
+        unsigned char *ptr;
+        SecBuffer *buffer;
+        int idx;
+
+        if (!pInput)
+            return SEC_E_INCOMPLETE_MESSAGE;
+
+        idx = schan_find_sec_buffer_idx(pInput, 0, SECBUFFER_TOKEN);
+        if (idx == -1)
+            return SEC_E_INCOMPLETE_MESSAGE;
+
+        buffer = &pInput->pBuffers[idx];
+        if (buffer->cbBuffer < 5)
+            return SEC_E_INCOMPLETE_MESSAGE;
+
+        ptr = buffer->pvBuffer;
+        expected_size = 5 + ((ptr[3] << 8) | ptr[4]);
+        if (buffer->cbBuffer < expected_size)
+        {
+            TRACE("Expected %u bytes, but buffer only contains %u bytes.\n", expected_size, buffer->cbBuffer);
+            return SEC_E_INCOMPLETE_MESSAGE;
+        }
+
         ctx = schan_get_object(phContext->dwLower, SCHAN_HANDLE_CTX);
     }
 
diff --git a/dlls/secur32/tests/schannel.c b/dlls/secur32/tests/schannel.c
index a729394..99394c3 100644
--- a/dlls/secur32/tests/schannel.c
+++ b/dlls/secur32/tests/schannel.c
@@ -640,21 +640,66 @@ static void test_communication(void)
 
     buffers[1].cBuffers = 1;
     buffers[1].pBuffers[0].BufferType = SECBUFFER_TOKEN;
-    data_size = buffers[0].pBuffers[0].cbBuffer;
     status = pInitializeSecurityContextA(&cred_handle, &context, (SEC_CHAR *)"localhost",
             ISC_REQ_CONFIDENTIALITY|ISC_REQ_STREAM,
             0, 0, &buffers[1], 0, NULL, &buffers[0], &attrs, NULL);
     ok(status == SEC_E_INVALID_TOKEN, "Expected SEC_E_INVALID_TOKEN, got %08x\n", status);
 
     buffers[0].pBuffers[0].cbBuffer = buf_size;
-    buffers[1].cBuffers = 4;
-    buffers[1].pBuffers[0].cbBuffer = buf_size;
 
     status = pInitializeSecurityContextA(&cred_handle, NULL, (SEC_CHAR *)"localhost",
             ISC_REQ_CONFIDENTIALITY|ISC_REQ_STREAM,
             0, 0, NULL, 0, &context, &buffers[0], &attrs, NULL);
     ok(status == SEC_I_CONTINUE_NEEDED, "Expected SEC_I_CONTINUE_NEEDED, got %08x\n", status);
 
+    buf = &buffers[0].pBuffers[0];
+    send(sock, buf->pvBuffer, buf->cbBuffer, 0);
+    buf->cbBuffer = buf_size;
+
+    status = pInitializeSecurityContextA(&cred_handle, &context, (SEC_CHAR *)"localhost",
+            ISC_REQ_CONFIDENTIALITY|ISC_REQ_STREAM,
+            0, 0, NULL, 0, NULL, &buffers[0], &attrs, NULL);
+    ok(status == SEC_E_INCOMPLETE_MESSAGE, "Got unexpected status %#x.\n", status);
+    ok(buffers[0].pBuffers[0].cbBuffer == buf_size, "Output buffer size changed.\n");
+    ok(buffers[0].pBuffers[0].BufferType == SECBUFFER_TOKEN, "Output buffer type changed.\n");
+
+    buffers[1].cBuffers = 4;
+    buffers[1].pBuffers[0].cbBuffer = 0;
+
+    status = pInitializeSecurityContextA(&cred_handle, &context, (SEC_CHAR *)"localhost",
+            ISC_REQ_CONFIDENTIALITY|ISC_REQ_STREAM,
+            0, 0, &buffers[1], 0, NULL, &buffers[0], &attrs, NULL);
+    ok(status == SEC_E_INCOMPLETE_MESSAGE, "Got unexpected status %#x.\n", status);
+    ok(buffers[0].pBuffers[0].cbBuffer == buf_size, "Output buffer size changed.\n");
+    ok(buffers[0].pBuffers[0].BufferType == SECBUFFER_TOKEN, "Output buffer type changed.\n");
+
+    buf = &buffers[1].pBuffers[0];
+    buf->cbBuffer = buf_size;
+    ret = receive_data(sock, buf);
+    if (ret == -1)
+        return;
+
+    buffers[1].pBuffers[0].cbBuffer = 4;
+    status = pInitializeSecurityContextA(&cred_handle, &context, (SEC_CHAR *)"localhost",
+            ISC_REQ_CONFIDENTIALITY|ISC_REQ_STREAM,
+            0, 0, &buffers[1], 0, NULL, &buffers[0], &attrs, NULL);
+    ok(status == SEC_E_INCOMPLETE_MESSAGE, "Got unexpected status %#x.\n", status);
+    ok(buffers[0].pBuffers[0].cbBuffer == buf_size, "Output buffer size changed.\n");
+    ok(buffers[0].pBuffers[0].BufferType == SECBUFFER_TOKEN, "Output buffer type changed.\n");
+
+    buffers[1].pBuffers[0].cbBuffer = 5;
+    status = pInitializeSecurityContextA(&cred_handle, &context, (SEC_CHAR *)"localhost",
+            ISC_REQ_CONFIDENTIALITY|ISC_REQ_STREAM,
+            0, 0, &buffers[1], 0, NULL, &buffers[0], &attrs, NULL);
+    ok(status == SEC_E_INCOMPLETE_MESSAGE, "Got unexpected status %#x.\n", status);
+    ok(buffers[0].pBuffers[0].cbBuffer == buf_size, "Output buffer size changed.\n");
+    ok(buffers[0].pBuffers[0].BufferType == SECBUFFER_TOKEN, "Output buffer type changed.\n");
+
+    buffers[1].pBuffers[0].cbBuffer = ret;
+    status = pInitializeSecurityContextA(&cred_handle, &context, (SEC_CHAR *)"localhost",
+            ISC_REQ_CONFIDENTIALITY|ISC_REQ_STREAM,
+            0, 0, &buffers[1], 0, NULL, &buffers[0], &attrs, NULL);
+    buffers[1].pBuffers[0].cbBuffer = buf_size;
     while (status == SEC_I_CONTINUE_NEEDED)
     {
         buf = &buffers[0].pBuffers[0];
-- 
1.7.3.4