websocket.h

Go to the documentation of this file.

// Program Name: websocket.h
// Created     : 18 Jan 2015
//
// Purpose     : Web Socket server
//
// Copyright (c) 2015 Stuart Morgan <smorgan@mythtv.org>
//
// Licensed under the GPL v2 or later, see COPYING for details
//

#ifndef __WEBSOCKET_H__
#define __WEBSOCKET_H__

#include "serverpool.h"
#include "upnpexp.h"
#include "mthreadpool.h"

#include <QRunnable>
#include <QSslConfiguration>
#include <QReadWriteLock>
#include <QEvent>
#include <QEventLoop>
#include <QTimer>

#include <cstdint>

class UPNP_PUBLIC WebSocketServer : public ServerPool
{
    Q_OBJECT

  public:
    WebSocketServer();
    virtual ~WebSocketServer();

    bool IsRunning(void) const
    {
        m_rwlock.lockForRead();
        bool tmp = m_running;
        m_rwlock.unlock();
        return tmp;
    }

  protected slots:
    void newTcpConnection(qt_socket_fd_t socket) override; // QTcpServer

  protected:
    mutable QReadWriteLock  m_rwlock;
    MThreadPool             m_threadPool;
    bool                    m_running {true}; // protected by m_rwlock

#ifndef QT_NO_OPENSSL
    QSslConfiguration       m_sslConfig;
#endif

  private:
     //void LoadSSLConfig();
};


class WebSocketFrame
{
  public:
    WebSocketFrame()
    {
        m_mask.reserve(4);
    }

   ~WebSocketFrame()
    {
        m_payload.clear();
        m_mask.clear();
    }

    void reset(void)
    {
        m_finalFrame = false;
        m_payload.clear();
        m_payload.resize(128);
        m_payload.squeeze();
        m_mask.clear();
        m_payloadSize = 0;
        m_opCode = kOpTextFrame;
        m_fragmented = false;
    }

    enum OpCode
    {
        kOpContinuation = 0x0,
        kOpTextFrame    = 0x1,
        kOpBinaryFrame  = 0x2,
        // Reserved
        kOpClose        = 0x8,
        kOpPing         = 0x9,
        kOpPong         = 0xA
        // Reserved
    };

    bool       m_finalFrame  {false};
    QByteArray m_payload;
    uint64_t   m_payloadSize {0};
    OpCode     m_opCode      {kOpTextFrame};
    bool       m_isMasked    {false};
    QByteArray m_mask;
    bool       m_fragmented  {false};
};

class WebSocketWorker;

class WebSocketExtension : public QObject
{
  Q_OBJECT

  public:
    WebSocketExtension() : QObject() { };
    virtual ~WebSocketExtension() = default;

    virtual bool HandleTextFrame(const WebSocketFrame &/*frame*/) { return false; }
    virtual bool HandleBinaryFrame(const WebSocketFrame &/*frame*/) { return false; }

  signals:
    void SendTextMessage(const QString &);
    void SendBinaryMessage(const QByteArray &);
};

class WebSocketWorkerThread : public QRunnable
{
  public:
    WebSocketWorkerThread(WebSocketServer &webSocketServer, qt_socket_fd_t sock,
                          PoolServerType type
#ifndef QT_NO_OPENSSL
                          , const QSslConfiguration& sslConfig
#endif
                        );
    virtual ~WebSocketWorkerThread() = default;

    void run(void) override; // QRunnable

  private:
    WebSocketServer  &m_webSocketServer;
    qt_socket_fd_t    m_socketFD;
    PoolServerType    m_connectionType;
#ifndef QT_NO_OPENSSL
    QSslConfiguration m_sslConfig;
#endif
};

class WebSocketWorker : public QObject
{
    Q_OBJECT

  public:

    WebSocketWorker(WebSocketServer &webSocketServer, qt_socket_fd_t sock,
                    PoolServerType type
#ifndef QT_NO_OPENSSL
                    , const QSslConfiguration& sslConfig
#endif
                    );
    virtual ~WebSocketWorker();

    void Exec();

    enum ErrorCode
    {
        kCloseNormal        = 1000,
        kCloseGoingAway     = 1001,
        kCloseProtocolError = 1002,
        kCloseUnsupported   = 1003,
        // Reserved - 1004
        // The following codes, 1005 and 1006 should not go over the wire,
        // they are internal codes generated if the client has disconnected
        // without sending a Close event or without indicating a reason for
        // closure.
        kCloseNoStatus      = 1005,
        kCloseAbnormal      = 1006,
        kCloseBadData       = 1007,
        kClosePolicy        = 1008, // Message violates server policy and 1003/1009 don't apply
        kCloseTooLarge      = 1009, // Message is larger than we support (32KB)
        kCloseNoExtensions  = 1010, // CLIENT ONLY
        kCloseUnexpectedErr = 1011, // SERVER ONLY
        // Reserved - 1012-1014
        kCloseNoTLS         = 1012  // Connection closed because it must use TLS
        // Reserved
    };

  public slots:
    void doRead();
    void CloseConnection();

    void SendHeartBeat();
    bool SendText(const QString &message);
    bool SendText(const QByteArray &message);
    bool SendBinary(const QByteArray &data);

  protected:
    bool ProcessHandshake(QTcpSocket *); 
    void ProcessFrames(QTcpSocket *); 

    void HandleControlFrame(const WebSocketFrame &frame);
    void HandleDataFrame(const WebSocketFrame &frame);

    void HandleCloseConnection(const QByteArray &payload);

    static QByteArray CreateFrame(WebSocketFrame::OpCode type, const QByteArray &payload);

    bool SendFrame(const QByteArray &frame);
    bool SendPing(const QByteArray &payload);
    bool SendPong(const QByteArray &payload);
    bool SendClose(ErrorCode errCode, const QString &message = QString());

    void SetupSocket();
    void CleanupSocket();

    void RegisterExtension(WebSocketExtension *extension);
    void DeregisterExtension(WebSocketExtension *extension);

    QEventLoop     *m_eventLoop    {nullptr};
    WebSocketServer &m_webSocketServer;
    qt_socket_fd_t m_socketFD;
    QTcpSocket    *m_socket        {nullptr};
    PoolServerType m_connectionType;

    // True if we've successfully upgraded from HTTP
    bool           m_webSocketMode {false};
    WebSocketFrame m_readFrame;

    uint8_t m_errorCount           {0};
    bool    m_isRunning            {false};

    QTimer *m_heartBeat            {nullptr};

#ifndef QT_NO_OPENSSL
    QSslConfiguration       m_sslConfig;
#endif

    bool    m_fuzzTesting          {false};

    QList<WebSocketExtension *> m_extensions;
};

#endif