linux_inet_diag: avoid unnecessary sockaddr initialization
[raindrops.git] / ext / raindrops / linux_inet_diag.c
blobd94dafd28452e4697c446aa8cb72e1dc5387ffe0
1 #include <ruby.h>
2 #ifdef HAVE_RUBY_ST_H
3 # include <ruby/st.h>
4 #else
5 # include <st.h>
6 #endif
7 #include "my_fileno.h"
8 #ifdef __linux__
10 /* Ruby 1.8.6+ macros (for compatibility with Ruby 1.9) */
11 #ifndef RSTRING_LEN
12 # define RSTRING_LEN(s) (RSTRING(s)->len)
13 #endif
15 /* partial emulation of the 1.9 rb_thread_blocking_region under 1.8 */
16 #ifndef HAVE_RB_THREAD_BLOCKING_REGION
17 # include <rubysig.h>
18 # define RUBY_UBF_IO ((rb_unblock_function_t *)-1)
19 typedef void rb_unblock_function_t(void *);
20 typedef VALUE rb_blocking_function_t(void *);
21 static VALUE
22 rb_thread_blocking_region(
23 rb_blocking_function_t *func, void *data1,
24 rb_unblock_function_t *ubf, void *data2)
26 VALUE rv;
28 TRAP_BEG;
29 rv = func(data1);
30 TRAP_END;
32 return rv;
34 #endif /* ! HAVE_RB_THREAD_BLOCKING_REGION */
36 #ifdef HAVE_RB_THREAD_IO_BLOCKING_REGION
37 VALUE rb_thread_io_blocking_region(rb_blocking_function_t *, void *, int);
38 #else
39 # define rb_thread_io_blocking_region(fn,data,fd) \
40 rb_thread_blocking_region((fn),(data),RUBY_UBF_IO,0)
41 #endif /* HAVE_RB_THREAD_IO_BLOCKING_REGION */
43 #include <assert.h>
44 #include <errno.h>
45 #include <sys/socket.h>
46 #include <sys/types.h>
47 #include <netdb.h>
48 #include <unistd.h>
49 #include <fcntl.h>
50 #include <string.h>
51 #include <asm/types.h>
52 #include <netinet/in.h>
53 #include <arpa/inet.h>
54 #include <netinet/tcp.h>
55 #include <linux/netlink.h>
56 #include <linux/rtnetlink.h>
57 #include <linux/inet_diag.h>
59 union any_addr {
60 struct sockaddr_storage ss;
61 struct sockaddr sa;
62 struct sockaddr_in in;
63 struct sockaddr_in6 in6;
66 static size_t page_size;
67 static unsigned g_seq;
68 static VALUE cListenStats, cIDSock;
69 static ID id_new;
71 struct listen_stats {
72 uint32_t active;
73 uint32_t listener_p:1;
74 uint32_t queued:31;
77 #define OPLEN (sizeof(struct inet_diag_bc_op) + \
78 sizeof(struct inet_diag_hostcond) + \
79 sizeof(struct sockaddr_storage))
81 struct nogvl_args {
82 st_table *table;
83 struct iovec iov[3]; /* last iov holds inet_diag bytecode */
84 struct listen_stats stats;
85 int fd;
88 #ifdef SOCK_CLOEXEC
89 # define my_SOCK_RAW (SOCK_RAW|SOCK_CLOEXEC)
90 # define FORCE_CLOEXEC(v) (v)
91 #else
92 # define my_SOCK_RAW SOCK_RAW
93 static VALUE FORCE_CLOEXEC(VALUE io)
95 int fd = my_fileno(io);
96 int flags = fcntl(fd, F_SETFD, FD_CLOEXEC);
97 if (flags == -1)
98 rb_sys_fail("fcntl(F_SETFD, FD_CLOEXEC)");
99 return io;
101 #endif
104 * call-seq:
105 * Raindrops::InetDiagSocket.new -> Socket
107 * Creates a new Socket object for the netlink inet_diag facility
109 static VALUE ids_s_new(VALUE klass)
111 VALUE argv[3];
113 argv[0] = INT2NUM(AF_NETLINK);
114 argv[1] = INT2NUM(my_SOCK_RAW);
115 argv[2] = INT2NUM(NETLINK_INET_DIAG);
117 return FORCE_CLOEXEC(rb_call_super(3, argv));
120 /* creates a Ruby ListenStats Struct based on our internal listen_stats */
121 static VALUE rb_listen_stats(struct listen_stats *stats)
123 VALUE active = UINT2NUM(stats->active);
124 VALUE queued = UINT2NUM(stats->queued);
126 return rb_struct_new(cListenStats, active, queued);
129 static int st_free_data(st_data_t key, st_data_t value, st_data_t ignored)
131 xfree((void *)key);
132 xfree((void *)value);
134 return ST_DELETE;
137 static int st_to_hash(st_data_t key, st_data_t value, VALUE hash)
139 struct listen_stats *stats = (struct listen_stats *)value;
141 if (stats->listener_p) {
142 VALUE k = rb_str_new2((const char *)key);
143 VALUE v = rb_listen_stats(stats);
145 OBJ_FREEZE(k);
146 rb_hash_aset(hash, k, v);
148 return st_free_data(key, value, 0);
151 static int st_AND_hash(st_data_t key, st_data_t value, VALUE hash)
153 struct listen_stats *stats = (struct listen_stats *)value;
155 if (stats->listener_p) {
156 VALUE k = rb_str_new2((const char *)key);
158 if (rb_hash_lookup(hash, k) == Qtrue) {
159 VALUE v = rb_listen_stats(stats);
160 OBJ_FREEZE(k);
161 rb_hash_aset(hash, k, v);
164 return st_free_data(key, value, 0);
167 static const char *addr_any(sa_family_t family)
169 static const char ipv4[] = "0.0.0.0";
170 static const char ipv6[] = "[::]";
172 if (family == AF_INET)
173 return ipv4;
174 assert(family == AF_INET6 && "unknown family");
175 return ipv6;
178 static void bug_warn(void)
180 fprintf(stderr, "Please report how you produced this at "\
181 "raindrops@librelist.org\n");
182 fflush(stderr);
185 static struct listen_stats *stats_for(st_table *table, struct inet_diag_msg *r)
187 char *key, *port, *old_key;
188 size_t alloca_len;
189 struct listen_stats *stats;
190 size_t keylen;
191 size_t portlen = sizeof("65535");
192 union any_addr sa;
193 socklen_t len = sizeof(struct sockaddr_storage);
194 int rc;
195 int flags = NI_NUMERICHOST | NI_NUMERICSERV;
197 switch ((sa.ss.ss_family = r->idiag_family)) {
198 case AF_INET: {
199 sa.in.sin_port = r->id.idiag_sport;
200 sa.in.sin_addr.s_addr = r->id.idiag_src[0];
201 keylen = INET_ADDRSTRLEN;
202 alloca_len = keylen + 1 + portlen;
203 key = alloca(alloca_len);
204 key[keylen] = 0; /* will be ':' later */
205 port = key + keylen + 1;
206 rc = getnameinfo(&sa.sa, len,
207 key, keylen, port, portlen, flags);
208 break;
210 case AF_INET6: {
211 sa.in6.sin6_port = r->id.idiag_sport;
212 memcpy(&sa.in6.sin6_addr, &r->id.idiag_src, sizeof(__be32[4]));
213 keylen = INET6_ADDRSTRLEN;
214 /* [ ] */
215 alloca_len = 1 + keylen + 1 + 1 + portlen;
216 key = alloca(alloca_len);
217 *key = '[';
218 key[1 + keylen + 1] = 0; /* will be ':' later */
219 port = 1 + key + keylen + 1 + 1;
220 rc = getnameinfo(&sa.sa, len,
221 key + 1, keylen, port, portlen, flags);
222 break;
224 default:
225 assert(0 && "unsupported address family, could that be IPv7?!");
227 if (rc != 0) {
228 fprintf(stderr, "BUG: getnameinfo: %s\n", gai_strerror(rc));
229 bug_warn();
230 *key = 0;
233 keylen = strlen(key);
234 portlen = strlen(port);
236 switch (sa.ss.ss_family) {
237 case AF_INET:
238 key[keylen] = ':';
239 memmove(key + keylen + 1, port, portlen + 1);
240 break;
241 case AF_INET6:
242 key[keylen] = ']';
243 key[keylen + 1] = ':';
244 memmove(key + keylen + 2, port, portlen + 1);
245 keylen++;
246 break;
247 default:
248 assert(0 && "unsupported address family, could that be IPv7?!");
251 if (st_lookup(table, (st_data_t)key, (st_data_t *)&stats))
252 return stats;
254 old_key = key;
256 if (r->idiag_state == TCP_ESTABLISHED) {
257 int n = snprintf(key, alloca_len, "%s:%u",
258 addr_any(sa.ss.ss_family),
259 ntohs(r->id.idiag_sport));
260 if (n <= 0) {
261 fprintf(stderr, "BUG: snprintf: %d\n", n);
262 bug_warn();
264 if (st_lookup(table, (st_data_t)key, (st_data_t *)&stats))
265 return stats;
266 if (n <= 0) {
267 key = xmalloc(1);
268 *key = '\0';
269 } else {
270 old_key = key;
271 key = xmalloc(n + 1);
272 memcpy(key, old_key, n + 1);
274 } else {
275 key = xmalloc(keylen + 1 + portlen + 1);
276 memcpy(key, old_key, keylen + 1 + portlen + 1);
278 stats = xcalloc(1, sizeof(struct listen_stats));
279 st_insert(table, (st_data_t)key, (st_data_t)stats);
280 return stats;
283 static void table_incr_active(st_table *table, struct inet_diag_msg *r)
285 struct listen_stats *stats = stats_for(table, r);
286 ++stats->active;
289 static void table_set_queued(st_table *table, struct inet_diag_msg *r)
291 struct listen_stats *stats = stats_for(table, r);
292 stats->listener_p = 1;
293 stats->queued = r->idiag_rqueue;
296 /* inner loop of inet_diag, called for every socket returned by netlink */
297 static inline void r_acc(struct nogvl_args *args, struct inet_diag_msg *r)
300 * inode == 0 means the connection is still in the listen queue
301 * and has not yet been accept()-ed by the server. The
302 * inet_diag bytecode cannot filter this for us.
304 if (r->idiag_inode == 0)
305 return;
306 if (r->idiag_state == TCP_ESTABLISHED) {
307 if (args->table)
308 table_incr_active(args->table, r);
309 else
310 args->stats.active++;
311 } else { /* if (r->idiag_state == TCP_LISTEN) */
312 if (args->table)
313 table_set_queued(args->table, r);
314 else
315 args->stats.queued = r->idiag_rqueue;
318 * we wont get anything else because of the idiag_states filter
322 static const char err_sendmsg[] = "sendmsg";
323 static const char err_recvmsg[] = "recvmsg";
324 static const char err_nlmsg[] = "nlmsg";
326 struct diag_req {
327 struct nlmsghdr nlh;
328 struct inet_diag_req r;
331 static void prep_msghdr(
332 struct msghdr *msg,
333 struct nogvl_args *args,
334 struct sockaddr_nl *nladdr,
335 size_t iovlen)
337 memset(msg, 0, sizeof(struct msghdr));
338 msg->msg_name = (void *)nladdr;
339 msg->msg_namelen = sizeof(struct sockaddr_nl);
340 msg->msg_iov = args->iov;
341 msg->msg_iovlen = iovlen;
344 static void prep_diag_args(
345 struct nogvl_args *args,
346 struct sockaddr_nl *nladdr,
347 struct rtattr *rta,
348 struct diag_req *req,
349 struct msghdr *msg)
351 memset(req, 0, sizeof(struct diag_req));
352 memset(nladdr, 0, sizeof(struct sockaddr_nl));
354 nladdr->nl_family = AF_NETLINK;
356 req->nlh.nlmsg_len = sizeof(struct diag_req) +
357 RTA_LENGTH(args->iov[2].iov_len);
358 req->nlh.nlmsg_type = TCPDIAG_GETSOCK;
359 req->nlh.nlmsg_flags = NLM_F_ROOT | NLM_F_MATCH | NLM_F_REQUEST;
360 req->nlh.nlmsg_pid = getpid();
361 req->r.idiag_states = (1<<TCP_ESTABLISHED) | (1<<TCP_LISTEN);
362 rta->rta_type = INET_DIAG_REQ_BYTECODE;
363 rta->rta_len = RTA_LENGTH(args->iov[2].iov_len);
365 args->iov[0].iov_base = req;
366 args->iov[0].iov_len = sizeof(struct diag_req);
367 args->iov[1].iov_base = rta;
368 args->iov[1].iov_len = sizeof(struct rtattr);
370 prep_msghdr(msg, args, nladdr, 3);
373 static void prep_recvmsg_buf(struct nogvl_args *args)
375 /* reuse buffer that was allocated for bytecode */
376 args->iov[0].iov_len = page_size;
377 args->iov[0].iov_base = args->iov[2].iov_base;
380 /* does the inet_diag stuff with netlink(), this is called w/o GVL */
381 static VALUE diag(void *ptr)
383 struct nogvl_args *args = ptr;
384 struct sockaddr_nl nladdr;
385 struct rtattr rta;
386 struct diag_req req;
387 struct msghdr msg;
388 const char *err = NULL;
389 unsigned seq = ++g_seq;
391 prep_diag_args(args, &nladdr, &rta, &req, &msg);
392 req.nlh.nlmsg_seq = seq;
394 if (sendmsg(args->fd, &msg, 0) < 0) {
395 err = err_sendmsg;
396 goto out;
399 prep_recvmsg_buf(args);
401 while (1) {
402 ssize_t readed;
403 size_t r;
404 struct nlmsghdr *h = (struct nlmsghdr *)args->iov[0].iov_base;
406 prep_msghdr(&msg, args, &nladdr, 1);
407 readed = recvmsg(args->fd, &msg, 0);
408 if (readed < 0) {
409 if (errno == EINTR)
410 continue;
411 err = err_recvmsg;
412 goto out;
414 if (readed == 0)
415 goto out;
416 r = (size_t)readed;
417 for ( ; NLMSG_OK(h, r); h = NLMSG_NEXT(h, r)) {
418 if (h->nlmsg_seq != seq)
419 continue;
420 if (h->nlmsg_type == NLMSG_DONE)
421 goto out;
422 if (h->nlmsg_type == NLMSG_ERROR) {
423 err = err_nlmsg;
424 goto out;
426 r_acc(args, NLMSG_DATA(h));
429 out:
431 int save_errno = errno;
432 if (err && args->table) {
433 st_foreach(args->table, st_free_data, 0);
434 st_free_table(args->table);
436 errno = save_errno;
438 return (VALUE)err;
441 /* populates sockaddr_storage struct by parsing +addr+ */
442 static void parse_addr(union any_addr *inet, VALUE addr)
444 char *host_ptr;
445 char *check;
446 char *colon = NULL;
447 char *rbracket = NULL;
448 void *dst;
449 long host_len;
450 int af, rc;
451 uint16_t *portdst;
452 unsigned long port;
454 Check_Type(addr, T_STRING);
455 host_ptr = StringValueCStr(addr);
456 host_len = RSTRING_LEN(addr);
457 if (*host_ptr == '[') { /* ipv6 address format (rfc2732) */
458 rbracket = memchr(host_ptr + 1, ']', host_len - 1);
460 if (rbracket == NULL)
461 rb_raise(rb_eArgError, "']' not found in IPv6 addr=%s",
462 host_ptr);
463 if (rbracket[1] != ':')
464 rb_raise(rb_eArgError, "':' not found in IPv6 addr=%s",
465 host_ptr);
466 colon = rbracket + 1;
467 host_ptr++;
468 *rbracket = 0;
469 inet->ss.ss_family = af = AF_INET6;
470 dst = &inet->in6.sin6_addr;
471 portdst = &inet->in6.sin6_port;
472 } else { /* ipv4 */
473 colon = memchr(host_ptr, ':', host_len);
474 inet->ss.ss_family = af = AF_INET;
475 dst = &inet->in.sin_addr;
476 portdst = &inet->in.sin_port;
479 if (!colon)
480 rb_raise(rb_eArgError, "port not found in: `%s'", host_ptr);
481 port = strtoul(colon + 1, &check, 10);
482 *colon = 0;
483 rc = inet_pton(af, host_ptr, dst);
484 *colon = ':';
485 if (rbracket) *rbracket = ']';
486 if (*check || ((uint16_t)port != port))
487 rb_raise(rb_eArgError, "invalid port: %s", colon + 1);
488 if (rc != 1)
489 rb_raise(rb_eArgError, "inet_pton failed for: `%s' with %d",
490 host_ptr, rc);
491 *portdst = ntohs((uint16_t)port);
494 /* generates inet_diag bytecode to match all addrs */
495 static void gen_bytecode_all(struct iovec *iov)
497 struct inet_diag_bc_op *op;
498 struct inet_diag_hostcond *cond;
500 /* iov_len was already set and base allocated in a parent function */
501 assert(iov->iov_len == OPLEN && iov->iov_base && "iov invalid");
502 op = iov->iov_base;
503 op->code = INET_DIAG_BC_S_COND;
504 op->yes = OPLEN;
505 op->no = sizeof(struct inet_diag_bc_op) + OPLEN;
506 cond = (struct inet_diag_hostcond *)(op + 1);
507 cond->family = AF_UNSPEC;
508 cond->port = -1;
509 cond->prefix_len = 0;
512 /* generates inet_diag bytecode to match a single addr */
513 static void gen_bytecode(struct iovec *iov, union any_addr *inet)
515 struct inet_diag_bc_op *op;
516 struct inet_diag_hostcond *cond;
518 /* iov_len was already set and base allocated in a parent function */
519 assert(iov->iov_len == OPLEN && iov->iov_base && "iov invalid");
520 op = iov->iov_base;
521 op->code = INET_DIAG_BC_S_COND;
522 op->yes = OPLEN;
523 op->no = sizeof(struct inet_diag_bc_op) + OPLEN;
525 cond = (struct inet_diag_hostcond *)(op + 1);
526 cond->family = inet->ss.ss_family;
527 switch (inet->ss.ss_family) {
528 case AF_INET: {
529 cond->port = ntohs(inet->in.sin_port);
530 cond->prefix_len = inet->in.sin_addr.s_addr == 0 ? 0 :
531 sizeof(inet->in.sin_addr.s_addr) * CHAR_BIT;
532 *cond->addr = inet->in.sin_addr.s_addr;
534 break;
535 case AF_INET6: {
536 cond->port = ntohs(inet->in6.sin6_port);
537 cond->prefix_len = memcmp(&in6addr_any, &inet->in6.sin6_addr,
538 sizeof(struct in6_addr)) == 0 ?
539 0 : sizeof(inet->in6.sin6_addr) * CHAR_BIT;
540 memcpy(&cond->addr, &inet->in6.sin6_addr,
541 sizeof(struct in6_addr));
543 break;
544 default:
545 assert(0 && "unsupported address family, could that be IPv7?!");
549 static void nl_errcheck(VALUE r)
551 const char *err = (const char *)r;
553 if (err) {
554 if (err == err_nlmsg)
555 rb_raise(rb_eRuntimeError, "NLMSG_ERROR");
556 else
557 rb_sys_fail(err);
561 static VALUE tcp_stats(struct nogvl_args *args, VALUE addr)
563 union any_addr query_addr;
565 parse_addr(&query_addr, addr);
566 gen_bytecode(&args->iov[2], &query_addr);
568 memset(&args->stats, 0, sizeof(struct listen_stats));
569 nl_errcheck(rb_thread_io_blocking_region(diag, args, args->fd));
571 return rb_listen_stats(&args->stats);
575 * call-seq:
576 * Raindrops::Linux.tcp_listener_stats([addrs[, sock]]) => hash
578 * If specified, +addr+ may be a string or array of strings representing
579 * listen addresses to filter for. Returns a hash with given addresses as
580 * keys and ListenStats objects as the values or a hash of all addresses.
582 * addrs = %w(0.0.0.0:80 127.0.0.1:8080)
584 * If +addr+ is nil or not specified, all (IPv4) addresses are returned.
585 * If +sock+ is specified, it should be a Raindrops::InetDiagSock object.
587 static VALUE tcp_listener_stats(int argc, VALUE *argv, VALUE self)
589 VALUE *ary;
590 long i;
591 VALUE rv = rb_hash_new();
592 struct nogvl_args args;
593 VALUE addrs, sock;
595 rb_scan_args(argc, argv, "02", &addrs, &sock);
598 * allocating page_size instead of OP_LEN since we'll reuse the
599 * buffer for recvmsg() later, we already checked for
600 * OPLEN <= page_size at initialization
602 args.iov[2].iov_len = OPLEN;
603 args.iov[2].iov_base = alloca(page_size);
604 args.table = NULL;
605 if (NIL_P(sock))
606 sock = rb_funcall(cIDSock, id_new, 0);
607 args.fd = my_fileno(sock);
609 switch (TYPE(addrs)) {
610 case T_STRING:
611 rb_hash_aset(rv, addrs, tcp_stats(&args, addrs));
612 return rv;
613 case T_ARRAY:
614 ary = RARRAY_PTR(addrs);
615 i = RARRAY_LEN(addrs);
616 if (i == 1) {
617 rb_hash_aset(rv, *ary, tcp_stats(&args, *ary));
618 return rv;
620 for (; --i >= 0; ary++) {
621 union any_addr check;
623 parse_addr(&check, *ary);
624 rb_hash_aset(rv, *ary, Qtrue);
626 /* fall through */
627 case T_NIL:
628 args.table = st_init_strtable();
629 gen_bytecode_all(&args.iov[2]);
630 break;
631 default:
632 rb_raise(rb_eArgError,
633 "addr must be an array of strings, a string, or nil");
636 nl_errcheck(rb_thread_io_blocking_region(diag, &args, args.fd));
638 st_foreach(args.table, NIL_P(addrs) ? st_to_hash : st_AND_hash, rv);
639 st_free_table(args.table);
641 /* let GC deal with corner cases */
642 if (argc < 2) rb_io_close(sock);
643 return rv;
646 void Init_raindrops_linux_inet_diag(void)
648 VALUE cRaindrops = rb_const_get(rb_cObject, rb_intern("Raindrops"));
649 VALUE mLinux = rb_define_module_under(cRaindrops, "Linux");
651 rb_require("socket");
652 cIDSock = rb_const_get(rb_cObject, rb_intern("Socket"));
653 id_new = rb_intern("new");
656 * Document-class: Raindrops::InetDiagSocket
658 * This is a subclass of +Socket+ specifically for talking
659 * to the inet_diag facility of Netlink.
661 cIDSock = rb_define_class_under(cRaindrops, "InetDiagSocket", cIDSock);
662 rb_define_singleton_method(cIDSock, "new", ids_s_new, 0);
664 cListenStats = rb_const_get(cRaindrops, rb_intern("ListenStats"));
666 rb_define_module_function(mLinux, "tcp_listener_stats",
667 tcp_listener_stats, -1);
669 page_size = getpagesize();
671 assert(OPLEN <= page_size && "bytecode OPLEN is not <= PAGE_SIZE");
673 #endif /* __linux__ */