chat.cpp

#include "chat.h"
#include <iostream>
#include <boost/asio.hpp>
#include <omp.h>
#include <cstring>
#ifdef _WIN32
#include <winsock2.h>
#else
#include <arpa/inet.h>
#endif

using boost::asio::ip::tcp;

const int CHUNK_SIZE = 32 * 1024; // 32KB
const int HEADER_SIZE = 4; // 4 bytes for message length
const int TAG_SIZE = 16; // AES-GCM tag size

Chat::Chat(ChatControl* chat_control, Crypto* crypto) : chat_control_(chat_control), crypto_(crypto) {}

std::vector<std::string> Chat::fragment_message(const std::string& message) {
    std::vector<std::string> chunks;
    // Always create at least one chunk, even for empty messages
    if (message.empty()) {
        chunks.push_back("");
    } else {
        for (size_t i = 0; i < message.length(); i += CHUNK_SIZE) {
            chunks.push_back(message.substr(i, CHUNK_SIZE));
        }
    }
    return chunks;
}

std::string Chat::reassemble_message(const std::vector<std::string>& chunks) {
    std::string message;
    for (const auto& chunk : chunks) {
        // Skip empty chunks (failed decryption)
        if (!chunk.empty()) {
            message += chunk;
        }
    }
    return message;
}

