
/*---------------------------------------------------------------------------*/
/* http_websocket_handler.cpp                                                */
/* copyright (c) innovaphone 2016                                            */
/*                                                                           */
/*---------------------------------------------------------------------------*/

#include "platform/platform.h"

#include "common/ilib/hash.h"
#include "common/ilib/base64.h"

#include "common/os/iomux.h"
#include "common/lib/httplib.h"
#include "common/interface/socket.h"
#include "common/interface/websocket_client.h"
#include "common/interface/webserver_plugin.h"
#include "common/lib/appservice.h"

#include "common/service/webserver_shared/webserver_internal.h"
#include "common/service/webserver_shared/inno_error_pages.h"

#include "types.h"
#include "apparguments.h"
#include "webserver.h"
#include "http_method_handler.h"
#include "http_connection.h"
#include "http_datasource.h"
#include "http_websocket_handler.h"

/*---------------------------------------------------------------------------*/
/* class HTTPWebsocketContext                                                */
/*---------------------------------------------------------------------------*/

class HTTPWebsocketContext : public ISocketContext {
public:
    virtual void EnableDTLS(bool useSrtp) {};
    virtual void SetServerCertificate(const byte * cert, size_t certLen, const char * hostName = NULL);
    virtual void SetClientCertificate(const byte * cert, size_t certLen, const char * hostName = NULL);
};

void HTTPWebsocketContext::SetServerCertificate(const byte * cert, size_t certLen, const char * hostName)
{
    // do nothing here for now
}

void HTTPWebsocketContext::SetClientCertificate(const byte * cert, size_t certLen, const char * hostName)
{
    // do nothing here for now
}

/*---------------------------------------------------------------------------*/
/* class TLSSocketContext	                                                 */
/*---------------------------------------------------------------------------*/

HTTPWebsocketHandler::HTTPWebsocketHandler(UHTTPMethodHandler * const user, class HTTPRegisteredPath * reqPathObj, class HTTPDataSource * dataSource, ISocket * const remoteSocket, IInstanceLog * const log, connection_flags_t connFlags)
    : HTTPMethodHandler("HTTPWebsocketHandler", user, reqPathObj, dataSource, remoteSocket, log, connFlags)
{
    this->proxySocketUser = NULL;
    this->websocketClient = NULL;
    this->websocketKey = NULL;
    this->ignoreSendResult = false;
    this->hasPendingPing = false;
}


HTTPWebsocketHandler::~HTTPWebsocketHandler()
{
    delete websocketClient;
    free(websocketKey);
}


