ext/raindrops/linux_inet_diag.c

   1 #include <ruby.h>
   2
   3 /* Ruby 1.8.6+ macros (for compatibility with Ruby 1.9) */
   4 #ifndef RSTRING_PTR
   5 #  define RSTRING_PTR(s) (RSTRING(s)->ptr)
   6 #endif
   7 #ifndef RSTRING_LEN
   8 #  define RSTRING_LEN(s) (RSTRING(s)->len)
   9 #endif
  10 #ifndef RSTRUCT_PTR
  11 #  define RSTRUCT_PTR(s) (RSTRUCT(s)->ptr)
  12 #endif
  13 #ifndef RSTRUCT_LEN
  14 #  define RSTRUCT_LEN(s) (RSTRUCT(s)->len)
  15 #endif
  16
  17 #ifndef HAVE_RB_STRUCT_ALLOC_NOINIT
  18 static ID id_new;
  19 static VALUE rb_struct_alloc_noinit(VALUE class)
  20 {
  21         return rb_funcall(class, id_new, 0, 0);
  22 }
  23 #endif /* !defined(HAVE_RB_STRUCT_ALLOC_NOINIT) */
  24
  25 /* partial emulation of the 1.9 rb_thread_blocking_region under 1.8 */
  26 #ifndef HAVE_RB_THREAD_BLOCKING_REGION
  27 #  include <rubysig.h>
  28 #  define RUBY_UBF_IO ((rb_unblock_function_t *)-1)
  29 typedef void rb_unblock_function_t(void *);
  30 typedef VALUE rb_blocking_function_t(void *);
  31 static VALUE
  32 rb_thread_blocking_region(
  33         rb_blocking_function_t *func, void *data1,
  34         rb_unblock_function_t *ubf, void *data2)
  35 {
  36         VALUE rv;
  37
  38         TRAP_BEG;
  39         rv = func(data1);
  40         TRAP_END;
  41
  42         return rv;
  43 }
  44 #endif /* ! HAVE_RB_THREAD_BLOCKING_REGION */
  45
  46 #include <assert.h>
  47 #include <errno.h>
  48 #include <sys/socket.h>
  49 #include <sys/types.h>
  50 #include <unistd.h>
  51 #include <string.h>
  52 #include <asm/types.h>
  53 #include <netinet/in.h>
  54 #include <arpa/inet.h>
  55 #include <netinet/tcp.h>
  56 #include <linux/netlink.h>
  57 #include <linux/rtnetlink.h>
  58 #include <linux/inet_diag.h>
  59
  60 static size_t page_size;
  61 static unsigned g_seq;
  62 static VALUE cListenStats;
  63
  64 struct my_addr {
  65         in_addr_t addr;
  66         uint16_t port;
  67 };
  68
  69 struct listen_stats {
  70         long active;
  71         long queued;
  72 };
  73
  74 #define OPLEN (sizeof(struct inet_diag_bc_op) + \
  75                sizeof(struct inet_diag_hostcond) + \
  76                sizeof(in_addr_t))
  77
  78 struct nogvl_args {
  79         struct iovec iov[3]; /* last iov holds inet_diag bytecode */
  80         struct my_addr addrs;
  81         struct listen_stats stats;
  82 };
  83
  84 /* creates a Ruby ListenStats Struct based on our internal listen_stats */
  85 static VALUE rb_listen_stats(struct listen_stats *stats)
  86 {
  87         VALUE rv = rb_struct_alloc_noinit(cListenStats);
  88         VALUE *ptr = RSTRUCT_PTR(rv);
  89
  90         ptr[0] = LONG2NUM(stats->active);
  91         ptr[1] = LONG2NUM(stats->queued);
  92
  93         return rv;
  94 }
  95
  96 /*
  97  * converts a base 10 string representing a port number into
  98  * an unsigned 16 bit integer.  Raises ArgumentError on failure
  99  */
 100 static uint16_t my_inet_port(const char *port)
 101 {
 102         char *err;
 103         unsigned long tmp = strtoul(port, &err, 10);
 104
 105         if (*err != 0 || tmp > 0xffff)
 106                 rb_raise(rb_eArgError, "port not parsable: `%s'\n", port);
 107
 108         return (uint16_t)tmp;
 109 }
 110
 111 /* inner loop of inet_diag, called for every socket returned by netlink */
 112 static inline void r_acc(struct nogvl_args *args, struct inet_diag_msg *r)
 113 {
 114         /*
 115          * inode == 0 means the connection is still in the listen queue
 116          * and has not yet been accept()-ed by the server.  The
 117          * inet_diag bytecode cannot filter this for us.
 118          */
 119         if (r->idiag_inode == 0)
 120                 return;
 121         if (r->idiag_state == TCP_ESTABLISHED)
 122                 args->stats.active++;
 123         else /* if (r->idiag_state == TCP_LISTEN) */
 124                 args->stats.queued = r->idiag_rqueue;
 125         /*
 126          * we wont get anything else because of the idiag_states filter
 127          */
 128 }
 129
 130 static const char err_socket[] = "socket";
 131 static const char err_sendmsg[] = "sendmsg";
 132 static const char err_recvmsg[] = "recvmsg";
 133 static const char err_nlmsg[] = "nlmsg";
 134
 135 /* does the inet_diag stuff with netlink(), this is called w/o GVL */
 136 static VALUE diag(void *ptr)
 137 {
 138         struct nogvl_args *args = ptr;
 139         struct sockaddr_nl nladdr;
 140         struct rtattr rta;
 141         struct {
 142                 struct nlmsghdr nlh;
 143                 struct inet_diag_req r;
 144         } req;
 145         struct msghdr msg;
 146         const char *err = NULL;
 147         unsigned seq = __sync_add_and_fetch(&g_seq, 1);
 148         int fd = socket(AF_NETLINK, SOCK_RAW, NETLINK_INET_DIAG);
 149
 150         if (fd < 0)
 151                 return (VALUE)err_socket;
 152
 153         memset(&args->stats, 0, sizeof(struct listen_stats));
 154
 155         memset(&nladdr, 0, sizeof(nladdr));
 156         nladdr.nl_family = AF_NETLINK;
 157
 158         memset(&req, 0, sizeof(req));
 159         req.nlh.nlmsg_len = sizeof(req) + RTA_LENGTH(args->iov[2].iov_len);
 160         req.nlh.nlmsg_type = TCPDIAG_GETSOCK;
 161         req.nlh.nlmsg_flags = NLM_F_ROOT | NLM_F_MATCH | NLM_F_REQUEST;
 162         req.nlh.nlmsg_pid = getpid();
 163         req.nlh.nlmsg_seq = seq;
 164         req.r.idiag_family = AF_INET;
 165         req.r.idiag_states = (1<<TCP_ESTABLISHED) | (1<<TCP_LISTEN);
 166         rta.rta_type = INET_DIAG_REQ_BYTECODE;
 167         rta.rta_len = RTA_LENGTH(args->iov[2].iov_len);
 168
 169         args->iov[0].iov_base = &req;
 170         args->iov[0].iov_len = sizeof(req);
 171         args->iov[1].iov_base = &rta;
 172         args->iov[1].iov_len = sizeof(rta);
 173
 174         memset(&msg, 0, sizeof(msg));
 175         msg.msg_name = (void *)&nladdr;
 176         msg.msg_namelen = sizeof(nladdr);
 177         msg.msg_iov = args->iov;
 178         msg.msg_iovlen = 3;
 179
 180         if (sendmsg(fd, &msg, 0) < 0) {
 181                 err = err_sendmsg;
 182                 goto out;
 183         }
 184
 185         /* reuse buffer that was allocated for bytecode */
 186         args->iov[0].iov_len = page_size;
 187         args->iov[0].iov_base = args->iov[2].iov_base;
 188
 189         while (1) {
 190                 ssize_t readed;
 191                 struct nlmsghdr *h = (struct nlmsghdr *)args->iov[0].iov_base;
 192
 193                 memset(&msg, 0, sizeof(msg));
 194                 msg.msg_name = (void *)&nladdr;
 195                 msg.msg_namelen = sizeof(nladdr);
 196                 msg.msg_iov = args->iov;
 197                 msg.msg_iovlen = 1;
 198
 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 inet my_addr struct by parsing +addr+ */
 231 static void parse_addr(struct my_addr *inet, VALUE addr)
 232 {
 233         char *host_port, *colon;
 234
 235         if (TYPE(addr) != T_STRING)
 236                 rb_raise(rb_eArgError, "addrs must be an Array of Strings");
 237
 238         host_port = RSTRING_PTR(addr);
 239         colon = memchr(host_port, ':', RSTRING_LEN(addr));
 240         if (!colon)
 241                 rb_raise(rb_eArgError, "port not found in: `%s'", host_port);
 242
 243         *colon = 0;
 244         inet->addr = inet_addr(host_port);
 245         *colon = ':';
 246         inet->port = htons(my_inet_port(colon + 1));
 247 }
 248
 249 /* generates inet_diag bytecode to match a single addr */
 250 static void gen_bytecode(struct iovec *iov, struct my_addr *inet)
 251 {
 252         struct inet_diag_bc_op *op;
 253         struct inet_diag_hostcond *cond;
 254
 255         /* iov_len was already set and base allocated in a parent function */
 256         assert(iov->iov_len == OPLEN && iov->iov_base && "iov invalid");
 257         op = iov->iov_base;
 258         op->code = INET_DIAG_BC_S_COND;
 259         op->yes = OPLEN;
 260         op->no = sizeof(struct inet_diag_bc_op) + OPLEN;
 261
 262         cond = (struct inet_diag_hostcond *)(op + 1);
 263         cond->family = AF_INET;
 264         cond->port = ntohs(inet->port);
 265         cond->prefix_len = inet->addr == 0 ? 0 : sizeof(in_addr_t) * CHAR_BIT;
 266         *cond->addr = inet->addr;
 267 }
 268
 269 static VALUE tcp_stats(struct nogvl_args *args, VALUE addr)
 270 {
 271         const char *err;
 272         VALUE verr;
 273
 274         parse_addr(&args->addrs, addr);
 275         gen_bytecode(&args->iov[2], &args->addrs);
 276
 277         verr = rb_thread_blocking_region(diag, args, RUBY_UBF_IO, 0);
 278         err = (const char *)verr;
 279         if (err) {
 280                 if (err == err_nlmsg)
 281                         rb_raise(rb_eRuntimeError, "NLMSG_ERROR");
 282                 else
 283                         rb_sys_fail(err);
 284         }
 285
 286         return rb_listen_stats(&args->stats);
 287 }
 288
 289 /*
 290  * call-seq:
 291  *      addrs = %w(0.0.0.0:80 127.0.0.1:8080)
 292  *      Raindrops::Linux.tcp_listener_stats(addrs) => hash
 293  *
 294  * Takes an array of strings representing listen addresses to filter for.
 295  * Returns a hash with given addresses as keys and ListenStats
 296  * objects as the values.
 297  */
 298 static VALUE tcp_listener_stats(VALUE obj, VALUE addrs)
 299 {
 300         VALUE *ary;
 301         long i;
 302         VALUE rv;
 303         struct nogvl_args args;
 304
 305         /*
 306          * allocating page_size instead of OP_LEN since we'll reuse the
 307          * buffer for recvmsg() later, we already checked for
 308          * OPLEN <= page_size at initialization
 309          */
 310         args.iov[2].iov_len = OPLEN;
 311         args.iov[2].iov_base = alloca(page_size);
 312
 313         if (TYPE(addrs) != T_ARRAY)
 314                 rb_raise(rb_eArgError, "addrs must be an Array or String");
 315
 316         rv = rb_hash_new();
 317         ary = RARRAY_PTR(addrs);
 318         for (i = RARRAY_LEN(addrs); --i >= 0; ary++)
 319                 rb_hash_aset(rv, *ary, tcp_stats(&args, *ary));
 320
 321         return rv;
 322 }
 323
 324 void Init_raindrops_linux_inet_diag(void)
 325 {
 326         VALUE cRaindrops = rb_const_get(rb_cObject, rb_intern("Raindrops"));
 327         VALUE mLinux = rb_define_module_under(cRaindrops, "Linux");
 328
 329         cListenStats = rb_const_get(cRaindrops, rb_intern("ListenStats"));
 330
 331         rb_define_module_function(mLinux, "tcp_listener_stats",
 332                                   tcp_listener_stats, 1);
 333
 334 #ifndef HAVE_RB_STRUCT_ALLOC_NOINIT
 335         id_new = rb_intern("new");
 336 #endif
 337         rb_require("raindrops/linux");
 338
 339         page_size = getpagesize();
 340
 341         assert(OPLEN <= page_size && "bytecode OPLEN is no <= PAGE_SIZE");
 342 }