bool Chat::send_message(const std::string& session_id, const std::string& message) {
    auto session = chat_control_->get_session(session_id);
    if (!session) {
        std::cerr << "Session not found: " << session_id << std::endl;
        return false;
    }

    try {
        std::cout << "Preparing to send message to " << session_id << std::endl;
        // Get or generate session AES key (reused for all messages in this session)
        std::vector<unsigned char> aes_key;
        {
            std::lock_guard<std::mutex> session_lock(session->session_mutex);
            if (!session->session_key_initialized) {
                std::cout << "Generating session AES key" << std::endl;
                // Generate session key once and encrypt it with peer's RSA public key
                session->session_aes_key = crypto_->generate_aes_key();
                session->session_key_initialized = true;
            } else {
                std::cout << "Reusing existing session AES key" << std::endl;
            }
            aes_key = session->session_aes_key;
        }
        
        // Generate a new IV for each message (required for GCM security)
        auto iv = crypto_->generate_unique_iv();
        if (iv.empty()) {
            std::cerr << "Error: Failed to generate IV" << std::endl;
            return false;
        }
        std::cout << "IV generated, length=" << iv.size() << std::endl;

        // Fragment message into chunks (will always return at least one chunk)
        auto chunks = fragment_message(message);
        std::cout << "Fragmented message into " << chunks.size() << " chunk(s)" << std::endl;

        // Encrypt chunks in parallel
        std::cout << "Encrypting chunks" << std::endl;
        auto encrypted_chunks = encrypt_chunks_parallel(chunks, aes_key, iv);
        std::cout << "Chunks encrypted: " << encrypted_chunks.size() << std::endl;

        // Create TCP connection
        boost::asio::io_context io_context;
        tcp::socket socket(io_context);
        tcp::resolver resolver(io_context);
        boost::system::error_code ec;
        boost::asio::connect(socket, resolver.resolve(session->peer.ip, "5545"), ec);

        if (ec) {
            std::cerr << "Failed to connect to " << session->peer.ip << ": " << ec.message() << std::endl;
            return false;
        }
        std::cout << "Connected to " << session->peer.ip << " for sending" << std::endl;

        // Send number of chunks
        uint32_t num_chunks = htonl((uint32_t)encrypted_chunks.size());
        boost::asio::write(socket, boost::asio::buffer(&num_chunks, sizeof(num_chunks)), ec);
        if (ec) { std::cerr << "Write chunk count failed: " << ec.message() << std::endl; return false; }
        std::cout << "Sending " << encrypted_chunks.size() << " chunks" << std::endl;

        // Send each encrypted chunk
        for (size_t i = 0; i < encrypted_chunks.size(); ++i) {
            const auto& chunk = encrypted_chunks[i];
            // Send chunk length (includes ciphertext + 16-byte GCM tag)
            // The tag is appended to ciphertext in encrypt_chunks_parallel, so chunk.length() includes both
            uint32_t chunk_len = htonl((uint32_t)chunk.length());
            boost::asio::write(socket, boost::asio::buffer(&chunk_len, sizeof(chunk_len)), ec);
            if (ec) { std::cerr << "Write chunk length failed: " << ec.message() << std::endl; return false; }
            std::cout << "Chunk " << i << " length=" << (uint32_t)chunk.length() << std::endl;

            // Send chunk data (ciphertext + 16-byte GCM authentication tag)
            boost::asio::write(socket, boost::asio::buffer(chunk), ec);
            if (ec) { std::cerr << "Write chunk data failed: " << ec.message() << std::endl; return false; }
            std::cout << "Chunk " << i << " data sent" << std::endl;
        }

        // Send IV
        uint32_t iv_len = htonl((uint32_t)iv.size());
        boost::asio::write(socket, boost::asio::buffer(&iv_len, sizeof(iv_len)), ec);
        if (ec) { std::cerr << "Write IV length failed: " << ec.message() << std::endl; return false; }
        boost::asio::write(socket, boost::asio::buffer(iv), ec);
        if (ec) { std::cerr << "Write IV data failed: " << ec.message() << std::endl; return false; }
        std::cout << "IV sent" << std::endl;

        // Send session AES key (encrypted with peer's RSA public key)
        // Note: For efficiency, we reuse the same session key for all messages in a session,
        // but still send it each time to maintain protocol compatibility with receivers
        auto peer_pub_key = crypto_->deserialize_rsa_public_key(session->peer.rsa_pub_key);
        if (!peer_pub_key) {
            std::cerr << "Failed to deserialize peer's public key" << std::endl;
            return false;
        }
        auto encrypted_aes_key = crypto_->rsa_encrypt(peer_pub_key, aes_key);
        EVP_PKEY_free(peer_pub_key);

        uint32_t key_len = htonl((uint32_t)encrypted_aes_key.size());
        boost::asio::write(socket, boost::asio::buffer(&key_len, sizeof(key_len)), ec);
        if (ec) { std::cerr << "Write key length failed: " << ec.message() << std::endl; return false; }
        boost::asio::write(socket, boost::asio::buffer(encrypted_aes_key), ec);
        if (ec) { std::cerr << "Write key data failed: " << ec.message() << std::endl; return false; }
        std::cout << "Encrypted session key sent (" << encrypted_aes_key.size() << " bytes)" << std::endl;

        socket.close(); // Explicitly close the socket
        std::cout << "Successfully sent " << encrypted_chunks.size() << " chunks to " << session->peer.ip << std::endl;
        return true;

    } catch (const std::exception& e) {
        std::cerr << "Exception in send_message: " << e.what() << std::endl;
        return false;
    }
}

