add block id

server/src/server.cpp

@@ -6,12 +6,94 @@
 #include <iostream>
 #include <map>
 #include <thread>
+#include <unordered_map>
+#include <mutex>
+#include <thread>
 
 // 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.

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
 {

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);
   }