ctdb-tools: Use db_hash in ctdb_killtcp
[samba.git] / ctdb / tools / ctdb_killtcp.c
1 /*
2    CTDB TCP connection killing utility
3
4    Copyright (C) Martin Schwenke <martin@meltin.net> 2016
5
6    This program is free software; you can redistribute it and/or modify
7    it under the terms of the GNU General Public License as published by
8    the Free Software Foundation; either version 3 of the License, or
9    (at your option) any later version.
10
11    This program is distributed in the hope that it will be useful,
12    but WITHOUT ANY WARRANTY; without even the implied warranty of
13    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
14    GNU General Public License for more details.
15
16    You should have received a copy of the GNU General Public License
17    along with this program; if not, see <http://www.gnu.org/licenses/>.
18 */
19
20 #include <talloc.h>
21 #include <tevent.h>
22
23 #include "replace.h"
24 #include "system/network.h"
25
26 #include "lib/util/debug.h"
27
28 #include "protocol/protocol.h"
29 #include "protocol/protocol_util.h"
30
31 #include "common/db_hash.h"
32 #include "common/system.h"
33 #include "common/logging.h"
34
35
36 /* Contains the listening socket and the list of TCP connections to
37  * kill */
38 struct ctdb_kill_tcp {
39         int capture_fd;
40         struct tevent_fd *fde;
41         struct db_hash_context *connections;
42         void *private_data;
43         void *destructor_data;
44         unsigned int attempts;
45         unsigned int max_attempts;
46         struct timeval retry_interval;
47         unsigned int batch_count;
48         unsigned int batch_size;
49 };
50
51 static const char *prog;
52
53 /*
54   called when we get a read event on the raw socket
55  */
56 static void capture_tcp_handler(struct tevent_context *ev,
57                                 struct tevent_fd *fde,
58                                 uint16_t flags, void *private_data)
59 {
60         struct ctdb_kill_tcp *killtcp = talloc_get_type(private_data, struct ctdb_kill_tcp);
61         /* 0 the parts that don't get set by ctdb_sys_read_tcp_packet */
62         struct ctdb_connection conn;
63         uint32_t ack_seq, seq;
64         int rst;
65         uint16_t window;
66         int ret;
67
68         if (ctdb_sys_read_tcp_packet(killtcp->capture_fd,
69                                      killtcp->private_data,
70                                      &conn.server, &conn.client,
71                                      &ack_seq, &seq, &rst, &window) != 0) {
72                 /* probably a non-tcp ACK packet */
73                 return;
74         }
75
76         if (window == htons(1234) && (rst || seq == 0)) {
77                 /* Ignore packets that we sent! */
78                 D_DEBUG("Ignoring packet: %s, "
79                         "seq=%"PRIu32", ack_seq=%"PRIu32", "
80                         "rst=%d, window=%"PRIu16"\n",
81                         ctdb_connection_to_string(killtcp, &conn, false),
82                         seq, ack_seq, rst, ntohs(window));
83                 return;
84         }
85
86         /* Check if this connection is one being reset, if found then delete */
87         ret = db_hash_delete(killtcp->connections,
88                              (uint8_t*)&conn, sizeof(conn));
89         if (ret == ENOENT) {
90                 /* Packet for some other connection, ignore */
91                 return;
92         }
93         if (ret != 0) {
94                 DBG_WARNING("Internal error (%s)\n", strerror(ret));
95                 return;
96         }
97
98         D_INFO("Sending a TCP RST to kill connection %s\n",
99                ctdb_connection_to_string(killtcp, &conn, true));
100
101         ctdb_sys_send_tcp(&conn.server, &conn.client, ack_seq, seq, 1);
102 }
103
104
105 static int tickle_connection_parser(uint8_t *keybuf, size_t keylen,
106                                     uint8_t *databuf, size_t datalen,
107                                     void *private_data)
108 {
109         struct ctdb_kill_tcp *killtcp = talloc_get_type_abort(
110                 private_data, struct ctdb_kill_tcp);
111         struct ctdb_connection *conn;
112
113         if (keylen != sizeof(*conn)) {
114                 DBG_WARNING("Unexpected data in connection hash\n");
115                 return 0;
116         }
117
118         conn = (struct ctdb_connection *)keybuf;
119
120         killtcp->batch_count++;
121         if (killtcp->batch_count > killtcp->batch_size) {
122                 /* Terminate the traverse */
123                 return 1;
124         }
125
126         ctdb_sys_send_tcp(&conn->server, &conn->client, 0, 0, 0);
127
128         return 0;
129 }
130
131 /*
132  * Called periodically until all sentenced connections have been reset
133  * or enough attempts have been made
134  */
135 static void ctdb_tickle_sentenced_connections(struct tevent_context *ev,
136                                               struct tevent_timer *te,
137                                               struct timeval t, void *private_data)
138 {
139         struct ctdb_kill_tcp *killtcp = talloc_get_type(private_data, struct ctdb_kill_tcp);
140         int count, ret;
141
142         /* loop over up to batch_size connections sending tickle ACKs */
143         killtcp->batch_count = 0;
144         ret = db_hash_traverse(killtcp->connections,
145                                tickle_connection_parser, killtcp, NULL);
146         if (ret != 0) {
147                 DBG_WARNING("Unexpected error traversing connections (%s)\n",
148                             strerror(ret));
149         }
150
151         killtcp->attempts++;
152
153         /*
154          * If there are no more connections to kill or we have tried
155          * too many times we can remove the entire killtcp structure
156          */
157         ret = db_hash_traverse(killtcp->connections, NULL, NULL, &count);
158         if (ret != 0) {
159                 /* What now?  Try again until max_attempts reached */
160                 DBG_WARNING("Unexpected error traversing connections (%s)\n",
161                             strerror(ret));
162                 count = 1;
163         }
164         if (count == 0 ||
165             killtcp->attempts >= killtcp->max_attempts) {
166                 talloc_free(killtcp);
167                 return;
168         }
169
170         /* try tickling them again in a seconds time
171          */
172         tevent_add_timer(ev, killtcp,
173                          tevent_timeval_current_ofs(
174                                  killtcp->retry_interval.tv_sec,
175                                  killtcp->retry_interval.tv_usec),
176                          ctdb_tickle_sentenced_connections, killtcp);
177 }
178
179 /* Add a TCP socket to the list of connections we want to RST.  The
180  * list is attached to *killtcp_arg.  If this is NULL then allocate
181  * the structure.  */
182 static int ctdb_killtcp(struct tevent_context *ev,
183                         TALLOC_CTX *mem_ctx,
184                         const char *iface,
185                         struct ctdb_connection *conn,
186                         struct ctdb_kill_tcp **killtcp_arg)
187 {
188         struct ctdb_kill_tcp *killtcp;
189         int ret;
190
191         if (killtcp_arg == NULL) {
192                 DEBUG(DEBUG_ERR, (__location__ " killtcp_arg is NULL!\n"));
193                 return -1;
194         }
195
196         killtcp = *killtcp_arg;
197
198         /* Allocate a new structure if necessary.  The structure is
199          * only freed when mem_ctx is freed. */
200         if (killtcp == NULL) {
201                 killtcp = talloc_zero(mem_ctx, struct ctdb_kill_tcp);
202                 if (killtcp == NULL) {
203                         DEBUG(DEBUG_ERR, (__location__ " out of memory\n"));
204                         return -1;
205                 }
206
207                 killtcp->capture_fd  = -1;
208                 ret = db_hash_init(killtcp, "connections", 2048, DB_HASH_SIMPLE,
209                                    &killtcp->connections);
210                 if (ret != 0) {
211                         D_ERR("Failed to initialise connection hash (%s)\n",
212                               strerror(ret));
213                         talloc_free(killtcp);
214                         return -1;
215                 }
216
217                 killtcp->attempts = 0;
218                 killtcp->max_attempts = 50;
219
220                 killtcp->retry_interval.tv_sec = 0;
221                 killtcp->retry_interval.tv_usec = 100 * 1000;
222
223                 killtcp->batch_count = 0;
224                 killtcp->batch_size = 300;
225
226                 *killtcp_arg = killtcp;
227         }
228
229         /* Connection is stored as a key in the connections hash */
230         ret = db_hash_add(killtcp->connections,
231                           (uint8_t *)conn, sizeof(*conn),
232                           NULL, 0);
233         if (ret != 0) {
234                 D_ERR("Error adding connection to hash (%s)\n", strerror(ret));
235                 return -1;
236         }
237
238         /*
239            If we don't have a socket to listen on yet we must create it
240          */
241         if (killtcp->capture_fd == -1) {
242                 killtcp->capture_fd =
243                         ctdb_sys_open_capture_socket(iface,
244                                                      &killtcp->private_data);
245                 if (killtcp->capture_fd == -1) {
246                         DEBUG(DEBUG_CRIT,(__location__ " Failed to open capturing "
247                                           "socket on iface '%s' for killtcp (%s)\n",
248                                           iface, strerror(errno)));
249                         return -1;
250                 }
251         }
252
253
254         if (killtcp->fde == NULL) {
255                 killtcp->fde = tevent_add_fd(ev, killtcp,
256                                              killtcp->capture_fd,
257                                              TEVENT_FD_READ,
258                                              capture_tcp_handler, killtcp);
259                 tevent_fd_set_auto_close(killtcp->fde);
260         }
261
262         return 0;
263 }
264
265 static int ctdb_killtcp_destructor(struct ctdb_kill_tcp *killtcp)
266 {
267         bool *done = killtcp->destructor_data;
268         *done = true;
269
270         return 0;
271 }
272
273 static void usage(void)
274 {
275         printf("usage: %s <interface> [ <srcip:port> <dstip:port> ]\n", prog);
276         exit(1);
277 }
278
279 int main(int argc, char **argv)
280 {
281         struct ctdb_connection conn;
282         struct ctdb_kill_tcp *killtcp = NULL;
283         struct tevent_context *ev = NULL;
284         struct TALLOC_CONTEXT *mem_ctx = NULL;
285         struct ctdb_connection_list *conn_list = NULL;
286         const char *t;
287         int debug_level;
288         bool done;
289         int i, ret;
290
291         /* Set the debug level */
292         t = getenv("CTDB_DEBUGLEVEL");
293         if (t != NULL) {
294                 if (debug_level_parse(t, &debug_level)) {
295                         DEBUGLEVEL = debug_level;
296                 } else {
297                         DEBUGLEVEL = DEBUG_ERR;
298                 }
299         }
300
301         prog = argv[0];
302
303         if (argc != 2 && argc != 4) {
304                 usage();
305         }
306
307         if (argc == 4) {
308                 ret = ctdb_sock_addr_from_string(argv[2], &conn.client, true);
309                 if (ret != 0) {
310                         D_ERR("Bad IP:port '%s'\n", argv[2]);
311                         goto fail;
312                 }
313
314                 ret = ctdb_sock_addr_from_string(argv[3], &conn.server, true);
315                 if (ret != 0) {
316                         D_ERR("Bad IP:port '%s'\n", argv[3]);
317                         goto fail;
318                 }
319
320
321                 conn_list = talloc_zero(mem_ctx, struct ctdb_connection_list);
322                 if (conn_list == NULL) {
323                         ret = ENOMEM;
324                         DBG_ERR("Internal error (%s)\n", strerror(ret));
325                         goto fail;
326                 }
327                 ret = ctdb_connection_list_add(conn_list, &conn);
328                 if (ret != 0) {
329                         DBG_ERR("Internal error (%s)\n", strerror(ret));
330                         goto fail;
331                 }
332         } else {
333                 ret = ctdb_connection_list_read(mem_ctx, true, &conn_list);
334                 if (ret != 0) {
335                         D_ERR("Unable to parse connections (%s)\n",
336                               strerror(ret));
337                         goto fail;
338                 }
339         }
340
341         mem_ctx = talloc_new(NULL);
342         if (mem_ctx == NULL) {
343                 DEBUG(DEBUG_ERR, (__location__ " out of memory\n"));
344                 goto fail;
345         }
346
347         ev = tevent_context_init(mem_ctx);
348         if (ev == NULL) {
349                 DEBUG(DEBUG_ERR, ("Failed to initialise tevent\n"));
350                 goto fail;
351         }
352
353         if (conn_list->num == 0) {
354                 /* No connections, done! */
355                 talloc_free(mem_ctx);
356                 return 0;
357         }
358
359         for (i = 0; i < conn_list->num; i++) {
360                 ret = ctdb_killtcp(ev, mem_ctx, argv[1],
361                                    &conn_list->conn[i], &killtcp);
362                 if (ret != 0) {
363                         DEBUG(DEBUG_ERR, ("Unable to killtcp\n"));
364                         goto fail;
365                 }
366         }
367
368         done = false;
369         killtcp->destructor_data = &done;
370         talloc_set_destructor(killtcp, ctdb_killtcp_destructor);
371
372         /* Do the initial processing of connections */
373         tevent_add_timer(ev, killtcp,
374                          tevent_timeval_current_ofs(0, 0),
375                          ctdb_tickle_sentenced_connections, killtcp);
376
377         while (!done) {
378                 tevent_loop_once(ev);
379         }
380
381         talloc_free(mem_ctx);
382
383         return 0;
384
385 fail:
386         TALLOC_FREE(mem_ctx);
387         return -1;
388 }