[Bug] Data loss in KtlsRecv on unexpected EOF due to unread mbedTLS buffers

[BadFr0ogy]

Description:
Hi, thanks for the great work on this project! I've been using it and encountered a data loss issue in KtlsRecv when the peer abruptly closes the connection. I have investigated the issue and made a local patch that seems to fix it. I'd like to discuss whether this is the best approach to merge upstream.

The Problem:
During TLS communication, if the underlying network layer (TDI) experiences a sudden disconnect or RST, mbedtls_ssl_read returns MBEDTLS_ERR_SSL_CONN_EOF.
In the original implementation of KtlsRecv, receiving this error code immediately returns STATUS_CONNECTION_DISCONNECTED. However, mbedTLS might have already decrypted some data that is currently sitting in its internal buffer. Returning immediately causes this pending decrypted data to be permanently lost.

Root Cause:
Because of the nature of the TDI layer and network buffering, the socket might report an EOF/disconnect to mbedTLS, prompting it to return MBEDTLS_ERR_SSL_CONN_EOF. But mbedTLS does not flush its internal available bytes automatically on EOF.

Proposed Solution:
I modified KtlsRecv to check for available bytes using mbedtls_ssl_get_bytes_avail before dropping the connection. If MBEDTLS_ERR_SSL_CONN_EOF is returned but there are still bytes left, we should force a read to salvage the remaining data before returning the disconnected status.

Additionally, I improved the Plain TCP path to properly handle STATUS_CONNECTION_RESET.

Here is the core logic of my patch for the TLS part:

    if (ret == MBEDTLS_ERR_SSL_CONN_EOF) {
        // Check if there is still decrypted data pending in mbedTLS internal buffers
        size_t available = mbedtls_ssl_get_bytes_avail(&Session->ssl);
        
        if (available > 0) {
            KTLS_LOG_WARN("Forcing read of remaining %zu bytes despite EOF", available);
            // Salvage the remaining data
            ret = mbedtls_ssl_read(&Session->ssl, (unsigned char*)Buffer, min(BufferSize, available));
            if (ret > 0) {
                *BytesReceived = (ULONG)ret;
                return STATUS_SUCCESS;
            }
        }
        
        KTLS_LOG_WARN("TLS Connection reset/closed unexpectedly (EOF).");
        return STATUS_CONNECTION_DISCONNECTED;
    }

Questions for Discussion:

1. Does this approach of using mbedtls_ssl_get_bytes_avail align with your design for KtlsRecv?

2. Is there any edge case in the TDI wrapper that might cause mbedtls_ssl_read to hang if we force read after an EOF? (In my tests, it works perfectly and saves the dropped packets).

3.I also added explicit handling for STATUS_CONNECTION_RESET in the Plain TCP section. Would you like me to submit a PR including both fixes?
Looking forward to your thoughts!

[BadFr0ogy]

Update: Another related issue when sending large payloads (~1MB)

While testing, I encountered another peculiar behavior related to MBEDTLS_ERR_SSL_CONN_EOF.

The Scenario:
I use KtlsSend in a loop to send a relatively large payload (around 1MB). The send operation completes successfully (100% sent). Immediately after sending, I call KtlsRecv to get the server's response.

However, KtlsRecv almost always fails immediately. It receives nothing, and hits the exact MBEDTLS_ERR_SSL_CONN_EOF block with available == 0. (Only once did it successfully receive the header and about 24% of the data before disconnecting).

Here are the driver logs:

[Driver][TLS 1MB] Sent: 1048576 / 1048580 bytes (99%)
[Driver][TLS] Sent: 1048580 / 1048580 bytes (100%)
[Driver][TLS] All data sent, waiting for response...
 
KTLS [DEBUG]: Calling mbedtls_ssl_read for 4 bytes...
KTLS [WARN]:  TLS Connection reset/closed unexpectedly (EOF). Pending bytes in mbedtls: 0
[Driver][TLS 1MB] Fatal Error receiving header: 0xC000020C
 
KTLS [INFO]:  Closing session
KTLS [INFO]:  Session closed successfully

My Question:
Could you please help clarify if this behavior is an issue with the underlying TDI implementation, or is it an issue with mbedTLS itself? I would really appreciate it if you could point me in the right direction or provide some advice on how to fix/debug this.

[4l3x777]

Description: Hi, thanks for the great work on this project! I've been using it and encountered a data loss issue in KtlsRecv when the peer abruptly closes the connection. I have investigated the issue and made a local patch that seems to fix it. I'd like to discuss whether this is the best approach to merge upstream.

The Problem: During TLS communication, if the underlying network layer (TDI) experiences a sudden disconnect or RST, mbedtls_ssl_read returns MBEDTLS_ERR_SSL_CONN_EOF. In the original implementation of KtlsRecv, receiving this error code immediately returns STATUS_CONNECTION_DISCONNECTED. However, mbedTLS might have already decrypted some data that is currently sitting in its internal buffer. Returning immediately causes this pending decrypted data to be permanently lost.

Root Cause: Because of the nature of the TDI layer and network buffering, the socket might report an EOF/disconnect to mbedTLS, prompting it to return MBEDTLS_ERR_SSL_CONN_EOF. But mbedTLS does not flush its internal available bytes automatically on EOF.

Proposed Solution: I modified KtlsRecv to check for available bytes using mbedtls_ssl_get_bytes_avail before dropping the connection. If MBEDTLS_ERR_SSL_CONN_EOF is returned but there are still bytes left, we should force a read to salvage the remaining data before returning the disconnected status.

