root/src/modules/lua_dns.c

Revision b12f6d214e3da63da73908974a3a36d24353b0c3, 15.1 kB (checked in by Theo Schlossnagle <jesus@omniti.com>, 7 months ago)

lua DNS resolution completion happen on originating thread

  • Property mode set to 100644
Line 
1 /*
2  * Copyright (c) 2010, OmniTI Computer Consulting, Inc.
3  * All rights reserved.
4  *
5  * Redistribution and use in source and binary forms, with or without
6  * modification, are permitted provided that the following conditions are
7  * met:
8  *
9  *     * Redistributions of source code must retain the above copyright
10  *       notice, this list of conditions and the following disclaimer.
11  *     * Redistributions in binary form must reproduce the above
12  *       copyright notice, this list of conditions and the following
13  *       disclaimer in the documentation and/or other materials provided
14  *       with the distribution.
15  *     * Neither the name OmniTI Computer Consulting, Inc. nor the names
16  *       of its contributors may be used to endorse or promote products
17  *       derived from this software without specific prior written
18  *       permission.
19  *
20  * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
21  * "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
22  * LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
23  * A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
24  * OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
25  * SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
26  * LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
27  * DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
28  * THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
29  * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
30  * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
31  */
32
33 #include "noit_defines.h"
34
35 #include <assert.h>
36 #include <math.h>
37 #include <errno.h>
38 #include <unistd.h>
39 #include <ctype.h>
40 #include <string.h>
41 #include <netinet/in.h>
42 #include <arpa/inet.h>
43 #ifdef HAVE_SYS_FILIO_H
44 #include <sys/filio.h>
45 #endif
46
47 #include "noit_conf.h"
48 #include "lua_noit.h"
49 #include "udns/udns.h"
50
51 static noit_hash_table dns_rtypes = NOIT_HASH_EMPTY;
52 static noit_hash_table dns_ctypes = NOIT_HASH_EMPTY;
53 static noit_hash_table dns_ctx_store = NOIT_HASH_EMPTY;
54 static pthread_mutex_t dns_ctx_store_lock;
55
56 typedef struct dns_ctx_handle {
57   char *ns;
58   struct dns_ctx *ctx;
59   noit_atomic32_t refcnt;
60   eventer_t e; /* eventer handling UDP traffic */
61   eventer_t timeout; /* the timeout managed by libudns */
62 } dns_ctx_handle_t;
63
64 typedef struct dns_lookup_ctx {
65   noit_lua_resume_info_t *ci;
66   dns_ctx_handle_t *h;
67   char *error;
68   char *results;
69   int results_len;
70   unsigned char dn[DNS_MAXDN];
71   enum dns_class query_ctype;
72   enum dns_type query_rtype;
73   int active;
74   noit_atomic32_t refcnt;
75 } dns_lookup_ctx_t;
76
77 static dns_ctx_handle_t *default_ctx_handle = NULL;
78
79 static int noit_lua_dns_eventer(eventer_t e, int mask, void *closure,
80                                 struct timeval *now) {
81   dns_ctx_handle_t *h = closure;
82   dns_ioevent(h->ctx, now->tv_sec);
83   return EVENTER_READ | EVENTER_EXCEPTION;
84 }
85
86 static int noit_lua_dns_timeouts(eventer_t e, int mask, void *closure,
87                                  struct timeval *now) {
88   dns_ctx_handle_t *h = closure;
89   dns_timeouts(h->ctx, 0, now->tv_sec);
90   return 0;
91 }
92
93 static void eventer_dns_utm_fn(struct dns_ctx *ctx, int timeout, void *data) {
94   dns_ctx_handle_t *h = data;
95   eventer_t e = NULL, newe = NULL;
96   if(h == NULL) return;
97   if(ctx == NULL) {
98     if(h->timeout) e = eventer_remove(h->timeout);
99   }
100   else {
101     assert(h->ctx == ctx);
102     if(timeout < 0) e = eventer_remove(h->timeout);
103     else {
104       newe = eventer_alloc();
105       newe->mask = EVENTER_TIMER;
106       newe->callback = noit_lua_dns_timeouts;
107       newe->closure = h;
108       gettimeofday(&newe->whence, NULL);
109       newe->whence.tv_sec += timeout;
110     }
111   }
112   if(e) eventer_free(e);
113   if(newe) eventer_add(newe);
114   h->timeout = newe;
115 }
116
117 static void dns_ctx_handle_free(void *vh) {
118   dns_ctx_handle_t *h = vh;
119   free(h->ns);
120   eventer_remove_fd(h->e->fd);
121   eventer_free(h->e);
122   h->e = NULL;
123   if(h->timeout) {
124     eventer_remove(h->timeout);
125     eventer_free(h->timeout);
126     h->timeout = NULL;
127   }
128   dns_close(h->ctx);
129   dns_free(h->ctx);
130   assert(h->timeout == NULL);
131   free(h);
132 }
133
134 static dns_ctx_handle_t *dns_ctx_alloc(const char *ns) {
135   void *vh;
136   dns_ctx_handle_t *h = NULL;
137   pthread_mutex_lock(&dns_ctx_store_lock);
138   if(ns == NULL && default_ctx_handle != NULL) {
139     /* special case -- default context */
140     h = default_ctx_handle;
141     noit_atomic_inc32(&h->refcnt);
142     goto bail;
143   }
144   if(ns &&
145      noit_hash_retrieve(&dns_ctx_store, ns, strlen(ns), &vh)) {
146     h = (dns_ctx_handle_t *)vh;
147     noit_atomic_inc32(&h->refcnt);
148   }
149   else {
150     int failed = 0;
151     h = calloc(1, sizeof(*h));
152     h->ns = ns ? strdup(ns) : NULL;
153     h->ctx = dns_new(NULL);
154     if(dns_init(h->ctx, 0) != 0) failed++;
155     if(ns) {
156       if(dns_add_serv(h->ctx, NULL) < 0) failed++;
157       if(dns_add_serv(h->ctx, ns) < 0) failed++;
158     }
159     if(dns_open(h->ctx) < 0) failed++;
160     if(failed) {
161       noitL(noit_error, "dns_open failed\n");
162       free(h->ns);
163       free(h);
164       h = NULL;
165       goto bail;
166     }
167     dns_set_tmcbck(h->ctx, eventer_dns_utm_fn, h);
168     h->e = eventer_alloc();
169     h->e->mask = EVENTER_READ | EVENTER_EXCEPTION;
170     h->e->closure = h;
171     h->e->callback = noit_lua_dns_eventer;
172     h->e->fd = dns_sock(h->ctx);
173     eventer_add(h->e);
174     h->refcnt = 1;
175     if(!ns)
176       default_ctx_handle = h;
177     else
178       noit_hash_store(&dns_ctx_store, h->ns, strlen(h->ns), h);
179   }
180  bail:
181   pthread_mutex_unlock(&dns_ctx_store_lock);
182   return h;
183 }
184
185 static void dns_ctx_release(dns_ctx_handle_t *h) {
186   if(h->ns == NULL) {
187     /* Special case for the default */
188     noit_atomic_dec32(&h->refcnt);
189     return;
190   }
191   pthread_mutex_lock(&dns_ctx_store_lock);
192   if(noit_atomic_dec32(&h->refcnt) == 0) {
193     /* I was the last one */
194     assert(noit_hash_delete(&dns_ctx_store, h->ns, strlen(h->ns),
195                             NULL, dns_ctx_handle_free));
196   }
197   pthread_mutex_unlock(&dns_ctx_store_lock);
198 }
199
200 void lookup_ctx_release(dns_lookup_ctx_t *v) {
201   if(!v) return;
202   if(v->results) free(v->results);
203   v->results = NULL;
204   if(v->error) free(v->error);
205   v->error = NULL;
206   if(noit_atomic_dec32(&v->refcnt) == 0) {
207     dns_ctx_release(v->h);
208     free(v);
209   }
210 }
211
212 int nl_dns_lookup(lua_State *L) {
213   dns_lookup_ctx_t *dlc, **holder;
214   const char *nameserver = NULL;
215   noit_lua_resume_info_t *ci;
216
217   ci = noit_lua_get_resume_info(L);
218   assert(ci);
219   if(lua_gettop(L) > 0)
220     nameserver = lua_tostring(L, 1);
221   holder = (dns_lookup_ctx_t **)lua_newuserdata(L, sizeof(*holder));
222   dlc = calloc(1, sizeof(*dlc));
223   dlc->refcnt = 1;
224   dlc->ci = ci;
225   dlc->h = dns_ctx_alloc(nameserver);
226   *holder = dlc;
227   luaL_getmetatable(L, "noit.dns");
228   lua_setmetatable(L, -2);
229   return 1;
230 }
231
232 static char *encode_txt(char *dst, const unsigned char *src, int len) {
233   int i;
234   for(i=0; i<len; i++) {
235     if(src[i] >= 127 || src[i] <= 31) {
236       snprintf(dst, 4, "\\%02x", src[i]);
237       dst += 3;
238     }
239     else if(src[i] == '\\') {
240       *dst++ = '\\';
241       *dst++ = '\\';
242     }
243     else {
244       *dst++ = (char)src[i];
245     }
246   }
247   *dst = '\0';
248   return dst;
249 }
250
251 static void dns_resume(dns_lookup_ctx_t *dlc) {
252   int r = dlc->results_len;
253   void *result = dlc->results;
254   struct dns_parse p;
255   struct dns_rr rr;
256   unsigned nrr = 0;
257   unsigned char dn[DNS_MAXDN];
258   const unsigned char *pkt, *cur, *end;
259   lua_State *L;
260
261   dlc->results = NULL; /* We will free 'result' in cleanup */
262   if(!dlc->active) goto cleanup;
263   if(!result) goto cleanup;
264
265   L = dlc->ci->coro_state;
266
267   pkt = result; end = pkt + r; cur = dns_payload(pkt);
268   dns_getdn(pkt, &cur, end, dn, sizeof(dn));
269   dns_initparse(&p, NULL, pkt, cur, end);
270   p.dnsp_qcls = 0;
271   p.dnsp_qtyp = 0;
272
273   while(dns_nextrr(&p, &rr) > 0) {
274     const char *fieldname = NULL;
275     char buff[DNS_MAXDN], *txt_str, *c;
276     int totalsize;
277     const unsigned char *pkt = p.dnsp_pkt;
278     const unsigned char *end = p.dnsp_end;
279     const unsigned char *dptr = rr.dnsrr_dptr;
280     const unsigned char *dend = rr.dnsrr_dend;
281     unsigned char *dn = rr.dnsrr_dn;
282     const unsigned char *tmp;
283
284     memset(buff, 0, sizeof(buff));
285
286     if (!dns_dnequal(dn, rr.dnsrr_dn)) continue;
287     if ((dlc->query_ctype == DNS_C_ANY || dlc->query_ctype == rr.dnsrr_cls) &&
288         (dlc->query_rtype == DNS_T_ANY || dlc->query_rtype == rr.dnsrr_typ)) {
289       lua_newtable(L);
290       lua_pushinteger(L, rr.dnsrr_ttl);
291       lua_setfield(L, -2, "ttl");
292
293       switch(rr.dnsrr_typ) {
294         case DNS_T_A:
295           if(rr.dnsrr_dsz == 4) {
296             snprintf(buff, sizeof(buff), "%d.%d.%d.%d",
297                      dptr[0], dptr[1], dptr[2], dptr[3]);
298             lua_pushstring(L, buff);
299             lua_setfield(L, -2, "a");
300           }
301           break;
302
303         case DNS_T_AAAA:
304           if(rr.dnsrr_dsz == 16) {
305             inet_ntop(AF_INET6, dptr, buff, 16);
306             lua_pushstring(L, buff);
307             lua_setfield(L, -2, "aaaa");
308           }
309           break;
310
311         case DNS_T_TXT:
312           totalsize = 0;
313           for(tmp = dptr; tmp < dend; totalsize += *tmp, tmp += *tmp + 1)
314             if(tmp + *tmp + 1 > dend) break;
315           /* worst case: every character escaped + '\0' */
316           txt_str = alloca(totalsize * 3 + 1);
317           if(!txt_str) break;
318           c = txt_str;
319           for(tmp = dptr; tmp < dend; tmp += *tmp + 1)
320             c = encode_txt(c, tmp+1, *tmp);
321           lua_pushstring(L, txt_str);
322           lua_setfield(L, -2, "txt");
323           break;
324
325         case DNS_T_MX:
326           lua_pushinteger(L, dns_get16(dptr));
327           lua_setfield(L, -2, "preference");
328           tmp = dptr + 2;
329           if(dns_getdn(pkt, &tmp, end, dn, DNS_MAXDN) <= 0 || tmp != dend)
330             break;
331           dns_dntop(dn, buff + strlen(buff), sizeof(buff) - strlen(buff));
332           lua_pushstring(L, buff);
333           lua_setfield(L, -2, "mx");
334           break;
335
336         case DNS_T_CNAME: if(!fieldname) fieldname = "cname";
337         case DNS_T_PTR: if(!fieldname) fieldname = "ptr";
338         case DNS_T_NS: if(!fieldname) fieldname = "ns";
339         case DNS_T_MB: if(!fieldname) fieldname = "mb";
340         case DNS_T_MD: if(!fieldname) fieldname = "md";
341         case DNS_T_MF: if(!fieldname) fieldname = "mf";
342         case DNS_T_MG: if(!fieldname) fieldname = "mg";
343         case DNS_T_MR: if(!fieldname) fieldname = "mr";
344          if(dns_getdn(pkt, &dptr, end, dn, DNS_MAXDN) <= 0) break;
345          dns_dntop(dn, buff, sizeof(buff));
346          lua_pushstring(L, buff);
347          lua_setfield(L, -2, fieldname);
348          break;
349
350         default:
351           break;
352       }
353       ++nrr;
354     }
355     else if (rr.dnsrr_typ == DNS_T_CNAME && !nrr) {
356       if (dns_getdn(pkt, &rr.dnsrr_dptr, end,
357                     p.dnsp_dnbuf, sizeof(p.dnsp_dnbuf)) <= 0 ||
358           rr.dnsrr_dptr != rr.dnsrr_dend) {
359         break;
360       }
361     }
362   }
363
364  cleanup:
365   if(result) free(result);
366   if(dlc->active) dlc->ci->lmc->resume(dlc->ci, nrr);
367   lookup_ctx_release(dlc);
368 }
369 static int dns_resume_event(eventer_t e, int mask, void *closure,
370                             struct timeval *now) {
371   dns_resume(closure);
372   return 0;
373 }
374
375 static void dns_cb(struct dns_ctx *ctx, void *result, void *data) {
376   eventer_t e;
377   int r = dns_status(ctx);
378   dns_lookup_ctx_t *dlc = data;
379   if(dlc->results) free(dlc->results);
380   dlc->results = NULL;
381   if(r >= 0) {
382     dlc->results_len = r;
383     dlc->results = malloc( r>0 ? r : r+1 );
384     memcpy(dlc->results, result, dlc->results_len);
385   }
386   if(pthread_equal(dlc->ci->bound_thread, pthread_self()))
387     return dns_resume(dlc);
388
389   /* We need to schedule this to run on another eventer thread */
390   e = eventer_alloc();
391   /* no whence, so epoch... or "right now!" */
392   e->thr_owner = dlc->ci->bound_thread;
393   e->closure = dlc;
394   e->callback = dns_resume_event;
395   e->mask = EVENTER_TIMER;
396   eventer_add(e);
397 }
398
399 static int noit_lua_dns_lookup(lua_State *L) {
400   dns_lookup_ctx_t *dlc, **holder;
401   const char *c, *query = "", *ctype = "IN", *rtype = "A";
402   char *ctype_up, *rtype_up, *d;
403   void *vnv_pair;
404   noit_lua_resume_info_t *ci;
405
406   ci = noit_lua_get_resume_info(L);
407   assert(ci);
408
409   holder = (dns_lookup_ctx_t **)lua_touserdata(L, lua_upvalueindex(1));
410   if(holder != lua_touserdata(L,1))
411     luaL_error(L, "Must be called as method\n");
412   dlc = *holder;
413
414   if(lua_gettop(L) > 1) query = lua_tostring(L, 2);
415   if(lua_gettop(L) > 2) rtype = lua_tostring(L, 3);
416   if(lua_gettop(L) > 3) ctype = lua_tostring(L, 4);
417
418   if(query == NULL || rtype == NULL || ctype == NULL) {
419     lua_pushnil(L);
420     return 1;
421   }
422
423   ctype_up = alloca(strlen(ctype)+1);
424   for(d = ctype_up, c = ctype; *c; d++, c++) *d = toupper(*c);
425   *d = '\0';
426   rtype_up = alloca(strlen(rtype)+1);
427   for(d = rtype_up, c = rtype; *c; d++, c++) *d = toupper(*c);
428   *d = '\0';
429
430   if(!noit_hash_retrieve(&dns_ctypes, ctype_up, strlen(ctype_up), &vnv_pair))
431     dlc->error = strdup("bad class");
432   else
433     dlc->query_ctype = (enum dns_class)((struct dns_nameval *)vnv_pair)->val;
434
435   if(!noit_hash_retrieve(&dns_rtypes, rtype_up, strlen(rtype_up), &vnv_pair))
436     dlc->error = strdup("bad rr type");
437   else
438     dlc->query_rtype = (enum dns_type)((struct dns_nameval *)vnv_pair)->val;
439
440   dlc->active = 1;
441   noit_atomic_inc32(&dlc->refcnt);
442   if(!dlc->error) {
443     int abs;
444     if(!dns_ptodn(query, strlen(query), dlc->dn, sizeof(dlc->dn), &abs) ||
445        !dns_submit_dn(dlc->h->ctx, dlc->dn, dlc->query_ctype, dlc->query_rtype,
446                       abs | DNS_NOSRCH, NULL, dns_cb, dlc)) {
447       dlc->error = strdup("submission error");
448       noit_atomic_dec32(&dlc->refcnt);
449     }
450     else {
451       struct timeval now;
452       gettimeofday(&now, NULL);
453       dns_timeouts(dlc->h->ctx, -1, now.tv_sec);
454     }
455   }
456   if(dlc->error) {
457     dlc->active = 0;
458     luaL_error(L, "dns: %s\n", dlc->error);
459   }
460   return noit_lua_yield(ci, 0);
461 }
462
463 int noit_lua_dns_gc(lua_State *L) {
464   dns_lookup_ctx_t **holder;
465   holder = (dns_lookup_ctx_t **)lua_touserdata(L,1);
466   (*holder)->active = 0;
467   lookup_ctx_release(*holder);
468   return 0;
469 }
470
471 int noit_lua_dns_index_func(lua_State *L) {
472   int n;
473   const char *k;
474   dns_lookup_ctx_t **udata;
475
476   n = lua_gettop(L);
477   assert(n == 2);
478   if(!luaL_checkudata(L, 1, "noit.dns"))
479     luaL_error(L, "metatable error, arg1 is not a noit.dns");
480   udata = lua_touserdata(L, 1);
481   if(!lua_isstring(L, 2))
482     luaL_error(L, "metatable error, arg2 is not a string");
483   k = lua_tostring(L, 2);
484   if(!strcmp(k, "lookup")) {
485     lua_pushlightuserdata(L, udata);
486     lua_pushcclosure(L, noit_lua_dns_lookup, 1);
487     return 1;
488   }
489   luaL_error(L, "noit.dns no such element: %s", k);
490   return 0;
491 }
492
493 void noit_lua_init_dns() {
494   int i;
495   const struct dns_nameval *nv;
496   struct dns_ctx *pctx;
497
498   /* HASH the rr types */
499   for(i=0, nv = dns_type_index(i); nv->name; nv = dns_type_index(++i))
500     noit_hash_store(&dns_rtypes,
501                     nv->name, strlen(nv->name),
502                     (void *)nv);
503   /* HASH the class types */
504   for(i=0, nv = dns_class_index(i); nv->name; nv = dns_class_index(++i))
505     noit_hash_store(&dns_ctypes,
506                     nv->name, strlen(nv->name),
507                     (void *)nv);
508
509   eventer_name_callback("lua/dns_eventer", noit_lua_dns_eventer);
510   eventer_name_callback("lua/dns_timeouts", noit_lua_dns_timeouts);
511
512   if (dns_init(NULL, 0) < 0 || (pctx = dns_new(NULL)) == NULL) {
513     noitL(noit_error, "Unable to initialize dns subsystem\n");
514   }
515   else
516     dns_free(pctx);
517 }
Note: See TracBrowser for help on using the browser.