diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 59530ee..fd570d5 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -42,3 +42,42 @@ jobs: - name: Test run: ctest --test-dir build --output-on-failure + + integration-test: + name: Integration Tests (Docker) + runs-on: ubuntu-24.04 + needs: build-linux + steps: + - uses: actions/checkout@v7 + + - name: Set up Python + uses: actions/setup-python@v5 + with: + python-version: "3.12" + + - name: Install test dependencies + run: pip install -r src/Tests/integration/requirements.txt + + - name: Build Docker image + run: docker build -t byteforge:ci . + + - name: Run ByteForge container + run: docker run -d --name byteforge -p 6625:6625 byteforge:ci + + - name: Wait for server to be ready + run: | + for i in {1..30}; do + (echo > /dev/tcp/127.0.0.1/6625) 2>/dev/null && break + sleep 1 + done + + - name: Run integration tests + run: pytest src/Tests/integration -v + + - name: Dump container logs on failure + if: failure() + run: docker logs byteforge + + - name: Stop container + if: always() + run: docker rm -f byteforge diff --git a/.gitignore b/.gitignore index e1402ff..b5747eb 100644 --- a/.gitignore +++ b/.gitignore @@ -8,6 +8,8 @@ CMakeFiles/ CMakeScripts/ cmake_install.cmake Makefile +.venv/ +__pycache__/ # ========================================== # Compiled Executables and Libraries and Logs diff --git a/src/Networking/EpollEventLoop.cpp b/src/Networking/EpollEventLoop.cpp index 35e7695..9fb8a0c 100644 --- a/src/Networking/EpollEventLoop.cpp +++ b/src/Networking/EpollEventLoop.cpp @@ -18,7 +18,7 @@ EpollEventLoop::~EpollEventLoop() void EpollEventLoop::add(SocketType sock) { epoll_event ev{}; - ev.events = EPOLLIN; // watch for "readable" (level-triggered by default) + ev.events = EPOLLIN | EPOLLONESHOT; // watch for "readable", (intentionally "mutes" the socket after the first notification) ev.data.fd = sock; epoll_ctl(epollFd, EPOLL_CTL_ADD, sock, &ev); @@ -67,4 +67,13 @@ int EpollEventLoop::wait(std::vector &out) return numReady; } +bool EpollEventLoop::rearm(SocketType sock) +{ + epoll_event ev{}; + ev.events = EPOLLIN | EPOLLONESHOT; + ev.data.fd = sock; + + return epoll_ctl(epollFd, EPOLL_CTL_MOD, sock, &ev) == 0; +} + #endif //_WIN32 \ No newline at end of file diff --git a/src/Networking/EpollEventLoop.h b/src/Networking/EpollEventLoop.h index 5bbc572..f43dc54 100644 --- a/src/Networking/EpollEventLoop.h +++ b/src/Networking/EpollEventLoop.h @@ -33,6 +33,7 @@ class EpollEventLoop : public IEventLoop void add(SocketType sock) override; void remove(SocketType sock) override; int wait(std::vector &out) override; + bool rearm(SocketType sock) override; private: /// File descriptor for the epoll instance itself. diff --git a/src/Networking/EventLoop.h b/src/Networking/EventLoop.h index 0f306e5..1c9524f 100644 --- a/src/Networking/EventLoop.h +++ b/src/Networking/EventLoop.h @@ -65,6 +65,9 @@ class IEventLoop * nothing ready, or -1 on error. */ virtual int wait(std::vector &out) = 0; + + // Explained inside the Server.h + virtual bool rearm(SocketType sock) = 0; }; #endif // KV_DATABASE_EVENTLOOP_H \ No newline at end of file diff --git a/src/Networking/IocpEventLoop.cpp b/src/Networking/IocpEventLoop.cpp index 42d6625..d4a0ff9 100644 --- a/src/Networking/IocpEventLoop.cpp +++ b/src/Networking/IocpEventLoop.cpp @@ -82,10 +82,17 @@ int IocpEventLoop::wait(std::vector &out) // recv() next; if THAT returns 0, that's how a graceful close is // detected - same pattern as epoll's EPOLLIN + recv()==0 on Linux. out.push_back(EventLoopEntry{sock, IOEvent::Readable}); - armRead(sock); // re-arm so we're notified again for the next batch of data + // We no longer auto-rearm here. The worker thread is now responsible + // for calling rearm() once it finishes processing fragmentation. } return 1; } +bool IocpEventLoop::rearm(SocketType sock) +{ + armRead(sock); + return true; // WSARecv failures handle themselves asynchronously in wait() +} + #endif //_WIN32 \ No newline at end of file diff --git a/src/Networking/IocpEventLoop.h b/src/Networking/IocpEventLoop.h index a5b2b76..517684f 100644 --- a/src/Networking/IocpEventLoop.h +++ b/src/Networking/IocpEventLoop.h @@ -49,6 +49,7 @@ class IocpEventLoop : public IEventLoop void add(SocketType sock) override; void remove(SocketType sock) override; int wait(std::vector &out) override; + bool rearm(SocketType sock) override; private: /// Posts (or re-posts) the zero-byte WSARecv that arms readiness notification for a socket. diff --git a/src/Server/Server.cpp b/src/Server/Server.cpp index fd0f367..4273e57 100644 --- a/src/Server/Server.cpp +++ b/src/Server/Server.cpp @@ -125,7 +125,7 @@ void Server::acceptClients() CloseSocket(AcceptSocket); continue; } - connections[AcceptSocket] = Connection{AcceptSocket, active_key}; + connections[AcceptSocket] = std::make_shared(AcceptSocket, active_key); eventLoop->add(AcceptSocket); std::cout << "[SERVER] Client connected and registered!\n"; } @@ -182,27 +182,14 @@ void Server::runEventLoop() if (entry.event == IOEvent::Readable) { auto it = connections.find(entry.socket); - if (it == connections.end()) - continue; // already cleaned up - - { - std::lock_guard lock(busyMutex); - // A job for this socket is already queued or running - - // skip this notification. Level-triggered epoll will - // notify us again next wait() if data is still unread. - if (busySockets.count(entry.socket)) - continue; - busySockets.insert(entry.socket); - } - - Connection conn = it->second; // small struct, cheap to copy - tpool.acceptJob([this, conn]() - { - messageHandler(conn.socket, conn.sessionKey); - std::lock_guard lock(busyMutex); - busySockets.erase(conn.socket); }); - } - else // HangUp or Error + if (it == connections.end()) continue; + + std::shared_ptr user_connection = it->second; // No Copy + // wont fire again until we re-arm it. + tpool.acceptJob([this, conn = std::move(user_connection)] { + messageHandler(conn); + }); + } else // HangUp or Error { closeConnection(entry.socket); } @@ -217,21 +204,22 @@ void Server::closeConnection(SocketType sock) auto it = connections.find(sock); if (it != connections.end()) { - userSessionManager.remove_session(it->second.sessionKey); + userSessionManager.remove_session(it->second->sessionKey); connections.erase(it); } CloseSocket(sock); } -void Server::messageHandler(SocketType clientSocket, const SessionKey &sessionKey) +void Server::messageHandler(std::shared_ptr clientConnection) { std::cout << "[CLIENT] Handling client message\n"; - char buffer[1024]; + char tempBuffer[1024]; + SocketType clientSocket = clientConnection->socket; // On Linux, the buffer is safely passed to standard recv - int bytesReceived = recv(clientSocket, buffer, sizeof(buffer) - 1, 0); + int bytesReceived = recv(clientSocket, tempBuffer, sizeof(tempBuffer), 0); if (bytesReceived == 0) { @@ -247,19 +235,37 @@ void Server::messageHandler(SocketType clientSocket, const SessionKey &sessionKe // Nothing to read right now - a different job for this socket // already drained it, or epoll notified us before this job // got scheduled. Not an error, just nothing to do. + this->rearmSocket(clientSocket); return; } std::cout << "[CLIENT] recv error, disconnecting\n"; closeConnection(clientSocket); return; } - userSessionManager.update_activity(sessionKey); + + clientConnection->commandBuffer.append(tempBuffer, bytesReceived); + + if (clientConnection->commandBuffer.length() > this->MAX_COMMAND_BUFFER_LENGTH) { + std::cout << "[CLIENT] Client is abusing the command buffer disconnecting them\n"; + closeConnection(clientSocket); + return; + } + + if (clientConnection->commandBuffer.find('\n') == std::string::npos) { + this->rearmSocket(clientSocket); + return; + } + + + + userSessionManager.update_activity(clientConnection->sessionKey); std::cout << "[CLIENT] Received " << bytesReceived << " bytes\n"; - std::string message(buffer, bytesReceived); + std::string message = clientConnection->commandBuffer; std::cout << "[CLIENT] Message: " << message << "\n"; std::istringstream iss(message); std::string command, key, value; + clientConnection->commandBuffer.clear(); iss >> command; iss >> key; @@ -277,6 +283,7 @@ void Server::messageHandler(SocketType clientSocket, const SessionKey &sessionKe { std::string response = "Empty Value Recieved, Try again\n"; send(clientSocket, response.c_str(), response.length(), 0); + this->rearmSocket(clientSocket); return; } @@ -317,6 +324,18 @@ void Server::messageHandler(SocketType clientSocket, const SessionKey &sessionKe std::string response = "No command was received\n"; send(clientSocket, response.c_str(), response.length(), 0); } + + //Re-enable epoll notifications for the next command from this client + this->rearmSocket(clientSocket); +} + +void Server::rearmSocket(SocketType clientSocket) +{ + if (!eventLoop->rearm(clientSocket)) + { + std::cout << "[SERVER] Failed to re-arm socket " << clientSocket << ", closing.\n"; + closeConnection(clientSocket); + } } void Server::onSessionExpired(SocketType sock) diff --git a/src/Server/Server.h b/src/Server/Server.h index 3b8a163..8b00dc1 100644 --- a/src/Server/Server.h +++ b/src/Server/Server.h @@ -37,11 +37,9 @@ class Server RateLimiter rt; /** Server owns an instance of RateLimter class */ UserSessionManager userSessionManager; /** Managing User Sessions */ UserSessionBackgroundWorker user_session_background_worker; /** Background worker that sweeps expired sessions*/ - std::unordered_map connections; /** Client connections, keyed by socket */ + std::unordered_map> connections; /** Client connections, keyed by socket */ std::unique_ptr eventLoop; /** Watches all client sockets for readiness */ std::thread eventLoopThread; /** Thread that runs runEventLoop() */ - std::mutex busyMutex; /** Guards busySockets */ - std::unordered_set busySockets; /** Sockets with a recv job already queued/running */ std::mutex expiredMutex; /** Mutex to guard the expiredSockets vector */ std::vector expiredSockets; /** Notifys the event loop of connections to remove */ @@ -89,16 +87,29 @@ class Server /** * @brief Handles one ready-to-read event for a client: one recv() call, * command parsing, and response. - * @param clientSocket The socket that has data available. - * @param sessionKey The session tied to this client. + * @param userConnection The user connection + * */ - void messageHandler(SocketType clientSocket, const SessionKey &sessionKey); + void messageHandler(std::shared_ptr userConnection); /** * @brief Adds sockets that are expired into the expired vector to be removed later by eventloop * @param sock which is the client socket that is to be removed */ void onSessionExpired(SocketType sock); + + /** + * What it is: An explicit call (epoll_ctl with EPOLL_CTL_MOD) executed when a thread finishes its work. + * What it does: Unmutes the socket so epoll can start listening for network activity again. + * Why we use it: Because EPOLLONESHOT completely mutes the socket, it will stay dead forever unless re-armed. + * Every exit path in your thread—whether it finished a full message or is waiting for more bytes (fragmentation)—must call rearm(). + * + * @param clientSocket + */ + void rearmSocket(SocketType clientSocket); + + + uint16_t MAX_COMMAND_BUFFER_LENGTH = 2024; }; #endif \ No newline at end of file diff --git a/src/Tests/integration/client.py b/src/Tests/integration/client.py new file mode 100644 index 0000000..364be2b --- /dev/null +++ b/src/Tests/integration/client.py @@ -0,0 +1,66 @@ +import socket +import time + + +class ByteForgeClient: + def __init__(self, host="127.0.0.1", port=6625, timeout=5.0, + max_retries=5, retry_delay=0.25): + self.host = host + self.port = port + self.timeout = timeout + self.max_retries = max_retries + self.retry_delay = retry_delay + self.sock = self._connect_with_retry() + + def _connect_with_retry(self): + last_error = None + for _ in range(self.max_retries): + try: + return socket.create_connection((self.host, self.port), timeout=self.timeout) + except (ConnectionRefusedError, ConnectionResetError, + ConnectionAbortedError, OSError) as e: + last_error = e + time.sleep(self.retry_delay) + raise ConnectionError( + f"Failed to connect to {self.host}:{self.port} after {self.max_retries} attempts" + ) from last_error + + def _send(self, command: str) -> str: + last_error = None + for _ in range(self.max_retries): + try: + self.sock.sendall((command + "\n").encode()) + return self._recv_line() + except (ConnectionResetError, ConnectionAbortedError, + BrokenPipeError, OSError) as e: + last_error = e + try: + self.sock.close() + except OSError: + pass + time.sleep(self.retry_delay) + self.sock = self._connect_with_retry() + raise ConnectionError( + f"Failed to send command after {self.max_retries} attempts" + ) from last_error + + def _recv_line(self) -> str: + data = b"" + while not data.endswith(b"\n"): + chunk = self.sock.recv(1024) + if not chunk: + break + data += chunk + return data.decode().strip() + + def insert(self, key: str, value: str) -> str: + return self._send(f"INSERT {key} {value}") + + def get(self, key: str) -> str: + return self._send(f"GET {key}") + + def delete(self, key: str) -> str: + return self._send(f"DELETE {key}") + + def close(self): + self.sock.close() \ No newline at end of file diff --git a/src/Tests/integration/requirements.txt b/src/Tests/integration/requirements.txt new file mode 100644 index 0000000..1c98737 --- /dev/null +++ b/src/Tests/integration/requirements.txt @@ -0,0 +1 @@ +pytest==8.3.4 \ No newline at end of file diff --git a/src/Tests/integration/test_advanced.py b/src/Tests/integration/test_advanced.py new file mode 100644 index 0000000..e69de29 diff --git a/src/Tests/integration/test_basic_commands.py b/src/Tests/integration/test_basic_commands.py new file mode 100644 index 0000000..194f643 --- /dev/null +++ b/src/Tests/integration/test_basic_commands.py @@ -0,0 +1,103 @@ +import pytest +from client import ByteForgeClient +import socket +import threading +import time + + +@pytest.fixture +def client(): + c = ByteForgeClient() + yield c + c.close() + + +def test_insert_get_delete(client): + assert client.insert("volt", "3333") == "Insert command was successful" + assert client.get("volt") == "Get command was successful: 3333" + assert client.delete("volt") == "Delete command was successful" + + +def test_get_after_delete(client): + client.insert("volt", "3333") + client.get("volt") + client.delete("volt") + + response = client.get("volt") + assert response == "Get command was successful but key dont exist" + + + +def test_unknown_command(client): + response = client._send("FOOBAR volt 3333") + assert response == "No command was received" + + +def test_concurrent_clients_distinct_keys(): + # Kept comfortably under the rate limiter's burst capacity (10), + # since earlier tests in the session may have already spent tokens + # that haven't fully refilled yet. + num_clients = 5 + results: list[str | None] = [None] * num_clients + errors: list[str | None] = [None] * num_clients + + def worker(i): + try: + c = ByteForgeClient() + key = f"vault{i}" + value = str(1000 + i) + c.insert(key, value) + results[i] = c.get(key) + c.delete(key) + c.close() + except OSError as e: + errors[i] = str(e) + + threads = [threading.Thread(target=worker, args=(i,)) for i in range(num_clients)] + for t in threads: + t.start() + for t in threads: + t.join() + + for i in range(num_clients): + assert errors[i] is None, f"client {i} failed: {errors[i]}" + expected_value = str(1000 + i) + assert results[i] == f"Get command was successful: {expected_value}" + + +def test_cross_client_visibility_after_delete(): + client_a = ByteForgeClient() + client_b = ByteForgeClient() + + try: + assert client_a.insert("shared", "42") == "Insert command was successful" + assert client_a.delete("shared") == "Delete command was successful" + + response = client_b.get("shared") + assert response == "Get command was successful but key dont exist" + finally: + client_a.close() + client_b.close() + + +def test_rate_limiter_blocks_excess_connections(): + host, port = "127.0.0.1", 6625 + max_attempts = 50 + blocked = False + + for _ in range(max_attempts): + try: + s = socket.create_connection((host, port), timeout=1.0) + s.sendall(b"GET volt\n") + data = s.recv(1024) + s.shutdown(socket.SHUT_RDWR) + s.close() + + if data == b"": + blocked = True + break + except (ConnectionResetError, ConnectionAbortedError, BrokenPipeError, OSError): + blocked = True + break + time.sleep(0.2) + assert blocked, f"Expected at least one blocked connection within {max_attempts} attempts" \ No newline at end of file diff --git a/src/UserSession/Connection.h b/src/UserSession/Connection.h index bd0caaa..32daa03 100644 --- a/src/UserSession/Connection.h +++ b/src/UserSession/Connection.h @@ -22,6 +22,12 @@ struct Connection /// The session key identifying this client's UserSession entry. SessionKey sessionKey; + + /// This is the command of the user (Used for fragmentation management) + std::string commandBuffer; + + Connection(SocketType sock, SessionKey key) + : socket(sock), sessionKey(key) {} }; #endif // CONNECTION_H