http_method_init_t HTTPWebsocketHandler::OnHandlerInit(const class HTTPParser * parser, byte * buffer, size_t size)
{
    // According to the RFC 6455, a websocket upgrade request must have
    // the following conditions
    // RequestGET   (already checked)
    // HTTP Version (at least 1.1 (for now (12.2015) only supported version)
    // Request URI  (well, we only can assume that it will be correct)
    // Host
    // Upgrade      (with Value "websocket", already checked)
    // Connect      (with value "upgrade")
    // Sec-WebSocket-Key
    // Sec-WebSocket-Version (must be 13)
    //
    // The rest of optional fields will be ignored...

    // Well, that never shoud happen with websocket...
    if (this->RequestIsChunkEncoded()) {
        this->log->Log(LOG_WEBSERVER, "#ERROR HTTPWebsocketHandler::Init(0x%x, 0x%x): Got WebsocketRequest for connection %u but transfer-encoding is chunked",
            this->DataSource(),
            this,
            User()->GetConnectionID());
        return HTTP_INIT_CRITICAL_ERROR;
    }

    size_t hnlen = (parser->GetHost() == NULL ? 0 : strlen(parser->GetHost()));

    if (!parser->IsConnectionType(HTTPLib::HTTP_CONNECTION_UPGRADE) || hnlen == 0 || parser->GetWebsocketVersion() != 13)
    {
        if (parser->GetWebsocketVersion() != 13) {
            // 426 Update Required, Sec-WebSocket-Version = 13
            ResultHeader()->SetHTTPResult(&HTTPLib::HTTP_426);
            ResultHeader()->SetCustomField("Sec-WebSocket-Version: 13");
        }
        else {
            ResultHeader()->SetHTTPResult(&HTTPLib::HTTP_400);
        }

        Socket()->Send(ResultHeader()->GetHeaderData(), ResultHeader()->GetHeaderSize());
        return HTTP_INIT_ERROR;
    }

    char * orgPath = const_cast<char *>(parser->GetRequestTarget());
    if (!this->GetRegisteredPathForRequest()->SupportsWebsocket()) {
        ResultHeader()->SetHTTPResult(&HTTPLib::HTTP_404);
        ResultHeader()->SetContentType("text/html; charset=utf-8");
        ResultHeader()->SetContentSize(INNO_404_LEN);
        Socket()->Send(ResultHeader()->GetHeaderData(), ResultHeader()->GetHeaderSize());
        Socket()->Send((void *)INNO_404, INNO_404_LEN);
        return HTTP_INIT_ERROR;
    }

    this->websocketKey = _strdup(parser->GetWebsocketKey());

    const char * host = parser->GetHost();
    const char * destPath = orgPath;
    destPath += this->GetRegisteredPathForRequest()->GetPathLen();
    if (!destPath || *destPath == '\0') destPath = "/";
    else if (destPath != orgPath && *destPath != '/') destPath--;

    // The base class (HTTPMethodHandler - init of it called above) already has the Registered Path Object for the current request.
    ExtendedHeader extHeader(this->log, this->DataSource()->ProtocolVersion());
    extHeader.SetRequestPath(destPath, strlen(destPath));
    extHeader.SetRegisteredPathForRequest(this->GetRegisteredPathForRequest()->GetPath(), this->GetRegisteredPathForRequest()->GetPathLen());
    extHeader.SetHost(host, strlen(host));
    extHeader.SetConnectionFlags(this->connFlags);
    extHeader.SetHeaderFields(parser);
    extHeader.SetRemoteAddress(this->Socket());
    extHeader.Encode(User()->GetConnectionID());

    DataSource()->SendMessage("HTTPWebsocketHandler", User()->GetConnectionID(), WSM_WEBSOCKET_LISTEN_RESULT, 0, extHeader.Buffer(), extHeader.BufferSize());

    return HTTP_INIT_DONE;
}


void HTTPWebsocketHandler::AppResponse(dword id, dword data, byte * buffer, size_t size)
{
    if (id == WSM_SEND_META_DATA && data == WSP_CANCEL_NOT_FOUND) {
        // That happens only if a websocket request will be done to an unkown resource. For that kind of response,
        // HandleControllMessage() always returns false, which means the socket should not be closed. So the handler
        // must be closed.
        this->User()->HandleControlMessage(ResultHeader(), WSP_CANCEL_NOT_FOUND, nullptr);
        this->User()->HTTPHandlerShutdown(DataSource(), false);
    }
    else {
        // For Websocket, the given ID is the message ID, data are the flags or other informations
        // (depending on mesage)...
        switch (id)
        {
        case WS_WEBSOCKET_CONNECT_ACCEPT:
        case WS_WEBSOCKET_CONNECT_REJECT:
            this->dataSendingActive = true; // From now on, the webserver won't send a message if the datasource closes.
            this->WebsocketAcceptCompleteMsg(id == WS_WEBSOCKET_CONNECT_ACCEPT);
            break;

        case WS_WEBSOCKET_SEND:
            websocketClient->Send(buffer, size, (ws_websocket_flags_t(data) & WS_WEBSOCKET_FLAG_TEXT) != 0);
            break;

        case WS_WEBSOCKET_RECV:
            // The size of the Recv() buffer is send in the flags part. If the given recvLen is 0,
            // the length of the buffer needed will be requested from tha app. In that case, there
            // is no need to give a buffer.
            websocketClient->Recv(data > 0 ? malloc(data) : NULL, data);
            break;

        case WS_WEBSOCKET_CLOSE:
            websocketClient->Close();
            break;

        default:
            this->log->Log(LOG_WEBSERVER, "#WARNING HTTPWebsocketHandler::AppResponse(0x%x): HTTPWebsocketHandler received unexpected message %u for connection %u",
                           this,
                           id,
                           User()->GetConnectionID());
            break;
        }
    }
}