std::string Chat::receive_message_from_socket(boost::asio::ip::tcp::socket& socket) {
    try {
        std::cout << "Receiving message on TCP chat socket" << std::endl;
        // Receive number of chunks
        uint32_t num_chunks;
        boost::system::error_code ec;
        boost::asio::read(socket, boost::asio::buffer(&num_chunks, sizeof(num_chunks)), ec);
        if (ec) {
            std::cerr << "Failed to receive chunk count: " << ec.message() << std::endl;
            return "";
        }
        num_chunks = ntohl(num_chunks);
        std::cout << "Chunk count received: " << num_chunks << std::endl;
        
        // Validate chunk count (prevent excessive memory allocation)
        if (num_chunks == 0) {
            std::cerr << "Error: Received zero chunks" << std::endl;
            return "";
        }
        if (num_chunks > 1000000) { // Reasonable upper limit (32KB * 1M = 32GB max message)
            std::cerr << "Error: Received excessive chunk count: " << num_chunks << std::endl;
            return "";
        }

        // Receive each encrypted chunk
        std::vector<std::string> encrypted_chunks(num_chunks);
        for (uint32_t i = 0; i < num_chunks; ++i) {
            // Receive chunk length (includes ciphertext + 16-byte GCM tag)
            uint32_t chunk_len;
            boost::asio::read(socket, boost::asio::buffer(&chunk_len, sizeof(chunk_len)), ec);
            if (ec) return "";
            chunk_len = ntohl(chunk_len);
             std::cout << "Chunk " << i << " length received: " << chunk_len << std::endl;
            
            // Validate chunk length is at least TAG_SIZE (16 bytes) and not excessive
            if (chunk_len < TAG_SIZE) {
                std::cerr << "Error: Received invalid chunk length " << chunk_len << " (minimum " << TAG_SIZE << " bytes required)" << std::endl;
                return "";
            }
            // Reasonable upper limit: 32KB + 16 bytes tag = 32832 bytes max per chunk
            if (chunk_len > CHUNK_SIZE + TAG_SIZE + 1024) { // Add 1KB tolerance
                std::cerr << "Error: Received excessive chunk length " << chunk_len << std::endl;
                return "";
            }

            // Receive chunk data (ciphertext + 16-byte GCM authentication tag)
            // Protocol: chunk_length = ciphertext_length + TAG_SIZE (16 bytes)
            // The tag will be extracted in decrypt_chunks_parallel
            encrypted_chunks[i].resize(chunk_len);
            boost::asio::read(socket, boost::asio::buffer(encrypted_chunks[i]), ec);
            if (ec) return "";
            std::cout << "Chunk " << i << " data received" << std::endl;
        }

        // Receive IV
        uint32_t iv_len;
        boost::asio::read(socket, boost::asio::buffer(&iv_len, sizeof(iv_len)), ec);
        if (ec) return "";
        iv_len = ntohl(iv_len);
        std::cout << "IV length received: " << iv_len << std::endl;
        
        // Validate IV length (AES-GCM requires 12-byte IV)
        if (iv_len != 12) {
            std::cerr << "Error: Invalid IV length " << iv_len << " (expected 12 bytes for AES-GCM)" << std::endl;
            return "";
        }

        std::vector<unsigned char> iv(iv_len);
        boost::asio::read(socket, boost::asio::buffer(iv), ec);
        if (ec) return "";
        std::cout << "IV data received" << std::endl;

        // Receive encrypted AES key
        uint32_t key_len;
        boost::asio::read(socket, boost::asio::buffer(&key_len, sizeof(key_len)), ec);
        if (ec) return "";
        key_len = ntohl(key_len);
        std::cout << "Encrypted key length received: " << key_len << std::endl;
        
        // Validate key length (RSA-2048 encrypted key should be 256 bytes, with some tolerance)
        if (key_len == 0 || key_len > 512) {
            std::cerr << "Error: Invalid encrypted key length " << key_len << std::endl;
            return "";
        }

        std::vector<unsigned char> encrypted_aes_key(key_len);
        boost::asio::read(socket, boost::asio::buffer(encrypted_aes_key), ec);
        if (ec) return "";
        std::cout << "Encrypted key data received" << std::endl;

        // Decrypt AES key with our private key
        std::cout << "Decrypting AES key" << std::endl;
        auto aes_key = crypto_->rsa_decrypt(encrypted_aes_key);
        if (aes_key.empty()) {
            std::cerr << "Failed to decrypt AES key" << std::endl;
            return "";
        }
        std::cout << "AES key decrypted, length=" << aes_key.size() << std::endl;

        // Decrypt chunks in parallel
        std::cout << "Decrypting chunks" << std::endl;
        auto decrypted_chunks = decrypt_chunks_parallel(encrypted_chunks, aes_key, iv);

        // Check if any chunks failed to decrypt
        bool all_decrypted = true;
        for (const auto& chunk : decrypted_chunks) {
            if (chunk.empty()) {
                all_decrypted = false;
                break;
            }
        }
        
        if (!all_decrypted) {
            std::cerr << "Error: Some chunks failed to decrypt" << std::endl;
            return "";
        }
        std::cout << "All chunks decrypted successfully" << std::endl;

        // Reassemble message
        std::string message = reassemble_message(decrypted_chunks);
        std::cout << "Message reassembled, size=" << message.size() << std::endl;
        return message;

    } catch (const std::exception& e) {
        std::cerr << "Exception in receive_message: " << e.what() << std::endl;
        return "";
    }
}

