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.