void HTTPWebsocketHandler::WebsocketAcceptCompleteMsg(bool accepted)
{
    if (!accepted) {
        ResultHeader()->SetHTTPResult(&HTTPLib::HTTP_501);
        Socket()->Send(ResultHeader()->GetHeaderData(), ResultHeader()->GetHeaderSize());
        User()->HTTPHandlerShutdown(DataSource(), true);
    }
    else {
        // HTTP/1.1 101 Switching Protocols
        // Upgrade: websocket
        // Connection: Upgrade
        // Sec-WebSocket-Accept: <Value-Of-wsServerSecKey>
        ResultHeader()->SetHTTPResult(&HTTPLib::HTTP_101);
        ResultHeader()->SetUpgrade(HTTPLib::HTTP_UPGRADE_WEBSOCKET);
        ResultHeader()->SetCustomField("Connection: Upgrade");

        // Sec-WebSocket-Accept = BASE64(Sec-WebSocket-Key + 258EAFA5-E914-47DA-95CA-C5AB0DC85B11)
        char tmp[61];
        _snprintf(tmp, 61, "%s258EAFA5-E914-47DA-95CA-C5AB0DC85B11", websocketKey);

        hash sha1;
        sha1.init(HASH_SHA1);

        unsigned char shaBuf[20];
        sha1.update(tmp, 60);
        sha1.final(shaBuf);

        char webserverKey[29];
        dword r = encode_base64((const byte *)shaBuf, webserverKey, 20);
        webserverKey[r] = 0;
        ResultHeader()->SetWebsocketAcceptKey(webserverKey);

        // We will receive one send result because of the header sent to the browser. That however
        // must not be a SendResult() of the Websocekt itself. Because of this, set a flag to ignore
        // the first Send-Result.
        ignoreSendResult = true;
        Socket()->Send(ResultHeader()->GetHeaderData(), ResultHeader()->GetHeaderSize());
        free(websocketKey);
        websocketKey = NULL;

        // Let the WebSocket Client work as server. This will be done by using a ProxySocket that
        // redirects socket calls (Recv(), Send(), Close()) and will be used to redirect USocket
        // calls (RecvResult(), SendResult(), SocketShutdown()). With this, the WebsocketClient
        // will be able to use the ISocket of the HTTPConnection object. By calling Connect() and
        // passing __ACTIVATE_SERVER_MODE as uri, the WebsocketClient will internally switch to
        // some kind of server mode (that is, the directon of the data encryption will be changed)
        // However, he needs to have a connection that already sucessfully did the websocket 
        // handshake, what already had been done here.
        websocketClient = IWebsocketClient::Create(NULL, this, this, NULL, NULL, this->log);
        websocketClient->Connect("__ACTIVATE_SERVER_MODE");

        if (DataSource() != nullptr) DataSource()->SendWebsocketMessage(User()->GetConnectionID(), WS_WEBSOCKET_CONNECT_COMPLETE, WS_WEBSOCKET_FLAG_NONE, NULL, 0);

        // Add to the list of Websocket-Connections to ping...
        this->User()->RegisterWebsocketForPing(this);
    }
}


void HTTPWebsocketHandler::SocketRecvResult(byte * buffer, size_t size)
{
    proxySocketUser->SocketRecvResult(this, buffer, size);
}


void HTTPWebsocketHandler::SocketSendResult()
{
    if (ignoreSendResult) ignoreSendResult = false;
    else if (proxySocketUser) proxySocketUser->SocketSendResult(this);
}


void HTTPWebsocketHandler::SocketShutdown(shutdownreason_t reason)
{
    // Race-Condition: If the TCP connection, that requested Websocket communication, will be
    // closed before the application will send a Websocket accept, proxySocketUser will be
    // nullptr when SocketShutdown() is called. In that case directly continue closing the
    // handler.
    if (proxySocketUser) proxySocketUser->SocketShutdown(this, reason);
    else this->User()->HTTPHandlerShutdown(this->DataSource(), true);
}

// Websocekt Callbacks
void HTTPWebsocketHandler::WebsocketConnectComplete(IWebsocketClient * const websocket)
{
}


