From 4fa74d0cfb3f492381a7236179d7a29e03e70b22 Mon Sep 17 00:00:00 2001 From: Nikola Petrov Date: Sat, 5 Apr 2025 23:20:35 +0200 Subject: [PATCH] add block id --- server/src/server.cpp | 87 ++++++++++++++++++++++++++++++++++++++++++- shared/inc/sql.hpp | 8 ++++ shared/src/sql.cpp | 11 ++++++ 3 files changed, 104 insertions(+), 2 deletions(-) diff --git a/server/src/server.cpp b/server/src/server.cpp index 16b544b..e338276 100644 --- a/server/src/server.cpp +++ b/server/src/server.cpp @@ -6,12 +6,94 @@ #include #include #include +#include +#include +#include // use pthread rw_lock to lock db so you can safy clone .db file +class BlockedList +{ +private: + std::mutex m; + + sqlite3 *db; + sqlite3_stmt *sel_stmt; + sqlite3_stmt *ins_stmt; + sqlite3_stmt *upd_stmt; + +public: + void init() + { + int rc = sql::open(DB_NAME, &db); + if (rc) + { + fprintf(stderr, "Can't open database: %s\n", sql::errmsg(db)); + return; + } + + sql::prepare_v2(db, get_warnings, -1, &sel_stmt, NULL); + sql::prepare_v2(db, insert_warnings, -1, &ins_stmt, NULL); + sql::prepare_v2(db, set_warnings, -1, &upd_stmt, NULL); + } + + bool isBlocked(uint32_t id) + { + std::scoped_lock lock(m); + sql::bind_int64(sel_stmt, 1, id); + int64_t found = 0; + if (sql::step(sel_stmt) == SQLITE_ROW) + { + int type = sql::column_type(sel_stmt, 0); + if (type != SQL_NULL) + found = sql::column_int64(sel_stmt, 0); + } + sql::reset(sel_stmt); + + return found >= 3; + } + + void addWarning(uint32_t id) + { + std::scoped_lock lock(m); + sql::bind_int64(sel_stmt, 1, id); + int64_t found = 0; + int res = sql::step(sel_stmt); + if (res == SQLITE_ROW) + { + int type = sql::column_type(sel_stmt, 0); + if (type != SQL_NULL) + { + found = sql::column_int64(sel_stmt, 0); + } + sql::reset(sel_stmt); + } + else if (res == SQL_DONE) + { + sql::bind_int64(ins_stmt, 1, id); + sql::step(ins_stmt); + sql::reset(ins_stmt); + sql::reset(sel_stmt); + return; + } + + sql::bind_int64(upd_stmt, 1, found + 1); + sql::bind_int64(upd_stmt, 2, id); + sql::step(upd_stmt); + sql::reset(upd_stmt); + } +}; + +BlockedList blockedList; + // When a new client connected: void call(int sock, sockaddr_in newSocketInfo) { + if (blockedList.isBlocked(newSocketInfo.sin_addr.s_addr)) + { + TcpSocket::closet(sock); + return; + } std::string add = TcpSocket::remoteAddress(newSocketInfo); printf("new: %s\n", add.c_str()); int64_t conf, id; @@ -24,6 +106,7 @@ void call(int sock, sockaddr_in newSocketInfo) { printf("StartHeader ERROR\n"); TcpSocket::closet(sock); + blockedList.addWarning(newSocketInfo.sin_addr.s_addr); return; } @@ -31,6 +114,7 @@ void call(int sock, sockaddr_in newSocketInfo) { printf("ID ERROR\n"); TcpSocket::closet(sock); + blockedList.addWarning(newSocketInfo.sin_addr.s_addr); return; } @@ -51,8 +135,6 @@ void call(int sock, sockaddr_in newSocketInfo) int64_t gen = 0; while (sql::step(max_gen_stmt) != SQL_DONE) { - int num_cols = sql::column_count(max_gen_stmt); - int type = sql::column_type(max_gen_stmt, 0); if (type == SQL_NULL) @@ -133,6 +215,7 @@ int main() { sql::init(); sql::create_tables(); + blockedList.init(); std::thread t(checker); std::thread l(listen_key); // Bind the server to a port. diff --git a/shared/inc/sql.hpp b/shared/inc/sql.hpp index baab27a..ea7590d 100644 --- a/shared/inc/sql.hpp +++ b/shared/inc/sql.hpp @@ -4,6 +4,12 @@ const char create_like_table[] = "CREATE TABLE IF NOT EXISTS like_table ( ID INT const char create_user_table[] = "CREATE TABLE IF NOT EXISTS user_table ( ID INTEGER PRIMARY KEY, USER_ID INTEGER, GEN INTEGER, CHECKED INTEGER);"; +const char create_blocked_ip_table[] = "CREATE TABLE IF NOT EXISTS blocked_ip_table ( ID INTEGER PRIMARY KEY, WARNINGS INTEGER);"; + +const char get_warnings[] = "SELECT WARNINGS FROM blocked_ip_table WHERE ID = ?;"; +const char insert_warnings[] = "INSERT INTO blocked_ip_table (ID, WARNINGS) VALUES(?, 1);"; +const char set_warnings[] = "UPDATE blocked_ip_table SET WARNINGS = ? WHERE ID = ?;"; + const char max_gen[] = "SELECT MAX(GEN) FROM user_table WHERE USER_ID = ?;"; const char insert_user_data[] = "INSERT INTO user_table (USER_ID, GEN, CHECKED) VALUES(?, ?, 0);"; @@ -27,8 +33,10 @@ constexpr char DB_NAME[] = "data.db"; struct sqlite3; struct sqlite3_stmt; +#define SQLITE_ROW 100 /* sqlite3_step() has another row ready */ #define SQL_DONE 101 #define SQL_NULL 5 +#define SQLITE_OK 0 namespace sql { diff --git a/shared/src/sql.cpp b/shared/src/sql.cpp index 3b97c7a..c60d1d9 100644 --- a/shared/src/sql.cpp +++ b/shared/src/sql.cpp @@ -50,6 +50,17 @@ namespace sql fprintf(stdout, "user_table created successfully\n"); } + rc = sqlite3_exec(db, create_blocked_ip_table, nullptr, 0, &zErrMsg); + if (rc != SQLITE_OK) + { + fprintf(stderr, "SQL error: %s\n", zErrMsg); + sqlite3_free(zErrMsg); + } + else + { + fprintf(stdout, "blocked_ip_table created successfully\n"); + } + sqlite3_close(db); }