diff --git a/core/src/agenthandler.cpp b/core/src/agenthandler.cpp index efbd924..1cd8470 100644 --- a/core/src/agenthandler.cpp +++ b/core/src/agenthandler.cpp @@ -1,6 +1,5 @@ #include "agenthandler.hpp" -// Constructor for the AgentHandler class AgentHandler::AgentHandler(boost::asio::streambuf &readBuffer, boost::asio::streambuf &writeBuffer, const std::shared_ptr &config, @@ -8,44 +7,35 @@ AgentHandler::AgentHandler(boost::asio::streambuf &readBuffer, const TCPClient::pointer &client, const std::string &clientConnStr, boost::uuids::uuid uuid) - : config_(config), // Initialize configuration - log_(log), // Initialize logging - client_(client), // Initialize TCP client - readBuffer_(readBuffer), // Reference to the read buffer - writeBuffer_(writeBuffer),// Reference to the write buffer + : config_(config), + log_(log), + client_(client), + readBuffer_(readBuffer), + writeBuffer_(writeBuffer), request_(HTTP::create(config, log, readBuffer, uuid)), - clientConnStr_(clientConnStr),// Initialize client connection string + clientConnStr_(clientConnStr), uuid_(uuid) {} -// Destructor for the AgentHandler class AgentHandler::~AgentHandler() {} -// Main handler function for processing requests void AgentHandler::handle() { - std::lock_guard lock(mutex_);// Lock the mutex for thread safety + std::lock_guard lock(mutex_); - // Initialize encryption status BoolStr encryption{false, std::string("FAILED")}; - // Encrypt the request data from the read buffer encryption = aes256Encrypt(hexStreambufToStr(readBuffer_), config_->agent().token); if (encryption.ok) { - // Log successful token validation log_->write("[" + to_string(uuid_) + "] [AgentHandler handle] [Encryption Done]", Log::Level::DEBUG); - // Generate a new HTTP POST request string with the encrypted message std::string newReq( request_->genHttpPostReqString(encode64(encryption.message))); - // Check if the request type is valid if (request_->detectType()) { - // Log the HTTP request details log_->write("[" + to_string(uuid_) + "] [AgentHandler handle] [Request] : " + request_->toString(), Log::Level::DEBUG); - // Log the client connention string and HTTP request target if (request_->parsedHttpRequest().target().length() > 0) { log_->write("[" + to_string(uuid_) + "] [CONNECT] [SRC " + clientConnStr_ + "]" + " [DST " + boost::lexical_cast( @@ -54,11 +44,9 @@ void AgentHandler::handle() { Log::Level::INFO); } - // If the client socket is not open or the request type is HTTP or CONNECT if (!client_->socket().is_open() || request_->httpType() == HTTP::HttpType::http || request_->httpType() == HTTP::HttpType::connect) { - // Connect the TCP client to the server boost::system::error_code ec; ; @@ -71,7 +59,6 @@ void AgentHandler::handle() { Log::Level::INFO); } - // Check for connection errors if (ec) { log_->write(std::string("[" + to_string(uuid_) + "] [AgentHandler handle] Connection error: ") + ec.message(), @@ -80,47 +67,44 @@ void AgentHandler::handle() { } } - // Copy the new request to the read buffer copyStringToStreambuf(newReq, readBuffer_); - // Log the request to be sent to the server log_->write("[" + to_string(uuid_) + "] [AgentHandler handle] [Request To Server] : \n" + newReq, Log::Level::DEBUG); - // Write the request to the client socket and initiate a read - // operation + client_->doWrite(readBuffer_); client_->doRead(); - // Check if there is data available in the read buffer + if (client_->readBuffer().size() > 0) { - // If the HTTP request type is not CONNECT + if (request_->httpType() != HTTP::HttpType::connect) { - // Create an HTTP response handler + HTTP::pointer response = HTTP::create(config_, log_, client_->readBuffer(), uuid_); - // Parse the HTTP response + if (response->parseHttpResp()) { - // Log the response details + log_->write( "[" + to_string(uuid_) + "] [AgentHandler handle] [Response] : " + response->restoString(), Log::Level::DEBUG); - // Decrypt the response body + BoolStr decryption{false, std::string("FAILED")}; decryption = aes256Decrypt(decode64(boost::lexical_cast( response->parsedHttpResponse().body())), config_->agent().token); - // Check if decryption was successful + if (decryption.ok) { - // Copy the decrypted message to the write buffer + copyStringToStreambuf(decryption.message, writeBuffer_); log_->write("[" + to_string(uuid_) + "] [AgentHandler handle] [Decryption Done]", Log::Level::DEBUG); } else { - // Log decryption failure and close the socket + log_->write("[" + to_string(uuid_) + "] [AgentHandler handle] [Decryption Failed] : [ " + decryption.message + "] ", Log::Level::DEBUG); @@ -130,7 +114,7 @@ void AgentHandler::handle() { client_->socket().close(); } } else { - // Log if the response is not an HTTP response + log_->write( "[AgentHandler handle] [NOT HTTP Response] " "[Response] : " + @@ -138,30 +122,30 @@ void AgentHandler::handle() { Log::Level::DEBUG); } } else { - // Log the response to a CONNECT request + log_->write("[" + to_string(uuid_) + "] [AgentHandler handle] [Response to connect] : \n" + streambufToString(client_->readBuffer()), Log::Level::DEBUG); - // Move the response from the read buffer to the write buffer + moveStreambuf(client_->readBuffer(), writeBuffer_); } } else { - // Close the socket if no data is available + client_->socket().close(); return; } } else { - // Log if the request is not a valid HTTP request + log_->write("[" + to_string(uuid_) + "] [AgentHandler handle] [NOT HTTP Request] [Request] : " + streambufToString(readBuffer_), Log::Level::DEBUG); - // Close the socket if no data is available + client_->socket().close(); return; } } else { - // Log encryption failure and close the socket + log_->write("[" + to_string(uuid_) + "] [AgentHandler handle] [Encryption Failed] : [ " + encryption.message + "] ", Log::Level::DEBUG); @@ -170,7 +154,7 @@ void AgentHandler::handle() { client_->socket().remote_endpoint().address().to_string() + ":" + std::to_string(client_->socket().remote_endpoint().port()) + "] ", Log::Level::INFO); - // Close the socket if no data is available + client_->socket().close(); return; } diff --git a/core/src/agenthandler.hpp b/core/src/agenthandler.hpp index 4098e4e..4812025 100644 --- a/core/src/agenthandler.hpp +++ b/core/src/agenthandler.hpp @@ -9,28 +9,11 @@ #include "log.hpp" #include "tcpclient.hpp" -/** - * @brief This class handles requests when the process is running in agent - * mode. - * - * The AgentHandler processes HTTP requests by encrypting/decrypting data, - * forwarding requests to a TCP server, and handling responses. - */ class AgentHandler : private Uncopyable { public: - using pointer = std::shared_ptr;///< Shared pointer type for - ///< AgentHandler + using pointer = std::shared_ptr; + - /** - * @brief Factory method to create an instance of AgentHandler. - * - * @param readBuffer The buffer to read incoming data from. - * @param writeBuffer The buffer to write outgoing data to. - * @param config Shared configuration object. - * @param log Shared logging object. - * @param client Shared pointer to a TCP client. - * @return pointer A shared pointer to the created AgentHandler instance. - */ static pointer create(boost::asio::streambuf &readBuffer, boost::asio::streambuf &writeBuffer, const std::shared_ptr &config, @@ -42,41 +25,15 @@ class AgentHandler : private Uncopyable { client, clientConnStr, uuid)); } - ~AgentHandler();///< Destructor for AgentHandler + ~AgentHandler(); - /** - * @brief Handles the HTTP request. - * - * This function is called from the TCPConnection::handleRead function. - * It processes the request, encrypts or decrypts data as needed, - * forwards requests to the server, and handles responses. - */ void handle(); - /** - * @brief Gets a reference to the HTTP request object. - * - * @return The HTTP request object. - */ inline const HTTP::pointer &request() & { return request_; } - /** - * @brief Gets a rvalue reference to the HTTP request object. - * - * @return The HTTP request object (rvalue). - */ inline const HTTP::pointer &&request() && { return std::move(request_); } private: - /** - * @brief Constructs an AgentHandler instance. - * - * @param readBuffer The buffer to read incoming data from. - * @param writeBuffer The buffer to write outgoing data to. - * @param config Shared configuration object. - * @param log Shared logging object. - * @param client Shared pointer to a TCP client. - */ AgentHandler(boost::asio::streambuf &readBuffer, boost::asio::streambuf &writeBuffer, const std::shared_ptr &config, @@ -85,16 +42,16 @@ class AgentHandler : private Uncopyable { const std::string &clientConnStr, boost::uuids::uuid uuid); - const std::shared_ptr &config_;///< Shared configuration object - const std::shared_ptr &log_; ///< Shared logging object - const TCPClient::pointer &client_; ///< Shared pointer to a TCP client + const std::shared_ptr &config_; + const std::shared_ptr &log_; + const TCPClient::pointer &client_; boost::asio::streambuf &readBuffer_, - &writeBuffer_; ///< References to buffers for reading and writing data - HTTP::pointer request_;///< HTTP request object + &writeBuffer_; + HTTP::pointer request_; const std::string - &clientConnStr_;///< socket client connection string "ip:port" + &clientConnStr_; boost::uuids::uuid uuid_; - std::mutex mutex_;///< Mutex to make the class thread-safe + std::mutex mutex_; }; diff --git a/core/src/config.cpp b/core/src/config.cpp index 8ebff93..34151d4 100644 --- a/core/src/config.cpp +++ b/core/src/config.cpp @@ -1,11 +1,5 @@ #include "config.hpp" -/** - * @brief Constructs a Config instance with specified mode and file path. - * - * @param mode The operational mode of the application (server or agent). - * @param filePath The path to the configuration file. - */ Config::Config(const RunMode &mode, const std::string &filePath) : runMode_(mode), filePath_(filePath), @@ -46,11 +40,7 @@ Config::Config(const RunMode &mode, const std::string &filePath) } } -/** - * @brief Copy constructor for Config. - * - * @param config A shared pointer to another Config instance. - */ + Config::Config(const Config::pointer &config) : runMode_(config->runMode()), configYaml_(YAML::LoadFile(config->filePath())), @@ -59,16 +49,9 @@ Config::Config(const Config::pointer &config) server_(config->server()), agent_(config->agent()) {} -/** - * @brief Destructor for Config. - */ + Config::~Config() = default; -/** - * @brief Returns a string representation of the current configuration. - * - * @return A string describing the configuration settings. - */ std::string Config::toString() const { std::lock_guard lock(configMutex_); std::stringstream ss; diff --git a/core/src/config.hpp b/core/src/config.hpp index 13e799e..7266d6a 100644 --- a/core/src/config.hpp +++ b/core/src/config.hpp @@ -3,7 +3,7 @@ #include #include -#include // Include mutex header for thread safety +#include #include #include "general.hpp" @@ -50,7 +50,7 @@ class Config : private Uncopyable { std::string listenIp_; unsigned short listenPort_; - mutable std::mutex configMutex_;// Mutex for thread safety + mutable std::mutex configMutex_; explicit Config(const RunMode &mode, const std::string &filePath); diff --git a/core/src/general.hpp b/core/src/general.hpp index c9d0652..7affc79 100644 --- a/core/src/general.hpp +++ b/core/src/general.hpp @@ -23,12 +23,6 @@ #include #include -/** - * @brief Base class to prevent copying of derived classes. - * - * This class is used to make derived classes uncopyable by deleting the - * copy constructor and copy assignment operator. - */ class Uncopyable { public: Uncopyable() = default; @@ -38,53 +32,26 @@ class Uncopyable { Uncopyable &operator=(const Uncopyable &) = delete; }; -/* - * FUCK function prints it on screen - */ inline void FUCK(const auto &message) { std::cout << "FUCK FUCK FUCK FUCK FUCK FUCK FUCK FUCK : " << message << std::endl; } -/** - * @brief Structure to return a boolean status and a message. - * - * This is used for returning status and messages from functions, particularly - * those dealing with encryption or configuration validation. - */ struct BoolStr { bool ok; ///< Indicates success or failure. std::string message;///< Message describing the result or error. }; -/** - * @brief Checks if a file exists. - * - * @param name The name of the file to check. - * @return true if the file exists, false otherwise. - */ inline bool fileExists(const std::string &name) { std::ifstream file(name.c_str()); return file.good(); } -/** - * @brief Converts a `boost::asio::streambuf` to a `std::string`. - * - * @param buff The streambuf to convert. - * @return The resulting string. - */ inline std::string streambufToString(const boost::asio::streambuf &buff) { return {boost::asio::buffers_begin(buff.data()), boost::asio::buffers_begin(buff.data()) + buff.size()}; } -/** - * @brief Copies a string to a `boost::asio::streambuf`. - * - * @param inputStr The string to copy. - * @param buff The target streambuf. - */ inline void copyStringToStreambuf(const std::string &inputStr, boost::asio::streambuf &buff) { buff.consume(buff.size()); @@ -92,12 +59,6 @@ inline void copyStringToStreambuf(const std::string &inputStr, os << inputStr; } -/** - * @brief Moves data from one `boost::asio::streambuf` to another. - * - * @param source The source streambuf. - * @param target The target streambuf. - */ inline void moveStreambuf(boost::asio::streambuf &source, boost::asio::streambuf &target) { auto bytes_copied = @@ -106,12 +67,6 @@ inline void moveStreambuf(boost::asio::streambuf &source, source.consume(source.size()); } -/** - * @brief Copies data from one `boost::asio::streambuf` to another. - * - * @param source The source streambuf. - * @param target The target streambuf. - */ inline void copyStreambuf(boost::asio::streambuf &source, boost::asio::streambuf &target) { auto bytes_copied = @@ -119,13 +74,6 @@ inline void copyStreambuf(boost::asio::streambuf &source, target.commit(bytes_copied); } -/** - * @brief Converts a hex array to a string. - * - * @param data Pointer to the data array. - * @param size The size of the data array. - * @return The resulting string. - */ inline std::string hexArrToStr(const unsigned char *data, std::size_t size) { std::stringstream tempStr; tempStr << std::hex << std::setfill('0'); @@ -135,12 +83,6 @@ inline std::string hexArrToStr(const unsigned char *data, std::size_t size) { return tempStr.str(); } -/** - * @brief Converts a hex string to an integer. - * - * @param hexString The hex string. - * @return The resulting integer. - */ inline unsigned short hexToInt(const std::string &hexString) { unsigned short result; std::stringstream hexStr(hexString); @@ -148,12 +90,6 @@ inline unsigned short hexToInt(const std::string &hexString) { return result; } -/** - * @brief Converts a hex string to ASCII. - * - * @param hex The hex string. - * @return The resulting ASCII string. - */ inline std::string hexToASCII(const std::string &hex) { std::string ascii; for (size_t i = 0; i < hex.length(); i += 2) { @@ -164,12 +100,6 @@ inline std::string hexToASCII(const std::string &hex) { return ascii; } -/** - * @brief Converts a character to its hex value. - * - * @param c The character. - * @return The hex value. - */ inline unsigned char charToHex(char c) { if ('0' <= c && c <= '9') return c - '0'; if ('A' <= c && c <= 'F') return c - 'A' + 10; @@ -177,12 +107,6 @@ inline unsigned char charToHex(char c) { return 0; } -/** - * @brief Converts a hex string to a hex array. - * - * @param hexStr The hex string. - * @return The resulting hex array. - */ inline std::vector strToHexArr(const std::string &hexStr) { std::vector result(hexStr.size() / 2); for (std::size_t i = 0; i < hexStr.size() / 2; ++i) @@ -190,25 +114,12 @@ inline std::vector strToHexArr(const std::string &hexStr) { return result; } -/** - * @brief Converts a hex `boost::asio::streambuf` to a string. - * - * @param buff The streambuf. - * @return The resulting string. - */ inline std::string hexStreambufToStr(const boost::asio::streambuf &buff) { return hexArrToStr( reinterpret_cast(streambufToString(buff).c_str()), buff.size()); } -/** - * @brief Splits a string by a specified token. - * - * @param str The string to split. - * @param token The token to split by. - * @return A vector of strings resulting from the split. - */ inline std::vector splitString(const std::string &str, const std::string &token) { std::vector result; @@ -227,12 +138,6 @@ inline std::vector splitString(const std::string &str, return result; } -/** - * @brief Decodes a Base64 encoded string. - * - * @param inputStr The Base64 encoded string. - * @return The decoded string. - */ inline std::string decode64(const std::string &inputStr) { using namespace boost::archive::iterators; @@ -249,12 +154,6 @@ inline std::string decode64(const std::string &inputStr) { } } -/** - * @brief Encodes a string to Base64. - * - * @param inputStr The string to encode. - * @return The Base64 encoded string. - */ inline std::string encode64(const std::string &inputStr) { using namespace boost::archive::iterators; @@ -269,13 +168,6 @@ inline std::string encode64(const std::string &inputStr) { return encoded; } -/** - * @brief Encrypts a plaintext string using AES-256-CBC. - * - * @param plaintext The plaintext string. - * @param key The encryption key. - * @return BoolStr containing success status and the encrypted message. - */ inline BoolStr aes256Encrypt(const std::string &plaintext, const std::string &key) { BoolStr result{false, "Encryption failed"}; @@ -333,13 +225,6 @@ inline BoolStr aes256Encrypt(const std::string &plaintext, return result; } -/** - * @brief Decrypts an AES-256-CBC encrypted string. - * - * @param ciphertext_with_iv The encrypted string with IV. - * @param key The decryption key. - * @return BoolStr containing success status and the decrypted message. - */ inline BoolStr aes256Decrypt(const std::string &ciphertext_with_iv, const std::string &key) { BoolStr result{false, "Decryption failed"}; @@ -394,13 +279,6 @@ inline BoolStr aes256Decrypt(const std::string &ciphertext_with_iv, return result; } -/** - * @brief Validates the command-line arguments and configuration file. - * - * @param argc The number of command-line arguments. - * @param argv The command-line arguments. - * @return BoolStr containing success status and validation message. - */ inline BoolStr validateConfig(int argc, const char *argv[]) { BoolStr result{false, "Validation failed"}; diff --git a/core/src/http.hpp b/core/src/http.hpp index 93a7e27..30402a9 100644 --- a/core/src/http.hpp +++ b/core/src/http.hpp @@ -10,12 +10,6 @@ #include "general.hpp" #include "log.hpp" -/* - * This class is for handling request. When a request comes to - * TCPConnection(see tcp.hpp), It calls the AgentHanler::handle function(see - * agenthandler.hpp) and object from this class will be created in - * AgentHanler::handle function to do all operations related to the request - */ class HTTP { public: using pointer = std::shared_ptr; @@ -27,60 +21,33 @@ class HTTP { return pointer(new HTTP(config, log, buffer, uuid)); } - /* - * To define the type of HTTP/HTTPS request - */ enum class HttpType { https, http, connect }; - /* - * To define the type of the TLS request - */ enum class TlsTypes { TLSHandshake, ChangeCipherSpec, ApplicationData }; - /* - * To define the TLS request, This struct is used while parsing a request - * after we detect that it is a TLS HTTP - */ struct TlsRequest { std::string sni; std::string body; TlsTypes type; }; - /* - * Copy constructor if you want to copy and initialize it - */ explicit HTTP(const HTTP &request); ~HTTP(); - /* - * Functions to handle private members - */ inline const HTTP::HttpType &httpType() const { return httpType_; } inline void httpType(const HTTP::HttpType &httpType) { httpType_ = httpType; } - /* - * This function will detect the type of request (HTTP|HTTPS) - * It checks the first byte of the body, in case of 16, 14, 17 it is HTTPS - * and else it may be HTTP - */ bool detectType(); - /* - * If the request is HTTP it will parse it and store in parsedHttpRequest_ - */ bool parseHttp(); bool parseHttpResp(); - /* - * If the request is HTTPS it will parse it and store in parsedTlsRequest_ - */ bool parseTls(); const std::string genHttpPostReqString(const std::string &body) const; @@ -99,9 +66,6 @@ class HTTP { inline const TlsRequest &parsedTlsRequest() & { return parsedTlsRequest_; } - /* - * This function returns the string of parsedTlsRequest_.type - */ const std::string tlsTypeToString() const; inline const std::string &dstIP() & { return dstIP_; } @@ -110,24 +74,15 @@ class HTTP { inline const unsigned short &dstPort() & { return dstPort_; } inline const unsigned short &&dstPort() && { return std::move(dstPort_); } - /* - * This function returns string of request based on the HttpType (HTTP|HTTPS) - */ const std::string toString() const; const std::string restoString() const; private: - /* - * default constructor - */ explicit HTTP(const std::shared_ptr &config, const std::shared_ptr &log, boost::asio::streambuf &buffer, boost::uuids::uuid uuid); - /* - * This function returns map of dst IP and Port for request - */ void setIPPort(); const std::shared_ptr &config_; diff --git a/core/src/log.cpp b/core/src/log.cpp index 705b6fd..1e403f9 100644 --- a/core/src/log.cpp +++ b/core/src/log.cpp @@ -7,7 +7,7 @@ Log::Log(const std::shared_ptr &config) else if (config_->runMode() == RunMode::agent) mode_ = std::string("AGENT"); - // Open the log file for appending. + std::ofstream logFile(config_->log().file, std::ios::out | std::ios::app); if (!logFile.is_open()) { std::cerr << "Error opening log file: " << config_->log().file @@ -33,7 +33,7 @@ Log::Log(const std::shared_ptr &log) Log::~Log() {} void Log::write(const std::string &message, Level level) const { - std::lock_guard lock(logMutex_);// Lock the mutex before writing + std::lock_guard lock(logMutex_); if (level <= level_ || level == Level::ERROR) { std::ofstream logFile(config_->log().file, std::ios::out | std::ios::app); diff --git a/core/src/log.hpp b/core/src/log.hpp index 1af5a6e..345f78a 100644 --- a/core/src/log.hpp +++ b/core/src/log.hpp @@ -10,14 +10,6 @@ #include "config.hpp" #include "general.hpp" -/** - * @brief A class for logging messages with different levels (INFO, TRACE, - * ERROR, DEBUG). - * - * The Log class handles logging messages to a file and/or console based on the - * configured log level. It provides facilities to create log instances, set - * the log level, and write log messages. - */ class Log : private Uncopyable { public: enum class Level { INFO, @@ -47,5 +39,5 @@ class Log : private Uncopyable { Level level_; std::string mode_; - mutable std::mutex logMutex_;// Mutex to protect write operations + mutable std::mutex logMutex_; }; diff --git a/core/src/main.cpp b/core/src/main.cpp index 68d36be..7efcc73 100644 --- a/core/src/main.cpp +++ b/core/src/main.cpp @@ -6,37 +6,28 @@ #include "runner.hpp" int main(int argc, char const *argv[]) { - // Ensure that at least two arguments are provided: the mode and the config - // file path. + if (argc < 3) { std::cerr << "Usage: " << argv[0] << " \n"; - return 1;// Exit with error code 1 indicating incorrect usage. + return 1; } - // Validate the configuration using the provided arguments. BoolStr configValidation = validateConfig(argc, argv); if (!configValidation.ok) { std::cerr << "Configuration validation failed: " << configValidation.message << "\n"; - return 1;// Exit with error code 1 indicating configuration validation - // failure. + return 1; } - // Determine the run mode based on the first argument. - // If "agent" is specified, set runMode_ to RunMode::agent; otherwise, set it - // to RunMode::server. std::string mode(argv[1]); RunMode runMode_ = (mode == "agent") ? RunMode::agent : RunMode::server; - // Initialize the main Config object with the run mode and config file path. Config::pointer config_ = Config::create(runMode_, std::string(argv[2])); - // Initialize the main Log object using the created Config object. Log::pointer log_ = Log::create(config_); - // Create and run the Runner object with the Config and Log objects. Runner runner_(config_, log_); runner_.run(); - return 0;// Exit with success code 0. + return 0; } diff --git a/core/src/runner.cpp b/core/src/runner.cpp index 452ece8..474b23b 100644 --- a/core/src/runner.cpp +++ b/core/src/runner.cpp @@ -1,65 +1,58 @@ #include "runner.hpp" -// Constructor for the Runner class Runner::Runner(const std::shared_ptr &config, const std::shared_ptr &log) - : config_(config),// Store the configuration object - log_(log), // Store the log object - io_context_(), // Initialize the I/O context for asynchronous operations + : config_(config), + log_(log), + io_context_(), work_guard_(boost::asio::make_work_guard( - io_context_)),// Prevents io_context from running out of work - running_(true), // Initialize running state to true + io_context_)), + running_(true), strand_(boost::asio::make_strand(io_context_)) {} -// Destructor for the Runner class Runner::~Runner() { running_.store( - false); // Set the running flag to false to stop the worker threads - io_context_.stop(); // Stop the I/O context, which will end the processing - for (auto &thread: threadPool_) {// Wait for all threads to complete + false); + io_context_.stop(); + for (auto &thread: threadPool_) { if (thread.joinable()) { thread.join(); } } } -// Main function to run the server void Runner::run() { try { - // Log the configuration mode and details + log_->write("Config initialized in " + config_->modeToString() + " mode", Log::Level::INFO); log_->write(config_->toString(), Log::Level::INFO); - // Create and initialize the TCP server auto tcpServer = TCPServer::create(io_context_, config_, log_); - // Start initial worker threads for (auto i = 0; i < config_->threads(); ++i) { threadPool_.emplace_back([this] { workerThread(); }); } - // Wait for all threads to finish for (auto &thread: threadPool_) { if (thread.joinable()) { thread.join(); } } } catch (const std::exception &error) { - // Log any exceptions thrown during execution + log_->write(std::string("[Runner run] ") + error.what(), Log::Level::ERROR); } catch (...) { - // Log any unknown exceptions + log_->write("[Runner run] Unknown error occurred", Log::Level::ERROR); } } -// Function for worker threads to process tasks void Runner::workerThread() { while (running_.load()) { try { io_context_.run(); - break;// Exit loop if io_context_.run() completes + break; } catch (const std::exception &e) { log_->write(std::string("[WorkerThread] Exception: ") + e.what(), Log::Level::ERROR); @@ -69,7 +62,6 @@ void Runner::workerThread() { } } -// Function to stop thread void Runner::stop() { running_.store(false); io_context_.stop(); diff --git a/core/src/runner.hpp b/core/src/runner.hpp index 0001ae0..365371d 100644 --- a/core/src/runner.hpp +++ b/core/src/runner.hpp @@ -11,64 +11,23 @@ #include "log.hpp" #include "tcpserver.hpp" -/* - * The Runner class is responsible for managing the application's main - * execution loop. It initializes the server, starts worker threads, and - * handles exceptions. This class is used in the main function to run the I/O - * context and manage the server lifecycle. - */ class Runner : private Uncopyable { public: - /* - * Constructor to initialize the Runner object with configuration and logging - * objects. - * @param config - Shared pointer to the configuration object. - * @param log - Shared pointer to the logging object. - */ explicit Runner(const std::shared_ptr &config, const std::shared_ptr &log); - - /* - * Destructor to clean up resources, stop the I/O context, and join all - * threads. - */ ~Runner(); - - /* - * Main function to start the server and run the I/O context. - * This function sets up the server, starts worker threads, and manages - * execution. - */ void run(); private: - /* - * Worker thread function that processes asynchronous operations. - * This function runs the I/O context and handles exceptions. - */ void workerThread(); - void stop(); - // Configuration object for the server. const std::shared_ptr &config_; - - // Logging object for recording messages and errors. const std::shared_ptr &log_; - - // Boost.Asio I/O context to manage asynchronous operations. boost::asio::io_context io_context_; - - // Work guard to prevent the I/O context from running out of work. boost::asio::executor_work_guard work_guard_; - - // Atomic flag to indicate if the server is still running. std::atomic running_; - - // Thread pool to manage worker threads. std::vector threadPool_; - - // Add strand to manage concurrency boost::asio::strand strand_; }; diff --git a/core/src/serverhandler.cpp b/core/src/serverhandler.cpp index 9f6453d..f70ccbf 100644 --- a/core/src/serverhandler.cpp +++ b/core/src/serverhandler.cpp @@ -1,13 +1,5 @@ #include "serverhandler.hpp" -/* - * Constructor to initialize the ServerHandler object. - * @param readBuffer - Buffer for reading incoming data. - * @param writeBuffer - Buffer for writing outgoing data. - * @param config - Shared pointer to the configuration object. - * @param log - Shared pointer to the logging object. - * @param client - Shared pointer to the TCP client object. - */ ServerHandler::ServerHandler(boost::asio::streambuf &readBuffer, boost::asio::streambuf &writeBuffer, const std::shared_ptr &config, @@ -24,34 +16,26 @@ ServerHandler::ServerHandler(boost::asio::streambuf &readBuffer, clientConnStr_(clientConnStr), uuid_(uuid) {} -/* - * Destructor for cleanup. No specific resources to release. - */ ServerHandler::~ServerHandler() {} -/* - * Handles the incoming request by processing it based on its type - * (HTTP/HTTPS). It performs decryption, handles connection, and processes the - * request. - */ void ServerHandler::handle() { - std::lock_guard lock(mutex_);// Lock the mutex for thread safety + std::lock_guard lock(mutex_); + - // Detects the type of request (HTTP/HTTPS). if (request_->detectType()) { - // Log the received request from the agent. + log_->write( "[" + to_string(uuid_) + "] [ServerHandler handle] [Request From Agent] : " + request_->toString(), Log::Level::DEBUG); - // Decrypt the request body using AES256. + BoolStr decryption{false, std::string("FAILED")}; decryption = aes256Decrypt(decode64(boost::lexical_cast( request_->parsedHttpRequest().body())), config_->agent().token); if (decryption.ok) { - // Log successful decryption and process the request. + log_->write( "[" + to_string(uuid_) + "] [ServerHandler handle] [Token Valid] : " + request_->toString(), Log::Level::DEBUG); @@ -59,14 +43,14 @@ void ServerHandler::handle() { std::string tempHexArrStr(tempHexArr.begin(), tempHexArr.end()); copyStringToStreambuf(tempHexArrStr, readBuffer_); - // Re-check the request type and process accordingly. + if (request_->detectType()) { log_->write( "[" + to_string(uuid_) + "] [ServerHandler handle] [Request] : " + request_->toString(), Log::Level::DEBUG); switch (request_->httpType()) { case HTTP::HttpType::connect: { - // Handle CONNECT request to establish a connection. + boost::asio::streambuf tempBuff; std::iostream os(&tempBuff); if (client_->doConnect(request_->dstIP(), request_->dstPort())) { @@ -93,7 +77,7 @@ void ServerHandler::handle() { } break; case HTTP::HttpType::http: case HTTP::HttpType::https: { - // Handle HTTP or HTTPS requests. + if (request_->httpType() == HTTP::HttpType::http) { if (client_->doConnect(request_->dstIP(), request_->dstPort())) { log_->write("[" + to_string(uuid_) + "] [CONNECT] [SRC " + clientConnStr_ + "] [DST " + @@ -114,7 +98,7 @@ void ServerHandler::handle() { client_->doWrite(readBuffer_); client_->doRead(); if (client_->readBuffer().size() > 0) { - // Encrypt the response and send it back. + BoolStr encryption{false, std::string("FAILED")}; encryption = aes256Encrypt(streambufToString(client_->readBuffer()), @@ -132,7 +116,7 @@ void ServerHandler::handle() { } } } else { - // Log encryption failure and close the connection. + log_->write( "[" + to_string(uuid_) + "] [ServerHandler handle] [Encryption " @@ -145,40 +129,40 @@ void ServerHandler::handle() { client_->socket().close(); } } else { - // Close the connection if no data was read. + client_->socket().close(); return; } } break; } } else { - // Log non-HTTP requests if detected. + log_->write("[" + to_string(uuid_) + "] [ServerHandler handle] [NOT HTTP Request] [Request] : " + streambufToString(readBuffer_), Log::Level::DEBUG); - // Close the socket if no data is available + client_->socket().close(); return; } } else { - // Log decryption failure and close the connection. + log_->write("[" + to_string(uuid_) + "] [ServerHandler handle] [Decryption Failed] : [ " + decryption.message + "] ", Log::Level::DEBUG); log_->write("[" + to_string(uuid_) + "] [ServerHandler handle] [Decryption Failed] : " + request_->toString(), Log::Level::INFO); - // Close the socket if no data is available + client_->socket().close(); return; } } else { - // Log if the request is not from an agent. + log_->write( "[" + to_string(uuid_) + "] [ServerHandler handle] [NOT HTTP Request From Agent] [Request] : " + streambufToString(readBuffer_), Log::Level::DEBUG); - // Close the socket if no data is available + client_->socket().close(); return; } diff --git a/core/src/serverhandler.hpp b/core/src/serverhandler.hpp index d6e1f73..9e94626 100644 --- a/core/src/serverhandler.hpp +++ b/core/src/serverhandler.hpp @@ -8,27 +8,12 @@ #include "log.hpp" #include "tcpclient.hpp" -/* - * The ServerHandler class is responsible for managing and handling HTTP/HTTPS - * requests when the process is running in server mode. It is initialized with - * buffers for reading and writing data, as well as configuration and logging - * objects, and a TCP client object for communication. - */ + class ServerHandler : private Uncopyable { public: - using pointer = std::shared_ptr;// Type alias for shared - // pointer to ServerHandler. + using pointer = std::shared_ptr; + - /* - * Static factory method to create a shared pointer to a ServerHandler - * instance. - * @param readBuffer - Buffer used for reading incoming data. - * @param writeBuffer - Buffer used for writing outgoing data. - * @param config - Shared pointer to the configuration object. - * @param log - Shared pointer to the logging object. - * @param client - Shared pointer to the TCP client object. - * @return A shared pointer to the created ServerHandler instance. - */ static pointer create(boost::asio::streambuf &readBuffer, boost::asio::streambuf &writeBuffer, const std::shared_ptr &config, @@ -40,39 +25,15 @@ class ServerHandler : private Uncopyable { client, clientConnStr, uuid)); } - /* - * Destructor to clean up resources. - */ ~ServerHandler(); - /* - * Handles incoming requests by processing them based on their type - * (HTTP/HTTPS). This method is invoked from the TCPConnection::handleRead - * function. - */ void handle(); - /* - * Accessor for the HTTP request object. - * @return A reference to the HTTP request object. - */ inline const HTTP::pointer &request() & { return request_; } - /* - * Move accessor for the HTTP request object. - * @return An rvalue reference to the HTTP request object. - */ inline const HTTP::pointer &&request() && { return std::move(request_); } private: - /* - * Constructor to initialize the ServerHandler object. - * @param readBuffer - Buffer used for reading incoming data. - * @param writeBuffer - Buffer used for writing outgoing data. - * @param config - Shared pointer to the configuration object. - * @param log - Shared pointer to the logging object. - * @param client - Shared pointer to the TCP client object. - */ explicit ServerHandler(boost::asio::streambuf &readBuffer, boost::asio::streambuf &writeBuffer, const std::shared_ptr &config, @@ -82,16 +43,16 @@ class ServerHandler : private Uncopyable { boost::uuids::uuid uuid); const std::shared_ptr - &config_; // Reference to the configuration object. - const std::shared_ptr &log_; // Reference to the logging object. - const TCPClient::pointer &client_;// Reference to the TCP client object. + &config_; + const std::shared_ptr &log_; + const TCPClient::pointer &client_; boost::asio::streambuf &readBuffer_, - &writeBuffer_; // References to the read and write buffers. - HTTP::pointer request_;// Shared pointer to the HTTP request object. + &writeBuffer_; + HTTP::pointer request_; const std::string - &clientConnStr_;///< socket client connection string "ip:port" + &clientConnStr_; boost::uuids::uuid uuid_; - std::mutex mutex_;///< Mutex to make the class thread-safe + std::mutex mutex_; }; diff --git a/core/src/tcpclient.cpp b/core/src/tcpclient.cpp index 1c75b12..06f71c3 100644 --- a/core/src/tcpclient.cpp +++ b/core/src/tcpclient.cpp @@ -2,16 +2,6 @@ #include -/* - * Constructor for the TCPClient class. - * Initializes the TCP client with given configuration and logging objects, - * and sets up the required I/O context, socket, resolver, and timer. - * - * @param io_context - The Boost Asio I/O context to be used for asynchronous - * operations. - * @param config - Shared pointer to the configuration object. - * @param log - Shared pointer to the logging object. - */ TCPClient::TCPClient(boost::asio::io_context &io_context, const std::shared_ptr &config, const std::shared_ptr &log) @@ -22,43 +12,26 @@ TCPClient::TCPClient(boost::asio::io_context &io_context, resolver_(io_context), timeout_(io_context) {} -/* - * Getter for the socket. - * - * @return Reference to the Boost Asio TCP socket. - */ boost::asio::ip::tcp::socket &TCPClient::socket() { std::lock_guard lock(mutex_); return socket_; } -/* - * Moves the contents of the given stream buffer to the internal write buffer. - * - * @param buffer - Stream buffer containing data to be written. - */ void TCPClient::writeBuffer(boost::asio::streambuf &buffer) { std::lock_guard lock(mutex_); moveStreambuf(buffer, writeBuffer_); } -/* - * Connects the TCP client to a specified destination IP and port. - * Logs the connection attempt and handles any connection errors. - * - * @param dstIP - Destination IP address as a string. - * @param dstPort - Destination port number. - */ bool TCPClient::doConnect(const std::string &dstIP, const unsigned short &dstPort) { std::lock_guard lock(mutex_); try { - // Log connection attempt + log_->write("[" + to_string(uuid_) + "] [TCPClient doConnect] [DST " + dstIP + ":" + std::to_string(dstPort) + "]", Log::Level::DEBUG); - // Resolve the endpoint and connect + boost::system::error_code error_code; auto endpoint = resolver_.resolve(dstIP.c_str(), std::to_string(dstPort).c_str(), error_code); if (error_code) { @@ -68,24 +41,18 @@ bool TCPClient::doConnect(const std::string &dstIP, return true; } } catch (std::exception &error) { - // Log connection errors + log_->write(std::string("[" + to_string(uuid_) + "] [TCPClient doConnect] ") + error.what(), Log::Level::ERROR); return false; } } -/* - * Writes data from the internal write buffer to the connected socket. - * Logs the write operation and handles any write errors. - * - * @param buffer - Stream buffer containing data to be written. - */ void TCPClient::doWrite(boost::asio::streambuf &buffer) { std::lock_guard lock(mutex_); try { moveStreambuf(buffer, writeBuffer_); - // Log details of the write operation + log_->write("[" + to_string(uuid_) + "] [TCPClient doWrite] [DST " + socket_.remote_endpoint().address().to_string() + ":" + std::to_string(socket_.remote_endpoint().port()) + "] " + @@ -96,7 +63,7 @@ void TCPClient::doWrite(boost::asio::streambuf &buffer) { std::to_string(socket_.remote_endpoint().port()) + "] " + "[Bytes " + std::to_string(writeBuffer_.size()) + "] ", Log::Level::TRACE); - // Perform the write operation + boost::system::error_code error; if (writeBuffer_.size() > 0) { resetTimeout(); @@ -118,30 +85,25 @@ void TCPClient::doWrite(boost::asio::streambuf &buffer) { return; } } catch (std::exception &error) { - // Log exceptions during the write operation + log_->write(std::string("[" + to_string(uuid_) + "] [TCPClient doWrite] [catch] ") + error.what(), Log::Level::DEBUG); socketShutdown(); } } -/* - * Reads data from the connected socket into the internal read buffer. - * Handles data reading, potential errors, and logs the read operation. - * Implements a retry mechanism based on configuration settings. - */ void TCPClient::doRead() { std::lock_guard lock(mutex_); try { - // Clear the read buffer + readBuffer_.consume(readBuffer_.size()); boost::system::error_code error; resetTimeout(); - // Read at least 39 bytes from the socket + boost::asio::read(socket_, readBuffer_, boost::asio::transfer_at_least(1), error); cancelTimeout(); - // Implement retry mechanism if data is still available + boost::asio::steady_timer timer(io_context_); for (auto i = 0; i <= config_->general().repeatWait; i++) { while (true) { @@ -160,13 +122,13 @@ void TCPClient::doRead() { log_->write("[" + to_string(uuid_) + "] [TCPClient doRead] [EOF] Connection closed by peer.", Log::Level::TRACE); socketShutdown(); - return;// Exit after closing the socket + return; } else if (error) { log_->write( std::string("[" + to_string(uuid_) + "] [TCPClient doRead] [error] ") + error.message(), Log::Level::ERROR); socketShutdown(); - return;// Exit after closing the socket + return; } } timer.expires_after(std::chrono::milliseconds(config_->general().timeWait)); @@ -175,7 +137,7 @@ void TCPClient::doRead() { if (readBuffer_.size() > 0) { try { - // Log the successful read operation + log_->write("[" + to_string(uuid_) + "] [TCPClient doRead] [SRC " + socket_.remote_endpoint().address().to_string() + ":" + std::to_string(socket_.remote_endpoint().port()) + @@ -188,30 +150,30 @@ void TCPClient::doRead() { "] ", Log::Level::TRACE); } catch (std::exception &error) { - // Log exceptions during logging + log_->write( std::string("[" + to_string(uuid_) + "] [TCPClient doRead] [catch log] ") + error.what(), Log::Level::DEBUG); } } else { socketShutdown(); - return;// Exit after closing the socket + return; } } catch (std::exception &error) { - // Log exceptions during the read operation + log_->write(std::string("[" + to_string(uuid_) + "] [TCPClient doRead] [catch read] ") + error.what(), Log::Level::DEBUG); socketShutdown(); - return;// Exit after closing the socket + return; } } void TCPClient::resetTimeout() { - // indicates no timeout + if (!config_->general().timeout) return; - // Start/Reset the timer and cancel old handlers + timeout_.expires_from_now(boost::posix_time::seconds(config_->general().timeout)); timeout_.async_wait(boost::bind(&TCPClient::onTimeout, shared_from_this(), @@ -219,16 +181,16 @@ void TCPClient::resetTimeout() { } void TCPClient::cancelTimeout() { - // If timeout is enabled, cancel the handlers + if (config_->general().timeout) timeout_.cancel(); } void TCPClient::onTimeout(const boost::system::error_code &error) { - // Ignore cancellation and only handle timer expiration. - if (error /* No Error */ || error == boost::asio::error::operation_aborted) return; - // Timeout has expired, do necessary actions + if (error || error == boost::asio::error::operation_aborted) return; + + log_->write( std::string("[" + to_string(uuid_) + "] [TCPClient onTimeout] [expiration] ") + std::to_string(+config_->general().timeout) + diff --git a/core/src/tcpclient.hpp b/core/src/tcpclient.hpp index 13d68f3..9773bbd 100644 --- a/core/src/tcpclient.hpp +++ b/core/src/tcpclient.hpp @@ -11,82 +11,34 @@ #include "http.hpp" #include "log.hpp" -/* - * TCPClient is a class for managing a TCP client connection. - * It handles connecting to an endpoint, sending data, and receiving data. - */ class TCPClient : public boost::enable_shared_from_this { public: - // Type alias for shared pointer to TCPClient using pointer = boost::shared_ptr; - /* - * Factory method to create a new TCPClient instance. - * - * @param io_context - The Boost Asio I/O context for asynchronous - * operations. - * @param config - Shared pointer to the configuration object. - * @param log - Shared pointer to the logging object. - * @return Shared pointer to a new TCPClient instance. - */ static pointer create(boost::asio::io_context &io_context, const std::shared_ptr &config, const std::shared_ptr &log) { return pointer(new TCPClient(io_context, config, log)); } - /* - * Getter for the socket. - * - * @return Reference to the Boost Asio TCP socket. - */ boost::asio::ip::tcp::socket &socket(); - /* - * Moves data from a given stream buffer to the internal write buffer. - * - * @param buffer - Stream buffer containing data to be written. - */ void writeBuffer(boost::asio::streambuf &buffer); - /* - * Accessors for the internal buffers. - * - * @return Reference or rvalue reference to the internal write buffer. - */ inline boost::asio::streambuf &writeBuffer() & { return writeBuffer_; } inline boost::asio::streambuf &&writeBuffer() && { return std::move(writeBuffer_); } - /* - * Accessors for the internal buffers. - * - * @return Reference or rvalue reference to the internal read buffer. - */ inline boost::asio::streambuf &readBuffer() & { return readBuffer_; } inline boost::asio::streambuf &&readBuffer() && { return std::move(readBuffer_); } - /* - * Initiates a connection to the specified destination IP and port. - * - * @param dstIP - Destination IP address as a string. - * @param dstPort - Destination port number. - */ bool doConnect(const std::string &dstIP, const unsigned short &dstPort); - /* - * Writes data from the internal write buffer to the connected socket. - * - * @param buffer - Stream buffer containing data to be written. - */ void doWrite(boost::asio::streambuf &buffer); - /* - * Reads data from the connected socket into the internal read buffer. - */ void doRead(); void socketShutdown(); @@ -94,38 +46,28 @@ class TCPClient : public boost::enable_shared_from_this { boost::uuids::uuid uuid_; private: - /* - * Private constructor for TCPClient. - * Initializes the TCP client with the given I/O context, configuration, and - * logging objects. - * - * @param io_context - The Boost Asio I/O context for asynchronous - * operations. - * @param config - Shared pointer to the configuration object. - * @param log - Shared pointer to the logging object. - */ explicit TCPClient(boost::asio::io_context &io_context, const std::shared_ptr &config, const std::shared_ptr &log); - // If the timeout is enabled, start/reset it + void resetTimeout(); - // Cancel the timeout + void cancelTimeout(); - // Handler in case of a timeout expiration + void onTimeout(const boost::system::error_code &error); - // Members + const std::shared_ptr - &config_; // Reference to the configuration object - const std::shared_ptr &log_; // Reference to the logging object - boost::asio::io_context &io_context_; // Reference to the I/O context - boost::asio::ip::tcp::socket socket_; // TCP socket for communication - boost::asio::streambuf writeBuffer_; // Buffer for data to be written - boost::asio::streambuf readBuffer_; // Buffer for data to be read - boost::asio::ip::tcp::resolver resolver_;// Resolver for endpoint resolution - boost::asio::deadline_timer timeout_; // Timer for managing timeouts - - // Mutex for protecting access to the socket and buffers + &config_; + const std::shared_ptr &log_; + boost::asio::io_context &io_context_; + boost::asio::ip::tcp::socket socket_; + boost::asio::streambuf writeBuffer_; + boost::asio::streambuf readBuffer_; + boost::asio::ip::tcp::resolver resolver_; + boost::asio::deadline_timer timeout_; + + mutable std::mutex mutex_; }; diff --git a/core/src/tcpconnection.cpp b/core/src/tcpconnection.cpp index 4f076fd..0486460 100644 --- a/core/src/tcpconnection.cpp +++ b/core/src/tcpconnection.cpp @@ -1,37 +1,34 @@ #include "tcpconnection.hpp" -/* - * TCPConnection - */ TCPConnection::TCPConnection(boost::asio::io_context &io_context, const std::shared_ptr &config, const std::shared_ptr &log, - const TCPClient::pointer &client) - : socket_(io_context), // Initialize the TCP socket with the I/O context - config_(config), // Store the configuration - log_(log), // Store the logging object - io_context_(io_context),// Store the I/O context - client_(client), // Store the TCP client pointer - timeout_(io_context), // Initialize the deadline timer with the I/O context + TCPClient::pointer client) + : socket_(io_context), + config_(config), + log_(log), + io_context_(io_context), + client_(client), + timeout_(io_context), strand_(boost::asio::make_strand(io_context_)) { uuid_ = boost::uuids::random_generator()(); } boost::asio::ip::tcp::socket &TCPConnection::socket() { - return socket_;// Return the TCP socket + return socket_; } void TCPConnection::start() { - std::lock_guard lock(mutex_);// Lock the mutex for thread safety + std::lock_guard lock(mutex_); client_->uuid_ = uuid_; - doRead();// Start reading from the socket + doRead(); } void TCPConnection::doRead() { try { - readBuffer_.consume(readBuffer_.size()); // Clear the read buffer - writeBuffer_.consume(writeBuffer_.size());// Clear the write buffer + readBuffer_.consume(readBuffer_.size()); + writeBuffer_.consume(writeBuffer_.size()); resetTimeout(); boost::asio::async_read( socket_, readBuffer_, boost::asio::transfer_exactly(1), @@ -40,11 +37,11 @@ void TCPConnection::doRead() { boost::bind(&TCPConnection::handleRead, shared_from_this(), boost::asio::placeholders::error, boost::asio::placeholders:: - bytes_transferred)));// Start an asynchronous - // read operation + bytes_transferred))); + } catch (std::exception &error) { log_->write(std::string("[" + to_string(uuid_) + "] [TCPConnection doRead] [catch] ") + error.what(), - Log::Level::ERROR);// Log any exceptions + Log::Level::ERROR); socketShutdown(); } } @@ -52,8 +49,6 @@ void TCPConnection::doRead() { void TCPConnection::handleRead(const boost::system::error_code &error, size_t) { try { - // Getting called means a successful read operation in doRead(), - // so cancel the timeout, and reset it before next I/O cancelTimeout(); if (error) { @@ -75,7 +70,7 @@ void TCPConnection::handleRead(const boost::system::error_code &error, size_t) { for (auto i = 0; i <= config_->general().repeatWait; i++) { while (true) { if (socket_.available() == 0) - break;// Exit if no more data is available + break; boost::system::error_code read_error; resetTimeout(); boost::asio::read(socket_, readBuffer_, @@ -97,10 +92,9 @@ void TCPConnection::handleRead(const boost::system::error_code &error, size_t) { } timer.expires_from_now( boost::posix_time::milliseconds(config_->general().timeWait)); - timer.wait();// Wait for the specified timeout duration + timer.wait(); } - // Log the successful read operation log_->write("[" + to_string(uuid_) + "] [TCPConnection handleRead] [SRC " + socket_.remote_endpoint().address().to_string() + ":" + std::to_string(socket_.remote_endpoint().port()) + @@ -111,26 +105,25 @@ void TCPConnection::handleRead(const boost::system::error_code &error, size_t) { std::to_string(socket_.remote_endpoint().port()) + "] " + "[Bytes " + std::to_string(readBuffer_.size()) + "] ", Log::Level::TRACE); - // Handle the request based on the running mode + if (config_->runMode() == RunMode::agent) { AgentHandler::pointer agentHandler_ = AgentHandler::create( readBuffer_, writeBuffer_, config_, log_, client_, socket_.remote_endpoint().address().to_string() + ":" + std::to_string(socket_.remote_endpoint().port()), uuid_); - agentHandler_->handle();// Process the request as an agent + agentHandler_->handle(); } else if (config_->runMode() == RunMode::server) { ServerHandler::pointer serverHandler_ = ServerHandler::create( readBuffer_, writeBuffer_, config_, log_, client_, socket_.remote_endpoint().address().to_string() + ":" + std::to_string(socket_.remote_endpoint().port()), uuid_); - serverHandler_->handle();// Process the request as a server + serverHandler_->handle(); } if (writeBuffer_.size() > 0) { - doWrite();// Write to the socket if there is data in the write - // buffer + doWrite(); } else { socketShutdown(); } @@ -144,45 +137,41 @@ void TCPConnection::handleRead(const boost::system::error_code &error, size_t) { } void TCPConnection::doWrite() { - std::lock_guard lock(mutex_);// Lock the mutex for thread safety + std::lock_guard lock(mutex_); try { log_->write("[" + to_string(uuid_) + "] [TCPConnection doWrite] [DST " + socket_.remote_endpoint().address().to_string() + ":" + std::to_string(socket_.remote_endpoint().port()) + "] [Bytes " + std::to_string(writeBuffer_.size()) + "] ", Log::Level::DEBUG); - log_->write("[" + to_string(uuid_) + "] [Write to] [DST " + - socket_.remote_endpoint().address().to_string() + ":" + - std::to_string(socket_.remote_endpoint().port()) + - "] [Bytes " + std::to_string(writeBuffer_.size()) + "] ", - Log::Level::TRACE);// Log the details of the write operation - boost::system::error_code error; resetTimeout(); - boost::asio::write(socket_, writeBuffer_, - error);// Write data to the socket - cancelTimeout(); - if (error) { - log_->write( - std::string("[" + to_string(uuid_) + "] [TCPConnection doWrite] [error] ") + error.message(), - Log::Level::ERROR);// Log any errors during the write operation - socketShutdown(); - return; - } - doRead();// Continue reading after writing + boost::asio::async_write(socket_, writeBuffer_, + boost::asio::bind_executor(strand_, + [self = shared_from_this()](const boost::system::error_code &error, std::size_t /*bytes_transferred*/) { + self->handleWrite(error); + })); } catch (std::exception &error) { log_->write( std::string("[" + to_string(uuid_) + "] [TCPConnection doWrite] [catch] ") + error.what(), - Log::Level::ERROR);// Log any exceptions during the write operation + Log::Level::ERROR); + socketShutdown(); + } +} + +void TCPConnection::handleWrite(const boost::system::error_code &error) { + cancelTimeout(); + if (!error) { + doRead(); + } else { + log_->write("[" + to_string(uuid_) + "] [TCPConnection handleWrite] [error] " + error.message(), + Log::Level::ERROR); socketShutdown(); } } void TCPConnection::resetTimeout() { - // indicates no timeout if (!config_->general().timeout) return; - - // Start/Reset the timer and cancel old handlers timeout_.expires_from_now(boost::posix_time::seconds(config_->general().timeout)); timeout_.async_wait(boost::bind(&TCPConnection::onTimeout, shared_from_this(), @@ -190,16 +179,12 @@ void TCPConnection::resetTimeout() { } void TCPConnection::cancelTimeout() { - // If timeout is enabled, cancel the handlers if (config_->general().timeout) timeout_.cancel(); } void TCPConnection::onTimeout(const boost::system::error_code &error) { - // Ignore cancellation and only handle timer expiration. if (error /* No Error */ || error == boost::asio::error::operation_aborted) return; - - // Timeout has expired, do necessary actions log_->write( std::string("[" + to_string(uuid_) + "] [TCPConnection onTimeout] [expiration] ") + std::to_string(+config_->general().timeout) + diff --git a/core/src/tcpconnection.hpp b/core/src/tcpconnection.hpp index 8069b74..b7350cf 100644 --- a/core/src/tcpconnection.hpp +++ b/core/src/tcpconnection.hpp @@ -9,34 +9,25 @@ #include "general.hpp" #include "serverhandler.hpp" -/* - * This class handles a single TCP connection. - * It is responsible for managing the TCP socket and its associated - * buffers, and it processes incoming and outgoing data. - */ class TCPConnection : public boost::enable_shared_from_this { public: using pointer = boost::shared_ptr;// Define a type alias for a shared // pointer to TCPConnection - // Factory method to create a new TCPConnection instance static pointer create(boost::asio::io_context &io_context, const std::shared_ptr &config, const std::shared_ptr &log, - const TCPClient::pointer &client) { + TCPClient::pointer client) { return pointer(new TCPConnection(io_context, config, log, client)); } - // Provides access to the TCP socket boost::asio::ip::tcp::socket &socket(); - // Moves the content of the given buffer to the write buffer inline void writeBuffer(boost::asio::streambuf &buffer) { moveStreambuf(buffer, writeBuffer_); } - // Accessors for the read buffer inline const boost::asio::streambuf &readBuffer() & { return readBuffer_; } inline const boost::asio::streambuf &&readBuffer() && { return std::move(readBuffer_); @@ -44,47 +35,32 @@ class TCPConnection : public boost::enable_shared_from_this { inline boost::uuids::uuid uuid() { return uuid_; } - // Starts the reading process void start(); - - // Initiates an asynchronous read operation void doRead(); - - // Handles the completion of an asynchronous read operation void handleRead(const boost::system::error_code &error, size_t bytes_transferred); - - // Initiates an asynchronous write operation void doWrite(); - + void handleWrite(const boost::system::error_code &error); void socketShutdown(); private: - // Private constructor to enforce use of the factory method explicit TCPConnection(boost::asio::io_context &io_context, const std::shared_ptr &config, const std::shared_ptr &log, - const TCPClient::pointer &client); - - // If the timeout is enabled, start/reset it + TCPClient::pointer client); void resetTimeout(); - // Cancel the timeout void cancelTimeout(); - // Handler in case of a timeout expiration void onTimeout(const boost::system::error_code &error); - boost::asio::ip::tcp::socket socket_; // The TCP socket for this connection - const std::shared_ptr &config_;// Configuration settings - const std::shared_ptr &log_; // Logging object - boost::asio::io_context &io_context_; // The I/O context - const TCPClient::pointer &client_; // Pointer to the associated TCP client + boost::asio::ip::tcp::socket socket_; + const std::shared_ptr &config_; + const std::shared_ptr &log_; + boost::asio::io_context &io_context_; + TCPClient::pointer client_; boost::asio::streambuf readBuffer_, - writeBuffer_; // Buffers for reading and writing data - boost::asio::deadline_timer timeout_;// Connections timeout - + writeBuffer_; + boost::asio::deadline_timer timeout_; boost::asio::strand strand_; - boost::uuids::uuid uuid_; - - std::mutex mutex_;///< Mutex to make the class thread-safe + std::mutex mutex_; }; diff --git a/core/src/tcpserver.cpp b/core/src/tcpserver.cpp index b6007fb..221ac23 100644 --- a/core/src/tcpserver.cpp +++ b/core/src/tcpserver.cpp @@ -1,54 +1,39 @@ #include "tcpserver.hpp" -// Constructor for TCPServer + TCPServer::TCPServer(boost::asio::io_context &io_context, const std::shared_ptr &config, const std::shared_ptr &log) - : config_(config),// Initialize configuration - log_(log), // Initialize logging - client_(TCPClient::create(io_context, config, - log)),// Create a TCPClient instance - io_context_(io_context), // Initialize the I/O context - acceptor_( - io_context, - boost::asio::ip::tcp::endpoint(boost::asio::ip::address::from_string( - config->listenIp()),// Listen IP - config->listenPort())), // Listen port - strand_(boost::asio::make_strand(io_context)) // Initialize the strand -{ - startAccept();// Begin accepting incoming connections + : config_(config), + log_(log), + io_context_(io_context), + acceptor_(io_context, + boost::asio::ip::tcp::endpoint(boost::asio::ip::address::from_string(config->listenIp()), + config->listenPort())) { + startAccept(); } -// Starts an asynchronous accept operation void TCPServer::startAccept() { - // Create a new TCP connection - auto connection = TCPConnection::create(io_context_, config_, log_, client_); + auto client = TCPClient::create(io_context_, config_, log_); + auto connection = TCPConnection::create(io_context_, config_, log_, client); - // Set socket option to allow address reuse acceptor_.set_option(boost::asio::socket_base::reuse_address(true)); - - // Initiate asynchronous accept operation using strand to synchronize acceptor_.async_accept( - connection->socket(),// Socket to be accepted - boost::asio::bind_executor( - strand_, [this, connection](const boost::system::error_code &error) { - // Handle the result of the accept operation - handleAccept(connection, error); - })); + connection->socket(), + [this, connection](const boost::system::error_code &error) { + handleAccept(connection, error); + }); } -// Handles the result of an asynchronous accept operation void TCPServer::handleAccept(TCPConnection::pointer connection, const boost::system::error_code &error) { if (!error) { - // No error, start the connection + connection->start(); } else { - // Log error details - log_->write("TCPServer handleAccept] " + std::string(error.message()), + + log_->write("[TCPServer handleAccept] " + std::string(error.message()), Log::Level::ERROR); } - - // Continue accepting new connections startAccept(); } diff --git a/core/src/tcpserver.hpp b/core/src/tcpserver.hpp index 57b7ba1..2385097 100644 --- a/core/src/tcpserver.hpp +++ b/core/src/tcpserver.hpp @@ -10,60 +10,27 @@ #include "serverhandler.hpp" #include "tcpconnection.hpp" -/** - * @brief This class is to create and handle a TCP server. - * It listens to the IP:Port and handles the socket. - */ class TCPServer { public: using pointer = boost::shared_ptr; - /** - * @brief Factory method to create a TCPServer instance. - * - * @param io_context The io_context to be used for asynchronous operations. - * @param config The shared configuration object. - * @param log The shared log object. - * @return pointer A shared pointer to the created TCPServer instance. - */ static pointer create(boost::asio::io_context &io_context, const std::shared_ptr &config, const std::shared_ptr &log) { return pointer(new TCPServer(io_context, config, log)); } + void startAccept(); private: - /** - * @brief Constructs a TCPServer instance. - * - * @param io_context The io_context to be used for asynchronous operations. - * @param config The shared configuration object. - * @param log The shared log object. - */ explicit TCPServer(boost::asio::io_context &io_context, const std::shared_ptr &config, const std::shared_ptr &log); - /** - * @brief Starts accepting incoming connections. - */ - void startAccept(); - - /** - * @brief Handles the acceptance of a new connection. - * - * @param connection The pointer to the new TCP connection. - * @param error The error code from the accept operation. - */ void handleAccept(TCPConnection::pointer connection, const boost::system::error_code &error); const std::shared_ptr &config_; const std::shared_ptr &log_; - TCPClient::pointer client_; boost::asio::io_context &io_context_; boost::asio::ip::tcp::acceptor acceptor_; - - // Strand to ensure that handler functions are not called concurrently. - boost::asio::strand strand_; }; diff --git a/nipovpn/etc/nipovpn/config.yaml b/nipovpn/etc/nipovpn/config.yaml index 70d4e05..85dada9 100644 --- a/nipovpn/etc/nipovpn/config.yaml +++ b/nipovpn/etc/nipovpn/config.yaml @@ -13,7 +13,7 @@ general: # timeout: unsigned short # Defines the timeout for I/O Operation in seconds. 0 indicates no timeout. # Useful to automatically close stalled connections. - timeout: 10 + timeout: 5 # repeatWait: unsigned short(1-65,635) # Defines the loop count which will try to repeat read from socket. # Same as timeWait @@ -22,7 +22,7 @@ general: # This block is to define log directives log: # logLevel: "INFO|TRACE|DEBUG" - logLevel: "TRACE" + logLevel: "INFO" # logFile: "/var/log/nipo/nipo.log" # Path of log file logFile: "/var/log/nipovpn/nipovpn.log" @@ -30,7 +30,7 @@ log: # This block is to define server directives server: # threads: number of threads in server - threads: 8 + threads: 16 # listenIp: which must be IPv4 which is used on server listenIp: "0.0.0.0" # listenPort: the port number which will be used on server