Add new iptables handle task
[rahunas] / src / rh-task-bandwidth.c
1 /**
2  * RahuNAS task bandwidth implementation 
3  * Author: Neutron Soutmun <neo.neutron@gmail.com>
4  * Date:   2008-11-20
5  */
6
7 #include <stdlib.h>
8 #include <syslog.h>
9 #include <sys/types.h>
10 #include <unistd.h>
11 #include <sys/wait.h>
12 #include <sys/stat.h>
13 #include <fcntl.h>
14
15 #include "rahunasd.h"
16 #include "rh-task.h"
17 #include "rh-task-bandwidth.h"
18 #include "rh-task-memset.h"
19 #include "rh-utils.h"
20
21
22
23 #define BANDWIDTH_WRAPPER "/etc/rahunas/bandwidth.sh"
24
25 static unsigned short slot_flags[MAX_SLOT_PAGE] = {1};
26 static unsigned short slot_count = 0;
27
28 unsigned short _get_slot_id()
29 {
30   unsigned short slot_id    = 0;
31   unsigned short page       = 0;
32   unsigned char  id_on_page = 0;
33   time_t random_time;
34
35   // Slot is full 
36   if (slot_count >= MAX_SLOT_ID)
37     return 0;
38
39   // Do a random slot_id 
40   while (slot_id == 0) {
41     time(&(random_time));
42     srandom(random_time);
43     slot_id = random()/(int)(((unsigned int)RAND_MAX + 1) / (MAX_SLOT_ID + 1));
44    
45     // Check validity
46     page = slot_id / PAGE_SIZE; 
47     id_on_page = slot_id % PAGE_SIZE;
48    
49     if (!(slot_flags[page] & (1 << id_on_page))) {
50       // Slot available
51       break;
52     }
53
54     // Slot not available, retry
55     slot_id = 0;
56   }
57
58   return slot_id;
59 }
60
61 void mark_reserved_slot_id(unsigned int slot_id)
62 {
63   unsigned short page       = 0;
64   unsigned char  id_on_page = 0;
65
66   page = slot_id / PAGE_SIZE; 
67   id_on_page = slot_id % PAGE_SIZE;
68
69   slot_count++;
70   slot_flags[page] |= 1 << id_on_page;
71 }
72
73 int bandwidth_exec(struct vserver *vs, char *const args[])
74 {
75   pid_t ws;
76   pid_t pid;
77   int status;
78   int exec_pipe[2];
79   char buffer[150];
80   char *endline = NULL;
81   int ret = 0;
82   int fd = 0;
83   
84   memset(buffer, '\0', sizeof(buffer));
85
86   if (pipe(exec_pipe) == -1) {
87     logmsg(RH_LOG_ERROR, "Error: pipe()");
88     return -1;
89   }
90   DP("pipe0=%d,pipe1=%d", exec_pipe[0], exec_pipe[1]);
91
92   pid = vfork();
93   dup2(exec_pipe[1], STDOUT_FILENO);
94
95   if (pid == 0) {
96     // Child
97     execv(BANDWIDTH_WRAPPER, args);
98   } else if (pid < 0) {
99     // Fork error
100     logmsg(RH_LOG_ERROR, "Error: vfork()"); 
101     ret = -1;
102   } else {
103     // Parent
104     ws = waitpid(pid, &status, 0);
105
106     DP("Bandwidth: Return (%d)", WEXITSTATUS(status));
107
108     // Return message log
109     DP("Read message");
110     read(exec_pipe[0], buffer, sizeof(buffer));
111
112     if (buffer != NULL) {
113       DP("Got message: %s", buffer);
114       endline = strstr(buffer, "\n");
115       if (endline != NULL) 
116         *endline = '\0';
117
118       if (vs != NULL) {
119         logmsg(RH_LOG_NORMAL, "[%s] Bandwidth: %s", 
120           vs->vserver_config->vserver_name, buffer);
121       } else {
122         logmsg(RH_LOG_NORMAL, "[main server] Bandwidth: %s", buffer);
123       }
124     }
125
126     if (WIFEXITED(status)) {
127       ret = WEXITSTATUS(status);
128     } else {
129       ret = -1;
130     } 
131   }
132
133   close(exec_pipe[0]);
134   close(exec_pipe[1]);
135   return ret;
136 }
137
138 int bandwidth_start(void)
139 {
140   char *args[3];
141
142   DP("Bandwidth: start");
143
144   args[0] = BANDWIDTH_WRAPPER;
145   args[1] = "start";
146   args[2] = (char *) 0;
147
148   return bandwidth_exec(NULL, args);
149 }
150
151 int bandwidth_stop(void)
152 {
153   char *args[3];
154
155   DP("Bandwidth: stop");
156
157   args[0] = BANDWIDTH_WRAPPER;
158   args[1] = "stop";
159   args[2] = (char *) 0;
160
161   return bandwidth_exec(NULL, args);
162 }
163
164 int bandwidth_add(struct vserver *vs, struct bandwidth_req *bw_req)
165 {
166   char *args[7];
167
168   DP("Bandwidth: request %s %s %s %s", bw_req->slot_id, 
169      bw_req->ip, bw_req->bandwidth_max_down, bw_req->bandwidth_max_up);
170
171   args[0] = BANDWIDTH_WRAPPER;
172   args[1] = "add";
173   args[2] = bw_req->slot_id;
174   args[3] = bw_req->ip;
175   args[4] = bw_req->bandwidth_max_down;
176   args[5] = bw_req->bandwidth_max_up;
177   args[6] = (char *) 0;
178
179   return bandwidth_exec(vs, args);
180 }
181
182 int bandwidth_del(struct vserver *vs, struct bandwidth_req *bw_req)
183 {
184   char *args[4];
185
186   DP("Bandwidth: request %s", bw_req->slot_id);
187
188   args[0] = BANDWIDTH_WRAPPER;
189   args[1] = "del";
190   args[2] = bw_req->slot_id;
191   args[3] = (char *) 0;
192
193   return bandwidth_exec(vs, args);
194 }
195
196 /* Start service task */
197 static int startservice (void)
198 {
199   return bandwidth_start();
200 }
201
202 /* Stop service task */
203 static int stopservice (void)
204 {
205   return bandwidth_stop();
206 }
207
208 /* Initialize */
209 static void init (struct vserver *vs)
210 {
211   /* Do nothing */
212 }
213
214 /* Cleanup */
215 static int cleanup (struct vserver *vs)
216 {
217   /* Do nothing */
218 }
219
220 /* Start session task */
221 static int startsess (struct vserver *vs, struct task_req *req)
222 {
223   struct bandwidth_req bw_req;
224   unsigned short slot_id;
225   unsigned char max_try = 3;
226   GList *member_node = NULL;
227   struct rahunas_member *member = NULL;
228
229   member_node = member_get_node_by_id(vs, req->id);
230   if (member_node == NULL)
231     return (-1);
232
233   member = (struct rahunas_member *) member_node->data;
234
235   if (member->bandwidth_max_down == 0 && member->bandwidth_max_up == 0)
236     return 0;
237
238   if (member->bandwidth_slot_id > 0)
239     return 0;
240   
241   // Formating the bandwidth request
242   sprintf(bw_req.ip, "%s", idtoip(vs->v_map, req->id));
243   sprintf(bw_req.bandwidth_max_down, "%lu", 
244     member->bandwidth_max_down);
245   sprintf(bw_req.bandwidth_max_up, "%lu", 
246     member->bandwidth_max_up);
247   
248   while (max_try > 0) { 
249     slot_id = _get_slot_id();
250     sprintf(bw_req.slot_id, "%d", slot_id);
251     if (bandwidth_add(vs, &bw_req) == 0)
252       break;
253     else
254       max_try--;
255   }
256
257
258   if (max_try == 0) {
259     logmsg(RH_LOG_ERROR, "[%s] Bandwidth: Maximum trying, failed!",
260            vs->vserver_config->vserver_name);
261     return -1;
262   }
263
264   mark_reserved_slot_id(slot_id);
265   member->bandwidth_slot_id = slot_id;
266  
267   return 0;
268 }
269
270 /* Stop session task */
271 static int stopsess  (struct vserver *vs, struct task_req *req)
272 {
273   struct bandwidth_req bw_req;
274   unsigned short slot_id = 0;
275   GList *member_node = NULL;
276   struct rahunas_member *member = NULL;
277   
278   member_node = member_get_node_by_id(vs, req->id);
279   if (member_node == NULL)
280     return (-1); 
281   
282   member = (struct rahunas_member *) member_node->data;
283   slot_id = member->bandwidth_slot_id;
284
285   if (slot_id < 1)
286     return 0;
287
288   sprintf(bw_req.slot_id, "%d", slot_id);
289
290   if (bandwidth_del(vs, &bw_req) == 0) {
291     member->bandwidth_slot_id = 0;
292
293     if (slot_count > 0)
294       slot_count--;
295   
296     return 0; 
297   } else {
298     return -1;
299   }
300 }
301
302 /* Commit start session task */
303 static int commitstartsess (struct vserver *vs, struct task_req *req)
304 {
305   /* Do nothing or need to implement */
306 }
307
308 /* Commit stop session task */
309 static int commitstopsess  (struct vserver *vs, struct task_req *req)
310 {
311   /* Do nothing or need to implement */
312 }
313
314 /* Rollback start session task */
315 static int rollbackstartsess (struct vserver *vs, struct task_req *req)
316 {
317   /* Do nothing or need to implement */
318 }
319
320 /* Rollback stop session task */
321 static int rollbackstopsess  (struct vserver *vs, struct task_req *req)
322 {
323   /* Do nothing or need to implement */
324 }
325
326 static struct task taskbandwidth = {
327   .taskname = "BANDWIDTH",
328   .taskprio = 20,
329   .init = &init,
330   .cleanup = &cleanup,
331   .startservice = &startservice,
332   .stopservice = &stopservice,
333   .startsess = &startsess,
334   .stopsess = &stopsess,
335   .commitstartsess = &commitstartsess,
336   .commitstopsess = &commitstopsess,
337   .rollbackstartsess = &rollbackstartsess,
338   .rollbackstopsess = &rollbackstopsess,
339 };
340
341 void rh_task_bandwidth_reg(struct main_server *ms) {
342   task_register(ms, &taskbandwidth);
343 }