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