SERVER-41587 Improve SECBUFFER_EXTRA handling

(cherry picked from commit 61c90181a842c8863df9e5685eb48e6fbc063ea4)

5 files changed, 75 insertions, 17 deletions

diff --git a/jstests/ssl/ssl_fragment.js b/jstests/ssl/ssl_fragment.js
index 6d08da8e834..bdb16559a0b 100644
--- a/jstests/ssl/ssl_fragment.js
+++ b/jstests/ssl/ssl_fragment.js
@@ -20,12 +20,27 @@
         assert.eq(s, read, "Did not receive value written");
     }
 
-    const options = {
+    let options = {
         sslMode: "requireSSL",
         sslPEMKeyFile: "jstests/libs/server.pem",
         networkMessageCompressors: 'disabled',
     };
 
+    let mongosOptions = {
+        sslMode: "requireSSL",
+        sslPEMKeyFile: "jstests/libs/server.pem",
+        networkMessageCompressors: 'disabled',
+    };
+
+    if (_isWindows()) {
+        // Force the ASIO stack to do small reads which will excerise the schannel buffering code
+        // and significantly slow down the test
+        options = Object.extend(options,
+                                {setParameter: {"failpoint.smallTLSReads": "{'mode':'alwaysOn'}"}});
+        mongosOptions = Object.extend(
+            mongosOptions, {setParameter: {"failpoint.smallTLSReads": "{'mode':'alwaysOn'}"}});
+    }
+
     const mongod = MongoRunner.runMongod(options);
     runTest(mongod);
     MongoRunner.stopMongod(mongod);
@@ -37,11 +52,13 @@
         config: 1,
         other: {
             configOptions: options,
-            mongosOptions: options,
+            mongosOptions: mongosOptions,
             shardOptions: options,
             shardAsReplicaSet: false,
         }
     });
+
     runTest(st.s0);
     st.stop();
+
 })();
