QAbstractSocket: wait with closing until all bytes have been written

only disconnect from host when all bytes have been written; i.e. not
only check whether the write buffer is empty, but also check whether
the socket engine has still bytes to write. This is necessary for
HTTP and SOCKS5 socket engine, because they both contain an inner TCP
socket which also does buffering. For the native socket engine, there
is no difference with this patch.

Reviewed-by: Markus Goetz
Reviewed-by: Thiago Macieira

9 files changed, 55 insertions, 7 deletions

diff --git a/src/network/socket/qabstractsocket.cpp b/src/network/socket/qabstractsocket.cpp
index 9fb0b47..955fee0 100644
--- a/src/network/socket/qabstractsocket.cpp
+++ b/src/network/socket/qabstractsocket.cpp
@@ -669,11 +669,11 @@ bool QAbstractSocketPrivate::canWriteNotification()
 
     if (socketEngine) {
 #if defined (Q_OS_WIN)
-	if (!writeBuffer.isEmpty())
-	    socketEngine->setWriteNotificationEnabled(true);
+        if (!writeBuffer.isEmpty())
+            socketEngine->setWriteNotificationEnabled(true);
 #else
-	if (writeBuffer.isEmpty())
-	    socketEngine->setWriteNotificationEnabled(false);
+        if (writeBuffer.isEmpty() && socketEngine->bytesToWrite() == 0)
+            socketEngine->setWriteNotificationEnabled(false);
 #endif
     }
 
@@ -710,11 +710,17 @@ void QAbstractSocketPrivate::connectionNotification()
 bool QAbstractSocketPrivate::flush()
 {
     Q_Q(QAbstractSocket);
-    if (!socketEngine || !socketEngine->isValid() || writeBuffer.isEmpty()) {
+    if (!socketEngine || !socketEngine->isValid() || (writeBuffer.isEmpty()
+        && socketEngine->bytesToWrite() == 0)) {
 #if defined (QABSTRACTSOCKET_DEBUG)
     qDebug("QAbstractSocketPrivate::flush() nothing to do: valid ? %s, writeBuffer.isEmpty() ? %s",
            socketEngine->isValid() ? "yes" : "no", writeBuffer.isEmpty() ? "yes" : "no");
 #endif
+
+        // this covers the case when the buffer was empty, but we had to wait for the socket engine to finish
+        if (state == QAbstractSocket::ClosingState)
+            q->disconnectFromHost();
+
         return false;
     }
 
@@ -751,7 +757,8 @@ bool QAbstractSocketPrivate::flush()
         }
     }
 
-    if (writeBuffer.isEmpty() && socketEngine && socketEngine->isWriteNotificationEnabled())
+    if (writeBuffer.isEmpty() && socketEngine && socketEngine->isWriteNotificationEnabled()
+        && !socketEngine->bytesToWrite())
         socketEngine->setWriteNotificationEnabled(false);
     if (state == QAbstractSocket::ClosingState)
         q->disconnectFromHost();
@@ -2347,7 +2354,8 @@ void QAbstractSocket::disconnectFromHostImplementation()
         }
 
         // Wait for pending data to be written.
