ext/raindrops/linux_inet_diag.c

   1 #include <ruby.h>
   2 #ifdef __linux__
   3
   4 /* Ruby 1.8.6+ macros (for compatibility with Ruby 1.9) */
   5 #ifndef RSTRING_PTR
   6 #  define RSTRING_PTR(s) (RSTRING(s)->ptr)
   7 #endif
   8 #ifndef RSTRING_LEN
   9 #  define RSTRING_LEN(s) (RSTRING(s)->len)
  10 #endif
  11 #ifdef RSTRUCT
  12 #  ifndef RSTRUCT_PTR
  13 #    define RSTRUCT_PTR(s) (RSTRUCT(s)->ptr)
  14 #   endif
  15 #  ifndef RSTRUCT_LEN
  16 #    define RSTRUCT_LEN(s) (RSTRUCT(s)->len)
  17 #  endif
  18 #endif
  19
  20 #ifndef HAVE_RB_STRUCT_ALLOC_NOINIT
  21 static ID id_new;
  22 static VALUE rb_struct_alloc_noinit(VALUE class)
  23 {
  24         return rb_funcall(class, id_new, 0, 0);
  25 }
  26 #endif /* !defined(HAVE_RB_STRUCT_ALLOC_NOINIT) */
  27
  28 /* partial emulation of the 1.9 rb_thread_blocking_region under 1.8 */
  29 #ifndef HAVE_RB_THREAD_BLOCKING_REGION
  30 #  include <rubysig.h>
  31 #  define RUBY_UBF_IO ((rb_unblock_function_t *)-1)
  32 typedef void rb_unblock_function_t(void *);
  33 typedef VALUE rb_blocking_function_t(void *);
  34 static VALUE
  35 rb_thread_blocking_region(
  36         rb_blocking_function_t *func, void *data1,
  37         rb_unblock_function_t *ubf, void *data2)
  38 {
  39         VALUE rv;
  40
  41         TRAP_BEG;
  42         rv = func(data1);
  43         TRAP_END;
  44
  45         return rv;
  46 }
  47 #endif /* ! HAVE_RB_THREAD_BLOCKING_REGION */
  48
  49 #include <assert.h>
  50 #include <errno.h>
  51 #include <sys/socket.h>
  52 #include <sys/types.h>
  53 #include <netdb.h>
  54 #include <unistd.h>
  55 #include <string.h>
  56 #include <asm/types.h>
  57 #include <netinet/in.h>
  58 #include <arpa/inet.h>
  59 #include <netinet/tcp.h>
  60 #include <linux/netlink.h>
  61 #include <linux/rtnetlink.h>
  62 #include <linux/inet_diag.h>
  63
  64 static size_t page_size;
  65 static unsigned g_seq;
  66 static VALUE cListenStats;
  67
  68 struct listen_stats {
  69         long active;
  70         long queued;
  71 };
  72
  73 #define OPLEN (sizeof(struct inet_diag_bc_op) + \
  74                sizeof(struct inet_diag_hostcond) + \
  75                sizeof(struct sockaddr_storage))
  76
  77 struct nogvl_args {
  78         struct iovec iov[3]; /* last iov holds inet_diag bytecode */
  79         struct sockaddr_storage query_addr;
  80         struct listen_stats stats;
  81 };
  82
  83 /* creates a Ruby ListenStats Struct based on our internal listen_stats */
  84 static VALUE rb_listen_stats(struct listen_stats *stats)
  85 {
  86         VALUE rv = rb_struct_alloc_noinit(cListenStats);
  87         VALUE active = LONG2NUM(stats->active);
  88         VALUE queued = LONG2NUM(stats->queued);
  89
  90 #ifdef RSTRUCT_PTR
  91         VALUE *ptr = RSTRUCT_PTR(rv);
  92         ptr[0] = active;
  93         ptr[1] = queued;
  94 #else /* Rubinius */
  95         rb_funcall(rv, rb_intern("active="), 1, active);
  96         rb_funcall(rv, rb_intern("queued="), 1, queued);
  97 #endif /* ! Rubinius */
  98         return rv;
  99 }
 100
 101 /* inner loop of inet_diag, called for every socket returned by netlink */
 102 static inline void r_acc(struct nogvl_args *args, struct inet_diag_msg *r)
 103 {
 104         /*
 105          * inode == 0 means the connection is still in the listen queue
 106          * and has not yet been accept()-ed by the server.  The
 107          * inet_diag bytecode cannot filter this for us.
 108          */
 109         if (r->idiag_inode == 0)
 110                 return;
 111         if (r->idiag_state == TCP_ESTABLISHED)
 112                 args->stats.active++;
 113         else /* if (r->idiag_state == TCP_LISTEN) */
 114                 args->stats.queued = r->idiag_rqueue;
 115         /*
 116          * we wont get anything else because of the idiag_states filter
 117          */
 118 }
 119
 120 static const char err_socket[] = "socket";
 121 static const char err_sendmsg[] = "sendmsg";
 122 static const char err_recvmsg[] = "recvmsg";
 123 static const char err_nlmsg[] = "nlmsg";
 124
 125 /* does the inet_diag stuff with netlink(), this is called w/o GVL */
 126 static VALUE diag(void *ptr)
 127 {
 128         struct nogvl_args *args = ptr;
 129         struct sockaddr_nl nladdr;
 130         struct rtattr rta;
 131         struct {
 132                 struct nlmsghdr nlh;
 133                 struct inet_diag_req r;
 134         } req;
 135         struct msghdr msg;
 136         const char *err = NULL;
 137         unsigned seq = ++g_seq; /* not atomic, rely on GVL for now */
 138         int fd = socket(AF_NETLINK, SOCK_RAW, NETLINK_INET_DIAG);
 139
 140         if (fd < 0)
 141                 return (VALUE)err_socket;
 142
 143         memset(&args->stats, 0, sizeof(struct listen_stats));
 144
 145         memset(&nladdr, 0, sizeof(nladdr));
 146         nladdr.nl_family = AF_NETLINK;
 147
 148         memset(&req, 0, sizeof(req));
 149         req.nlh.nlmsg_len = sizeof(req) + RTA_LENGTH(args->iov[2].iov_len);
 150         req.nlh.nlmsg_type = TCPDIAG_GETSOCK;
 151         req.nlh.nlmsg_flags = NLM_F_ROOT | NLM_F_MATCH | NLM_F_REQUEST;
 152         req.nlh.nlmsg_pid = getpid();
 153         req.nlh.nlmsg_seq = seq;
 154         req.r.idiag_family = AF_INET | AF_INET6;
 155         req.r.idiag_states = (1<<TCP_ESTABLISHED) | (1<<TCP_LISTEN);
 156         rta.rta_type = INET_DIAG_REQ_BYTECODE;
 157         rta.rta_len = RTA_LENGTH(args->iov[2].iov_len);
 158
 159         args->iov[0].iov_base = &req;
 160         args->iov[0].iov_len = sizeof(req);
 161         args->iov[1].iov_base = &rta;
 162         args->iov[1].iov_len = sizeof(rta);
 163
 164         memset(&msg, 0, sizeof(msg));
 165         msg.msg_name = (void *)&nladdr;
 166         msg.msg_namelen = sizeof(nladdr);
 167         msg.msg_iov = args->iov;
 168         msg.msg_iovlen = 3;
 169
 170         if (sendmsg(fd, &msg, 0) < 0) {
 171                 err = err_sendmsg;
 172                 goto out;
 173         }
 174
 175         /* reuse buffer that was allocated for bytecode */
 176         args->iov[0].iov_len = page_size;
 177         args->iov[0].iov_base = args->iov[2].iov_base;
 178
 179         while (1) {
 180                 ssize_t readed;
 181                 struct nlmsghdr *h = (struct nlmsghdr *)args->iov[0].iov_base;
 182
 183                 memset(&msg, 0, sizeof(msg));
 184                 msg.msg_name = (void *)&nladdr;
 185                 msg.msg_namelen = sizeof(nladdr);
 186                 msg.msg_iov = args->iov;
 187                 msg.msg_iovlen = 1;
 188
 189                 readed = recvmsg(fd, &msg, 0);
 190                 if (readed < 0) {
 191                         if (errno == EINTR)
 192                                 continue;
 193                         err = err_recvmsg;
 194                         goto out;
 195                 }
 196                 if (readed == 0)
 197                         goto out;
 198
 199                 for ( ; NLMSG_OK(h, readed); h = NLMSG_NEXT(h, readed)) {
 200                         if (h->nlmsg_seq != seq)
 201                                 continue;
 202                         if (h->nlmsg_type == NLMSG_DONE)
 203                                 goto out;
 204                         if (h->nlmsg_type == NLMSG_ERROR) {
 205                                 err = err_nlmsg;
 206                                 goto out;
 207                         }
 208                         r_acc(args, NLMSG_DATA(h));
 209                 }
 210         }
 211 out:
 212         {
 213                 int save_errno = errno;
 214                 close(fd);
 215                 errno = save_errno;
 216         }
 217         return (VALUE)err;
 218 }
 219
 220 /* populates sockaddr_storage struct by parsing +addr+ */
 221 static void parse_addr(struct sockaddr_storage *inet, VALUE addr)
 222 {
 223         char *host_ptr;
 224         char *colon = NULL;
 225         char *rbracket = NULL;
 226         long host_len;
 227         struct addrinfo hints;
 228         struct addrinfo *res;
 229         int rc;
 230
 231         if (TYPE(addr) != T_STRING)
 232                 rb_raise(rb_eArgError, "addrs must be an Array of Strings");
 233
 234         host_ptr = StringValueCStr(addr);
 235         host_len = RSTRING_LEN(addr);
 236         if (*host_ptr == '[') { /* ipv6 address format (rfc2732) */
 237                 rbracket = memchr(host_ptr + 1, ']', host_len - 1);
 238
 239                 if (rbracket) {
 240                         if (rbracket[1] == ':') {
 241                                 colon = rbracket + 1;
 242                                 host_ptr++;
 243                                 *rbracket = 0;
 244                         } else {
 245                                 rbracket = NULL;
 246                         }
 247                 }
 248         } else { /* ipv4 */
 249                 colon = memchr(host_ptr, ':', host_len);
 250         }
 251
 252         if (!colon)
 253                 rb_raise(rb_eArgError, "port not found in: `%s'", host_ptr);
 254
 255         hints.ai_family = AF_UNSPEC;
 256         hints.ai_socktype = SOCK_STREAM;
 257         hints.ai_protocol = IPPROTO_TCP;
 258         hints.ai_flags = AI_NUMERICHOST | AI_NUMERICSERV;
 259
 260         *colon = 0;
 261         if (rbracket) *rbracket = 0;
 262         rc = getaddrinfo(host_ptr, colon + 1, &hints, &res);
 263         *colon = ':';
 264         if (rbracket) *rbracket = ']';
 265         if (rc != 0)
 266                 rb_raise(rb_eArgError, "getaddrinfo(%s): %s",
 267                          host_ptr, gai_strerror(rc));
 268
 269         memcpy(inet, res->ai_addr, res->ai_addrlen);
 270         freeaddrinfo(res);
 271 }
 272
 273 /* generates inet_diag bytecode to match a single addr */
 274 static void gen_bytecode(struct iovec *iov, struct sockaddr_storage *inet)
 275 {
 276         struct inet_diag_bc_op *op;
 277         struct inet_diag_hostcond *cond;
 278
 279         /* iov_len was already set and base allocated in a parent function */
 280         assert(iov->iov_len == OPLEN && iov->iov_base && "iov invalid");
 281         op = iov->iov_base;
 282         op->code = INET_DIAG_BC_S_COND;
 283         op->yes = OPLEN;
 284         op->no = sizeof(struct inet_diag_bc_op) + OPLEN;
 285
 286         cond = (struct inet_diag_hostcond *)(op + 1);
 287         cond->family = inet->ss_family;
 288         switch (inet->ss_family) {
 289         case AF_INET: {
 290                 struct sockaddr_in *in = (struct sockaddr_in *)inet;
 291
 292                 cond->port = ntohs(in->sin_port);
 293                 cond->prefix_len = in->sin_addr.s_addr == 0 ? 0 :
 294                                    sizeof(in->sin_addr.s_addr) * CHAR_BIT;
 295                 *cond->addr = in->sin_addr.s_addr;
 296                 }
 297                 break;
 298         case AF_INET6: {
 299                 struct sockaddr_in6 *in6 = (struct sockaddr_in6 *)inet;
 300
 301                 cond->port = ntohs(in6->sin6_port);
 302                 cond->prefix_len = memcmp(&in6addr_any, &in6->sin6_addr,
 303                                           sizeof(struct in6_addr)) == 0 ?
 304                                   0 : sizeof(in6->sin6_addr) * CHAR_BIT;
 305                 memcpy(&cond->addr, &in6->sin6_addr, sizeof(struct in6_addr));
 306                 }
 307                 break;
 308         default:
 309                 assert("unsupported address family, could that be IPv7?!");
 310         }
 311 }
 312
 313 static VALUE tcp_stats(struct nogvl_args *args, VALUE addr)
 314 {
 315         const char *err;
 316         VALUE verr;
 317
 318         parse_addr(&args->query_addr, addr);
 319         gen_bytecode(&args->iov[2], &args->query_addr);
 320
 321         verr = rb_thread_blocking_region(diag, args, RUBY_UBF_IO, 0);
 322         err = (const char *)verr;
 323         if (err) {
 324                 if (err == err_nlmsg)
 325                         rb_raise(rb_eRuntimeError, "NLMSG_ERROR");
 326                 else
 327                         rb_sys_fail(err);
 328         }
 329
 330         return rb_listen_stats(&args->stats);
 331 }
 332
 333 /*
 334  * call-seq:
 335  *      addrs = %w(0.0.0.0:80 127.0.0.1:8080)
 336  *      Raindrops::Linux.tcp_listener_stats(addrs) => hash
 337  *
 338  * Takes an array of strings representing listen addresses to filter for.
 339  * Returns a hash with given addresses as keys and ListenStats
 340  * objects as the values.
 341  */
 342 static VALUE tcp_listener_stats(VALUE obj, VALUE addrs)
 343 {
 344         VALUE *ary;
 345         long i;
 346         VALUE rv;
 347         struct nogvl_args args;
 348
 349         /*
 350          * allocating page_size instead of OP_LEN since we'll reuse the
 351          * buffer for recvmsg() later, we already checked for
 352          * OPLEN <= page_size at initialization
 353          */
 354         args.iov[2].iov_len = OPLEN;
 355         args.iov[2].iov_base = alloca(page_size);
 356
 357         if (TYPE(addrs) != T_ARRAY)
 358                 rb_raise(rb_eArgError, "addrs must be an Array of Strings");
 359
 360         rv = rb_hash_new();
 361         ary = RARRAY_PTR(addrs);
 362         for (i = RARRAY_LEN(addrs); --i >= 0; ary++)
 363                 rb_hash_aset(rv, *ary, tcp_stats(&args, *ary));
 364
 365         return rv;
 366 }
 367
 368 void Init_raindrops_linux_inet_diag(void)
 369 {
 370         VALUE cRaindrops = rb_const_get(rb_cObject, rb_intern("Raindrops"));
 371         VALUE mLinux = rb_define_module_under(cRaindrops, "Linux");
 372
 373         cListenStats = rb_const_get(cRaindrops, rb_intern("ListenStats"));
 374
 375         rb_define_module_function(mLinux, "tcp_listener_stats",
 376                                   tcp_listener_stats, 1);
 377
 378 #ifndef HAVE_RB_STRUCT_ALLOC_NOINIT
 379         id_new = rb_intern("new");
 380 #endif
 381         page_size = getpagesize();
 382
 383         assert(OPLEN <= page_size && "bytecode OPLEN is not <= PAGE_SIZE");
 384 }
 385 #endif /* __linux__ */