diff --git a/src/mongo/util/net/ssl/detail/impl/schannel.ipp b/src/mongo/util/net/ssl/detail/impl/schannel.ipp
index 81f8ffc21c5..0fcbf547c39 100644
--- a/src/mongo/util/net/ssl/detail/impl/schannel.ipp
+++ b/src/mongo/util/net/ssl/detail/impl/schannel.ipp
@@ -327,10 +327,13 @@ ssl_want SSLHandshakeManager::doServerHandshake(asio::error_code& ec,
     // ASC_RET_MUTUAL_AUTH is not set since we do our own certificate validation later.
     invariant(attribs == (retAttribs | ASC_RET_EXTENDED_ERROR | ASC_RET_MUTUAL_AUTH));
 
-    if (inputBuffers[1].BufferType == SECBUFFER_EXTRA && inputBuffers[1].pvBuffer != nullptr &&
-        inputBuffers[1].cbBuffer > 0) {
+    if (inputBuffers[1].BufferType == SECBUFFER_EXTRA && inputBuffers[1].cbBuffer > 0) {
+        // SECBUFFER_EXTRA do not set pvBuffer, just cbBuffer.
+        // cbBuffer tells us how much remaining in the buffer is extra
         _pExtraEncryptedBuffer->reset();
-        _pExtraEncryptedBuffer->append(inputBuffers[1].pvBuffer, inputBuffers[1].cbBuffer);
+        _pExtraEncryptedBuffer->append(_pInBuffer->data() +
+                                           (_pInBuffer->size() - inputBuffers[1].cbBuffer),
+                                       inputBuffers[1].cbBuffer);
     }
 
 
@@ -350,12 +353,14 @@ ssl_want SSLHandshakeManager::doServerHandshake(asio::error_code& ec,
     // Reset the input buffer
     _pInBuffer->reset();
 
-    // Check if we have any additional encrypted data
+    // Check if we have any additional data
     if (!_pExtraEncryptedBuffer->empty()) {
         _pInBuffer->swap(*_pExtraEncryptedBuffer);
         _pExtraEncryptedBuffer->reset();
 
-        setState(State::HaveEncryptedData);
+        // When doing the handshake and we have extra data, this means we have an incomplete tls
+        // record and need more bytes to complete the tls record.
+        setState(State::NeedMoreHandshakeData);
     }
 
     if (needOutput) {
@@ -467,10 +472,13 @@ ssl_want SSLHandshakeManager::doClientHandshake(asio::error_code& ec) {
 
     if (_pInBuffer->size()) {
         // Locate (optional) extra buffer
-        if (inputBuffers[1].BufferType == SECBUFFER_EXTRA && inputBuffers[1].pvBuffer != nullptr &&
-            inputBuffers[1].cbBuffer > 0) {
+        if (inputBuffers[1].BufferType == SECBUFFER_EXTRA && inputBuffers[1].cbBuffer > 0) {
+            // SECBUFFER_EXTRA do not set pvBuffer, just cbBuffer
+            // cbBuffer tells us how much remaining in the buffer is extra
             _pExtraEncryptedBuffer->reset();
-            _pExtraEncryptedBuffer->append(inputBuffers[1].pvBuffer, inputBuffers[1].cbBuffer);
+            _pExtraEncryptedBuffer->append(_pInBuffer->data() +
+                                               (_pInBuffer->size() - inputBuffers[1].cbBuffer),
+                                           inputBuffers[1].cbBuffer);
         }
     }
 
@@ -495,7 +503,9 @@ ssl_want SSLHandshakeManager::doClientHandshake(asio::error_code& ec) {
         _pInBuffer->swap(*_pExtraEncryptedBuffer);
         _pExtraEncryptedBuffer->reset();
 
-        setState(State::HaveEncryptedData);
+        // When doing the handshake and we have extra data, this means we have an incomplete tls
+        // record and need more bytes to complete the tls record.
+        setState(State::NeedMoreHandshakeData);
     }
 
     if (needOutput) {
diff --git a/src/mongo/util/net/ssl/detail/io.hpp b/src/mongo/util/net/ssl/detail/io.hpp
index b943b1fb16f..8a702abc9dd 100644
--- a/src/mongo/util/net/ssl/detail/io.hpp
+++ b/src/mongo/util/net/ssl/detail/io.hpp
@@ -18,6 +18,7 @@
 #include "asio/detail/config.hpp"
 
 #include "asio/write.hpp"
+#include "mongo/util/fail_point_service.h"
 #include "mongo/util/net/ssl/detail/engine.hpp"
 #include "mongo/util/net/ssl/detail/stream_core.hpp"
 
@@ -27,6 +28,9 @@ namespace asio {
 namespace ssl {
 namespace detail {
 
+// Failpoint to force small reads of data to exercise the SChannel buffering code
+MONGO_FAIL_POINT_DECLARE(smallTLSReads);
+
 template <typename Stream, typename Operation>
 std::size_t io(Stream& next_layer, stream_core& core, const Operation& op, asio::error_code& ec) {
     std::size_t bytes_transferred = 0;
@@ -36,9 +40,22 @@ std::size_t io(Stream& next_layer, stream_core& core, const Operation& op, asio:
 
                 // If the input buffer is empty then we need to read some more data from
                 // the underlying transport.
-                if (core.input_.size() == 0)
-                    core.input_ = asio::buffer(core.input_buffer_,
-                                               next_layer.read_some(core.input_buffer_, ec));
+                if (core.input_.size() == 0) {
+                    // Read tiny amounts of TLS data to test the SChannel buffering code
+                    if (MONGO_FAIL_POINT(smallTLSReads)) {
+                        core.input_ =
+                            asio::buffer(core.input_buffer_,
+                                         next_layer.read_some(
+                                             mutable_buffer(core.input_buffer_.data(),
+                                                            std::min(core.input_buffer_.size(),
+                                                                     static_cast<size_t>(10UL))),
+                                             ec));
+
+                    } else {
+                        core.input_ = asio::buffer(core.input_buffer_,
+                                                   next_layer.read_some(core.input_buffer_, ec));
+                    }
+                }
 
                 // Pass the new input data to the engine.
                 core.input_ = core.engine_.put_input(core.input_);
diff --git a/src/mongo/util/net/ssl/detail/schannel.hpp b/src/mongo/util/net/ssl/detail/schannel.hpp
index 5b72e84b2de..8f1780cb4c6 100644
--- a/src/mongo/util/net/ssl/detail/schannel.hpp
+++ b/src/mongo/util/net/ssl/detail/schannel.hpp
@@ -361,9 +361,12 @@ private:
      * +----------------+     +-----------------------+     +-------------------+     +------+
      * | HandshakeStart | --> | NeedMoreHandshakeData | --> | HaveEncryptedData | --> | Done |
      * +----------------+     +-----------------------+     +-------------------+     +------+
+     *                          ^                   |
+     *                          +-------------------+
      *
-     * "[ HandshakeStart ] --> [ NeedMoreHandshakeData ] --> [HaveEncryptedData] -> [
-     * NeedMoreHandshakeData], [Done] " | graph-easy
+     * echo "[ HandshakeStart ] --> [ NeedMoreHandshakeData ] --> [ NeedMoreHandshakeData ] -->
+     * [HaveEncryptedData] -> [NeedMoreHandshakeData], [Done]" | graph-easy "[ HandshakeStart ] -->
+     * [ NeedMoreHandshakeData ] --> [HaveEncryptedData] -> [
      */
     enum class State {
         // Initial state
@@ -385,7 +388,7 @@ private:
     void setState(State s) {
         ASSERT_STATE_TRANSITION(_state == State::HandshakeStart, s == State::NeedMoreHandshakeData);
         ASSERT_STATE_TRANSITION(_state == State::NeedMoreHandshakeData,
-                                s == State::HaveEncryptedData);
+                                s == State::HaveEncryptedData || s == State::NeedMoreHandshakeData);
         ASSERT_STATE_TRANSITION(_state == State::HaveEncryptedData,
                                 s == State::NeedMoreHandshakeData || s == State::Done);
         _state = s;
diff --git a/src/mongo/util/net/ssl_stream.cpp b/src/mongo/util/net/ssl_stream.cpp
index a99723d4a32..6bbe58864b0 100644
--- a/src/mongo/util/net/ssl_stream.cpp
+++ b/src/mongo/util/net/ssl_stream.cpp
@@ -11,6 +11,17 @@
 
 #include "mongo/config.h"
 
+#include "mongo/util/fail_point_service.h"
+
 #ifdef MONGO_CONFIG_SSL
 #include "mongo/util/net/ssl/impl/src.hpp"
+
+namespace asio {
+namespace ssl {
+namespace detail {
+MONGO_FAIL_POINT_DEFINE(smallTLSReads);
+}  // namespce detail
+}  // namespce ssl
+}  // namespce asio
+
 #endif