-        if (d->socketEngine && d->socketEngine->isValid() && d->writeBuffer.size() > 0) {
+        if (d->socketEngine && d->socketEngine->isValid() && (d->writeBuffer.size() > 0
+            || d->socketEngine->bytesToWrite() > 0)) {
             d->socketEngine->setWriteNotificationEnabled(true);
 
 #if defined(QABSTRACTSOCKET_DEBUG)
diff --git a/src/network/socket/qabstractsocketengine_p.h b/src/network/socket/qabstractsocketengine_p.h
index c639092..14b3c81 100644
--- a/src/network/socket/qabstractsocketengine_p.h
+++ b/src/network/socket/qabstractsocketengine_p.h
@@ -126,6 +126,8 @@ public:
     virtual qint64 pendingDatagramSize() const = 0;
 #endif
 
+    virtual qint64 bytesToWrite() const = 0;
+
     virtual int option(SocketOption option) const = 0;
     virtual bool setOption(SocketOption option, int value) = 0;
 
diff --git a/src/network/socket/qhttpsocketengine.cpp b/src/network/socket/qhttpsocketengine.cpp
index fb61dbf..5c28318 100644
--- a/src/network/socket/qhttpsocketengine.cpp
+++ b/src/network/socket/qhttpsocketengine.cpp
@@ -276,6 +276,16 @@ qint64 QHttpSocketEngine::pendingDatagramSize() const
 }
 #endif // QT_NO_UDPSOCKET
 
+qint64 QHttpSocketEngine::bytesToWrite() const
+{
+    Q_D(const QHttpSocketEngine);
+    if (d->socket) {
+        return d->socket->bytesToWrite();
+    } else {
+        return 0;
+    }
+}
+
 int QHttpSocketEngine::option(SocketOption option) const
 {
     Q_D(const QHttpSocketEngine);
diff --git a/src/network/socket/qhttpsocketengine_p.h b/src/network/socket/qhttpsocketengine_p.h
index a423116..76430db 100644
--- a/src/network/socket/qhttpsocketengine_p.h
+++ b/src/network/socket/qhttpsocketengine_p.h
@@ -110,6 +110,8 @@ public:
     qint64 pendingDatagramSize() const;
 #endif // QT_NO_UDPSOCKET
 
+    qint64 bytesToWrite() const;
+
     int option(SocketOption option) const;
     bool setOption(SocketOption option, int value);
 
diff --git a/src/network/socket/qnativesocketengine.cpp b/src/network/socket/qnativesocketengine.cpp
index e7f8401..a150b26 100644
--- a/src/network/socket/qnativesocketengine.cpp
+++ b/src/network/socket/qnativesocketengine.cpp
@@ -754,6 +754,12 @@ qint64 QNativeSocketEngine::write(const char *data, qint64 size)
     return d->nativeWrite(data, size);
 }
 
+
+qint64 QNativeSocketEngine::bytesToWrite() const
+{
+    return 0;
+}
+
 /*!
     Reads up to \a maxSize bytes into \a data from the socket.
     Returns the number of bytes read, or -1 if an error occurred.
diff --git a/src/network/socket/qnativesocketengine_p.h b/src/network/socket/qnativesocketengine_p.h
index 1f6a243..a03d8f1 100644
--- a/src/network/socket/qnativesocketengine_p.h
+++ b/src/network/socket/qnativesocketengine_p.h
@@ -135,6 +135,8 @@ public:
     bool hasPendingDatagrams() const;
     qint64 pendingDatagramSize() const;
 
+    qint64 bytesToWrite() const;
+
     qint64 receiveBufferSize() const;
     void setReceiveBufferSize(qint64 bufferSize);
 
diff --git a/src/network/socket/qsocks5socketengine.cpp b/src/network/socket/qsocks5socketengine.cpp
index 30074cf..bd60ad1 100644
--- a/src/network/socket/qsocks5socketengine.cpp
+++ b/src/network/socket/qsocks5socketengine.cpp
@@ -1235,6 +1235,9 @@ void QSocks5SocketEnginePrivate::_q_controlSocketError(QAbstractSocket::SocketEr
         if (!readNotificationPending)
             connectData->readBuffer.clear();
         emitReadNotification();
+        data->controlSocket->close();
+        // cause a disconnect in the outer socket
+        emitWriteNotification();
     } else if (socks5State == Uninitialized
                || socks5State == AuthenticationMethodsSent
                || socks5State == Authenticating
@@ -1245,6 +1248,7 @@ void QSocks5SocketEnginePrivate::_q_controlSocketError(QAbstractSocket::SocketEr
     } else {
         q_func()->setError(data->controlSocket->error(), data->controlSocket->errorString());
         emitReadNotification();
+        emitWriteNotification();
     }
 }
 
@@ -1623,6 +1627,16 @@ qint64 QSocks5SocketEngine::pendingDatagramSize() const
 }
 #endif // QT_NO_UDPSOCKET
 
+qint64 QSocks5SocketEngine::bytesToWrite() const
+{
+    Q_D(const QSocks5SocketEngine);
+    if (d->data && d->data->controlSocket) {
+        return d->data->controlSocket->bytesToWrite();
+    } else {
+        return 0;
+    }
+}
+
 int QSocks5SocketEngine::option(SocketOption option) const
 {
     Q_D(const QSocks5SocketEngine);
diff --git a/src/network/socket/qsocks5socketengine_p.h b/src/network/socket/qsocks5socketengine_p.h
index 7cb0920..2402517 100644
--- a/src/network/socket/qsocks5socketengine_p.h
+++ b/src/network/socket/qsocks5socketengine_p.h
@@ -100,6 +100,8 @@ public:
     qint64 pendingDatagramSize() const;
 #endif // QT_NO_UDPSOCKET
 
+    qint64 bytesToWrite() const;
+
     int option(SocketOption option) const;
     bool setOption(SocketOption option, int value);
 
diff --git a/src/network/ssl/qsslsocket.cpp b/src/network/ssl/qsslsocket.cpp
index 2c88130..608d772 100644
--- a/src/network/ssl/qsslsocket.cpp
+++ b/src/network/ssl/qsslsocket.cpp
@@ -707,6 +707,8 @@ void QSslSocket::close()
     qDebug() << "QSslSocket::close()";
 #endif
     Q_D(QSslSocket);
+    if (d->plainSocket)
+        d->plainSocket->close();
     QTcpSocket::close();
 
     // must be cleared, reading/writing not possible on closed socket: