smsub/SMSubServer.cpp

/*****************************************************************************
* smsub Server Manager Subprocess
* Copyright (C) 2020 Syping
*
* Redistribution and use in source and binary forms, with or without modification,
* are permitted provided that the following conditions are met:
*
* 1. Redistributions of source code must retain the above copyright notice,
* this list of conditions and the following disclaimer.
*
* 2. Redistributions in binary form must reproduce the above copyright notice,
* this list of conditions and the following disclaimer in the documentation
* and/or other materials provided with the distribution.
*
* This software is provided as-is, no warranties are given to you, we are not
* responsible for anything with use of the software, you are self responsible.
*****************************************************************************/

#include <QCoreApplication>
#include <QTimer>
#include <QUuid>
#include "SMSubServer.h"
#include "smsub.h"

SMSubServer::SMSubServer(SMSubServerSettings *serverSettings, const QString &socket) : serverSettings(serverSettings)
{
    QLocalServer *localServer = new QLocalServer(this);
    localServer->setSocketOptions(QLocalServer::UserAccessOption | QLocalServer::GroupAccessOption);
    localServer->listen(socket);

    QObject::connect(localServer, &QLocalServer::newConnection, this, &SMSubServer::newConnection);

    type = ServerType::Local;
    server = localServer;
}

SMSubServer::SMSubServer(SMSubServerSettings *serverSettings, const QString &serverName, const quint16 &port) : serverSettings(serverSettings)
{
    QWebSocketServer *webSocketServer = new QWebSocketServer(serverName, QWebSocketServer::NonSecureMode, this);
    webSocketServer->listen(QHostAddress::LocalHost, port);

    QObject::connect(webSocketServer, &QWebSocketServer::newConnection, this, &SMSubServer::newConnection);

    type = ServerType::WebSocket;
    server = webSocketServer;
}

bool SMSubServer::isListening()
{
    if (likely(type == ServerType::Local)) {
        return static_cast<QLocalServer*>(server)->isListening();
    }
    else if (type == ServerType::WebSocket) {
        return static_cast<QWebSocketServer*>(server)->isListening();
    }
    return false;
}

void SMSubServer::newConnection()
{
    QObject *socket;
    if (likely(type == ServerType::Local)) {
        QLocalSocket *localSocket = static_cast<QLocalServer*>(server)->nextPendingConnection();
        QObject::connect(localSocket, &QLocalSocket::readyRead, this, &SMSubServer::lsReadyRead);
        QObject::connect(localSocket, &QLocalSocket::disconnected, this, &SMSubServer::deleteSocket);
        localSocket->write(QString("SMSub Version %1\n").arg(QCoreApplication::applicationVersion()).toUtf8());
        socket = localSocket;
    }
    else if (type == ServerType::WebSocket) {
        QWebSocket *webSocket = static_cast<QWebSocketServer*>(server)->nextPendingConnection();
        QObject::connect(webSocket, &QWebSocket::binaryMessageReceived, this, &SMSubServer::wsMessageReceived);
        QObject::connect(webSocket, &QWebSocket::disconnected, this, &SMSubServer::deleteSocket);
        webSocket->sendBinaryMessage(QString("SMSub Version %1\n").arg(QCoreApplication::applicationVersion()).toUtf8());
        socket = webSocket;
    }
    else {
        // Just for being sure
        return;
    }

    // Set authentication state
    if (serverSettings->isLocal) {
        socket->setProperty("Authenticated", true);
        sockets << socket;
    }
    else {
        socket->setProperty("Authenticated", false);
    }
}

