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