websocket.cpp

#include <iostream>
#include <string>
#include <vector>
#include <cstring>
#include <sstream>
#include <iomanip>
#include <map>
#include <thread>
#include <mutex>
#include <algorithm>

// POSIX sockets
#include <sys/socket.h>
#include <netinet/in.h>
#include <arpa/inet.h>
#include <unistd.h>
#include <fcntl.h>

// For SHA1 and Base64
#include <openssl/sha.h>
#include <openssl/bio.h>
#include <openssl/evp.h>
#include <openssl/buffer.h>

class WebSocketServer {
private:
    int serverSocket;
    std::vector<int> clients;
    std::mutex clientsMutex;
    bool running;

    // Base64 encode
    std::string base64Encode(const unsigned char* buffer, size_t length) {
        BIO *bio, *b64;
        BUF_MEM *bufferPtr;

        b64 = BIO_new(BIO_f_base64());
        bio = BIO_new(BIO_s_mem());
        bio = BIO_push(b64, bio);

        BIO_set_flags(bio, BIO_FLAGS_BASE64_NO_NL);
        BIO_write(bio, buffer, length);
        BIO_flush(bio);
        BIO_get_mem_ptr(bio, &bufferPtr);

        std::string result(bufferPtr->data, bufferPtr->length);
        BIO_free_all(bio);

        return result;
    }

    // Generate WebSocket accept key
    std::string generateAcceptKey(const std::string& key) {
        std::string magicString = "258EAFA5-E914-47DA-95CA-C5AB0DC85B11";
        std::string combined = key + magicString;
        
        unsigned char hash[SHA_DIGEST_LENGTH];
        SHA1(reinterpret_cast<const unsigned char*>(combined.c_str()), 
             combined.length(), hash);
        
        return base64Encode(hash, SHA_DIGEST_LENGTH);
    }

    // Parse HTTP headers
    std::map<std::string, std::string> parseHeaders(const std::string& request) {
        std::map<std::string, std::string> headers;
        std::istringstream stream(request);
        std::string line;
        
        // Skip first line (GET / HTTP/1.1)
        std::getline(stream, line);
        
        while (std::getline(stream, line) && line != "\r") {
            size_t colonPos = line.find(':');
            if (colonPos != std::string::npos) {
                std::string key = line.substr(0, colonPos);
                std::string value = line.substr(colonPos + 2); // Skip ": "
                
                // Remove \r if present
                if (!value.empty() && value.back() == '\r') {
                    value.pop_back();
                }
                
                headers[key] = value;
            }
        }
        
        return headers;
    }

    // Perform WebSocket handshake
    bool performHandshake(int clientSocket) {
        char buffer[4096];
        ssize_t bytesRead = recv(clientSocket, buffer, sizeof(buffer) - 1, 0);
        
        if (bytesRead <= 0) {
            return false;
        }
        
        buffer[bytesRead] = '\0';
        std::string request(buffer);
        
        // Parse headers
        auto headers = parseHeaders(request);
        
        // Check for WebSocket upgrade
        if (headers.find("Upgrade") == headers.end() || 
            headers["Upgrade"].find("websocket") == std::string::npos) {
            return false;
        }
        
        // Get WebSocket key
        if (headers.find("Sec-WebSocket-Key") == headers.end()) {
            return false;
        }
        
        std::string key = headers["Sec-WebSocket-Key"];
        std::string acceptKey = generateAcceptKey(key);
        
        // Send handshake response
        std::ostringstream response;
        response << "HTTP/1.1 101 Switching Protocols\r\n";
        response << "Upgrade: websocket\r\n";
        response << "Connection: Upgrade\r\n";
        response << "Sec-WebSocket-Accept: " << acceptKey << "\r\n";
        response << "\r\n";
        
        std::string responseStr = response.str();
        send(clientSocket, responseStr.c_str(), responseStr.length(), 0);
        
        std::cout << "WebSocket handshake completed with client " << clientSocket << std::endl;
        return true;
    }