bool SMSubServer::messageReceived(QObject *socket, const QByteArray &message)
{
    // Only allow commands being sent if authenticated
    bool isAuthenticated = socket->property("Authenticated").toBool();
    if (likely(isAuthenticated)) {
        if (message.startsWith("+dbg")) {
            socket->setProperty("ReceiveDbgMsg", true);
            sendMessage(socket, "Debug messages enabled!\n");
        }
        else if (message.startsWith("-dbg")) {
            socket->setProperty("ReceiveDbgMsg", false);
            sendMessage(socket, "Debug messages disabled!\n");
        }
        else if (message.startsWith("+log")) {
            socket->setProperty("ReceiveLog", true);
            debugOutput(socket, "Log output enabled!");
        }
        else if (message.startsWith("-log")) {
            socket->setProperty("ReceiveLog", false);
            debugOutput(socket, "Log output disabled!");
        }
        else if (message.startsWith("+reg")) {
            if (likely(serverSettings->canRegister)) {
                QByteArray authUuid = QUuid::createUuid().toByteArray(QUuid::Id128);
                authUuid = QByteArray::fromHex(authUuid).toBase64(QByteArray::OmitTrailingEquals);
                emit tokenRegistered(QString::fromUtf8(authUuid));
                sendMessage(socket, "Token: " + authUuid + '\n');
            }
            else {
                sendMessage(socket, "Permission denied!\n");
            }
        }
        else if (message.startsWith("kill")) {
            emit killRequested();
            debugOutput(socket, "Killing server!");
        }
        else if (message.startsWith("stop")) {
            emit stopRequested();
            debugOutput(socket, "Stopping server!");
        }
        else if (message.startsWith("wl")) {
            const QByteArray writeData = message.mid(3);
            emit inputWritten(writeData + '\n');
            debugOutput(socket, "Write line \"" + writeData + "\"!");
        }
        else if (message.startsWith("w")) {
            const QByteArray writeData = message.mid(2);
            emit inputWritten(writeData);
            debugOutput(socket, "Write \"" + writeData + "\"!");
        }
    }
    else {
        // Authenticate when token is valid, otherwise disconnect
        if (unlikely(tokens.contains(QString::fromUtf8(message)))) {
            // Set client as authenticated and add it to vector
            socket->setProperty("Authenticated", true);
            sendMessage(socket, "Login successful!\n");
            isAuthenticated = true;
            sockets << socket;
        }
        else {
            // Stop receiving data and disconnect socket
            if (likely(type == ServerType::Local)) {
                QLocalSocket *localSocket = static_cast<QLocalSocket*>(socket);
                QObject::disconnect(localSocket, &QLocalSocket::readyRead, this, &SMSubServer::lsReadyRead);
                localSocket->write("Incorrect token!\n");
                localSocket->disconnectFromServer();
                return false;
            }
            else if (type == ServerType::WebSocket) {
                QWebSocket *webSocket = static_cast<QWebSocket*>(socket);
                QObject::disconnect(webSocket, &QWebSocket::binaryMessageReceived, this, &SMSubServer::wsMessageReceived);
                webSocket->sendBinaryMessage("Incorrect token!\n");
                webSocket->close(QWebSocketProtocol::CloseCodeNormal);
                return false;
            }
        }
    }
    return true;
}

void SMSubServer::wsMessageReceived(const QByteArray &message)
{
    QWebSocket *socket = static_cast<QWebSocket*>(sender());
    messageReceived(socket, message.trimmed());
}

void SMSubServer::lsReadyRead()
{
    QLocalSocket *socket = static_cast<QLocalSocket*>(sender());
    while (socket->canReadLine()) {
        const QByteArray message = socket->readLine().trimmed();
        if (!messageReceived(socket, message))
            return;
    }
}

void SMSubServer::deleteSocket()
{
    // Delete socket and remove from index
    QObject *socket = sender();
    sockets.removeAll(socket);
    socket->deleteLater();
}

void SMSubServer::debugOutput(QObject *socket, const QByteArray &message)
{
    // Only send debug messages when the client opted-in
    const QVariant variant = socket->property("ReceiveDbgMsg");
    if (unlikely(variant.type() == QVariant::Bool)) {
        bool receiveDbgMsg = variant.toBool();
        if (likely(receiveDbgMsg)) {
            sendMessage(socket, message + '\n');
        }
    }
}

void SMSubServer::writeOutput(const QByteArray &output)
{
    // Read process output when client opted-in for log
    QVector<QObject*>::const_iterator it = sockets.constBegin();
    QVector<QObject*>::const_iterator end = sockets.constEnd();
    while (it != end) {
        const QVariant variant = (*it)->property("ReceiveLog");
        if (unlikely(variant.type() == QVariant::Bool)) {
            bool receiveLog = variant.toBool();
            if (likely(receiveLog)) {
                sendMessage(*it, output);
            }
        }
        it++;
    }
}

void SMSubServer::sendMessage(QObject *socket, const QByteArray &message)
{
    if (likely(type == ServerType::Local)) {
        QLocalSocket *localSocket = static_cast<QLocalSocket*>(socket);
        localSocket->write(message);
    }
    else if (type == ServerType::WebSocket) {
        QWebSocket *webSocket = static_cast<QWebSocket*>(socket);
        webSocket->sendBinaryMessage(message);
    }
}

void SMSubServer::registerToken(const QString &token)
{
    // Register temporary token for a secure remote connection
    tokens << token;
    QTimer::singleShot(30000, [this, token]() {
        tokens.removeAll(token);
    });
}