]> git.stg.codes - stg.git/blob - projects/rscriptd/main.cpp
В плагинах захвата трафика для Divert-сокетов и BPF реализована
[stg.git] / projects / rscriptd / main.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 : Boris Mikhailenko <stg34@stargazer.dp.ua>
19  *    Author : Maxim Mamontov <faust@stargazer.dp.ua>
20  */
21
22  /*
23  $Revision: 1.19 $
24  $Author: faust $
25  $Date: 2010/09/10 06:37:45 $
26  */
27
28 #include <sys/types.h>
29 #include <sys/ipc.h>
30 #include <sys/msg.h>
31 #include <sys/stat.h>
32 #include <sys/wait.h>
33 #include <fcntl.h> // creat
34 #include <unistd.h>
35
36 #include <cstdlib>
37 #include <csignal>
38 #include <set>
39
40 #include "common.h"
41 #include "stg_logger.h"
42 #include "script_executer.h"
43 #include "conffiles.h"
44 #include "listener.h"
45 #include "pidfile.h"
46 #include "version.h"
47
48 using namespace std;
49
50 #ifdef DEBUG
51 # define MAIN_DEBUG  1
52 # define NO_DAEMON    1
53 #endif
54
55 #define START_FILE "/._ST_ART_ED_"
56
57 static bool childExited = false;
58 set<pid_t> executersPid;
59 static pid_t stgChildPid;
60 volatile time_t stgTime = time(NULL);
61
62 class STG_STOPPER
63 {
64 public:
65     STG_STOPPER() { nonstop = true; }
66     bool GetStatus() const { return nonstop; };
67     void Stop(const char * __file__, int __line__)
68         {
69         #ifdef NO_DAEMON
70         printfd(__FILE__, "rscriptd stopped at %s:%d\n", __file__, __line__);
71         #endif
72         nonstop = false;
73         }
74 private:
75     bool nonstop;
76 };
77 //-----------------------------------------------------------------------------
78 STG_STOPPER nonstop;
79 //-----------------------------------------------------------------------------
80 void CatchPROF(int)
81 {
82 }
83 //-----------------------------------------------------------------------------
84 void CatchUSR1(int)
85 {
86 }
87 //-----------------------------------------------------------------------------
88 void CatchTERM(int sig)
89 {
90 /*
91  *Function Name:CatchINT
92  *Parameters: sig_num - signal number
93  *Description: INT signal handler
94  *Returns: Nothing
95  */
96 STG_LOGGER & WriteServLog = GetStgLogger();
97 WriteServLog("Shutting down... %d", sig);
98
99 nonstop.Stop(__FILE__, __LINE__);
100
101 struct sigaction newsa, oldsa;
102 sigset_t sigmask;
103
104 sigemptyset(&sigmask);
105 sigaddset(&sigmask, SIGTERM);
106 newsa.sa_handler = SIG_IGN;
107 newsa.sa_mask = sigmask;
108 newsa.sa_flags = 0;
109 sigaction(SIGTERM, &newsa, &oldsa);
110
111 sigemptyset(&sigmask);
112 sigaddset(&sigmask, SIGINT);
113 newsa.sa_handler = SIG_IGN;
114 newsa.sa_mask = sigmask;
115 newsa.sa_flags = 0;
116 sigaction(SIGINT, &newsa, &oldsa);
117 }
118 //-----------------------------------------------------------------------------
119 void CatchPIPE(int)
120 {
121 STG_LOGGER & WriteServLog = GetStgLogger();
122 WriteServLog("Broken pipe!");
123 }
124 //-----------------------------------------------------------------------------
125 void CatchHUP(int)
126 {
127 }
128 //-----------------------------------------------------------------------------
129 void CatchCHLD(int)
130 {
131 int status;
132 pid_t childPid;
133 childPid = waitpid(-1, &status, WNOHANG);
134
135 set<pid_t>::iterator pid;
136 pid = executersPid.find(childPid);
137 if (pid != executersPid.end())
138     {
139     executersPid.erase(pid);
140     }
141 if (childPid == stgChildPid)
142     {
143     childExited = true;
144     }
145 }
146 //-----------------------------------------------------------------------------
147 void SetSignalHandlers()
148 {
149 struct sigaction newsa, oldsa;
150 sigset_t sigmask;
151
152 sigemptyset(&sigmask);
153 sigaddset(&sigmask, SIGTERM);
154 newsa.sa_handler = CatchTERM;
155 newsa.sa_mask = sigmask;
156 newsa.sa_flags = 0;
157 sigaction(SIGTERM, &newsa, &oldsa);
158
159 sigemptyset(&sigmask);
160 sigaddset(&sigmask, SIGUSR1);
161 newsa.sa_handler = CatchUSR1;
162 newsa.sa_mask = sigmask;
163 newsa.sa_flags = 0;
164 sigaction(SIGUSR1, &newsa, &oldsa);
165
166 sigemptyset(&sigmask);
167 sigaddset(&sigmask, SIGINT);
168 newsa.sa_handler = CatchTERM;
169 newsa.sa_mask = sigmask;
170 newsa.sa_flags = 0;
171 sigaction(SIGINT, &newsa, &oldsa);
172
173 sigemptyset(&sigmask);
174 sigaddset(&sigmask, SIGPIPE);
175 newsa.sa_handler = CatchPIPE;
176 newsa.sa_mask = sigmask;
177 newsa.sa_flags = 0;
178 sigaction(SIGPIPE, &newsa, &oldsa);
179
180 sigemptyset(&sigmask);
181 sigaddset(&sigmask, SIGHUP);
182 newsa.sa_handler = CatchHUP;
183 newsa.sa_mask = sigmask;
184 newsa.sa_flags = 0;
185 sigaction(SIGHUP, &newsa, &oldsa);
186
187 sigemptyset(&sigmask);
188 sigaddset(&sigmask, SIGCHLD);
189 newsa.sa_handler = CatchCHLD;
190 newsa.sa_mask = sigmask;
191 newsa.sa_flags = 0;
192 sigaction(SIGCHLD, &newsa, &oldsa);
193 }
194 //-----------------------------------------------------------------------------
195 int StartScriptExecuter(char * procName, int msgKey, int * msgID)
196 {
197 STG_LOGGER & WriteServLog = GetStgLogger();
198
199 if (*msgID == -11)   // If msgID == -11 - first call. Create queue
200     {
201     for (int i = 0; i < 2; i++)
202         {
203         *msgID = msgget(msgKey, IPC_CREAT | IPC_EXCL | 0600);
204
205         if (*msgID == -1)
206             {
207             *msgID = msgget(msgKey, 0);
208             if (*msgID == -1)
209                 {
210                 WriteServLog("Message queue not created.");
211                 return -1;
212                 }
213             else
214                 {
215                 msgctl(*msgID, IPC_RMID, NULL);
216                 }
217             }
218         else
219             {
220             WriteServLog("Message queue created successfully. msgKey=%d msgID=%d", msgKey, *msgID);
221             break;
222             }
223         }
224     }
225
226 pid_t executerPid = fork();
227
228 switch (executerPid)
229     {
230     case -1:    // Failure
231         WriteServLog("Fork error!");
232         return -1;
233
234     case 0:     // Child
235         //close(0);
236         //close(1);
237         //close(2);
238         //setsid();
239         Executer(msgKey, *msgID, executerPid, procName);
240         return 1;
241
242     default:    // Parent
243         if (executersPid.empty())
244             Executer(msgKey, *msgID, executerPid, NULL);
245         executersPid.insert(executerPid);
246     }
247 return 0;
248 }
249 //-----------------------------------------------------------------------------
250 #ifdef NO_DAEMON
251 int ForkAndWait(const string &)
252 #else
253 int ForkAndWait(const string & confDir)
254 #endif
255 {
256 #ifndef NO_DAEMON
257 stgChildPid = fork();
258 string startFile = confDir + START_FILE;
259 unlink(startFile.c_str());
260
261 switch (stgChildPid)
262     {
263     case -1:    // Failure
264         return -1;
265         break;
266
267     case 0:     // Child
268         //close(0);
269         close(1);
270         close(2);
271         setsid();
272         break;
273
274     default:    // Parent
275         for (int i = 0; i < 120 * 5; i++)
276             {
277             if (access(startFile.c_str(), F_OK) == 0)
278                 {
279                 //printf("Fork successfull. Exit.\n");
280                 unlink(startFile.c_str());
281                 exit(0);
282                 }
283
284             if (childExited)
285                 {
286                 unlink(startFile.c_str());
287                 exit(1);
288                 }
289             usleep(200000);
290             }
291         unlink(startFile.c_str());
292         exit(1);
293         break;
294     }
295 #endif
296 return 0;
297 }
298 //-----------------------------------------------------------------------------
299 void KillExecuters()
300 {
301 set<pid_t>::iterator pid;
302 pid = executersPid.begin();
303 while (pid != executersPid.end())
304     {
305     printfd(__FILE__, "KillExecuters pid=%d\n", *pid);
306     kill(*pid, SIGUSR1);
307     ++pid;
308     }
309 }
310 //-----------------------------------------------------------------------------
311 int main(int argc, char * argv[])
312 {
313
314 /*
315   Initialization order:
316   - Logger
317   - Config
318   - Set signal nandlers
319   - Fork and exit
320  */
321
322 CONFIGFILE * cfg = NULL;
323 LISTENER * listener = NULL;
324 int msgID = -11;
325 int execNum = 0;
326 int execMsgKey = 0;
327
328 string logFileName;
329 string confDir;
330 string password;
331 string onConnect;
332 string onDisconnect;
333 int port;
334 int userTimeout;
335
336 if (getuid())
337     {
338     printf("You must be root. Exit.\n");
339     exit(1);
340     }
341
342 if (argc == 2)
343     cfg = new CONFIGFILE(argv[1]);
344 else
345     cfg = new CONFIGFILE("/etc/rscriptd/rscriptd.conf");
346
347 if (cfg->Error())
348     {
349     STG_LOGGER & WriteServLog = GetStgLogger();
350     WriteServLog.SetLogFileName("/var/log/rscriptd.log");
351     WriteServLog("Error reading config file!");
352     delete cfg;
353     return EXIT_FAILURE;
354     }
355
356 cfg->ReadString("LogFileName", &logFileName, "/var/log/rscriptd.log");
357 cfg->ReadInt("ExecutersNum", &execNum, 1);
358 cfg->ReadInt("ExecMsgKey", &execMsgKey, 5555);
359 cfg->ReadString("ConfigDir", &confDir, "/etc/rscriptd");
360 cfg->ReadString("Password", &password, "");
361 cfg->ReadInt("Port", &port, 5555);
362 cfg->ReadInt("UserTimeout", &userTimeout, 60);
363 cfg->ReadString("ScriptOnConnect", &onConnect, "/etc/rscriptd/OnConnect");
364 cfg->ReadString("ScriptOnDisconnect", &onDisconnect, "/etc/rscriptd/OnDisconnect");
365
366 SetSignalHandlers();
367 if (ForkAndWait(confDir) < 0)
368     {
369     STG_LOGGER & WriteServLog = GetStgLogger();
370     WriteServLog("Fork error!");
371     delete cfg;
372     return EXIT_FAILURE;
373     }
374
375 STG_LOGGER & WriteServLog = GetStgLogger();
376 PIDFile pidFile("/var/run/rscriptd.pid");
377 WriteServLog.SetLogFileName(logFileName);
378 WriteServLog("rscriptd v. %s", SERVER_VERSION);
379
380 for (int i = 0; i < execNum; i++)
381     {
382     int ret = StartScriptExecuter(argv[0], execMsgKey, &msgID);
383     if (ret < 0)
384         {
385         STG_LOGGER & WriteServLog = GetStgLogger();
386         WriteServLog("Start Script Executer error!");
387         delete cfg;
388         return EXIT_FAILURE;
389         }
390     if (ret == 1)
391         {
392         delete cfg;
393         return EXIT_SUCCESS;
394         }
395     }
396
397 listener = new LISTENER();
398 listener->SetPort(port);
399 listener->SetPassword(password);
400 listener->SetUserTimeout(userTimeout);
401 listener->SetScriptOnConnect(onConnect);
402 listener->SetScriptOnDisconnect(onDisconnect);
403
404 listener->Start();
405
406 WriteServLog("rscriptd started successfully.");
407 WriteServLog("+++++++++++++++++++++++++++++++++++++++++++++");
408
409 #ifndef NO_DAEMON
410 string startFile(confDir + START_FILE);
411 creat(startFile.c_str(), S_IRUSR);
412 #endif
413
414 while (nonstop.GetStatus())
415     {
416     usleep(100000);
417     }
418
419 listener->Stop();
420
421 WriteServLog("+++++++++++++++++++++++++++++++++++++++++++++");
422
423 int res = msgctl(msgID, IPC_RMID, NULL);
424 if (res)
425     WriteServLog("Queue was not removed. id=%d", msgID);
426 else
427     WriteServLog("Queue removed successfully.");
428
429 KillExecuters();
430
431 WriteServLog("rscriptd stopped successfully.");
432 WriteServLog("---------------------------------------------");
433
434 delete listener;
435 delete cfg;
436 return EXIT_SUCCESS;
437 }
438 //-----------------------------------------------------------------------------
439