projects/stargazer/traffcounter_impl.h

   1 /*
   2  *    This program is free software; you can redistribute it and/or modify
   3  *    it under the terms of the GNU General Public License as published by
   4  *    the Free Software Foundation; either version 2 of the License, or
   5  *    (at your option) any later version.
   6  *
   7  *    This program is distributed in the hope that it will be useful,
   8  *    but WITHOUT ANY WARRANTY; without even the implied warranty of
   9  *    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
  10  *    GNU General Public License for more details.
  11  *
  12  *    You should have received a copy of the GNU General Public License
  13  *    along with this program; if not, write to the Free Software
  14  *    Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA
  15  */
  16
  17 /*
  18  *    Author : Boris Mikhailenko <stg34@stargazer.dp.ua>
  19  */
  20
  21 #pragma once
  22
  23 #include "stg/traffcounter.h"
  24 #include "stg/logger.h"
  25 #include "stg/raw_ip_packet.h"
  26 #include "stg/subscriptions.h"
  27 #include "user_impl.h"
  28
  29 #include <list>
  30 #include <vector>
  31 #include <tuple>
  32 #include <map>
  33 #include <string>
  34 #include <mutex>
  35 #pragma GCC diagnostic push
  36 #pragma GCC diagnostic ignored "-Wshadow"
  37 #include <jthread.hpp>
  38 #pragma GCC diagnostic pop
  39 #include <cstdint>
  40 #include <ctime>
  41
  42 #define PROTOMAX    (5)
  43
  44 namespace STG
  45 {
  46
  47 class UsersImpl;
  48
  49 //-----------------------------------------------------------------------------
  50 struct Rule {
  51 uint32_t    ip;             // IP
  52 uint32_t    mask;           // Network mask
  53 uint16_t    port1;          // Min port
  54 uint16_t    port2;          // Max port
  55 uint8_t     proto;          // Protocol
  56 uint32_t    dir;            // Direction
  57 };
  58 //-----------------------------------------------------------------------------
  59 struct PacketExtraData {
  60     PacketExtraData()
  61         : flushTime(0),
  62           updateTime(0),
  63           userU(NULL),
  64           userUPresent(false),
  65           userD(NULL),
  66           userDPresent(false),
  67           dirU(DIR_NUM),
  68           dirD(DIR_NUM),
  69           lenU(0),
  70           lenD(0)
  71     {}
  72
  73     time_t      flushTime;          // Last flush time
  74     time_t      updateTime;         // Last update time
  75     UserImpl * userU;              // Uploader
  76     bool        userUPresent;       // Uploader is registered user
  77     UserImpl * userD;              // Downloader
  78     bool        userDPresent;       // Downloader is registered user
  79     int         dirU;               // Upload direction
  80     int         dirD;               // Download direction
  81     uint32_t    lenU;               // Upload length
  82     uint32_t    lenD;               // Download length
  83 };
  84 //-----------------------------------------------------------------------------
  85 class TraffCounterImpl : public TraffCounter {
  86     public:
  87         TraffCounterImpl(UsersImpl * users, const std::string & rulesFileName);
  88         ~TraffCounterImpl();
  89
  90         int         Reload();
  91         int         Start();
  92         int         Stop();
  93
  94         void        process(const RawPacket & rawPacket) override;
  95         void        SetMonitorDir(const std::string & monitorDir);
  96
  97         size_t      rulesCount() const override { return rules.size(); }
  98
  99     private:
 100         bool        ParseAddress(const char * ta, Rule * rule) const;
 101         uint32_t    CalcMask(uint32_t msk) const;
 102         void        FreeRules();
 103         bool        ReadRules(bool test = false);
 104
 105         void        Run(std::stop_token token);
 106
 107         void        DeterminateDir(const RawPacket & packet,
 108                                    int * dirU, // Direction for upload
 109                                    int * dirD) const; // Direction for download
 110
 111         void        FlushAndRemove();
 112
 113         void        AddUser(UserImpl * user);
 114         void        DelUser(uint32_t uip);
 115         void        SetUserNotifiers(UserImpl* user);
 116         void        UnSetUserNotifiers(UserImpl* user);
 117
 118         typedef std::list<Rule>::iterator rule_iter;
 119
 120         std::list<Rule>          rules;
 121
 122         typedef std::map<RawPacket, PacketExtraData> Packets;
 123         typedef Packets::iterator pp_iter;
 124         typedef std::multimap<uint32_t, pp_iter> Index;
 125         typedef Index::iterator ip2p_iter;
 126         typedef Index::const_iterator ip2p_citer;
 127
 128         Packets packets; // Packets tree
 129
 130         Index ip2packets; // IP-to-Packet index
 131
 132         std::string              dirName[DIR_NUM + 1];
 133
 134         Logger &             WriteServLog;
 135         std::string              rulesFileName;
 136
 137         std::string              monitorDir;
 138         bool                     monitoring;
 139         time_t                   touchTimeP;
 140
 141         UsersImpl *             users;
 142
 143         bool                     stopped;
 144         std::mutex               m_mutex;
 145         std::jthread             m_thread;
 146
 147         ScopedConnection m_onAddUserConn;
 148         ScopedConnection m_onDelUserConn;
 149
 150         using OnIPConns = std::tuple<int, ScopedConnection, ScopedConnection>;
 151         std::vector<OnIPConns> m_onIPConns;
 152         void beforeIPChange(uint32_t oldVal);
 153         void afterIPChange(UserImpl* user, uint32_t newVal);
 154 };
 155
 156 }
 157 //-----------------------------------------------------------------------------