Send and parse headers for the handshake request

Task-number: QTBUG-40842
Change-Id: I02787a0e55701966e564a312d51297f4160bd89a
Reviewed-by: Liang Qi <liang.qi@theqtcompany.com>

4 files changed, 34 insertions, 4 deletions

diff --git a/src/websockets/qwebsocket.cpp b/src/websockets/qwebsocket.cpp
index 3d948a5..3635003 100644
--- a/src/websockets/qwebsocket.cpp
+++ b/src/websockets/qwebsocket.cpp
@@ -398,6 +398,8 @@ void QWebSocket::open(const QUrl &url)
     \since 5.6
 
     The \a request url will be used to open the WebSocket connection.
+    Headers present in the request will be sent to the server in the upgrade request,
+    together with the ones needed for the websocket handshake.
  */
 void QWebSocket::open(const QNetworkRequest &request)
 {
diff --git a/src/websockets/qwebsocket_p.cpp b/src/websockets/qwebsocket_p.cpp
index 28c24ae..3ed8d8d 100644
--- a/src/websockets/qwebsocket_p.cpp
+++ b/src/websockets/qwebsocket_p.cpp
@@ -282,6 +282,11 @@ QWebSocket *QWebSocketPrivate::upgradeFrom(QTcpSocket *pTcpSocket,
     QWebSocket *pWebSocket = new QWebSocket(pTcpSocket, response.acceptedVersion(), parent);
     if (Q_LIKELY(pWebSocket)) {
         QNetworkRequest netRequest(request.requestUrl());
+        QMapIterator<QString, QString> headerIter(request.headers());
+        while (headerIter.hasNext()) {
+            headerIter.next();
+            netRequest.setRawHeader(headerIter.key().toLatin1(), headerIter.value().toLatin1());
+        }
         pWebSocket->d_func()->setExtension(response.acceptedExtension());
         pWebSocket->d_func()->setOrigin(request.origin());
         pWebSocket->d_func()->setRequest(netRequest);
@@ -1044,6 +1049,12 @@ void QWebSocketPrivate::processStateChanged(QAbstractSocket::SocketState socketS
     case QAbstractSocket::ConnectedState:
         if (webSocketState == QAbstractSocket::ConnectingState) {
             m_key = generateKey();
+
+            QList<QPair<QString, QString> > headers;
+            foreach (const QByteArray &key, m_request.rawHeaderList())
+                headers << qMakePair(QString::fromLatin1(key),
+                                     QString::fromLatin1(m_request.rawHeader(key)));
+
             const QString handshake =
                     createHandShakeRequest(m_resourceName,
                                            m_request.url().host()
@@ -1052,7 +1063,8 @@ void QWebSocketPrivate::processStateChanged(QAbstractSocket::SocketState socketS
                                            origin(),
                                            QString(),
                                            QString(),
-                                           m_key);
+                                           m_key,
+                                           headers);
             if (handshake.isEmpty()) {
                 m_pSocket->abort();
                 Q_EMIT q->error(QAbstractSocket::ConnectionRefusedError);
@@ -1151,7 +1163,8 @@ QString QWebSocketPrivate::createHandShakeRequest(QString resourceName,
                                                   QString origin,
                                                   QString extensions,
                                                   QString protocols,
-                                                  QByteArray key)
+                                                  QByteArray key,
+                                                  QList<QPair<QString, QString> > headers)
 {
     QStringList handshakeRequest;
     if (resourceName.contains(QStringLiteral("\r\n"))) {
@@ -1193,6 +1206,12 @@ QString QWebSocketPrivate::createHandShakeRequest(QString resourceName,
         handshakeRequest << QStringLiteral("Sec-WebSocket-Extensions: ") % extensions;
     if (protocols.length() > 0)
         handshakeRequest << QStringLiteral("Sec-WebSocket-Protocol: ") % protocols;
+
+    QListIterator<QPair<QString, QString> > headerIter(headers);
+    while (headerIter.hasNext()) {
+        const QPair<QString,QString> &header = headerIter.next();
+        handshakeRequest << header.first % QStringLiteral(": ") % header.second;
+    }
     handshakeRequest << QStringLiteral("\r\n");
 
     return handshakeRequest.join(QStringLiteral("\r\n"));
diff --git a/src/websockets/qwebsocket_p.h b/src/websockets/qwebsocket_p.h
index 615bcfc..260e8e8 100644
--- a/src/websockets/qwebsocket_p.h
+++ b/src/websockets/qwebsocket_p.h
@@ -184,7 +184,8 @@ private:
                                    QString origin,
                                    QString extensions,
                                    QString protocols,
-                                   QByteArray key);
+                                   QByteArray key,
+                                   QList<QPair<QString, QString> > headers);
 
     static QWebSocket *upgradeFrom(QTcpSocket *tcpSocket,
                                    const QWebSocketHandshakeRequest &request,
diff --git a/tests/auto/qwebsocket/tst_qwebsocket.cpp b/tests/auto/qwebsocket/tst_qwebsocket.cpp
index 2422012..4a0e06b 100644
--- a/tests/auto/qwebsocket/tst_qwebsocket.cpp
+++ b/tests/auto/qwebsocket/tst_qwebsocket.cpp
@@ -52,6 +52,7 @@ public:
 
 Q_SIGNALS:
     void newConnection(QUrl requestUrl);
+    void newConnection(QNetworkRequest request);
 
 private Q_SLOTS:
     void onNewConnection();
@@ -87,6 +88,7 @@ void EchoServer::onNewConnection()
     QWebSocket *pSocket = m_pWebSocketServer->nextPendingConnection();
 
     Q_EMIT newConnection(pSocket->requestUrl());
+    Q_EMIT newConnection(pSocket->request());
 
     connect(pSocket, SIGNAL(textMessageReceived(QString)), this, SLOT(processTextMessage(QString)));
     connect(pSocket, SIGNAL(binaryMessageReceived(QByteArray)), this, SLOT(processBinaryMessage(QByteArray)));
@@ -589,17 +591,23 @@ void tst_QWebSocket::tst_openRequest()
     QWebSocket socket;
 
     QSignalSpy socketConnectedSpy(&socket, SIGNAL(connected()));
+    QSignalSpy serverRequestSpy(&echoServer, SIGNAL(newConnection(QNetworkRequest)));
 
     QUrl url = QUrl(QStringLiteral("ws://") + echoServer.hostAddress().toString() +
                     QLatin1Char(':') + QString::number(echoServer.port()));
     url.addQueryItem("queryitem", "with encoded characters");
     QNetworkRequest req(url);
+    req.setRawHeader("X-Custom-Header", "A custom header");
     socket.open(req);
 
     if (socketConnectedSpy.count() == 0)
         QVERIFY(socketConnectedSpy.wait(500));
     QCOMPARE(socket.state(), QAbstractSocket::ConnectedState);
-
+    QCOMPARE(serverRequestSpy.count(), 1);
+    QList<QVariant> arguments = serverRequestSpy.takeFirst();
+    QNetworkRequest requestConnected = arguments.at(0).value<QNetworkRequest>();
+    QCOMPARE(requestConnected.url(), req.url());
+    QCOMPARE(requestConnected.rawHeader("X-Custom-Header"), req.rawHeader("X-Custom-Header"));
     socket.close();
 }