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