daemon/worker: estimate pktbuf sizes for tasks

diff --git a/daemon/worker.c b/daemon/worker.c
index a4d1b9de01b09796a60edc44408aa703ab62e6c6..d5aa0e0d36b2caad41009216249b7cc5dc9c9209 100644 (file)
--- a/daemon/worker.c
+++ b/daemon/worker.c
@@ -61,7 +61,7 @@ static int parse_query(knot_pkt_t *query)
        return kr_ok();
 }
 
-static struct qr_task *qr_task_create(struct worker_ctx *worker, uv_handle_t *handle, const struct sockaddr *addr)
+static struct qr_task *qr_task_create(struct worker_ctx *worker, uv_handle_t *handle, knot_pkt_t *query, const struct sockaddr *addr)
 {
        mm_ctx_t pool;
        mm_ctx_mempool(&pool, MM_DEFAULT_BLKSIZE);
@@ -80,15 +80,28 @@ static struct qr_task *qr_task_create(struct worker_ctx *worker, uv_handle_t *ha
                memcpy(&task->source.addr, addr, sockaddr_len(addr));
        }
 
+       /* How much can client handle? */
+       size_t answer_max = KNOT_WIRE_MIN_PKTSIZE;
+       if (!addr) { /* TCP */
+               answer_max = KNOT_WIRE_MAX_PKTSIZE;
+       } else if (knot_pkt_has_edns(query)) { /* EDNS */
+               answer_max = knot_edns_get_payload(query->opt_rr);
+       }
+       /* How much space do we need for intermediate packets? */
+       size_t pktbuf_max = KNOT_EDNS_MAX_UDP_PAYLOAD;
+       if (pktbuf_max < answer_max) {
+               pktbuf_max = answer_max;
+       }
+
        /* Create buffers */
-       knot_pkt_t *next_query = knot_pkt_new(NULL, KNOT_EDNS_MAX_UDP_PAYLOAD, &task->req.pool);
-       knot_pkt_t *answer = knot_pkt_new(NULL, KNOT_WIRE_MAX_PKTSIZE, &task->req.pool);
-       if (!next_query || !answer) {
+       knot_pkt_t *pktbuf = knot_pkt_new(NULL, pktbuf_max, &task->req.pool);
+       knot_pkt_t *answer = knot_pkt_new(NULL, answer_max, &task->req.pool);
+       if (!pktbuf || !answer) {
                mp_delete(pool.ctx);
                return NULL;
        }
        task->req.answer = answer;
-       task->next_query = next_query;
+       task->next_query = pktbuf;
 
        /* Start resolution */
        uv_timer_init(handle->loop, &task->timeout);
@@ -240,7 +253,7 @@ int worker_exec(struct worker_ctx *worker, uv_handle_t *handle, knot_pkt_t *quer
                if (ret != 0 || knot_wire_get_qr(query->wire)) {
                        return kr_error(EINVAL); /* Ignore. */
                }
-               task = qr_task_create(worker, handle, addr);
+               task = qr_task_create(worker, handle, query, addr);
                if (!task) {
                        return kr_error(ENOMEM);
                }