    // Decode WebSocket frame
    std::string decodeFrame(const std::vector<uint8_t>& frame, size_t& frameSize) {
        if (frame.size() < 2) {
            frameSize = 0;
            return "";
        }

        uint8_t byte1 = frame[0];
        uint8_t byte2 = frame[1];
        
        bool fin = (byte1 & 0x80) != 0;
        uint8_t opcode = byte1 & 0x0F;
        bool masked = (byte2 & 0x80) != 0;
        uint64_t payloadLength = byte2 & 0x7F;
        
        size_t pos = 2;
        
        // Extended payload length
        if (payloadLength == 126) {
            if (frame.size() < pos + 2) {
                frameSize = 0;
                return "";
            }
            payloadLength = (frame[pos] << 8) | frame[pos + 1];
            pos += 2;
        } else if (payloadLength == 127) {
            if (frame.size() < pos + 8) {
                frameSize = 0;
                return "";
            }
            payloadLength = 0;
            for (int i = 0; i < 8; i++) {
                payloadLength = (payloadLength << 8) | frame[pos + i];
            }
            pos += 8;
        }
        
        // Masking key
        uint8_t maskingKey[4];
        if (masked) {
            if (frame.size() < pos + 4) {
                frameSize = 0;
                return "";
            }
            for (int i = 0; i < 4; i++) {
                maskingKey[i] = frame[pos + i];
            }
            pos += 4;
        }
        
        // Check if we have complete frame
        if (frame.size() < pos + payloadLength) {
            frameSize = 0;
            return "";
        }
        
        frameSize = pos + payloadLength;
        
        // Decode payload
        std::string payload;
        for (uint64_t i = 0; i < payloadLength; i++) {
            uint8_t byte = frame[pos + i];
            if (masked) {
                byte ^= maskingKey[i % 4];
            }
            payload += static_cast<char>(byte);
        }
        
        // Handle opcodes
        if (opcode == 0x8) { // Close frame
            return "";
        } else if (opcode == 0x9) { // Ping frame
            return "PING";
        } else if (opcode == 0x1) { // Text frame
            return payload;
        }
        
        return payload;
    }

    // Encode WebSocket frame
    std::vector<uint8_t> encodeFrame(const std::string& message) {
        std::vector<uint8_t> frame;
        
        // Byte 1: FIN (1) + RSV (000) + Opcode (0001 for text)
        frame.push_back(0x81);
        
        size_t length = message.length();
        
        // Byte 2: MASK (0) + Payload length
        if (length <= 125) {
            frame.push_back(static_cast<uint8_t>(length));
        } else if (length <= 65535) {
            frame.push_back(126);
            frame.push_back((length >> 8) & 0xFF);
            frame.push_back(length & 0xFF);
        } else {
            frame.push_back(127);
            for (int i = 7; i >= 0; i--) {
                frame.push_back((length >> (i * 8)) & 0xFF);
            }
        }
        
        // Payload
        for (char c : message) {
            frame.push_back(static_cast<uint8_t>(c));
        }
        
        return frame;
    }

    // Handle client connection
    void handleClient(int clientSocket) {
        // Perform handshake
        if (!performHandshake(clientSocket)) {
            std::cout << "Handshake failed for client " << clientSocket << std::endl;
            close(clientSocket);
            return;
        }
        
        // Add to clients list
        {
            std::lock_guard<std::mutex> lock(clientsMutex);
            clients.push_back(clientSocket);
        }
        
        std::vector<uint8_t> buffer;
        uint8_t tempBuffer[4096];
        
        while (running) {
            ssize_t bytesRead = recv(clientSocket, tempBuffer, sizeof(tempBuffer), 0);
            
            if (bytesRead <= 0) {
                std::cout << "Client " << clientSocket << " disconnected" << std::endl;
                break;
            }
            
            // Append to buffer
            for (ssize_t i = 0; i < bytesRead; i++) {
                buffer.push_back(tempBuffer[i]);
            }
            
            // Try to decode frames
            while (!buffer.empty()) {
                size_t frameSize = 0;
                std::string message = decodeFrame(buffer, frameSize);
                
                if (frameSize == 0) {
                    // Incomplete frame, wait for more data
                    break;
                }
                
                // Remove processed frame from buffer
                buffer.erase(buffer.begin(), buffer.begin() + frameSize);
                
                if (message.empty()) {
                    // Close or empty frame
                    running = false;
                    break;
                }
                
                if (message == "PING") {
                    // Respond to ping with pong
                    std::vector<uint8_t> pong = {0x8A, 0x00}; // Pong frame
                    send(clientSocket, pong.data(), pong.size(), 0);
                    continue;
                }
                
                std::cout << "Received from client " << clientSocket << ": " << message << std::endl;
                
                // Broadcast to all clients
                broadcast(message, clientSocket);
            }
        }
        
        // Remove from clients list
        {
            std::lock_guard<std::mutex> lock(clientsMutex);
            clients.erase(std::remove(clients.begin(), clients.end(), clientSocket), clients.end());
        }
        
        close(clientSocket);
    }

