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/noncopyable.h"
  27 #include "stg/notifer.h"
  28 #include "user_impl.h"
  29
  30 #include <ctime>
  31 #include <list>
  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
  41 #define PROTOMAX    (5)
  42
  43 namespace STG
  44 {
  45
  46 class UsersImpl;
  47
  48 //-----------------------------------------------------------------------------
  49 struct Rule {
  50 uint32_t    ip;             // IP
  51 uint32_t    mask;           // Network mask
  52 uint16_t    port1;          // Min port
  53 uint16_t    port2;          // Max port
  54 uint8_t     proto;          // Protocol
  55 uint32_t    dir;            // Direction
  56 };
  57 //-----------------------------------------------------------------------------
  58 struct PacketExtraData {
  59     PacketExtraData()
  60         : flushTime(0),
  61           updateTime(0),
  62           userU(NULL),
  63           userUPresent(false),
  64           userD(NULL),
  65           userDPresent(false),
  66           dirU(DIR_NUM),
  67           dirD(DIR_NUM),
  68           lenU(0),
  69           lenD(0)
  70     {}
  71
  72     time_t      flushTime;          // Last flush time
  73     time_t      updateTime;         // Last update time
  74     UserImpl * userU;              // Uploader
  75     bool        userUPresent;       // Uploader is registered user
  76     UserImpl * userD;              // Downloader
  77     bool        userDPresent;       // Downloader is registered user
  78     int         dirU;               // Upload direction
  79     int         dirD;               // Download direction
  80     uint32_t    lenU;               // Upload length
  81     uint32_t    lenD;               // Download length
  82 };
  83 //-----------------------------------------------------------------------------
  84 class TraffCounterImpl;
  85 //-----------------------------------------------------------------------------
  86 class TRF_IP_BEFORE: public PropertyNotifierBase<uint32_t> {
  87 public:
  88                 TRF_IP_BEFORE(TraffCounterImpl & t, UserImpl * u)
  89                     : PropertyNotifierBase<uint32_t>(),
  90                       traffCnt(t),
  91                       user(u)
  92                 {}
  93                 TRF_IP_BEFORE(const TRF_IP_BEFORE & rvalue)
  94                     : PropertyNotifierBase<uint32_t>(),
  95                       traffCnt(rvalue.traffCnt),
  96                       user(rvalue.user)
  97                 {}
  98     void        notify(const uint32_t & oldValue, const uint32_t & newValue) override;
  99     void        SetUser(UserImpl * u) { user = u; }
 100     UserImpl * GetUser() const { return user; }
 101
 102 private:
 103     TRF_IP_BEFORE & operator=(const TRF_IP_BEFORE & rvalue);
 104
 105     TraffCounterImpl & traffCnt;
 106     UserImpl * user;
 107 };
 108 //-----------------------------------------------------------------------------
 109 class TRF_IP_AFTER: public PropertyNotifierBase<uint32_t> {
 110 public:
 111                 TRF_IP_AFTER(TraffCounterImpl & t, UserImpl * u)
 112                     : PropertyNotifierBase<uint32_t>(),
 113                       traffCnt(t),
 114                       user(u)
 115                 {}
 116                 TRF_IP_AFTER(const TRF_IP_AFTER & rvalue)
 117                     : PropertyNotifierBase<uint32_t>(),
 118                       traffCnt(rvalue.traffCnt),
 119                       user(rvalue.user)
 120                 {}
 121     void        notify(const uint32_t & oldValue, const uint32_t & newValue) override;
 122     void        SetUser(UserImpl * u) { user = u; }
 123     UserImpl * GetUser() const { return user; }
 124 private:
 125     TRF_IP_AFTER & operator=(const TRF_IP_AFTER & rvalue);
 126
 127     TraffCounterImpl & traffCnt;
 128     UserImpl * user;
 129 };
 130
 131 using UserImplPtr = UserImpl*;
 132 //-----------------------------------------------------------------------------
 133 class ADD_USER_NONIFIER: public NotifierBase<UserImplPtr> {
 134 public:
 135             explicit ADD_USER_NONIFIER(TraffCounterImpl & t) :
 136                 NotifierBase<UserImplPtr>(),
 137                 traffCnt(t)
 138             {}
 139     virtual ~ADD_USER_NONIFIER() {}
 140     void    notify(const UserImplPtr & user) override;
 141
 142 private:
 143     ADD_USER_NONIFIER(const ADD_USER_NONIFIER & rvalue);
 144     ADD_USER_NONIFIER & operator=(const ADD_USER_NONIFIER & rvalue);
 145
 146     TraffCounterImpl & traffCnt;
 147 };
 148 //-----------------------------------------------------------------------------
 149 class DEL_USER_NONIFIER: public NotifierBase<UserImplPtr> {
 150 public:
 151             explicit DEL_USER_NONIFIER(TraffCounterImpl & t) :
 152                 NotifierBase<UserImplPtr>(),
 153                 traffCnt(t)
 154             {}
 155     virtual ~DEL_USER_NONIFIER() {}
 156     void    notify(const UserImplPtr & user) override;
 157
 158 private:
 159     DEL_USER_NONIFIER(const DEL_USER_NONIFIER & rvalue);
 160     DEL_USER_NONIFIER & operator=(const DEL_USER_NONIFIER & rvalue);
 161
 162     TraffCounterImpl & traffCnt;
 163 };
 164 //-----------------------------------------------------------------------------
 165 class TraffCounterImpl : public TraffCounter {
 166     friend class ADD_USER_NONIFIER;
 167     friend class DEL_USER_NONIFIER;
 168     friend class TRF_IP_BEFORE;
 169     friend class TRF_IP_AFTER;
 170     public:
 171         TraffCounterImpl(UsersImpl * users, const std::string & rulesFileName);
 172         ~TraffCounterImpl();
 173
 174         int         Reload();
 175         int         Start();
 176         int         Stop();
 177
 178         void        process(const RawPacket & rawPacket) override;
 179         void        SetMonitorDir(const std::string & monitorDir);
 180
 181         size_t      rulesCount() const override { return rules.size(); }
 182
 183     private:
 184         bool        ParseAddress(const char * ta, Rule * rule) const;
 185         uint32_t    CalcMask(uint32_t msk) const;
 186         void        FreeRules();
 187         bool        ReadRules(bool test = false);
 188
 189         void        Run(std::stop_token token);
 190
 191         void        DeterminateDir(const RawPacket & packet,
 192                                    int * dirU, // Direction for upload
 193                                    int * dirD) const; // Direction for download
 194
 195         void        FlushAndRemove();
 196
 197         void        AddUser(UserImpl * user);
 198         void        DelUser(uint32_t uip);
 199         void        SetUserNotifiers(UserImpl * user);
 200         void        UnSetUserNotifiers(UserImpl * user);
 201
 202         typedef std::list<Rule>::iterator rule_iter;
 203
 204         std::list<Rule>          rules;
 205
 206         typedef std::map<RawPacket, PacketExtraData> Packets;
 207         typedef Packets::iterator pp_iter;
 208         typedef std::multimap<uint32_t, pp_iter> Index;
 209         typedef Index::iterator ip2p_iter;
 210         typedef Index::const_iterator ip2p_citer;
 211
 212         Packets packets; // Packets tree
 213
 214         Index ip2packets; // IP-to-Packet index
 215
 216         std::string              dirName[DIR_NUM + 1];
 217
 218         Logger &             WriteServLog;
 219         std::string              rulesFileName;
 220
 221         std::string              monitorDir;
 222         bool                     monitoring;
 223         time_t                   touchTimeP;
 224
 225         UsersImpl *             users;
 226
 227         bool                     stopped;
 228         std::mutex               m_mutex;
 229         std::jthread             m_thread;
 230
 231         std::list<TRF_IP_BEFORE> ipBeforeNotifiers;
 232         std::list<TRF_IP_AFTER>  ipAfterNotifiers;
 233
 234         ADD_USER_NONIFIER        addUserNotifier;
 235         DEL_USER_NONIFIER        delUserNotifier;
 236 };
 237
 238 }
 239 //-----------------------------------------------------------------------------