fix: session not disconnected properly

perf: use server sockfd directly instead of a dupe of it

8 files changed, 86 insertions, 46 deletions

diff --git a/src/format.hpp b/src/format.hpp
index 6dc5e4d..cb9ca17 100644
--- a/src/format.hpp
+++ b/src/format.hpp
@@ -13,22 +13,34 @@ enum eFormatType : int8_t {
 
 namespace NFormatter {
 inline std::string fmt(eFormatType type, std::string str) {
-    std::string prefix;
-    switch (type) {
-        case LOG:  prefix = "[LOG] "; break;
-        case WARN: prefix = "[WARN] "; break;
-        case ERR:  prefix = "[ERROR] "; break;
-        case SYS:  prefix = "[SYSTEM] "; break;
-        case CHAT: prefix = "[CHAT] "; break;
-        default:   prefix = ""; break;
-    }
-    if (!str.empty() && str.back() != '\n' && !prefix.empty())
-        str += '\n';
-    return prefix + str;
+	std::string prefix;
+	switch (type) {
+	case LOG:
+		prefix = "[LOG] ";
+		break;
+	case WARN:
+		prefix = "[WARN] ";
+		break;
+	case ERR:
+		prefix = "[ERROR] ";
+		break;
+	case SYS:
+		prefix = "[SYSTEM] ";
+		break;
+	case CHAT:
+		prefix = "[CHAT] ";
+		break;
+	default:
+		prefix = "";
+		break;
+	}
+	if (!str.empty() && str.back() != '\n' && !prefix.empty())
+		str += '\n';
+	return prefix + str;
 }
 
 template <typename... Args>
 std::string fmt(eFormatType type, std::format_string<Args...> fmt, Args &&...args) {
-    return NFormatter::fmt(type, std::format(fmt, std::forward<Args>(args)...));
+	return NFormatter::fmt(type, std::format(fmt, std::forward<Args>(args)...));
 }
 } // namespace NFormatter
diff --git a/src/log.hpp b/src/log.hpp
index 8cfc917..5e5f1d3 100644
--- a/src/log.hpp
+++ b/src/log.hpp
@@ -4,20 +4,20 @@
 #include <print>
 #include <string>
 
-inline void log(FILE* _stream, eFormatType type, std::string str) {
-    std::println(_stream, "{}", NFormatter::fmt(type, str));
+inline void log(FILE *_stream, eFormatType type, std::string str) {
+	std::println(_stream, "{}", NFormatter::fmt(type, str));
 }
 
 inline void log(eFormatType type, std::string str) {
-    log(stdout, type, str);
+	log(stdout, type, str);
 }
 
 template <typename... Args>
-void log(FILE* _stream, eFormatType type, std::format_string<Args...> fmt, Args &&...args) {
-    log(_stream, type, std::format(fmt, std::forward<Args>(args)...));
+void log(FILE *_stream, eFormatType type, std::format_string<Args...> fmt, Args &&...args) {
+	log(_stream, type, std::format(fmt, std::forward<Args>(args)...));
 }
 
 template <typename... Args>
 void log(eFormatType type, std::format_string<Args...> fmt, Args &&...args) {
-    log(stdout, type, std::format(fmt, std::forward<Args>(args)...));
+	log(stdout, type, std::format(fmt, std::forward<Args>(args)...));
 }
diff --git a/src/server.cpp b/src/server.cpp
index 7893893..ac4506f 100644
--- a/src/server.cpp
+++ b/src/server.cpp
@@ -27,6 +27,6 @@ CServer::~CServer() {
 	m_sockfd.reset();
 }
 
-Hyprutils::OS::CFileDescriptor CServer::getSocket() const {
-	return m_sockfd.duplicate();
+Hyprutils::OS::CFileDescriptor *CServer::getSocket() {
+	return &m_sockfd;
 }
diff --git a/src/server.hpp b/src/server.hpp
index 2cd15b3..e440a23 100644
--- a/src/server.hpp
+++ b/src/server.hpp
@@ -8,7 +8,7 @@ class CServer {
   public:
 	CServer(uint16_t port);
 	~CServer();
-	Hyprutils::OS::CFileDescriptor getSocket() const;
+	Hyprutils::OS::CFileDescriptor *getSocket();
 
   private:
 	Hyprutils::OS::CFileDescriptor m_sockfd;
diff --git a/src/session.cpp b/src/session.cpp
index d18dbd6..7ce2238 100644
--- a/src/session.cpp
+++ b/src/session.cpp
@@ -1,5 +1,7 @@
 #include "session.hpp"
+#include "format.hpp"
 #include "log.hpp"
+#include "sessionManager.hpp"
 #include <cerrno>
 #include <cstring>
 #include <sys/socket.h>
@@ -17,34 +19,34 @@ CSession::~CSession() {
 }
 
 void CSession::onConnect() const {
-	log(INFO, "Client connected");
+	log(LOG, "Client connected");
 }
 
 void CSession::onDisconnect() const {
-	log(INFO, "Client disconnected");
+	log(LOG, "Client disconnected");
 }
 
 void CSession::onErrno(eEventType eventType) const {
 	const auto err = strerror(errno);
 	switch (eventType) {
 	case READ:
-		log(ERROR, "Error while reading: {}", err);
+		log(LOG, "Error while reading: {}", err);
 	case WRITE:
-		log(ERROR, "Error while writing: {}", err);
+		log(LOG, "Error while writing: {}", err);
 	}
 }
 
 void CSession::onRecv(const SRecvData &data) const {
-	log(INFO, "Received: {}", data.data, data.dataSize);
+	log(LOG, "Received: {}", data.data, data.dataSize);
 }
 
 void CSession::onSend(const std::string &msg) const {
-	log(INFO, "Sent: {}", msg);
+	log(LOG, "Sent: {}", msg);
 }
 
 void CSession::recvManager() {
 	while (true) {
-		CSession::SRecvData recvData = read();
+		SRecvData recvData = read();
 		if (!recvData.good)
 			break;
 	}
@@ -59,7 +61,7 @@ CSession::SRecvData CSession::read() {
 	}
 	if (dataSize == 0) {
 		recvData.good = false;
-		onDisconnect();
+		// onDisconnect();
 	}
 
 	onRecv(recvData);
@@ -67,7 +69,7 @@ CSession::SRecvData CSession::read() {
 }
 
 CSession::SRecvData CSession::read(const std::string &msg) {
-	write(eFormatType::STANDARD, "{}", msg, false);
+	write({}, "{}", msg, false);
 	return read();
 }
 
@@ -75,6 +77,8 @@ void CSession::run() {
 	onConnect();
 	registerSession();
 	recvManager();
+
+	g_pSessionManager->removeSession(self);
 }
 
 bool CSession::registerSession() {
@@ -83,13 +87,13 @@ bool CSession::registerSession() {
 		return false;
 
 	if (recvData.data[0] == '\n' || recvData.data[0] == ' ') {
-		write("Name cannot be empty", eFormatType::ERROR);
+		write("Name cannot be empty", eFormatType::ERR);
 		registerSession();
 	}
 
 	m_name = recvData.data;
 
-	std::println("Client registered as: {}", m_name);
+	log(LOG, "Client registered as: {}", m_name);
 
 	return true;
 }
@@ -106,7 +110,7 @@ bool CSession::isValid() const {
 
 template <typename... Args>
 bool CSession::write(eFormatType type, std::format_string<Args...> fmt, Args &&...args) const {
-	std::string msg = CFormatter::fmt(type, fmt, std::forward<Args>(args)...);
+	std::string msg = NFormatter::fmt(type, fmt, std::forward<Args>(args)...);
 
 	if (send(m_sockfd.get(), msg.c_str(), msg.size(), 0) < 0) {
 		onErrno(WRITE);
@@ -123,3 +127,7 @@ bool CSession::write(const std::string &msg, eFormatType type) const {
 std::string CSession::getName() const {
 	return m_name;
 }
+
+void CSession::setSelf(std::pair<std::jthread, std::shared_ptr<CSession>> *self) {
+	this->self = self;
+}
diff --git a/src/session.hpp b/src/session.hpp
index 0c5e4f9..d5fb7e6 100644
--- a/src/session.hpp
+++ b/src/session.hpp
@@ -2,7 +2,9 @@
 
 #include "FileDescriptor.hpp"
 #include "format.hpp"
+#include <memory>
 #include <string>
+#include <thread>
 
 class CSession {
   public:
@@ -21,8 +23,10 @@ class CSession {
 	template <typename... Args>
 	bool write(eFormatType type, std::format_string<Args...> fmt, Args &&...args) const;
 	void run();
+	void setSelf(std::pair<std::jthread, std::shared_ptr<CSession>> *self);
 
   private:
+	std::pair<std::jthread, std::shared_ptr<CSession>> *self;
 	enum eEventType {
 		READ,
 		WRITE
diff --git a/src/sessionManager.cpp b/src/sessionManager.cpp
index d960aeb..290b8c6 100644
--- a/src/sessionManager.cpp
+++ b/src/sessionManager.cpp
@@ -1,7 +1,9 @@
 #include "sessionManager.hpp"
 #include "server.hpp"
 #include "session.hpp"
-#include <print>
+#include <algorithm>
+#include <unistd.h>
+#include <utility>
 
 CSessionManager::CSessionManager() = default;
 
@@ -10,17 +12,29 @@ CSessionManager::~CSessionManager() {
 		thread.join();
 }
 
-void CSessionManager::newSession(Hyprutils::OS::CFileDescriptor sockfd) {
-	std::shared_ptr session = std::make_shared<CSession>(std::move(sockfd));
-	m_vSessions.emplace_back(std::jthread(&CSession::run, session), std::move(session));
+std::pair<std::jthread, std::shared_ptr<CSession>> *CSessionManager::newSession(Hyprutils::OS::CFileDescriptor sockfd) {
+	std::shared_ptr session	 = std::make_shared<CSession>(std::move(sockfd));
+	const auto		instance = &m_vSessions.emplace_back(std::jthread(&CSession::run, session), std::move(session));
+	instance->second->setSelf(instance);
+	return instance;
 }
 
 void CSessionManager::enterLoop() {
 	while (true)
-		newSession(Hyprutils::OS::CFileDescriptor{accept(g_pServer->getSocket().take(), nullptr, nullptr) /* awaits new client/session */});
+		newSession(Hyprutils::OS::CFileDescriptor{accept(g_pServer->getSocket()->get(), nullptr, nullptr) /* awaits new client/session */});
 }
 
-void CSessionManager::broadcast(const std::string &msg) const {
+void CSessionManager::broadcast(const std::string &msg, eFormatType type) const {
 	for (const auto &[thread, session] : m_vSessions)
-		session->write(msg);
+		session->write(msg, type);
+}
+
+void CSessionManager::removeSession(std::pair<std::jthread, std::shared_ptr<CSession>> *session) {
+	auto it = std::ranges::find_if(m_vSessions, [session](const auto &s) { return s.second.get() == session->second.get(); });
+
+	if (it != m_vSessions.end()) {
+		if (it->first.joinable())
+			it->first.detach(); // .detach the thread since it's removing itself
+		m_vSessions.erase(it);
+	}
 }
diff --git a/src/sessionManager.hpp b/src/sessionManager.hpp
index 2dd825a..0b447b5 100644
--- a/src/sessionManager.hpp
+++ b/src/sessionManager.hpp
@@ -1,5 +1,6 @@
 #pragma once
 
+#include "format.hpp"
 #include "session.hpp"
 #include <memory>
 #include <thread>
@@ -9,11 +10,12 @@ class CSessionManager {
   public:
 	CSessionManager();
 	~CSessionManager();
-	void enterLoop();
-	void newSession(Hyprutils::OS::CFileDescriptor sockfd);
-	void broadcast(const std::string &message) const;
-	void kick(const CSession *session) const;
-	void addSession(const CSession &session);
+	void												enterLoop();
+	std::pair<std::jthread, std::shared_ptr<CSession>> *newSession(Hyprutils::OS::CFileDescriptor sockfd);
+	void												broadcast(const std::string &message, eFormatType type = {}) const;
+	void												kick(const CSession *session) const;
+	void												addSession(const CSession &session);
+	void												removeSession(std::pair<std::jthread, std::shared_ptr<CSession>> *session);
 
   private:
 	std::vector<std::pair<std::jthread, std::shared_ptr<CSession>>> m_vSessions;