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.
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.
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
22 * Author : Boris Mikhailenko <stg34@stargazer.dp.ua>
27 $Date: 2010/11/03 11:28:07 $
32 #include <sys/types.h>
33 #include <sys/socket.h>
34 #include <netinet/in.h>
35 #include <arpa/inet.h>
39 #include <cstdio> // fopen and similar
40 #include <cstdlib> // strtol
42 #include "stg/common.h"
43 #include "stg/locker.h"
44 #include "stg/const.h" // MONITOR_TIME_DELAY_SEC
45 #include "traffcounter_impl.h"
46 #include "stg_timer.h"
47 #include "users_impl.h"
48 #include "async_pool.h"
50 #define FLUSH_TIME (10)
51 #define REMOVE_TIME (31)
53 using STG::TraffCounterImpl;
55 namespace AsyncPoolST = STG::AsyncPoolST;
57 const char protoName[PROTOMAX][8] =
58 {"TCP", "UDP", "ICMP", "TCP_UDP", "ALL"};
62 tcp = 0, udp, icmp, tcp_udp, all
65 //-----------------------------------------------------------------------------
66 TraffCounterImpl::TraffCounterImpl(UsersImpl * u, const std::string & fn)
67 : WriteServLog(Logger::get()),
70 touchTimeP(stgTime - MONITOR_TIME_DELAY_SEC),
74 for (int i = 0; i < DIR_NUM; i++)
75 strprintf(&dirName[i], "DIR%d", i);
77 dirName[DIR_NUM] = "NULL";
79 m_onAddUserConn = users->onImplAdd([this](auto user){
80 AsyncPoolST::enqueue([this, user](){ SetUserNotifiers(user); });
82 m_onDelUserConn = users->onImplDel([this](auto user){
83 AsyncPoolST::enqueue([this, user](){ UnSetUserNotifiers(user); DelUser(user->GetCurrIP()); });
86 //-----------------------------------------------------------------------------
87 TraffCounterImpl::~TraffCounterImpl()
90 //-----------------------------------------------------------------------------
91 int TraffCounterImpl::Start()
93 std::lock_guard<std::mutex> lock(m_mutex);
100 printfd(__FILE__, "TraffCounterImpl::Start() - Cannot read rules\n");
101 WriteServLog("TraffCounter: Cannot read rules.");
105 printfd(__FILE__, "TraffCounter::Start()\n");
106 int h = users->OpenSearch();
107 assert(h && "USERS::OpenSearch is always correct");
110 while (users->SearchNext(h, &u) == 0)
112 users->CloseSearch(h);
114 m_thread = std::jthread([this](auto token){ Run(std::move(token)); });
117 //-----------------------------------------------------------------------------
118 int TraffCounterImpl::Stop()
123 m_thread.request_stop();
125 int h = users->OpenSearch();
126 assert(h && "USERS::OpenSearch is always correct");
130 //5 seconds to thread stops itself
131 for (int i = 0; i < 25 && !stopped; i++)
132 std::this_thread::sleep_for(std::chrono::milliseconds(200));
142 printfd(__FILE__, "TraffCounter::Stop()\n");
146 //-----------------------------------------------------------------------------
147 void TraffCounterImpl::Run(std::stop_token token)
150 sigfillset(&signalSet);
151 pthread_sigmask(SIG_BLOCK, &signalSet, NULL);
156 time_t touchTime = stgTime - MONITOR_TIME_DELAY_SEC;
157 while (!token.stop_requested())
159 std::this_thread::sleep_for(std::chrono::milliseconds(500));
160 if (token.stop_requested())
166 if (monitoring && (touchTime + MONITOR_TIME_DELAY_SEC <= stgTime))
168 std::string monFile(monitorDir + "/traffcounter_r");
169 printfd(__FILE__, "Monitor=%d file TraffCounter %s\n", monitoring, monFile.c_str());
174 if (++c % FLUSH_TIME == 0)
180 //-----------------------------------------------------------------------------
181 void TraffCounterImpl::process(const RawPacket & rawPacket)
183 if (monitoring && (touchTimeP + MONITOR_TIME_DELAY_SEC <= stgTime))
185 std::string monFile = monitorDir + "/traffcounter_p";
186 printfd(__FILE__, "Monitor=%d file TraffCounter %s\n", monitoring, monFile.c_str());
187 touchTimeP = stgTime;
191 std::lock_guard<std::mutex> lock(m_mutex);
193 //printfd(__FILE__, "TraffCounter::Process()\n");
194 //TODO replace find with lower_bound.
196 // Searching a new packet in a tree.
197 pp_iter pi = packets.find(rawPacket);
199 // Packet found - update length and time
200 if (pi != packets.end())
202 pi->second.lenU += rawPacket.GetLen();
203 pi->second.lenD += rawPacket.GetLen();
204 pi->second.updateTime = stgTime;
205 /*printfd(__FILE__, "=============================\n");
206 printfd(__FILE__, "Packet found!\n");
207 printfd(__FILE__, "Version=%d\n", rawPacket.GetIPVersion());
208 printfd(__FILE__, "HeaderLen=%d\n", rawPacket.GetHeaderLen());
209 printfd(__FILE__, "PacketLen=%d\n", rawPacket.GetLen());
210 printfd(__FILE__, "SIP=%s\n", inet_ntostring(rawPacket.GetSrcIP()).c_str());
211 printfd(__FILE__, "DIP=%s\n", inet_ntostring(rawPacket.GetDstIP()).c_str());
212 printfd(__FILE__, "src port=%d\n", rawPacket.GetSrcPort());
213 printfd(__FILE__, "pst port=%d\n", rawPacket.GetDstPort());
214 printfd(__FILE__, "len=%d\n", rawPacket.GetLen());
215 printfd(__FILE__, "proto=%d\n", rawPacket.GetProto());
216 printfd(__FILE__, "PacketDirU=%d\n", pi->second.dirU);
217 printfd(__FILE__, "PacketDirD=%d\n", pi->second.dirD);
218 printfd(__FILE__, "=============================\n");*/
224 // Packet not found - add new packet
226 ed.updateTime = stgTime;
227 ed.flushTime = stgTime;
230 userU is that whose user_ip == packet_ip_src
231 userD is that whose user_ip == packet_ip_dst
234 uint32_t ipU = rawPacket.GetSrcIP();
235 uint32_t ipD = rawPacket.GetDstIP();
237 // Searching users with such IP
238 if (users->FindByIPIdx(ipU, &ed.userU) == 0)
240 ed.userUPresent = true;
243 if (users->FindByIPIdx(ipD, &ed.userD) == 0)
245 ed.userDPresent = true;
248 if (ed.userUPresent ||
251 DeterminateDir(rawPacket, &ed.dirU, &ed.dirD);
253 ed.lenD = ed.lenU = rawPacket.GetLen();
255 //TODO use result of lower_bound to inserting new record
257 // Adding packet to a tree.
258 std::pair<pp_iter, bool> insertResult = packets.insert(std::make_pair(rawPacket, ed));
259 pp_iter newPacket = insertResult.first;
261 // Adding packet reference to an IP index.
262 ip2packets.insert(std::make_pair(ipU, newPacket));
263 ip2packets.insert(std::make_pair(ipD, newPacket));
266 //-----------------------------------------------------------------------------
267 void TraffCounterImpl::FlushAndRemove()
269 std::lock_guard<std::mutex> lock(m_mutex);
271 Packets::size_type oldPacketsSize = packets.size();
272 Index::size_type oldIp2packetsSize = ip2packets.size();
275 pi = packets.begin();
277 ip2packets.erase(ip2packets.begin(), ip2packets.end());
278 while (pi != packets.end())
281 if (stgTime - pi->second.flushTime > FLUSH_TIME)
283 if (pi->second.userUPresent)
285 //printfd(__FILE__, "+++ Flushing U user %s (%s:%d) of length %d\n", pi->second.userU->GetLogin().c_str(), inet_ntostring(pi->first.GetSrcIP()).c_str(), pi->first.GetSrcPort(), pi->second.lenU);
288 if (pi->second.dirU < DIR_NUM)
290 #ifdef TRAFF_STAT_WITH_PORTS
291 pi->second.userU->AddTraffStatU(pi->second.dirU,
292 pi->first.GetDstIP(),
293 pi->first.GetDstPort(),
296 pi->second.userU->AddTraffStatU(pi->second.dirU,
297 pi->first.GetDstIP(),
303 pi->second.flushTime = stgTime;
306 if (pi->second.userDPresent)
308 //printfd(__FILE__, "+++ Flushing D user %s (%s:%d) of length %d\n", pi->second.userD->GetLogin().c_str(), inet_ntostring(pi->first.GetDstIP()).c_str(), pi->first.GetDstPort(), pi->second.lenD);
311 if (pi->second.dirD < DIR_NUM)
313 #ifdef TRAFF_STAT_WITH_PORTS
314 pi->second.userD->AddTraffStatD(pi->second.dirD,
315 pi->first.GetSrcIP(),
316 pi->first.GetSrcPort(),
319 pi->second.userD->AddTraffStatD(pi->second.dirD,
320 pi->first.GetSrcIP(),
326 pi->second.flushTime = stgTime;
330 if (stgTime - pi->second.updateTime < REMOVE_TIME)
332 std::pair<pp_iter, bool> res = newPackets.insert(*pi);
335 ip2packets.insert(std::make_pair(pi->first.GetSrcIP(), res.first));
336 ip2packets.insert(std::make_pair(pi->first.GetDstIP(), res.first));
341 swap(packets, newPackets);
342 printfd(__FILE__, "FlushAndRemove() packets: %d(rem %d) ip2packets: %d(rem %d)\n",
344 oldPacketsSize - packets.size(),
346 oldIp2packetsSize - ip2packets.size());
349 //-----------------------------------------------------------------------------
350 void TraffCounterImpl::AddUser(UserImpl * user)
352 printfd(__FILE__, "AddUser: %s\n", user->GetLogin().c_str());
353 uint32_t uip = user->GetCurrIP();
354 std::pair<ip2p_iter, ip2p_iter> pi;
356 std::lock_guard<std::mutex> lock(m_mutex);
357 // Find all packets with IP belongs to this user
358 pi = ip2packets.equal_range(uip);
360 while (pi.first != pi.second)
362 if (pi.first->second->first.GetSrcIP() == uip)
364 assert((!pi.first->second->second.userUPresent ||
365 pi.first->second->second.userU == user) &&
366 "U user present but it's not current user");
368 pi.first->second->second.lenU = 0;
369 pi.first->second->second.userU = user;
370 pi.first->second->second.userUPresent = true;
373 if (pi.first->second->first.GetDstIP() == uip)
375 assert((!pi.first->second->second.userDPresent ||
376 pi.first->second->second.userD == user) &&
377 "D user present but it's not current user");
379 pi.first->second->second.lenD = 0;
380 pi.first->second->second.userD = user;
381 pi.first->second->second.userDPresent = true;
387 //-----------------------------------------------------------------------------
388 void TraffCounterImpl::DelUser(uint32_t uip)
390 printfd(__FILE__, "DelUser: %s \n", inet_ntostring(uip).c_str());
391 std::pair<ip2p_iter, ip2p_iter> pi;
393 std::lock_guard<std::mutex> lock(m_mutex);
394 pi = ip2packets.equal_range(uip);
396 while (pi.first != pi.second)
398 if (pi.first->second->first.GetSrcIP() == uip)
400 if (pi.first->second->second.dirU < DIR_NUM && pi.first->second->second.userUPresent)
402 #ifdef TRAFF_STAT_WITH_PORTS
403 pi.first->second->second.userU->AddTraffStatU(pi.first->second->second.dirU,
404 pi.first->second->first.GetDstIP(),
405 pi.first->second->first.GetDstPort(),
406 pi.first->second->second.lenU);
408 pi.first->second->second.userU->AddTraffStatU(pi.first->second->second.dirU,
409 pi.first->second->first.GetDstIP(),
410 pi.first->second->second.lenU);
413 pi.first->second->second.userUPresent = false;
416 if (pi.first->second->first.GetDstIP() == uip)
418 if (pi.first->second->second.dirD < DIR_NUM && pi.first->second->second.userDPresent)
420 #ifdef TRAFF_STAT_WITH_PORTS
421 pi.first->second->second.userD->AddTraffStatD(pi.first->second->second.dirD,
422 pi.first->second->first.GetSrcIP(),
423 pi.first->second->first.GetSrcPort(),
424 pi.first->second->second.lenD);
426 pi.first->second->second.userD->AddTraffStatD(pi.first->second->second.dirD,
427 pi.first->second->first.GetSrcIP(),
428 pi.first->second->second.lenD);
432 pi.first->second->second.userDPresent = false;
438 ip2packets.erase(pi.first, pi.second);
440 //-----------------------------------------------------------------------------
441 void TraffCounterImpl::SetUserNotifiers(UserImpl* user)
443 // Adding user. Adding notifiers to user.
444 m_onIPConns.emplace_back(
446 user->beforeCurrIPChange([this](auto oldVal, auto /*newVal*/){ beforeIPChange(oldVal); }),
447 user->afterCurrIPChange([this, user](auto /*oldVal*/, auto newVal){ afterIPChange(user, newVal); })
450 //-----------------------------------------------------------------------------
451 void TraffCounterImpl::UnSetUserNotifiers(UserImpl * user)
453 // Removing user. Removing notifiers from user.
454 m_onIPConns.erase(std::remove_if(m_onIPConns.begin(), m_onIPConns.end(),
455 [user](const auto& cs){ return std::get<0>(cs) == user->GetID(); }),
458 //-----------------------------------------------------------------------------
459 void TraffCounterImpl::DeterminateDir(const RawPacket & packet,
460 int * dirU, // Direction for incoming packet
461 int * dirD) const // Direction for outgoing packet
463 bool addrMatchU = false;
464 bool portMatchU = false;
465 bool addrMatchD = false;
466 bool portMatchD = false;
467 bool foundU = false; // Was rule for U found ?
468 bool foundD = false; // Was rule for D found ?
469 //printfd(__FILE__, "foundU=%d, foundD=%d\n", foundU, foundD);
471 enum { ICMP_RPOTO = 1, TCP_PROTO = 6, UDP_PROTO = 17 };
473 std::list<Rule>::const_iterator ln;
476 while (ln != rules.end())
489 portMatchU = (packet.GetProto() == ICMP_RPOTO);
493 if (packet.GetProto() == TCP_PROTO || packet.GetProto() == UDP_PROTO)
494 portMatchU = (packet.GetDstPort() >= ln->port1) && (packet.GetDstPort() <= ln->port2);
498 if (packet.GetProto() == TCP_PROTO)
499 portMatchU = (packet.GetDstPort() >= ln->port1) && (packet.GetDstPort() <= ln->port2);
503 if (packet.GetProto() == UDP_PROTO)
504 portMatchU = (packet.GetDstPort() >= ln->port1) && (packet.GetDstPort() <= ln->port2);
508 printfd(__FILE__, "Error! Incorrect rule!\n");
512 addrMatchU = (packet.GetDstIP() & ln->mask) == ln->ip;
514 if (!foundU && addrMatchU && portMatchU)
518 //printfd(__FILE__, "Up rule ok! %d\n", ln->dir);
534 portMatchD = (packet.GetProto() == ICMP_RPOTO);
538 if (packet.GetProto() == TCP_PROTO || packet.GetProto() == UDP_PROTO)
539 portMatchD = (packet.GetSrcPort() >= ln->port1) && (packet.GetSrcPort() <= ln->port2);
543 if (packet.GetProto() == TCP_PROTO)
544 portMatchD = (packet.GetSrcPort() >= ln->port1) && (packet.GetSrcPort() <= ln->port2);
548 if (packet.GetProto() == UDP_PROTO)
549 portMatchD = (packet.GetSrcPort() >= ln->port1) && (packet.GetSrcPort() <= ln->port2);
553 printfd(__FILE__, "Error! Incorrect rule!\n");
557 addrMatchD = (packet.GetSrcIP() & ln->mask) == ln->ip;
559 if (!foundD && addrMatchD && portMatchD)
563 //printfd(__FILE__, "Down rule ok! %d\n", ln->dir);
568 } //while (ln != rules.end())
576 //-----------------------------------------------------------------------------
577 bool TraffCounterImpl::ReadRules(bool test)
579 //printfd(__FILE__, "TraffCounter::ReadRules()\n");
584 char tp[100]; // protocol
585 char ta[100]; // address
586 char td[100]; // target direction
589 f = fopen(rulesFileName.c_str(), "rt");
593 printfd(__FILE__, "TraffCounterImpl::ReadRules() - File '%s' cannot be opened.\n", rulesFileName.c_str());
594 WriteServLog("File '%s' cannot be oppened.", rulesFileName.c_str());
598 while (fgets(str, 1023, f))
601 if (str[strspn(str," \t")] == '#' || str[strspn(str," \t")] == '\n')
606 r = sscanf(str,"%99s %99s %99s", tp, ta, td);
609 printfd(__FILE__, "TraffCounterImpl::ReadRules() - Error in file '%s' at line %d. There must be 3 parameters.\n", rulesFileName.c_str(), lineNumber);
610 WriteServLog("Error in file '%s' at line %d. There must be 3 parameters.", rulesFileName.c_str(), lineNumber);
618 for (uint8_t i = 0; i < PROTOMAX; i++)
620 if (strcasecmp(tp, protoName[i]) == 0)
624 for (uint32_t i = 0; i < DIR_NUM + 1; i++)
626 if (td == dirName[i])
630 if (rul.dir == 0xff || rul.proto == 0xff)
632 printfd(__FILE__, "TraffCounterImpl::ReadRules() - Error in file '%s' at line %d.\n", rulesFileName.c_str(), lineNumber);
633 WriteServLog("Error in file %s. Line %d.",
634 rulesFileName.c_str(), lineNumber);
639 if (ParseAddress(ta, &rul) != 0)
641 printfd(__FILE__, "TraffCounterImpl::ReadRules() - Error in file '%s' at line %d. Error in adress.\n", rulesFileName.c_str(), lineNumber);
642 WriteServLog("Error in file %s. Error in adress. Line %d.",
643 rulesFileName.c_str(), lineNumber);
648 rules.push_back(rul);
653 // Adding lastest rule: ALL 0.0.0.0/0 NULL
654 rul.dir = DIR_NUM; //NULL
655 rul.ip = 0; //0.0.0.0
662 rules.push_back(rul);
666 //-----------------------------------------------------------------------------
667 int TraffCounterImpl::Reload()
669 std::lock_guard<std::mutex> lock(m_mutex);
673 printfd(__FILE__, "TraffCounterImpl::Reload() - Failed to reload rules.\n");
674 WriteServLog("TraffCounter: Cannot reload rules. Errors found.");
680 printfd(__FILE__, "TraffCounterImpl::Reload() - Reloaded rules successfully.\n");
681 WriteServLog("TraffCounter: Reloaded rules successfully.");
684 //-----------------------------------------------------------------------------
685 bool TraffCounterImpl::ParseAddress(const char * ta, Rule * rule) const
687 char addr[50], mask[20], port1[20], port2[20], ports[40];
689 size_t len = strlen(ta);
692 memset(addr, 0, sizeof(addr));
693 for (i = 0; i < len; i++)
695 if (ta[i] == '/' || ta[i] == ':')
733 if (!(rule->proto == tcp || rule->proto == udp || rule->proto == tcp_udp))
735 printfd(__FILE__, "TraffCounterImpl::ParseAddress() - No ports specified for this protocol.\n");
736 WriteServLog("No ports specified for this protocol.");
741 ports[i - p] = ta[i];
747 strcpy(ports, "0-65535");
754 if ((sss = strchr(ports, '-')) != NULL)
756 strncpy(port1, ports, int(sss-ports));
757 port1[int(sss - ports)] = 0;
758 strcpy(port2, sss + 1);
762 strcpy(port1, ports);
763 strcpy(port2, ports);
766 // Convert strings to mask, ports and IP
767 uint16_t prt1, prt2, msk;
768 struct in_addr ipaddr;
771 msk = static_cast<uint16_t>(strtol(mask, &res, 10));
775 prt1 = static_cast<uint16_t>(strtol(port1, &res, 10));
779 prt2 = static_cast<uint16_t>(strtol(port2, &res, 10));
783 int r = inet_aton(addr, &ipaddr);
787 rule->ip = ipaddr.s_addr;
788 rule->mask = CalcMask(msk);
790 if ((ipaddr.s_addr & rule->mask) != ipaddr.s_addr)
792 printfd(__FILE__, "TraffCounterImpl::ParseAddress() - Address does'n match mask.\n");
793 WriteServLog("Address does'n match mask.");
802 //-----------------------------------------------------------------------------
803 uint32_t TraffCounterImpl::CalcMask(uint32_t msk) const
805 if (msk >= 32) return 0xFFffFFff;
806 if (msk == 0) return 0;
807 return htonl(0xFFffFFff << (32 - msk));
809 //---------------------------------------------------------------------------
810 void TraffCounterImpl::FreeRules()
814 //-----------------------------------------------------------------------------
815 void TraffCounterImpl::SetMonitorDir(const std::string & dir)
818 monitoring = !monitorDir.empty();
820 //-----------------------------------------------------------------------------
821 void TraffCounterImpl::beforeIPChange(uint32_t oldVal)
823 // User changes his address. Remove old IP
827 AsyncPoolST::enqueue([this, oldVal](){ DelUser(oldVal); });
829 //-----------------------------------------------------------------------------
830 void TraffCounterImpl::afterIPChange(UserImpl* user, uint32_t newVal)
832 // User changes his address. Add new IP
836 AsyncPoolST::enqueue([this, user](){ AddUser(user); });
838 //-----------------------------------------------------------------------------