    // Broadcast message to all clients
    void broadcast(const std::string& message, int senderSocket) {
        std::string fullMessage = "Client " + std::to_string(senderSocket) + ": " + message;
        std::vector<uint8_t> frame = encodeFrame(fullMessage);
        
        std::lock_guard<std::mutex> lock(clientsMutex);
        for (int clientSocket : clients) {
            send(clientSocket, frame.data(), frame.size(), 0);
        }
    }

public:
    WebSocketServer() : serverSocket(-1), running(false) {}

    ~WebSocketServer() {
        stop();
    }

    bool start(int port) {
        // Create socket
        serverSocket = socket(AF_INET, SOCK_STREAM, 0);
        if (serverSocket < 0) {
            std::cerr << "Failed to create socket" << std::endl;
            return false;
        }

        // Set socket options
        int opt = 1;
        if (setsockopt(serverSocket, SOL_SOCKET, SO_REUSEADDR, &opt, sizeof(opt)) < 0) {
            std::cerr << "Failed to set socket options" << std::endl;
            return false;
        }

        // Bind socket
        sockaddr_in serverAddr;
        serverAddr.sin_family = AF_INET;
        serverAddr.sin_addr.s_addr = INADDR_ANY;
        serverAddr.sin_port = htons(port);

        if (bind(serverSocket, (struct sockaddr*)&serverAddr, sizeof(serverAddr)) < 0) {
            std::cerr << "Failed to bind socket" << std::endl;
            return false;
        }

        // Listen
        if (listen(serverSocket, 10) < 0) {
            std::cerr << "Failed to listen on socket" << std::endl;
            return false;
        }

        running = true;
        std::cout << "WebSocket server started on port " << port << std::endl;

        // Accept connections
        while (running) {
            sockaddr_in clientAddr;
            socklen_t clientAddrLen = sizeof(clientAddr);
            
            int clientSocket = accept(serverSocket, (struct sockaddr*)&clientAddr, &clientAddrLen);
            
            if (clientSocket < 0) {
                if (running) {
                    std::cerr << "Failed to accept connection" << std::endl;
                }
                continue;
            }

            std::cout << "New connection from " << inet_ntoa(clientAddr.sin_addr) 
                      << ":" << ntohs(clientAddr.sin_port) << std::endl;

            // Handle client in new thread
            std::thread clientThread(&WebSocketServer::handleClient, this, clientSocket);
            clientThread.detach();
        }

        return true;
    }

    void stop() {
        running = false;
        
        // Close all client sockets
        {
            std::lock_guard<std::mutex> lock(clientsMutex);
            for (int clientSocket : clients) {
                close(clientSocket);
            }
            clients.clear();
        }
        
        // Close server socket
        if (serverSocket >= 0) {
            close(serverSocket);
            serverSocket = -1;
        }
    }
};

int main() {
    WebSocketServer server;
    
    std::cout << "Starting WebSocket server on port 8080..." << std::endl;
    std::cout << "Connect with: ws://localhost:8080" << std::endl;
    std::cout << "\nTest with HTML client:\n" << std::endl;
    std::cout << "<!DOCTYPE html>\n";
    std::cout << "<html><body>\n";
    std::cout << "<input id=\"msg\" type=\"text\" placeholder=\"Enter message\">\n";
    std::cout << "<button onclick=\"send()\">Send</button>\n";
    std::cout << "<div id=\"messages\"></div>\n";
    std::cout << "<script>\n";
    std::cout << "const ws = new WebSocket('ws://localhost:8080');\n";
    std::cout << "ws.onmessage = e => {\n";
    std::cout << "  document.getElementById('messages').innerHTML += '<p>' + e.data + '</p>';\n";
    std::cout << "};\n";
    std::cout << "function send() {\n";
    std::cout << "  ws.send(document.getElementById('msg').value);\n";
    std::cout << "  document.getElementById('msg').value = '';\n";
    std::cout << "}\n";
    std::cout << "</script>\n";
    std::cout << "</body></html>\n" << std::endl;
    
    server.start(8080);
    
    return 0;
}

/**
 g++ websocket_server.cpp -o websocket_server -lpthread -lssl -lcrypto
./websocket_server
 */