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