]> git.stg.codes - stg.git/blob - projects/traffcounter/traffcounter.cpp
Додано патч від Alexey Osipov <lion-simba@pridelands.ru>
[stg.git] / projects / traffcounter / traffcounter.cpp
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 : Maxim Mamontov <faust@stargazer.dp.ua>
19  */
20
21 /*
22  $Revision: 1.5 $
23  $Date: 2009/10/12 08:43:32 $
24  $Author: faust $
25  */
26
27 #include <csignal>
28 #include <cassert>
29 #include <algorithm>
30
31 #include "traffcounter.h"
32 #include "logger.h"
33 #include "lock.h"
34 #include "utils.h"
35
36 //-----------------------------------------------------------------------------
37 STG::TRAFFCOUNTER::TRAFFCOUNTER()
38     : rulesFinder(),
39       pendingPackets(),
40       sessions(),
41       cacheHits(0),
42       cacheMisses(0),
43       pendingCount(0),
44       maxPending(0),
45       ip2sessions(),
46       stopped(true),
47       running(false)
48 {
49 LOG_IT << "TRAFFCOUNTER::TRAFFCOUNTER()\n";
50 pthread_mutex_init(&sessionMutex, NULL);
51 pthread_mutex_init(&pendingMutex, NULL);
52 pthread_mutex_init(&ipMutex, NULL);
53 pthread_mutex_init(&rulesMutex, NULL);
54 pthread_cond_init(&pendingCond, NULL);
55 }
56 //-----------------------------------------------------------------------------
57 STG::TRAFFCOUNTER::~TRAFFCOUNTER()
58 {
59 LOG_IT << "TRAFFCOUNTER::~TRAFFCOUNTER()\n";
60 pthread_cond_destroy(&pendingCond);
61 pthread_mutex_destroy(&rulesMutex);
62 pthread_mutex_destroy(&ipMutex);
63 pthread_mutex_destroy(&pendingMutex);
64 pthread_mutex_destroy(&sessionMutex);
65 }
66 //-----------------------------------------------------------------------------
67 //  Starting processing thread
68 bool STG::TRAFFCOUNTER::Start()
69 {
70 LOG_IT << "TRAFFCOUNTER::Start()\n";
71
72 if (running)
73     return false;
74
75 running = true;
76 stopped = true;
77
78 if (pthread_create(&thread, NULL, Run, this))
79     {
80     LOG_IT << "TRAFFCOUNTER::Start() Error: Cannot start thread!\n";
81     return true;
82     }
83
84 return false;
85 }
86 //-----------------------------------------------------------------------------
87 bool STG::TRAFFCOUNTER::Stop()
88 {
89 LOG_IT << "TRAFFCOUNTER::Stop()\n";
90 LOG_IT << "maxPending: " << maxPending << std::endl;
91
92 if (!running)
93     return false;
94
95 running = false;
96 // Awake thread
97 pthread_cond_signal(&pendingCond);
98
99 //5 seconds to thread stops itself
100 for (int i = 0; i < 25 && !stopped; ++i)
101     {
102     usleep(200000);
103     }
104
105 //after 5 seconds waiting thread still running. now kill it
106 if (!stopped)
107     {
108     LOG_IT << "TRAFFCOUNTER::Stop() Killing thread\n";
109     if (pthread_kill(thread, SIGINT))
110         {
111         return true;
112         }
113     LOG_IT << "TRAFFCOUNTER::Stop() Thread killed\n";
114     }
115
116 return false;
117 }
118 //-----------------------------------------------------------------------------
119 double STG::TRAFFCOUNTER::StreamQuality() const
120 {
121 if (!cacheHits && !cacheMisses)
122     {
123     return 0;
124     }
125
126 double quality = cacheHits;
127 return quality / (quality + cacheMisses);
128 }
129 //-----------------------------------------------------------------------------
130 void STG::TRAFFCOUNTER::AddPacket(const iphdr & ipHdr, uint16_t sport, uint16_t dport)
131 {
132 /*
133  *  Intersects with AddIP (from user thread), DeleteIP (from user thread) and
134  * Process (internal thread). AddPacket is calling from capturer's thread
135  *
136  *  ips is affected by AddIP (logarithmic lock time) and
137  *  DeleteIP (from user thread)
138  *
139  *  May be locked by AddIP or DeleteIP (from user thread)
140  *
141  *  Lock AddIP (user thread) or DeleteIP (user thread)
142  *  Logarithmic lock time
143  */
144
145 bool srcExists;
146 bool dstExists;
147
148     {
149     SCOPED_LOCK lock(ipMutex);
150     srcExists = std::binary_search(ips.begin(), ips.end(), ipHdr.saddr);
151     dstExists = std::binary_search(ips.begin(), ips.end(), ipHdr.daddr);
152     }
153
154 if (!srcExists &&
155     !dstExists)
156     {
157     // Just drop the packet
158     return;
159     }
160
161 STG::PENDING_PACKET p(ipHdr, sport, dport);
162
163 // Packet classification
164 if (srcExists)
165     {
166     if (dstExists)
167         {
168         // Both src and dst are countable
169         p.direction = PENDING_PACKET::LOCAL;
170         }
171     else
172         {
173         // Src is countable
174         p.direction = PENDING_PACKET::OUTGOING;
175         }
176     }
177 else
178     {
179     if (dstExists)
180         {
181         // Dst is countable
182         p.direction = PENDING_PACKET::INCOMING;
183         }
184     else
185         {
186         assert(0);
187         // Not src nor dst are countable
188         p.direction = PENDING_PACKET::FOREIGN;
189         }
190     }
191
192 /*
193  *  pendingPackets is affected by Process (from internal thread)
194  *
195  *  May be locked by Process (internal thread)
196  *
197  *  Lock Process (internal thread)
198  *  Constant lock time
199  */
200 SCOPED_LOCK lock(pendingMutex);
201 pendingPackets.push_back(p);
202 pendingCount++;
203 #ifdef STATISTIC
204 if (pendingCount > maxPending)
205     maxPending = pendingCount;
206 #endif
207 pthread_cond_signal(&pendingCond);
208
209 }
210 //-----------------------------------------------------------------------------
211 void STG::TRAFFCOUNTER::AddIP(uint32_t ip)
212 {
213 /*
214  *  AddIP is calling from users and affect DeleteIP and AddPacket.
215  * DeleteIP cannot be called concurrently with AddIP - it's the same
216  * thread. AddPacket is calling from capturer's thread - concurrently
217  * with AddIP.
218  *
219  * May be locked by AddPacket (from capturer's thread)
220  * Logarithmic lock time
221  *
222  * Lock AddPacket (capturer's thread)
223  * Logarithmic lock time
224  */
225 SCOPED_LOCK lock(ipMutex);
226 IP_ITER it(std::lower_bound(ips.begin(), ips.end(), ip));
227
228 if (it != ips.end() && *it == ip)
229     {
230     return;
231     }
232 // Insertion
233 ips.insert(it, ip);
234 }
235 //-----------------------------------------------------------------------------
236 void STG::TRAFFCOUNTER::DeleteIP(uint32_t ip, STG::TRAFF_DATA * traff)
237 {
238 /*
239  *  DeleteIP is calling from users and affect AddIP, AddPacket, GetIP and
240  * Process. AddIP and GetIP cannot be called concurrently with DeleteIP - it's
241  * the same thread. AddPacket is calling from capturer's thread - concurrently
242  * with DeleteIP. Process is calling from internal thread - concurrently with
243  * DeleteIP.
244  *
245  * May be locked by AddPacket (from capturer's thread)
246  * Logarithmic lock time
247  *
248  * Lock AddPacket (capturer's thread)
249  * Logarithmic lock time
250  */
251
252     {
253     SCOPED_LOCK lock(ipMutex);
254
255     IP_ITER it(std::lower_bound(ips.begin(), ips.end(), ip));
256     if (it == ips.end())
257         {
258         return;
259         }
260     if (*it != ip)
261         {
262         return;
263         }
264
265     ips.erase(it);
266     }
267
268 // Get sessions for this ip
269 std::pair<INDEX_ITER,
270           INDEX_ITER> range;
271
272 SCOPED_LOCK lock(sessionMutex);
273 range = ip2sessions.equal_range(ip);
274 std::list<INDEX_ITER> toDelete;
275
276 // Lock session growing
277 for (INDEX_ITER it = range.first; it != range.second; ++it)
278     {
279     traff->push_back(STG::TRAFF_ITEM(it->second->first, it->second->second));
280
281     // Include self
282     toDelete.push_back(it);
283
284     /*if (ip == it->second->first.saddr)
285         {
286         toDelete.push_back(it->second->second.dIdx);
287         }
288     else
289         {
290         toDelete.push_back(it->second->second.sIdx);
291         }*/
292
293     --it->second->second.refCount;
294
295     // Remove session
296     /*
297      *  Normally we will lock here only in case of session between
298      *  two users from ips list
299      */
300     if (!it->second->second.refCount)
301         {
302         sessions.erase(it->second);
303         }
304     }
305
306 // Remove indexes
307 for (std::list<INDEX_ITER>::iterator it = toDelete.begin();
308      it != toDelete.end();
309      ++it)
310     {
311     ip2sessions.erase(*it);
312     }
313 }
314 //-----------------------------------------------------------------------------
315 void STG::TRAFFCOUNTER::GetIP(uint32_t ip, STG::TRAFF_DATA * traff)
316 {
317 /*
318  *  Normally we will lock here only in case of session between
319  *  two users from ips list
320  */
321 std::pair<INDEX_ITER,
322           INDEX_ITER> range;
323
324 SCOPED_LOCK lock(sessionMutex);
325 range = ip2sessions.equal_range(ip);
326 std::list<INDEX_ITER> toDelete;
327
328 // TODO: replace with foreach
329 for (SESSION_INDEX::iterator it = range.first;
330      it != range.second;
331      ++it)
332     {
333     traff->push_back(STG::TRAFF_ITEM(it->second->first, it->second->second));
334     toDelete.push_back(it);
335     --it->second->second.refCount;
336     if (!it->second->second.refCount)
337         {
338         sessions.erase(it->second);
339         }
340     }
341
342 for (std::list<INDEX_ITER>::iterator it = toDelete.begin();
343      it != toDelete.end();
344      ++it)
345     {
346     ip2sessions.erase(*it);
347     }
348 }
349 //-----------------------------------------------------------------------------
350 void * STG::TRAFFCOUNTER::Run(void * data)
351 {
352 STG::TRAFFCOUNTER * tc = static_cast<STG::TRAFFCOUNTER *>(data);
353 tc->stopped = false;
354
355 while (tc->running)
356     {
357     STG::PENDING_PACKET packet;
358         {
359         SCOPED_LOCK lock(tc->pendingMutex);
360         if (tc->pendingPackets.empty())
361             {
362             pthread_cond_wait(&tc->pendingCond, &tc->pendingMutex);
363             }
364         if (!tc->running)
365             {
366             break;
367             }
368         packet = *tc->pendingPackets.begin();
369         tc->pendingPackets.pop_front();
370         --tc->pendingCount;
371         }
372     tc->Process(packet);
373     }
374
375 tc->stopped = true;
376 return NULL;
377 }
378 //-----------------------------------------------------------------------------
379 void STG::TRAFFCOUNTER::Process(const STG::PENDING_PACKET & p)
380 {
381 // Bypass on stop
382 if (!running)
383     return;
384
385 // Fail on foreign packets
386 if (p.direction == PENDING_PACKET::FOREIGN) {
387     assert(0);
388 }
389
390 // Searching a new packet in a tree.
391 SESSION_ITER si;
392     {
393     SCOPED_LOCK lock(sessionMutex);
394     si = sessions.find(STG::SESSION_ID(p));
395     }
396
397 // Packet found - update length and time
398 if (si != sessions.end())
399     {
400     // Grow session
401     SCOPED_LOCK lock(sessionMutex);
402     si->second.length += p.length;
403     ++cacheHits;
404     return;
405     }
406
407 ++cacheMisses;
408
409 // Packet not found - add new packet
410
411 // This packet is alowed to create session
412 STG::SESSION_ID sid(p);
413 SESSION_FULL_DATA sd;
414
415 // Identify a packet
416     {
417     SCOPED_LOCK lock(rulesMutex);
418     sd.dir = rulesFinder.GetDir(p);
419     }
420
421 sd.length = p.length;
422
423 if (p.direction == PENDING_PACKET::LOCAL)
424     {
425     sd.refCount = 2;
426     }
427 else
428     {
429     sd.refCount = 1;
430     }
431
432 // Create a session
433 std::pair<SESSION_ITER,
434           bool> sIt(sessions.insert(std::make_pair(sid, sd)));
435     {
436     SCOPED_LOCK lock(sessionMutex);
437     std::pair<SESSION_ITER,
438               bool> sIt(sessions.insert(std::make_pair(sid, sd)));
439
440     // Create an indexes
441     sIt.first->second.sIdx = ip2sessions.insert(std::make_pair(p.saddr, sIt.first));
442     sIt.first->second.dIdx = ip2sessions.insert(std::make_pair(p.daddr, sIt.first));
443     }
444
445 }
446 //-----------------------------------------------------------------------------
447 void STG::TRAFFCOUNTER::SetRules(const STG::RULES & data)
448 {
449 /*
450  *  SetRules is calling from outside internel thread. Process is calling
451  * from internal thread and calls DeterminateDir which use rules data.
452  *
453  * May be locked by DeterminateDir (Process) from internal thread.
454  *
455  * Lock DeterminateDir (Process) - internal thread.
456  * Linear lock time
457  */
458 SCOPED_LOCK lock(rulesMutex);
459 rulesFinder.SetRules(data);
460 }