inet_diag: force the use of 32-bit counters
[raindrops.git] / ext / raindrops / linux_inet_diag.c
blob0eff08d459cfabae4877e6d85976e6a3df03ac19
1 #include <ruby.h>
2 #ifdef __linux__
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
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)
20 VALUE rv;
22 TRAP_BEG;
23 rv = func(data1);
24 TRAP_END;
26 return rv;
28 #endif /* ! HAVE_RB_THREAD_BLOCKING_REGION */
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>
45 static size_t page_size;
46 static unsigned g_seq;
47 static VALUE cListenStats;
49 struct listen_stats {
50 uint32_t active;
51 uint32_t queued;
54 #define OPLEN (sizeof(struct inet_diag_bc_op) + \
55 sizeof(struct inet_diag_hostcond) + \
56 sizeof(struct sockaddr_storage))
58 struct nogvl_args {
59 struct iovec iov[3]; /* last iov holds inet_diag bytecode */
60 struct listen_stats stats;
63 /* creates a Ruby ListenStats Struct based on our internal listen_stats */
64 static VALUE rb_listen_stats(struct listen_stats *stats)
66 VALUE active = UINT2NUM(stats->active);
67 VALUE queued = UINT2NUM(stats->queued);
69 return rb_struct_new(cListenStats, active, queued);
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)
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.
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;
87 * we wont get anything else because of the idiag_states filter
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";
96 struct diag_req {
97 struct nlmsghdr nlh;
98 struct inet_diag_req r;
101 static void prep_msghdr(
102 struct msghdr *msg,
103 struct nogvl_args *args,
104 struct sockaddr_nl *nladdr,
105 size_t iovlen)
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;
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)
121 struct inet_diag_bc_op *op = args->iov[2].iov_base;
122 struct inet_diag_hostcond *cond = (struct inet_diag_hostcond *)(op + 1);
124 memset(req, 0, sizeof(struct diag_req));
125 memset(nladdr, 0, sizeof(struct sockaddr_nl));
127 nladdr->nl_family = AF_NETLINK;
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);
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);
144 prep_msghdr(msg, args, nladdr, 3);
147 static void prep_recvmsg_buf(struct nogvl_args *args)
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;
154 /* does the inet_diag stuff with netlink(), this is called w/o GVL */
155 static VALUE diag(void *ptr)
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);
166 if (fd < 0)
167 return (VALUE)err_socket;
169 prep_diag_args(args, &nladdr, &rta, &req, &msg);
170 req.nlh.nlmsg_seq = seq;
172 if (sendmsg(fd, &msg, 0) < 0) {
173 err = err_sendmsg;
174 goto out;
177 prep_recvmsg_buf(args);
179 while (1) {
180 ssize_t readed;
181 struct nlmsghdr *h = (struct nlmsghdr *)args->iov[0].iov_base;
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;
191 if (readed == 0)
192 goto out;
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;
203 r_acc(args, NLMSG_DATA(h));
206 out:
208 int save_errno = errno;
209 close(fd);
210 errno = save_errno;
212 return (VALUE)err;
215 /* populates sockaddr_storage struct by parsing +addr+ */
216 static void parse_addr(struct sockaddr_storage *inet, VALUE addr)
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;
226 if (TYPE(addr) != T_STRING)
227 rb_raise(rb_eArgError, "addrs must be an Array of Strings");
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);
234 if (rbracket) {
235 if (rbracket[1] == ':') {
236 colon = rbracket + 1;
237 host_ptr++;
238 *rbracket = 0;
239 } else {
240 rbracket = NULL;
243 } else { /* ipv4 */
244 colon = memchr(host_ptr, ':', host_len);
247 if (!colon)
248 rb_raise(rb_eArgError, "port not found in: `%s'", host_ptr);
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;
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));
264 memcpy(inet, res->ai_addr, res->ai_addrlen);
265 freeaddrinfo(res);
268 /* generates inet_diag bytecode to match a single addr */
269 static void gen_bytecode(struct iovec *iov, struct sockaddr_storage *inet)
271 struct inet_diag_bc_op *op;
272 struct inet_diag_hostcond *cond;
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;
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;
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;
292 break;
293 case AF_INET6: {
294 struct sockaddr_in6 *in6 = (struct sockaddr_in6 *)inet;
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));
302 break;
303 default:
304 assert(0 && "unsupported address family, could that be IPv7?!");
308 static VALUE tcp_stats(struct nogvl_args *args, VALUE addr)
310 const char *err;
311 VALUE verr;
312 struct sockaddr_storage query_addr;
314 parse_addr(&query_addr, addr);
315 gen_bytecode(&args->iov[2], &query_addr);
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);
327 return rb_listen_stats(&args->stats);
331 * call-seq:
332 * addrs = %w(0.0.0.0:80 127.0.0.1:8080)
333 * Raindrops::Linux.tcp_listener_stats(addrs) => hash
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.
339 static VALUE tcp_listener_stats(VALUE obj, VALUE addrs)
341 VALUE *ary;
342 long i;
343 VALUE rv;
344 struct nogvl_args args;
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
351 args.iov[2].iov_len = OPLEN;
352 args.iov[2].iov_base = alloca(page_size);
354 if (TYPE(addrs) != T_ARRAY)
355 rb_raise(rb_eArgError, "addrs must be an Array of Strings");
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));
362 return rv;
365 void Init_raindrops_linux_inet_diag(void)
367 VALUE cRaindrops = rb_const_get(rb_cObject, rb_intern("Raindrops"));
368 VALUE mLinux = rb_define_module_under(cRaindrops, "Linux");
370 cListenStats = rb_const_get(cRaindrops, rb_intern("ListenStats"));
372 rb_define_module_function(mLinux, "tcp_listener_stats",
373 tcp_listener_stats, 1);
375 page_size = getpagesize();
377 assert(OPLEN <= page_size && "bytecode OPLEN is not <= PAGE_SIZE");
379 #endif /* __linux__ */