Teaspeak-Server/flooder/src/ProxiedClient.cpp

#include <mutex>
#include <netinet/tcp.h>
#include <arpa/inet.h>
#include "ProxiedClient.h"
#include "TSClient.h"

using namespace std;
using namespace ts::flood;

ProxiedClient::ProxiedClient(event_base* base, const sockaddr_in &proxyAddr, const sockaddr_in &remoteAddr) : evBase(base) {
    this->proxyAddr = new sockaddr_in{};
    this->remoteAddr = new sockaddr_in{};

    memcpy(this->proxyAddr, &proxyAddr, sizeof(proxyAddr));
    memcpy(this->remoteAddr, &remoteAddr, sizeof(proxyAddr));

    this->client = new TSClient(this);
}

ProxiedClient::~ProxiedClient() {
    delete this->proxyAddr;
    delete this->remoteAddr;
}

#define CERR(msg) \
do { \
    cerr << "Could not connect: " << msg << "(" << errno << "/" << strerror(errno) << ")" << endl; \
    return false; \
} while(0)

#if defined(TCP_CORK) && !defined(TCP_NOPUSH)
    #define TCP_NOPUSH TCP_CORK
#endif

static int enabled = 1;
static int disabled = 0;

bool ProxiedClient::connect() {
    assert(this->state == ProxyState::PROXY_UNCONNECTED);
    TAILQ_INIT(&this->writeQueue);

    this->fileDescriptor = socket(AF_INET, SOCK_STREAM, 0);
    if(this->fileDescriptor < 0) CERR("Socket setup failed");
    if(::connect(this->fileDescriptor, reinterpret_cast<const sockaddr *>(this->proxyAddr), sizeof(*this->proxyAddr)) < 0) CERR("connect() failed");
    if(setsockopt(this->fileDescriptor, SOL_SOCKET, SO_REUSEADDR, &enabled, sizeof(enabled)) < 0) CERR("could not set reuse addr");
    if(setsockopt(this->fileDescriptor, IPPROTO_TCP, TCP_NOPUSH, &disabled, sizeof(disabled)) < 0) CERR("could not set no push");

    cout << "Connected to " << inet_ntoa(this->proxyAddr->sin_addr) << endl;

    this->rEvent = event_new(this->evBase, this->fileDescriptor, EV_READ | EV_PERSIST, ProxiedClient::handleEventRead, this);
    this->wEvent = event_new(this->evBase, this->fileDescriptor, EV_WRITE, ProxiedClient::handleEventWrite, this);

    event_add(rEvent, nullptr);
    this->state = ProxyState::PROXY_INIT_METHODS;

    this->requestProxyConnection();
    return true;
}

void ProxiedClient::disconnect() {
    this->closeConnection();
}

void ProxiedClient::closeConnection() {
    {
        lock_guard<threads::Mutex> lock(this->stateLock);
        if(this->state == PROXY_UNCONNECTED) return;
        this->state = PROXY_UNCONNECTED;
    }

    event_del(this->wEvent);
    event_del(this->rEvent);

    this->wEvent = nullptr;
    this->rEvent = nullptr;
}

void ProxiedClient::sendMessage(const std::string& message) {
    buffer::RawBuffer* buffer;
    if(this->state != PROXY_CONNECTED){
        buffer = new buffer::RawBuffer{message.length()};
        memcpy(buffer->buffer, message.data(), message.length());
    } else {
        cout << "Send " << message.length() << " bytes with relay" << endl;
        int relayHeaderLength = 2 + 1 + 1 + 4 + 2;
        buffer = new buffer::RawBuffer{relayHeaderLength + message.length()};
        buffer->index = 0;
        char preBuffer[relayHeaderLength];
        preBuffer[0] = 0x00;
        preBuffer[1] = 0x00;
        preBuffer[2] = 0x00;
        preBuffer[3] = 0x01;

        IPv4 addr{this->relayAddr->sin_addr.s_addr};
        preBuffer[4] = addr._1;
        preBuffer[5] = addr._2;
        preBuffer[6] = addr._3;
        preBuffer[7] = addr._4;

        preBuffer[8] = (ntohs(this->relayAddr->sin_port) >> 8) & 0xFF;
        preBuffer[9] = (ntohs(this->relayAddr->sin_port) >> 0) & 0xFF;
        //memset(&preBuffer[4], 0, 6);

        memcpy(&buffer->buffer[0], preBuffer, relayHeaderLength);
        memcpy(&buffer->buffer[relayHeaderLength], message.data(), message.length());
    }

    {
        lock_guard<threads::Mutex> lock(this->queueLock);
        TAILQ_INSERT_TAIL(&this->writeQueue, buffer, tail);
    }
    event_add(this->wEvent, nullptr);
}

void ProxiedClient::handleMessage(const std::string &message) {
    if(this->state == PROXY_UNCONNECTED) return;
    if(this->state == PROXY_CONNECTED) return; //TODO
    this->handleProxyMessage(message);
}

extern void hexdump(std::ostream& outs, const std::string& s, size_t line_len = 16);
void ProxiedClient::handleEventWrite(int fd, short, void* ptrClient) {
    auto* client = static_cast<ProxiedClient *>(ptrClient);

    buffer::RawBuffer* buffer = nullptr;
    {
        lock_guard<threads::Mutex> lock(client->queueLock);
        buffer = TAILQ_FIRST(&client->writeQueue);
        if(!buffer) return;

        ssize_t writtenBytes = 0;
        if(client->state == PROXY_CONNECTED){
            cout << "Write bytes to relay - " << fd << " - " << inet_ntoa(client->relayAddr->sin_addr) << ":" << ntohs(client->relayAddr->sin_port) << endl;
            hexdump(cout, string((const char*) buffer->buffer, buffer->length));
            writtenBytes = sendto(fd, buffer->buffer, buffer->length, 0, (const sockaddr *) client->relayAddr, sizeof(*client->relayAddr));
        } else
            writtenBytes = send(fd, &buffer->buffer[buffer->index], buffer->length - buffer->index, 0);
        buffer->index += writtenBytes;
        cout << "Written: " << writtenBytes << endl;

        if(buffer->index >= buffer->length || client->state == PROXY_CONNECTED) {
            TAILQ_REMOVE(&client->writeQueue, buffer, tail);
            delete buffer;
        }
        if(!TAILQ_EMPTY(&client->writeQueue))
            event_add(client->wEvent, nullptr);
    }
}

void ProxiedClient::handleEventRead(int fd, short, void* ptrClient) {
    auto* client = static_cast<ProxiedClient *>(ptrClient);

    char buffer[1024];
    sockaddr_in remoteAddr{};
    socklen_t remoteAddrSize = sizeof(remoteAddr);
    auto read = recvfrom(fd, buffer, 1024, MSG_DONTWAIT, reinterpret_cast<sockaddr *>(&remoteAddr), &remoteAddrSize);
    cout << "Read " << read << " bytes" << endl;
    if(read < 0){
        if(errno == EWOULDBLOCK) return;
        cerr << "Invalid read: " << errno << "/" << strerror(errno) << endl;
        client->disconnect();
        return;
    } else if(read == 0){
        cerr << "Client hangs up!" << endl;
        client->closeConnection();
        return;
    }

    hexdump(cout, string(buffer, read));
    client->handleMessage(string(buffer, read));
}

void ProxiedClient::proxyInizalisized() {
    this->client->startConnect();
}