]> git.stg.codes - stg.git/blob - projects/stargazer/traffcounter_impl.h
ed639f5c3ddcffb2fc1db1d669678a41a673add6
[stg.git] / 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);
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);
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);
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);
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 }