bool Chat::send_image(const std::string& session_id, const std::vector<unsigned char>& image_bytes) {
    // Re-use logic from send_message, simplified here for brevity but following same protocol
    std::string image_data(image_bytes.begin(), image_bytes.end());
    return send_message(session_id, image_data);
}

std::vector<unsigned char> Chat::receive_image_from_socket(boost::asio::ip::tcp::socket& socket) {
    std::string data = receive_message_from_socket(socket);
    return std::vector<unsigned char>(data.begin(), data.end());
}

std::vector<std::string> Chat::encrypt_chunks_parallel(const std::vector<std::string>& chunks, const std::vector<unsigned char>& aes_key, const std::vector<unsigned char>& iv) {
    std::vector<std::string> encrypted_chunks(chunks.size());

    #pragma omp parallel for
    for (int i = 0; i < (int)chunks.size(); ++i) {
        std::vector<unsigned char> plaintext(chunks[i].begin(), chunks[i].end());
        std::vector<unsigned char> tag;
        auto ciphertext = crypto_->aes_gcm_encrypt(aes_key, iv, plaintext, tag);
        
        // Append tag to ciphertext
        ciphertext.insert(ciphertext.end(), tag.begin(), tag.end());
        
        encrypted_chunks[i] = std::string(ciphertext.begin(), ciphertext.end());
    }

    return encrypted_chunks;
}

std::vector<std::string> Chat::decrypt_chunks_parallel(const std::vector<std::string>& chunks, const std::vector<unsigned char>& aes_key, const std::vector<unsigned char>& iv) {
    std::vector<std::string> decrypted_chunks(chunks.size());

    #pragma omp parallel for
    for (int i = 0; i < (int)chunks.size(); ++i) {
        // Validate chunk size: must be at least TAG_SIZE (16 bytes) to contain both ciphertext and tag
        if (chunks[i].size() < TAG_SIZE) {
            std::cerr << "Error: Chunk " << i << " too small (" << chunks[i].size() << " bytes), minimum " << TAG_SIZE << " bytes required (ciphertext + tag)" << std::endl;
            decrypted_chunks[i] = ""; // Mark as failed
            continue; 
        }

        std::vector<unsigned char> chunk_data(chunks[i].begin(), chunks[i].end());
        
        // Extract GCM authentication tag (last 16 bytes)
        // The chunk contains: [ciphertext][16-byte tag]
        std::vector<unsigned char> tag(chunk_data.end() - TAG_SIZE, chunk_data.end());
        
        // Extract ciphertext (everything before the tag)
        std::vector<unsigned char> ciphertext(chunk_data.begin(), chunk_data.end() - TAG_SIZE);
        
        // Decrypt and verify authentication tag
        auto plaintext = crypto_->aes_gcm_decrypt(aes_key, iv, ciphertext, tag);
        if (plaintext.empty()) {
            std::cerr << "Error: Failed to decrypt chunk " << i << " (authentication tag verification may have failed)" << std::endl;
            decrypted_chunks[i] = "";
        } else {
            decrypted_chunks[i] = std::string(plaintext.begin(), plaintext.end());
        }
    }

    return decrypted_chunks;
}