void HTTPWebsocketHandler::WebsocketCloseComplete(IWebsocketClient * const websocket, closereason_t reason)
{
    // Only send the data, if the handler on the app side not hat been closed
    word temp;
    switch (reason)
    {
    case UWebsocketClient::WSCR_NORMAL_CLOSE:
        temp = ws_close_reason_t::WSCR_NORMAL_CLOSE;
        break;

    case UWebsocketClient::WSCR_BUFFER_OVERFLOW:
        temp = ws_close_reason_t::WSCR_BUFFER_OVERFLOW;
        break;

    case UWebsocketClient::WSCR_SOCKET_LOST:
        temp = ws_close_reason_t::WSCR_SOCKET_LOST;
        break;

    case UWebsocketClient::WSCR_DOUBLE_RECV_CALL:
        this->log->Log(LOG_WEBSERVER, "ERROR: Received multiple websocket recv calls from app with registered path 's' - closing connection.", this->DataSource()->GetRoot());
        temp = ws_close_reason_t::WSCR_DOUBLE_RECV_CALL;
        break;

    default:
        temp = WSCR_ERROR;
        break;
    }

    byte tmpbuf[sizeof(word)];
    WRITE_WORD_NET(tmpbuf, temp);

    if (DataSource() != nullptr) DataSource()->SendWebsocketMessage(User()->GetConnectionID(), WS_WEBSOCKET_CLOSE_COMPLETE, WS_WEBSOCKET_FLAG_NONE, tmpbuf, sizeof(word));

    User()->UnregisterWebsocketForPing(this);

    if (websocket == nullptr) {
        // if websocket is nullptr, the function had been called not by the websocket client. Instead, the HTTPWebsocketHandler() itself called
        // the function to simulate a close of the datasource because of an internal error (e. g. the datasource had been closed because the
        // application crashed before sending the accept complete for the websocket message). So let the connectio send an error message before
        // closing the socket.
        User()->HTTPHandlerShutdown(DataSource(), true, &HTTPLib::HTTP_500);
    }
    else User()->HTTPHandlerShutdown(DataSource(), false); 
}


void HTTPWebsocketHandler::WebsocketSendResult(IWebsocketClient * const websocket)
{
    if (DataSource() != nullptr) DataSource()->SendWebsocketMessage(User()->GetConnectionID(), WS_WEBSOCKET_SEND_RESULT, WS_WEBSOCKET_FLAG_NONE, NULL, 0);
}


void HTTPWebsocketHandler::WebsocketRecvResult(IWebsocketClient * const websocket, void * buf, size_t len, bool text, bool isFragment)
{
    ws_websocket_flags_t flags = (text ? WS_WEBSOCKET_FLAG_TEXT : WS_WEBSOCKET_FLAG_NONE);
    if (isFragment) flags = ws_websocket_flags_t(flags | WS_WEBSOCKET_FLAG_FRAGMENTED);

    if (DataSource() != nullptr) DataSource()->SendWebsocketMessage(User()->GetConnectionID(), WS_WEBSOCKET_RECV_RESULT, flags, static_cast<byte *>(buf), len);
    free(buf);
}


void * HTTPWebsocketHandler::WebsocketRecvBuffer(size_t len)
{
    return malloc(len);
}


void HTTPWebsocketHandler::WebsocketRecvCanceled(IWebsocketClient * const websocket, void * buf)
{
    free(buf);
}


void HTTPWebsocketHandler::WebsocketPongReceived(IWebsocketClient * const websocket)
{
    this->hasPendingPing = false;
}


ISocketContext * HTTPWebsocketHandler::CreateSocketContext(class IInstanceLog * const log)
{
    return new class HTTPWebsocketContext();
}

// The HTTPWebsocketHandler is also a socket proxy
ISocket * HTTPWebsocketHandler::CreateSocket(class IIoMux * const iomux, USocket * const user, class IInstanceLog * const log, bool useIPv6, class ISocketContext * socketContext)
{
    proxySocketUser = user;
    return this;
}


void HTTPWebsocketHandler::Shutdown()
{
    Socket()->Shutdown();
}


void HTTPWebsocketHandler::Send(const void * buf, size_t len)
{
    Socket()->Send(buf, len);
}


void HTTPWebsocketHandler::Recv(void * buf, size_t len, bool recvPartial)
{
    Socket()->Recv(buf, len, recvPartial);
}


void HTTPWebsocketHandler::Ping()
{
    if (this->hasPendingPing) {
        this->log->Log(LOG_WEBSERVER, "#WARNING HTTPWebsocketHandler::Ping(0x%x): PING unanswered, closing WebSocket connection", this);
        this->websocketClient->Cancel();
    }
    else {
        this->hasPendingPing = true;
        this->websocketClient->Ping();
    }
}


void HTTPWebsocketHandler::ShutdownHandler()
{
    if (this->websocketClient != nullptr) this->websocketClient->Cancel();
    else WebsocketCloseComplete(nullptr, WSCR_SOCKET_LOST);
}