Additionally, I improved the Plain TCP path to properly handle STATUS_CONNECTION_RESET.

Here is the core logic of my patch for the TLS part:

    if (ret == MBEDTLS_ERR_SSL_CONN_EOF) {
        // Check if there is still decrypted data pending in mbedTLS internal buffers
        size_t available = mbedtls_ssl_get_bytes_avail(&Session->ssl);
        
        if (available > 0) {
            KTLS_LOG_WARN("Forcing read of remaining %zu bytes despite EOF", available);
            // Salvage the remaining data
            ret = mbedtls_ssl_read(&Session->ssl, (unsigned char*)Buffer, min(BufferSize, available));
            if (ret > 0) {
                *BytesReceived = (ULONG)ret;
                return STATUS_SUCCESS;
            }
        }
        
        KTLS_LOG_WARN("TLS Connection reset/closed unexpectedly (EOF).");
        return STATUS_CONNECTION_DISCONNECTED;
    }

Questions for Discussion:

1. Does this approach of using mbedtls_ssl_get_bytes_avail align with your design for KtlsRecv?
2. Is there any edge case in the TDI wrapper that might cause mbedtls_ssl_read to hang if we force read after an EOF? (In my tests, it works perfectly and saves the dropped packets).

3.I also added explicit handling for STATUS_CONNECTION_RESET in the Plain TCP section. Would you like me to submit a PR including both fixes? Looking forward to your thoughts!

Hi! Thank you for your attention to the project, it's a pleasure!

Regarding the first point, I agree, the approach with mbedtls_ssl_get_bytes_avail is absolutely correct and corresponds to the best practices for mbedTLS, you can send your version to the pull request.

According to the second point, theoretically there are borderline cases, but with the correct implementation of mbedtls transport callbacks, there should be no hanging.
If the recv callback (TDI wrapper) returns EOF, mbedtls_ssl_read() eventually returns 0 as a sign of a closed connection. The hangup does not occur in mbedtls itself, but due to the fact that the transport callback does not signal the final EOF and always returns, for example (MBEDTLS_ERR_SSL_WANT_READ). Alternatively, you can see how the underlying TDI layer gives EOF (in my opinion, this would be correct)

For the third one, of course, send it, I will be glad of your solutions and comments!

[4l3x777]

Update: Another related issue when sending large payloads (~1MB)

While testing, I encountered another peculiar behavior related to MBEDTLS_ERR_SSL_CONN_EOF.

The Scenario: I use KtlsSend in a loop to send a relatively large payload (around 1MB). The send operation completes successfully (100% sent). Immediately after sending, I call KtlsRecv to get the server's response.

However, KtlsRecv almost always fails immediately. It receives nothing, and hits the exact MBEDTLS_ERR_SSL_CONN_EOF block with available == 0. (Only once did it successfully receive the header and about 24% of the data before disconnecting).

Here are the driver logs:

[Driver][TLS 1MB] Sent: 1048576 / 1048580 bytes (99%)
[Driver][TLS] Sent: 1048580 / 1048580 bytes (100%)
[Driver][TLS] All data sent, waiting for response...
 
KTLS [DEBUG]: Calling mbedtls_ssl_read for 4 bytes...
KTLS [WARN]:  TLS Connection reset/closed unexpectedly (EOF). Pending bytes in mbedtls: 0
[Driver][TLS 1MB] Fatal Error receiving header: 0xC000020C
 
KTLS [INFO]:  Closing session
KTLS [INFO]:  Session closed successfully

My Question: Could you please help clarify if this behavior is an issue with the underlying TDI implementation, or is it an issue with mbedTLS itself? I would really appreciate it if you could point me in the right direction or provide some advice on how to fix/debug this.

I also tried to transfer a large file via https chunked streaming (about 100MB in size), but I checked only localy, i.e. there was an https server with a self-signed certificate.
The file was transmitted correctly and completely (the checksums matched), the response from the server was checked through the client and server logs, they also matched. Here I propose to take the following section as a basis:

// Stream file from disk without loading into memory
UNICODE_STRING FilePath;
RtlInitUnicodeString(&FilePath, L"\\??\\C:\\test_file.bin");

KHTTP_FILE File = {
    .FieldName = "file",
    .FileName = "test_file.bin",
    .ContentType = "application/octet-stream",
    .UseFileStream = TRUE,
    .FilePath = &FilePath,
    .Data = NULL,
    .DataLength = 0
};

KHTTP_CONFIG Config = {
    .UseHttps = TRUE,
    .TimeoutMs = 300000,        // 5 minutes
    .UseChunkedTransfer = TRUE,
    .ChunkSize = 256 * 1024,    // 256KB chunks
    .ProgressCallback = ProgressCallback
};

PKHTTP_RESPONSE Response = NULL;
NTSTATUS Status = KhttpPostMultipartChunked(
    "https://192.168.56.1:8443/upload",
    NULL,
    NULL, 0,
    &File, 1,
    &Config,
    &Response
);

As a research:

1. Try to look at the logs, you can put them after KtlsSend, before KtlsRecv, to see what happens in the session.
2. Since TCP ACKs from the server may be delayed, try to put small delays of about 100-500ms between KtlsSend and KtlsRecv. Here it is possible that the server at the KtlsRecv moment may either not have received all the data yet, or has already sent a response and closed the connection.
   PS: In general, it seems to me that the main problem is the out-of-sync between the client code and the real state of the TCP connection on the TDI layer. This requires deeper integration with TDI to monitor the status of the send buffer before attempting to read, if I understood the problem correctly.