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